perf(moe): iq2_s c32 gathers run 16 k-lanes per row on m5 and later

CHANGELOG.md

@@ -11,6 +11,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   `KQ_QMV_SPLITK=0` restores the previous kernels.
 - The iq2_s fused MoE gate/up and down gathers at one activation row take less
   time on every GPU. `KQ_MOE_C32=0` restores the previous kernels.
+- On M5 GPUs the iq2_s MoE gathers at one activation row take a few percent
+  less time again. `KQ_MOE_NX=8` restores the earlier width.
 
 ## [0.5.1]
 

README.md

@@ -305,8 +305,9 @@ All optional; the defaults are right for normal use.
   A/B lever.
 - `KQ_MOE_C32` - on Metal, the 32-weight decode step of the iq2_s fused MoE gate/up and down
   gathers. By default it runs at one activation row on every GPU. `0` restores the 16-weight
-  kernels, and `1` also runs it at wider rows where the dedupe kernels do not apply. `KQ_MOE_NX=16`
-  selects its 16-lane form, and `KQ_MOE_NX=32` keeps the 16-weight kernels. Read per call. The
+  kernels, and `1` also runs it at wider rows where the dedupe kernels do not apply. It runs 16
+  K-lanes per row on M5 and later GPUs and 8 on earlier ones. `KQ_MOE_NX=8` or `16` forces a
+  width, and `KQ_MOE_NX=32` keeps the 16-weight kernels. Read per call. The
   outputs differ from the 16-weight kernels only in summation order. A/B lever.
 
 - `KQ_SDPA_IDX_NAX=0` - run `sdpa_fa_indexed` on its simdgroup kernel even where the NAX tile

bindings.cpp

@@ -127,6 +127,13 @@ NB_MODULE(_ext, m) {
       "codec, rows of x, N and K, as [min_m, min_k, max_k, min_n, max_n]; "
       "empty where the qmm_t kernel runs.");
 
+  m.def(
+      "_moe_c32_nx",
+      &mlx_kquant::moe_c32_nx,
+      "arch"_a,
+      "The K-lanes per output row (8 or 16) of the iq2_s c32 MoE gathers "
+      "for an architecture string, unless KQ_MOE_NX forces one.");
+
   m.def(
       "_dequantize_route_takes",
       &mlx_kquant::dequantize_route_takes,

docs/kernels.md

@@ -1041,8 +1041,9 @@ term; requires the expert biases and is instantiated for `mxfp4`/`nvfp4` only).
 At one activation row the iq2_s gate/up and down gathers decode 32 weights per lane step, one
 32-weight group of a block, where the other codecs decode 16. The block scale, the high-bit byte
 and the scale byte then load once per 32 weights, and the grid indices and sign bytes load as
-16-bit pairs. Each lane covers one 32-weight group per step at 8 K-lanes per output row, or 16
-under `KQ_MOE_NX=16`. Outputs differ from the 16-weight kernels only in summation order.
+16-bit pairs. Each lane covers one 32-weight group per step, with 16 K-lanes per output row on
+M5 and later GPUs and 8 on earlier ones. `KQ_MOE_NX=8` or `16` forces a width. Outputs differ from
+the 16-weight kernels only in summation order.
 `KQ_MOE_C32=0` restores the 16-weight kernels. `KQ_MOE_C32=1` also runs the 32-weight kernels at
 wider rows where the dedupe kernels do not apply. A forced `KQ_MOE_NX=32` keeps the 16-weight
 kernels.
@@ -1052,6 +1053,9 @@ On M3 Max at E=256, N=K=4096 and top 8, back to back, the gate/up gather took 0.
 step of eight Muse Spark UD-IQ2_S layers, the MoE blocks took 8.3% less time. With the `q8_0`
 split-K kernel, the step went from 1.414 to 1.335 ms per layer (-5.6%).
 
+On M5 Max the 16-lane form is faster. Back to back at the same shapes, a MoE block replica with
+both gathers at 16 K-lanes took 2.3% to 5.2% less time than at 8, in both arm orders.
+
 The grid lookup is the next lever. With the grid read removed (wrong output), the 32-weight
 gathers took about 25% less time, and 37% less than the 16-weight gathers. That bounds the gain of
 a cheaper lookup. Two other grid sources gave no gain. A 64-bit grid in a device buffer matched the

