perf(kernels): faster q8_0 mat-vec and iq2_s MoE gathers

CHANGELOG.md

@@ -6,6 +6,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Changed
+- `q8_0` `quantized_matmul` at one activation row takes less time on every GPU.
+  `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.
+
 ## [0.5.1]
 
 ### Added

README.md

@@ -303,6 +303,11 @@ All optional; the defaults are right for normal use.
   shared-expert slots stay on the float path). Read per call. Default off: the outputs differ from
   the float kernels at half rounding level, and an activation past 65504 in magnitude turns to inf.
   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
+  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
   kernel is available. Read once per process. The simdgroup kernel matches `sdpa_fa_verify` over

docs/kernels.md

@@ -207,6 +207,15 @@ and one FMA. The row width is a Metal function constant, so each width gets its
 the superblock count and row stride folded. Under sustained load that ran about 1% faster than
 reading the width at run time.
 
+A `q8_0` `quantized_matmul` with one activation row and no batch dimensions runs a split-K kernel
+too, when N is even and 256 divides K. The eight simdgroups of a threadgroup share two output rows.
+They stride the 256-weight steps between them and add their partial sums in threadgroup memory. A
+lane keeps the per-row kernels' share of a step, eight weights of one block. It reads them as four
+16-bit words, since a quarter block sits 2 bytes past a 2-byte boundary. The split puts more
+simdgroups in flight at layer-size N. On M3 Max, back to back over ten shapes, it took 1.6% to 6.6%
+less time than the per-row kernels: 2.7% less at K 8192 and N 4096, and 3.2% less at K 4096 and N
+8192. In a synced decode step of eight Muse Spark UD-IQ2_S layers, it saved 1.1% per layer.
+
 Under sustained decode the GPU runs below its top clock, where the per-row kernels turn ALU-bound.
 Measured on M5 Max over 90 seconds of back-to-back mat-vecs at the gemma-4 31B MLP shapes, the
 split-K kernel read 450-465 GB/s at the lowest clock, against 385-395 GB/s for the per-row kernels
@@ -339,10 +348,10 @@ Tuning levers (defaults are right for normal use):
 - `KQ_QMV_FINE` - tiling of the per-row mat-vec. `1` forces two output rows per threadgroup, `0`
   forces the coarse tiling, and unset takes the fine tiling at one activation row up to the
   per-codec N ceiling that `kq.qmv_fine_max_n(codec)` reports. Read live per call.
-- `KQ_QMV_SPLITK` - the `q6_k` split-K kernel for one activation row. `0` restores the per-row
-  kernels, as does a set `KQ_QMV_FINE`, so a tiling A/B keeps measuring the tiling it names. A
-  lever that forces another route at one row, such as `KQ_VERIFY_NAX=1`, takes precedence. Read
-  live per call.
+- `KQ_QMV_SPLITK` - the `q6_k` and `q8_0` split-K kernels for one activation row. `0` restores
+  the per-row kernels, as does a set `KQ_QMV_FINE`, so a tiling A/B keeps measuring the tiling it
+  names. A lever that forces another route at one row, such as `KQ_VERIFY_NAX=1`, takes
+  precedence. Read live per call.
 - `KQ_VERIFY_EXT` - the mat-vec route on the M 2-12 band. `1` forces `mv_ext` for every codec with
   the kernel, `0` forces `verify_qmv` where the codec has it and per-row qmv elsewhere, and unset
   takes the per-codec default. On NAX GPUs a set value also keeps per-row qmv and NAX split-K off
@@ -1029,6 +1038,26 @@ The GLU activation is selected per model: plain SwiGLU/GELU, the clamped `silu_l
 `swiglu_clamp` (gpt-oss clamped SwiGLU: biases added, sigmoid slope `alpha`, and a `(u + 1)` linear
 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.
+`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.
+
+On M3 Max at E=256, N=K=4096 and top 8, back to back, the gate/up gather took 0.303 ms against
+0.340 ms (-11.1%), and the down gather took 0.166 ms against 0.191 ms (-12.5%). In a synced decode
+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%).
+
+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
+constant table. A 2-bit code table, in threadgroup memory or in a device buffer, was about 20%
+slower.
+
 ## Attention
 
 Scaled-dot-product variants for shapes stock MLX's fused allowlist excludes, plus the sparse

metal/kq_moe_glu_kq.metal

@@ -314,6 +314,29 @@ instantiate_kq_ext_sx2_lowbit(iq1_m, KqIq1_mExt)
 instantiate_kq_ext_sx_mix(q2_0, KqQ2_0Ext, iq4_nl, KqIq4_nlExt, bfloat16_t)
 instantiate_kq_ext_sx_mix(q2_0, KqQ2_0Ext, iq4_nl, KqIq4_nlExt, float16_t)
 
+// 32-weight decode step (KqExtC32) for the uniform gathers at decode.
+#define instantiate_kq_ext_c32_nx(codec, traits, type, nx, sfx)               \
+  instantiate_kernel(                                                         \
+      "kq_" #codec "_moe_glu_gather_silu_c32" sfx "_" #type,                  \
+      kq_ext_moe_glu_gather, type, traits, KQ_GLU_ACT_SILU, nx, true)          \
+  instantiate_kernel(                                                         \
+      "kq_" #codec "_moe_glu_gather_gelu_c32" sfx "_" #type,                  \
+      kq_ext_moe_glu_gather, type, traits, KQ_GLU_ACT_GELU, nx, true)          \
+  instantiate_kernel(                                                         \
+      "kq_" #codec "_moe_glu_gather_silu_limit_c32" sfx "_" #type,            \
+      kq_ext_moe_glu_gather, type, traits, KQ_GLU_ACT_SILU_LIMIT, nx, true)    \
+  instantiate_kernel(                                                         \
+      "kq_" #codec "_gather_qmv_c32" sfx "_" #type,                           \
+      kq_ext_gather_qmv, type, traits, nx, true)
+
+#define instantiate_kq_ext_c32(codec, traits)                                 \
+  instantiate_kq_ext_c32_nx(codec, traits, bfloat16_t, 8, "")                  \
+  instantiate_kq_ext_c32_nx(codec, traits, bfloat16_t, 16, "_nx16")            \
+  instantiate_kq_ext_c32_nx(codec, traits, float16_t, 8, "")                   \
+  instantiate_kq_ext_c32_nx(codec, traits, float16_t, 16, "_nx16")
+
+instantiate_kq_ext_c32(iq2_s, KqIq2_sExt)
+
 // q6_k/q8_0 mix_ns (generic; no tuned ns kernels exist).
 instantiate_kq_ext_mix_ns(q6_k, KqQ6_KExt)
 instantiate_kq_ext_mix_ns(q8_0, KqQ8_0Ext)

metal/kq_quantized.metal

@@ -158,6 +158,7 @@
   instantiate_kquant_batched(qmv_fast, type, 32, 8, 0, q8_0)            \
   instantiate_kquant_batched(qmv_fast, type, 32, 8, 1, q8_0)            \
   instantiate_kquant_batched(qmv_fast_fine, type, 32, 8, 0, q8_0)       \
+  instantiate_kquant_batched(qmv_splitk, type, 32, 8, 0, q8_0)          \
   instantiate_kquant_batched(qmv_fine, type, 32, 8, 0, q8_0)            \
   instantiate_kquant_batched(qmv,      type, 32, 8, 0, q8_0)            \
   instantiate_kquant_batched(qmv,      type, 32, 8, 1, q8_0)            \

metal/mlx/backend/metal/kernels/kq_moe_glu_kq.h

@@ -780,11 +780,32 @@ template <typename T, int results_per_simdgroup = 4>
 
 #define KQ_EXT_NXPSG 8
 
+// Dot of a 32-weight group (l0, l1 unscaled, sc their two scales) with
+// x[0..31].
+template <typename T>
+METAL_FUNC float kq_ext_dot32(
+    const thread float4x4& l0,
+    const thread float4x4& l1,
+    float2 sc,
+    const device T* xp) {
+  return sc.x *
+      (dot(l0[0], float4(*(const device vec<T, 4>*)(xp + 0))) +
+       dot(l0[1], float4(*(const device vec<T, 4>*)(xp + 4))) +
+       dot(l0[2], float4(*(const device vec<T, 4>*)(xp + 8))) +
+       dot(l0[3], float4(*(const device vec<T, 4>*)(xp + 12)))) +
+      sc.y *
+      (dot(l1[0], float4(*(const device vec<T, 4>*)(xp + 16))) +
+       dot(l1[1], float4(*(const device vec<T, 4>*)(xp + 20))) +
+       dot(l1[2], float4(*(const device vec<T, 4>*)(xp + 24))) +
+       dot(l1[3], float4(*(const device vec<T, 4>*)(xp + 28))));
+}
+
 // Partial dot of x (one activation row) against wire-byte weight row `row`,
 // this thread's K-stripe only (chunks tx, tx + NX, ...). `luts` is the
 // codec's threadgroup-staged table block (KqTgLuts; unused for table-free
-// codecs).
-template <typename T, typename Codec, int NX>
+// codecs). With C32 a lane steps one 32-weight group at a time through
+// KqExtC32 (the "_c32" kernels).
+template <typename T, typename Codec, int NX, bool C32 = false>
 METAL_FUNC float kq_ext_row_partial(
     const device uint8_t* w,
     const device T* x,
@@ -796,6 +817,19 @@ METAL_FUNC float kq_ext_row_partial(
   const int nb = K / Codec::superblock;
   const device uint8_t* w_row = w + row * (int64_t)nb * Codec::block_bytes;
   float acc = 0.0f;
+  if (C32 && KqExtC32<Codec>::ok) {
+    constexpr short gpb = Codec::superblock / 32;
+    for (int ig = tx; 32 * ig < K; ig += NX) {
+      float4x4 l0, l1;
+      const float2 sc = KqExtC32<Codec>::deq_group32s(
+          w_row + (int64_t)(ig / gpb) * Codec::block_bytes,
+          short(ig % gpb),
+          l0,
+          l1);
+      acc += kq_ext_dot32(l0, l1, sc, x + ig * 32);
+    }
+    return acc;
+  }
   for (int ich = tx; 16 * ich < K; ich += NX) {
     const device uint8_t* block =
         w_row + (int64_t)(ich / chpb) * Codec::block_bytes;
@@ -813,7 +847,7 @@ METAL_FUNC float kq_ext_row_partial(
 }
 
 // GLU pair variant: gate and up rows share each activation chunk load.
-template <typename T, typename Codec, int NX>
+template <typename T, typename Codec, int NX, bool C32 = false>
 METAL_FUNC float2 kq_ext_glu_row_partial(
     const device uint8_t* gw,
     const device uint8_t* uw,
@@ -828,6 +862,35 @@ METAL_FUNC float2 kq_ext_glu_row_partial(
   const device uint8_t* g_row = gw + row_off;
   const device uint8_t* u_row = uw + row_off;
   float2 acc = float2(0.0f);
+  if (C32 && KqExtC32<Codec>::ok) {
+    constexpr short gpb = Codec::superblock / 32;
+    for (int ig = tx; 32 * ig < K; ig += NX) {
+      const int64_t boff = (int64_t)(ig / gpb) * Codec::block_bytes;
+      const short g8 = short(ig % gpb);
+      const device T* xp = x + ig * 32;
+      float4 a[8];
+#pragma unroll
+      for (int i = 0; i < 8; ++i) {
+        a[i] = float4(*(const device vec<T, 4>*)(xp + 4 * i));
+      }
+      float4x4 l0, l1;
+      float2 sc = KqExtC32<Codec>::deq_group32s(g_row + boff, g8, l0, l1);
+      acc.x += sc.x *
+              (dot(l0[0], a[0]) + dot(l0[1], a[1]) + dot(l0[2], a[2]) +
+               dot(l0[3], a[3])) +
+          sc.y *
+              (dot(l1[0], a[4]) + dot(l1[1], a[5]) + dot(l1[2], a[6]) +
+               dot(l1[3], a[7]));
+      sc = KqExtC32<Codec>::deq_group32s(u_row + boff, g8, l0, l1);
+      acc.y += sc.x *
+              (dot(l0[0], a[0]) + dot(l0[1], a[1]) + dot(l0[2], a[2]) +
+               dot(l0[3], a[3])) +
+          sc.y *
+              (dot(l1[0], a[4]) + dot(l1[1], a[5]) + dot(l1[2], a[6]) +
+               dot(l1[3], a[7]));
+    }
+    return acc;
+  }
   for (int ich = tx; 16 * ich < K; ich += NX) {
     const int64_t boff = (int64_t)(ich / chpb) * Codec::block_bytes;
     const short cch = short(ich % chpb);
@@ -873,7 +936,12 @@ METAL_FUNC float kq_ext_reduce(float v) {
   return v;
 }
 
-template <typename T, typename Codec, int ACT, int NX = KQ_EXT_NXPSG>
+template <
+    typename T,
+    typename Codec,
+    int ACT,
+    int NX = KQ_EXT_NXPSG,
+    bool C32 = false>
 [[kernel]] void kq_ext_moe_glu_gather(
     const device uint8_t* gw [[buffer(0)]],
     const device uint8_t* uw [[buffer(1)]],
@@ -899,7 +967,7 @@ template <typename T, typename Codec, int ACT, int NX = KQ_EXT_NXPSG>
   out += ((int64_t)tid.z * R + tid.y) * N;
 
   KQ_EXT_STAGE_LUTS(Codec, kq_luts)
-  const float2 gu = kq_ext_glu_row_partial<T, Codec, NX>(
+  const float2 gu = kq_ext_glu_row_partial<T, Codec, NX, C32>(
       gw, uw, x, (int64_t)expert * N + out_row, K, tx, kq_luts);
   const float g = kq_ext_reduce<NX>(gu.x);
   const float u = kq_ext_reduce<NX>(gu.y);
@@ -997,7 +1065,7 @@ template <typename T, typename Codec, int NX = KQ_EXT_NXPSG>
   }
 }
 
-template <typename T, typename Codec, int NX = KQ_EXT_NXPSG>
+template <typename T, typename Codec, int NX = KQ_EXT_NXPSG, bool C32 = false>
 [[kernel]] void kq_ext_gather_qmv(
     const device uint8_t* w [[buffer(0)]],
     const device T* x [[buffer(1)]],
@@ -1022,7 +1090,7 @@ template <typename T, typename Codec, int NX = KQ_EXT_NXPSG>
   out += row_idx * N;
 
   KQ_EXT_STAGE_LUTS(Codec, kq_luts)
-  const float r = kq_ext_reduce<NX>(kq_ext_row_partial<T, Codec, NX>(
+  const float r = kq_ext_reduce<NX>(kq_ext_row_partial<T, Codec, NX, C32>(
       w, x, (int64_t)expert * N + out_row, K, tx, kq_luts));
   if (tx == 0) {
     out[out_row] = static_cast<T>(r);

metal/mlx/backend/metal/kernels/kq_quantized.h

@@ -1711,6 +1711,104 @@ template <typename T, int group_size, int bits, bool batched>
       w, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
 }
 
+// Split-K M=1 mat-vec. The nsg simdgroups of a threadgroup share rps rows
+// and stride the 256-weight steps by nsg, so the grid keeps N / rps
+// threadgroups while each simdgroup walks only K / nsg of a row. Lane l
+// owns 8 weights (block l / 4, quarter l % 4) of every 8-block step, the
+// lane mapping of kq_q8_0_qmv_fast_impl. The quarter sits 2 bytes past a
+// 2-byte boundary, so it loads as four ushorts. Partial sums reduce through
+// `red` (nsg * rps floats). Needs N % rps == 0 and K % 256 == 0.
+template <typename T, int rps, int nsg>
+METAL_FUNC void kq_q8_0_qmv_splitk_impl(
+    const device uint8_t* w,
+    const device T* x,
+    device T* y,
+    const int in_vec_size,
+    threadgroup float* red,
+    uint3 tid,
+    uint simd_gid,
+    uint simd_lid) {
+  typedef float U;
+  constexpr int step = 8 * KQ_Q8_0_GROUP;
+  thread U result[rps] = {0};
+
+  const int row_bytes = in_vec_size * KQ_Q8_0_BLOCK_BYTES / KQ_Q8_0_GROUP;
+  const int out_row = tid.y * rps;
+  const int lane_k = simd_lid * 8;
+  const int quarter = simd_lid % 4;
+  const device uint8_t* w0 = w + static_cast<int64_t>(out_row) * row_bytes +
+      (simd_lid / 4) * KQ_Q8_0_BLOCK_BYTES;
+
+  for (int k = int(simd_gid) * step; k < in_vec_size; k += nsg * step) {
+    const vec<T, 4> xa = *(const device vec<T, 4>*)(x + k + lane_k);
+    const vec<T, 4> xb = *(const device vec<T, 4>*)(x + k + lane_k + 4);
+    const U x0 = U(xa[0]), x1 = U(xa[1]), x2 = U(xa[2]), x3 = U(xa[3]);
+    const U x4 = U(xb[0]), x5 = U(xb[1]), x6 = U(xb[2]), x7 = U(xb[3]);
+    const int boff = (k / KQ_Q8_0_GROUP) * KQ_Q8_0_BLOCK_BYTES;
+#pragma unroll
+    for (int row = 0; row < rps; row++) {
+      const device uint8_t* b = w0 + row * row_bytes + boff;
+      const U d = U(*(const device half*)b);
+      const device ushort* q =
+          (const device ushort*)(b + KQ_Q8_0_Q_OFFSET + quarter * 8);
+      const float2 q0 = float2(as_type<char2>(q[0]));
+      const float2 q1 = float2(as_type<char2>(q[1]));
+      const float2 q2 = float2(as_type<char2>(q[2]));
+      const float2 q3 = float2(as_type<char2>(q[3]));
+      const U p = q0.x * x0 + q0.y * x1 + q1.x * x2 + q1.y * x3 + q2.x * x4 +
+          q2.y * x5 + q3.x * x6 + q3.y * x7;
+      result[row] += d * p;
+    }
+  }
+
+#pragma unroll
+  for (int row = 0; row < rps; row++) {
+    const U s = simd_sum(result[row]);
+    if (simd_lid == 0) {
+      red[simd_gid * rps + row] = s;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_gid == 0 && simd_lid < uint(rps)) {
+    U t = U(0);
+    for (int g = 0; g < nsg; g++) {
+      t += red[g * rps + simd_lid];
+    }
+    y[out_row + simd_lid] = static_cast<T>(t);
+  }
+}
+
+#define KQ_Q8_0_DEFINE_QMV_SPLITK(name, RPS, NSG)                           \
+  template <typename T, int group_size, int bits, bool batched>             \
+  [[kernel, max_total_threads_per_threadgroup(32 * NSG)]] void name(        \
+      const device uint8_t* w,                                              \
+      const device uint8_t* /* scales */,                                   \
+      const device T* x,                                                    \
+      device T* y,                                                          \
+      const constant int& in_vec_size,                                      \
+      const constant int& /* out_vec_size */,                               \
+      const constant int& /* x_batch_ndims */,                              \
+      const constant int* /* x_shape */,                                    \
+      const constant int64_t* /* x_strides */,                              \
+      const constant int& /* w_batch_ndims */,                              \
+      const constant int* /* w_shape */,                                    \
+      const constant int64_t* /* w_strides */,                              \
+      const constant int64_t* /* s_strides */,                              \
+      uint3 tid [[threadgroup_position_in_grid]],                           \
+      uint simd_gid [[simdgroup_index_in_threadgroup]],                     \
+      uint simd_lid [[thread_index_in_simdgroup]]) {                        \
+    static_assert(                                                          \
+        group_size == KQ_Q8_0_GROUP, "Q8_0 kernel requires group_size=32"); \
+    static_assert(bits == 8, "Q8_0 kernel requires bits=8");                \
+    threadgroup float red[NSG * RPS];                                       \
+    kq_q8_0_qmv_splitk_impl<T, RPS, NSG>(                                   \
+        w, x, y, in_vec_size, red, tid, simd_gid, simd_lid);                \
+  }
+
+// Split-K M=1 q8_0 mat-vec. The host copy of the tile is
+// kq_q8_0_splitk_rps/nsg in src/kquant_metal_internal.h.
+KQ_Q8_0_DEFINE_QMV_SPLITK(kq_q8_0_qmv_splitk, 2, 8)
+
 // Bias-fused variant of kq_q8_0_qmv_fast: decode-only (B=1) fast path for a
 // KQuantLinear whose GGUF weight carries a real linear bias (e.g. gpt-oss
 // attention QKVO). Fuses the post-matmul add into this kernel's own output