feat(sdpa): starts + q8-KV operands for batched decode; env-gated small-M qmm kernels

CHANGELOG.md

@@ -10,6 +10,16 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - `route_shed(indices, scores, slot_table)`: GPU-side routed-expert slot
   remap plus residency shed for streamed MoE decode; non-resident experts
   are shed and reported (miss ids and scores) without a host sync.
+- `sdpa_decode_gqa` optional `starts` (int32 [B]): per-batch-row key start
+  offsets for left-padded batched KV caches; padded-out key chunks are
+  skipped, not staged.
+- `sdpa_decode_gqa` q8 KV operands (affine wire, bits 8, group 64): batched
+  decode attends over quantized KV directly, dequantizing on the staged
+  tile; up to 1.9x/call at depth vs dequantize-then-attend.
+- Env-gated small-M qmm experiment kernels (`KQ_QMM_SPLITK`,
+  `KQ_QMM_SPLITK_NAX`, `KQ_MV_EXT_SB`, `KQ_MV_EXT_NX`, `KQ_MV_EXT_HD`):
+  the NAX split-K path lifts the collapsed M9-16 band 65-76%; the rest
+  measured flat to negative on M5 and stay off by default.
 
 ## [0.3.7]
 

bindings.cpp

@@ -207,6 +207,11 @@ NB_MODULE(_ext, m) {
       "sinks"_a = nb::none(),
       "splits"_a = 0,
       "tile_c"_a = 0,
+      "starts"_a = nb::none(),
+      "k_scales"_a = nb::none(),
+      "k_biases"_a = nb::none(),
+      "v_scales"_a = nb::none(),
+      "v_biases"_a = nb::none(),
       nb::kw_only(),
       "stream"_a = nb::none(),
       R"(
@@ -215,7 +220,12 @@ NB_MODULE(_ext, m) {
         chunk is streamed through threadgroup-staged K/V tiles shared by the
         whole GQA group, so device memory reads the KV once per kv-head. At
         qL 2..4 (speculative-verify width) every query also shares the staged
-        tiles, causally clamped to its own trailing position.
+        tiles, causally clamped to its own trailing position. With `starts`,
+        batch row b attends keys [starts[b], kL) -- a left-padded batched KV
+        cache -- and fully padded-out key chunks are skipped, not staged.
+        With k_scales/k_biases/v_scales/v_biases (all four), k and v are
+        mlx affine-quantized wire (uint32, bits 8, group 64) and dequant is
+        fused into the tile stage.
 
         Args:
             q (array): queries [B, n_q_heads, qL, D], float16/bfloat16;
@@ -229,6 +239,9 @@ NB_MODULE(_ext, m) {
             splits (int): key-axis split count; 0 picks the default.
             tile_c (int): staged tile height, 8/16/32; 0 (default) picks by
                 head_dim (32 up to D=128, 16 at D=256, 8 at D=512).
+            starts (array, optional): per-batch-row key start offsets,
+                int32 [B], each in [0, kL - qL]; row b attends [starts[b],
+                kL). Out-of-range values read as an empty row (zero output).
 
         Returns:
             array: attention output [B, n_q_heads, qL, D].

metal/kq_quantized.metal

@@ -397,6 +397,109 @@ instantiate_mv_ext_all(iq1_m, 256, 1)
   instantiate_mv_ext_nr2_for_type(codec, gs, bits, float16_t)
 instantiate_mv_ext_nr2_all(q6_k, 256, 6)
 
+// Shuffle-broadcast experiment (KQ_MV_EXT_SB=1): the four ty-lanes sharing
+// a tx column exchange float4 quarters of the activation window over
+// simd_shuffle instead of each loading all 16 elements -- activation cache
+// traffic / 4, no barriers, no extra accumulators. q6_k only, M 4-12 where
+// the nr0=1 activation decay bites. Suffix _sb.
+#define instantiate_mv_ext_sb(codec, type, gs, bits, m)                 \
+  instantiate_kernel(                                                   \
+      "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m  \
+      "_sb",                                                            \
+      kq_ ## codec ## _mv_ext_sb, type, m, 2, 8)
+#define instantiate_mv_ext_sb_for_type(codec, gs, bits, type)           \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 4)                       \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 5)                       \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 6)                       \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 7)                       \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 8)                       \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 9)                       \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 10)                      \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 11)                      \
+  instantiate_mv_ext_sb(codec, type, gs, bits, 12)
+#define instantiate_mv_ext_sb_all(codec, gs, bits)                      \
+  instantiate_mv_ext_sb_for_type(codec, gs, bits, float)                \
+  instantiate_mv_ext_sb_for_type(codec, gs, bits, bfloat16_t)           \
+  instantiate_mv_ext_sb_for_type(codec, gs, bits, float16_t)
+instantiate_mv_ext_sb_all(q6_k, 256, 6)
+
+// Wide-nxpsg experiment (KQ_MV_EXT_NX=16|32): more K lanes per simdgroup
+// means fewer output rows per simdgroup, so each activation element has
+// nypsg*nsg = 4 (x16) or 2 (x32) redundant readers instead of 8 -- the
+// L1-capacity pressure behind the M4->8 decay -- and the grid gains
+// threadgroups. Same impl, different template params. q6_k M 4-12.
+#define instantiate_mv_ext_nx(codec, type, gs, bits, m, nx)             \
+  instantiate_kernel(                                                   \
+      "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m  \
+      "_x" #nx,                                                         \
+      kq_ ## codec ## _mv_ext, type, m, 2, nx)
+#define instantiate_mv_ext_nx_for_type(codec, gs, bits, type, nx)       \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 4, nx)                   \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 5, nx)                   \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 6, nx)                   \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 7, nx)                   \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 8, nx)                   \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 9, nx)                   \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 10, nx)                  \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 11, nx)                  \
+  instantiate_mv_ext_nx(codec, type, gs, bits, 12, nx)
+#define instantiate_mv_ext_nx_all(codec, gs, bits, nx)                  \
+  instantiate_mv_ext_nx_for_type(codec, gs, bits, float, nx)            \
+  instantiate_mv_ext_nx_for_type(codec, gs, bits, bfloat16_t, nx)       \
+  instantiate_mv_ext_nx_for_type(codec, gs, bits, float16_t, nx)
+instantiate_mv_ext_nx_all(q6_k, 256, 6, 16)
+instantiate_mv_ext_nx_all(q6_k, 256, 6, 32)
+
+// T-precision-dot experiment (KQ_MV_EXT_HD=1): the FMA-issue-bound band's
+// ALU lever. Dequanted chunk converts float->T once (amortized over M rows),
+// activations load at native T width with no per-row convert, and the
+// 16-term chunk dot runs at half/bfloat issue rate before an f32 fold.
+// q6_k only, M 4-12; no float x variant (no rate advantage). Suffix _hd.
+#define instantiate_mv_ext_hd(codec, type, gs, bits, m)                 \
+  instantiate_kernel(                                                   \
+      "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m  \
+      "_hd",                                                            \
+      kq_ ## codec ## _mv_ext_hd, type, m, 2, 8)
+#define instantiate_mv_ext_hd_for_type(codec, gs, bits, type)           \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 4)                       \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 5)                       \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 6)                       \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 7)                       \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 8)                       \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 9)                       \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 10)                      \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 11)                      \
+  instantiate_mv_ext_hd(codec, type, gs, bits, 12)
+#define instantiate_mv_ext_hd_all(codec, gs, bits)                      \
+  instantiate_mv_ext_hd_for_type(codec, gs, bits, bfloat16_t)           \
+  instantiate_mv_ext_hd_for_type(codec, gs, bits, float16_t)
+instantiate_mv_ext_hd_all(q6_k, 256, 6)
+
+// Staged-activation experiment (KQ_MV_EXT_TS=1): the M x 128 activation
+// window stages into threadgroup memory once per K-step and 8 simdgroups
+// (32 output rows) share it -- the activation-path lever the sb/nr2/nx
+// falsifications never isolated. Dot math identical to base. q6_k M 4-12.
+#define instantiate_mv_ext_ts(codec, type, gs, bits, m)                 \
+  instantiate_kernel(                                                   \
+      "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m  \
+      "_ts",                                                            \
+      kq_ ## codec ## _mv_ext_ts, type, m, 8, 8)
+#define instantiate_mv_ext_ts_for_type(codec, gs, bits, type)           \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 4)                       \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 5)                       \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 6)                       \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 7)                       \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 8)                       \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 9)                       \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 10)                      \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 11)                      \
+  instantiate_mv_ext_ts(codec, type, gs, bits, 12)
+#define instantiate_mv_ext_ts_all(codec, gs, bits)                      \
+  instantiate_mv_ext_ts_for_type(codec, gs, bits, float)                \
+  instantiate_mv_ext_ts_for_type(codec, gs, bits, bfloat16_t)           \
+  instantiate_mv_ext_ts_for_type(codec, gs, bits, float16_t)
+instantiate_mv_ext_ts_all(q6_k, 256, 6)
+
 #define instantiate_kquant_q3_k_for_type(type)                          \
   instantiate_kquant_batched(verify_qmv, type, 256, 3, 0, q3_k)          \
   instantiate_kquant_batched(qmv_fast, type, 256, 3, 0, q3_k)            \
@@ -863,4 +966,37 @@ instantiate_kquant_gather_qmm_rhs_codec(256, 1, iq1_m)
   map[3u * slot + 1u] = r;
   map[3u * slot + 2u] = min(64u, seg_len - rank);
 }
-// clang-format on
+
+// Split-K partial fold for qmm_t_splitk: y[i] = (T)(f32 sum over the
+// splits axis of partials[z * stride + i]). Partials carry one T rounding
+// per slice; the fold accumulates in f32 and rounds once more. One thread
+// per output element; cost is noise next to the GEMM pass.
+template <typename T>
+[[kernel]] void kq_qmm_splitk_accum(
+    const device T* partials [[buffer(0)]],
+    device T* y [[buffer(1)]],
+    const constant int& n_elems [[buffer(2)]],
+    const constant int& splits [[buffer(3)]],
+    const constant int& stride [[buffer(4)]],
+    uint gid [[thread_position_in_grid]]) {
+  if (gid >= static_cast<uint>(n_elems)) {
+    return;
+  }
+  float acc = 0.0f;
+  for (int z = 0; z < splits; ++z) {
+    acc += static_cast<float>(
+        partials[static_cast<int64_t>(z) * stride + gid]);
+  }
+  y[gid] = static_cast<T>(acc);
+}
+
+#define instantiate_kq_qmm_splitk_accum(type)   \
+  instantiate_kernel(                           \
+      "kquant_qmm_splitk_accum_" #type,         \
+      kq_qmm_splitk_accum,                      \
+      type)
+
+instantiate_kq_qmm_splitk_accum(float)
+instantiate_kq_qmm_splitk_accum(bfloat16_t)
+instantiate_kq_qmm_splitk_accum(float16_t)
+    // clang-format on

metal/kq_quantized_nax.metal

@@ -116,6 +116,26 @@
   instantiate_kquant_nax_qmm_t_smallbm(codec, float16_t,  gs, bits)          \
   instantiate_kquant_nax_qmm_t_smallbm(codec, bfloat16_t, gs, bits)
 
+// Split-K qmm_t on the NAX BM=32 tile (KQ_QMM_SPLITK_NAX experiment):
+// grid.z K-slices into T partials + shared accum fold. The plain small-M
+// grid is TG-count starved (ceil(N/64) x 1 threadgroups at decode shapes);
+// splitting K multiplies occupancy without touching the fragment shape.
+// q6_k + q8_0 only; no batched or float x variants (route gates match
+// qmm_nax and non_batched).
+#define instantiate_kquant_nax_qmm_t_splitk(type, gs, bits, aligned_N, codec) \
+  instantiate_kernel(                                                        \
+      "kquant_" #codec "_qmm_t_nax_splitk_" #type "_gs_" #gs "_b_" #bits     \
+          "_bm32_bn64_bk64_wm2_wn2_alN_" #aligned_N,                         \
+      kq_ ## codec ## _qmm_t_nax_splitk,                                     \
+      type, gs, bits, aligned_N, 32, 64, 2, 2)
+#define instantiate_kquant_nax_splitk(codec, gs, bits)                       \
+  instantiate_kquant_nax_qmm_t_splitk(float16_t,  gs, bits, true,  codec)    \
+  instantiate_kquant_nax_qmm_t_splitk(float16_t,  gs, bits, false, codec)    \
+  instantiate_kquant_nax_qmm_t_splitk(bfloat16_t, gs, bits, true,  codec)    \
+  instantiate_kquant_nax_qmm_t_splitk(bfloat16_t, gs, bits, false, codec)
+instantiate_kquant_nax_splitk(q6_k, 256, 6)
+instantiate_kquant_nax_splitk(q8_0, 32, 8)
+
 // Double-buffered BM=64 qmm_t, name-suffixed _db: dispatched by the host
 // solely for the M33-64 decode band (kq_smallbm_policy db64 + KQ_NAX_DB64).
 // As a blanket BM=64 default the doubled Ws cut occupancy (M96+ -3-15%,