src/kquant.h

@@ -84,6 +84,11 @@ bool dequantize_route_takes(
     int N,
     int K);
 
+// The K-lanes per output row (8 or 16) of the iq2_s c32 MoE gathers on a GPU
+// with this architecture string (as get_architecture gives it), when
+// KQ_MOE_NX does not force one.
+int moe_c32_nx(const std::string& arch);
+
 // Every kernel name the split-K dispatchers can build (qmm_splitk,
 // qmm_nax_splitk, verify_mma, verify_nax and their fold), for the kernel
 // name tests; empty on non-Metal builds.

src/kquant_moe_glu.cpp

@@ -5,6 +5,7 @@
 //   K-quant wire bytes: moe_glu_gather_kq (gate + up + GLU epilogue;
 //     optional biases with act "swiglu_clamp", mxfp4/nvfp4 only) and
 //     gather_qmv_kq, per-codec kernels (full codec matrix).
+#include <cctype>
 #include <cstdio>
 #include <cstdlib>
 #include <set>
@@ -358,8 +359,19 @@ inline bool kq_moe_c32(const std::string& codec, int T) {
   return v == -1 ? T == 1 : v != 0;
 }
 
+// A forced KQ_MOE_NX of 8 or 16, else this GPU's class width (moe_c32_nx).
 inline int kq_moe_c32_nx() {
-  return kq_moe_nx_env() == 16 ? 16 : 8;
+  const int v = kq_moe_nx_env();
+  if (v == 8 || v == 16) {
+    return v;
+  }
+#ifdef _METAL_
+  static const int here =
+      moe_c32_nx(mx::metal::device(mx::Device::gpu).get_architecture());
+  return here;
+#else
+  return 8;
+#endif
 }
 
 inline bool kq_gather_fine(int64_t coarse_tgs) {
@@ -1317,6 +1329,26 @@ std::vector<mx::Shape> KQuantGatherQMVMixBias::output_shapes(
   return {{inputs[4].shape(0), inputs[0].shape(1)}};
 }
 
+// applegpu_g<gen><suffix>, two digit gen, a lower case suffix other than p:
+// NX 16 from gen 17 (M5), where it beat NX 8 by 2 to 6% on a Muse Spark MoE
+// block (M5 Max, E=256, top 8, N=K=4096). NX 8 elsewhere.
+int moe_c32_nx(const std::string& arch) {
+  const std::string head = "applegpu_g";
+  if (arch.size() != head.size() + 3 || arch.compare(0, head.size(), head)) {
+    return 8;
+  }
+  const char d0 = arch[head.size()];
+  const char d1 = arch[head.size() + 1];
+  const char suffix = arch[head.size() + 2];
+  if (!std::isdigit(static_cast<unsigned char>(d0)) ||
+      !std::isdigit(static_cast<unsigned char>(d1)) ||
+      !std::islower(static_cast<unsigned char>(suffix)) || suffix == 'p') {
+    return 8;
+  }
+  const int gen = (d0 - '0') * 10 + (d1 - '0');
+  return gen >= 17 ? 16 : 8;
+}
+
 // Codecs with the fused kq GLU/gather kernels wired (kq_moe_glu_kq.h):
 // tuned kernels for q6_k/q8_0, generic Ext-trait kernels for the rest of the
 // codec matrix. Unsupported codecs must fall back to the stock path

tests/test_moe_c32_class.py

@@ -0,0 +1,31 @@
+"""The K-lanes per row of the iq2_s c32 MoE gathers by GPU class.
+
+``_moe_c32_nx`` answers for any architecture string, so these cases run on
+every machine, CPU builds included.
+"""
+
+import pytest
+from mlx_kquant._ext import _moe_c32_nx
+
+WIDE = [f"applegpu_g{g}{c}" for g in (17, 18, 19) for c in "gsdc"]
+NARROW = [f"applegpu_g{g}{c}" for g in (13, 14, 15, 16) for c in "gsdc"] + [
+    "applegpu_g17p",
+    "applegpu_g18p",
+    "applegpu_g12s",
+    "applegpu_g17",
+    "applegpu_g17S",
+    "applegpu_g1xs",
+    "applegpu_g170s",
+    "gpu_g17s",
+    "",
+]
+
+
+@pytest.mark.parametrize("arch", WIDE)
+def test_m5_and_later_take_nx16(arch):
+    assert _moe_c32_nx(arch) == 16
+
+
+@pytest.mark.parametrize("arch", NARROW)
+def test_earlier_phones_and_malformed_keep_nx8(arch):
+    assert _moe_c32_nx(arch) == 8

tests/test_moe_glu_variants.py

@@ -552,24 +552,37 @@ def _c32_run(T, dtype, env, tmp_path, arm):
     return np.load(f), r.stderr
 
 
+def _c32_class_suffix():
+    import mlx.core as mx
+    from mlx_kquant._ext import _moe_c32_nx
+
+    nx = _moe_c32_nx(mx.device_info().get("architecture", ""))
+    return "" if nx == 8 else f"_nx{nx}"
+
+
 @pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
 def test_c32_dispatch(dtype, tmp_path):
     """iq2_s gate/up and down gathers take the 32-weight-step kernels at one
-    row by default, at NX = 8 or a forced 16. KQ_MOE_C32=0 and a forced
-    KQ_MOE_NX=32 keep the 16-weight kernels. Two rows keep them too unless
-    KQ_MOE_C32=1. The forms differ only in summation order."""
-    want_glu = f"kq_iq2_s_moe_glu_gather_silu_c32_{dtype}_t"
-    want_qmv = f"kq_iq2_s_gather_qmv_c32_{dtype}_t"
+    row by default, at this GPU class's NX (8, or 16 from M5), or a forced 8
+    or 16. KQ_MOE_C32=0 and a forced KQ_MOE_NX=32 keep the 16-weight kernels.
+    Two rows keep them too unless KQ_MOE_C32=1. The forms differ only in
+    summation order."""
+    sfx = _c32_class_suffix()
+    want_glu = f"kq_iq2_s_moe_glu_gather_silu_c32{sfx}_{dtype}_t"
+    want_qmv = f"kq_iq2_s_gather_qmv_c32{sfx}_{dtype}_t"
     new, names = _c32_run(1, dtype, {}, tmp_path, "default")
     assert want_glu in names and want_qmv in names, names
     old, names = _c32_run(1, dtype, {"KQ_MOE_C32": "0"}, tmp_path, "old")
     assert "_c32" not in names, names
+    narrow, names = _c32_run(1, dtype, {"KQ_MOE_NX": "8"}, tmp_path, "nx8")
+    assert f"kq_iq2_s_moe_glu_gather_silu_c32_{dtype}_t" in names, names
+    assert f"kq_iq2_s_gather_qmv_c32_{dtype}_t" in names, names
     wide, names = _c32_run(1, dtype, {"KQ_MOE_NX": "16"}, tmp_path, "nx16")
     assert "kq_iq2_s_moe_glu_gather_silu_c32_nx16_" in names, names
     assert "kq_iq2_s_gather_qmv_c32_nx16_" in names, names
     for key in ("glu", "qmv"):
         scale = np.abs(old[key]).max()
-        for got in (new, wide):
+        for got in (new, narrow, wide):
             rel = np.abs(got[key] - old[key]).max() / scale
             assert rel < 2e-3, (key, rel)
     _, names = _c32_run(1, dtype, {"KQ_MOE_NX": "32"}, tmp_path, "nx32")
@@ -610,8 +623,9 @@ def test_c32_reads_live():
         {"KQ_MOE_NX_LOG": "1"},
         "c32 live read",
     )
+    sfx = _c32_class_suffix()
     assert "kq_iq2_s_gather_qmv_bfloat16_t" in r.stderr, r.stderr
-    assert "kq_iq2_s_gather_qmv_c32_bfloat16_t" in r.stderr, r.stderr
+    assert f"kq_iq2_s_gather_qmv_c32{sfx}_bfloat16_t" in r.stderr, r.stderr
 
 
 _C32_REF_SNIPPET = r"""