Add a multi-row 2-pass SDPA vector kernel for 8 < qL <= 16

[katlun-lgtm]

Proposed changes

Addresses #3826 (2.3–2.8× latency jump from L=8 to L=12 query rows, flat plateau until L≥48).

The 2-pass vector kernel packs every (query head, query position) pair into its own simdgroup, which caps the threadgroup at gqa_factor * qL <= 32 — for GQA 32/8 the wall sits exactly at qL=8. Anything above leaves the decode-shaped kernel and lands on the full-attention kernel, which launches one query tile per head and streams all S keys serially, so L=12 pays ~L=48 cost (mechanism detailed in the issue thread).

This revision addresses the review: the change is purely additive and does one thing — give 8 < qL <= 16 (the speculative-decoding verify window) a 2-pass vector path.

• New kernel sdpa_vector_2pass_1_multirow: each simdgroup owns up to 4 consecutive query rows (2 for head dims > 128, bounding per-lane register state), reading each K/V row into registers once and reusing it across its rows. The block-parallel reduction over the key sequence is retained and pass 2 is unchanged.
• New dispatch function sdpa_vector_2pass_multirow (same structure as the existing one) and one additive else if in eval_gpu routing 8 < qL <= 16 there when kL >= 1024 (the same threshold the existing decode routing uses for its 2-pass split) or when the full kernel does not support the head dim.
• sdpa_vector_2pass_1, its dispatch, and the qL <= 8 routing are byte-identical to main: the diff is +587/−2, and the only removed line is the use_fallback return statement picking up the new branch. Decode cannot regress — no code it runs was touched.

Benchmarks (M3 Max g15s, fp16, D=128, GQA 32/8, B=1; latency per call — mx.eval each step, as a verify loop runs — min of 10 trials × 20 reps; baseline = mlx 0.32.0, whose behavior at these shapes is identical to current main on this hardware: the only SDPA changes since are M5-only #3843 and env-gated #3875):

S	mask	L=8	L=9	L=12	L=16	L=17 / 24
16384	none	1.11 → 1.10	2.38 → 1.24 (1.92×)	2.38 → 1.43 (1.66×)	2.40 → 1.94 (1.23×)	parity
16384	causal	1.15 → 1.15	2.50 → 1.26 (1.98×)	2.49 → 1.51 (1.65×)	2.50 → 1.97 (1.27×)	parity
32768	none	1.94 → 1.99	4.51 → 2.10 (2.14×)	4.51 → 2.51 (1.79×)	4.51 → 3.52 (1.28×)	parity
32768	causal	2.03 → 2.11	4.71 → 2.21 (2.14×)	4.71 → 2.61 (1.80×)	4.72 → 3.54 (1.33×)	parity

qL ≤ 8 and qL ≥ 17 keep their existing routes and measure parity within noise (L 1–8 scatter is ±3–4% in both directions across the sweep).

Correctness: test_fast_sdpa.py passes in full (17 tests, 602 subtests). The new additive test_sdpa_multirow_query sweeps Lq {9,12,16} × Nkv {1,2} at D=128/kL=4096 through the new route, plus D=256 (the 2-rows-per-simdgroup bound) and a D=96/kL=512 case (the unsupported-full-head-dim route), each against masks {none, scalar-bool, bool-vector, two broadcast bool layouts, additive float, causal}, plus the transposed-query layout and sinks. The pre-existing tests are untouched.

Independent validation on M5 Max (nax path) — confirmed by pierre427 on the previous revision of this branch. The multi-row kernel body and the routing thresholds are unchanged by the restructure — only the packaging changed (in-place branch in the existing kernel → separate additive kernel):

S	L=12	L=16
16384, mask=None	2.00 → 1.03 ms (1.95×)	1.98 → 1.26 ms (1.56×)
32768, mask=None	3.83 → 1.75 ms (2.19×)	3.83 → 2.23 ms (1.72×)
16384, causal	1.98 → 1.11 ms (1.78×)	1.96 → 1.37 ms (1.42×)
32768, causal	3.72 → 1.90 ms (1.96×)	3.75 → 2.48 ms (1.51×)

Boundary behavior on nax matched the design: qL ≤ 8 decode unchanged, qL ≥ 24 returns to the existing full-attention route with no regression, correctness at fp16 noise (max abs err 1.4e-4).

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes (clang-format v21.1.8 as pinned)
• I have added tests that prove my fix is effective or that my feature works (test_sdpa_multirow_query, additive)
• I have updated the necessary documentation (if needed)

[pierre427]

Confirmed on M5 Max — the win transfers to the nax path, and it's a bit larger than your M3 Max numbers.

Method: built the merge-base (4367c73b) and this branch head (b09124a1) from source, each with MACOSX_DEPLOYMENT_TARGET=26.2 so the NAX kernels are compiled in (verified active in both builds: bf16 4096³ GEMM 2.4–2.6 ms, vs ~9 ms when NAX is gated out — so the baseline L>8 shapes are genuinely taking sdpa_full_self_attention_nax, bq=64). Original repro, GQA 32/8, D=128, fp16, mask=None and mask="causal", 40-iter, one warmup pass then a measured pass per build.

mask=None:

S	L	base (4367c73)	#3838	speedup
16384	12	1.998 ms	1.027 ms	1.95×
16384	16	1.975 ms	1.262 ms	1.56×
32768	12	3.834 ms	1.751 ms	2.19×
32768	16	3.832 ms	2.233 ms	1.72×

mask="causal":

S	L	base	#3838	speedup
16384	12	1.979 ms	1.113 ms	1.78×
16384	16	1.955 ms	1.373 ms	1.42×
32768	12	3.721 ms	1.902 ms	1.96×
32768	16	3.748 ms	2.476 ms	1.51×

Boundaries behave exactly as intended on the nax path: qL≤8 decode is unchanged (16k L=8 0.755→0.802, 32k L=8 1.292→1.300 — within noise), and qL≥24 snaps back to the full-attn plateau (16k L=24 1.983→1.911, 32k L=24 3.814→3.638), so the routing cut at 16 is clean and nothing above it regresses. Correctness vs an fp32 reference stays at fp16 noise (max abs err 1.4e-4 at L=12/16, S=2048).

So the mechanism you describe — replacing "one query tile per head, serial over S" with the 2-pass block-parallel-over-S reduction — carries over to the nax full-attn baseline unchanged, and the window recovers near-linear scaling in L (16k: L=16 patched 1.262 ms ≈ 1.67× the L=8 cost, vs 2.6× before). Your framing that L=24–32 is better left to steel split-K holds here too: on nax those shapes sit right back on the ~1.9/3.6 ms plateau, so split-K is still the missing piece for that range.

One note for the description: the M5 gains at L=16 (1.56–1.72×) are meaningfully above the M3 L=16 numbers (1.28–1.31×), so the nax path benefits across the whole 9–16 window rather than tapering near 16 the way it does on M3.

[katlun-lgtm]

Thanks for the rigorous run — building both sides from source with NAX verified active (via the GEMM canary) makes this the comparison I couldn't produce, and I've folded your tables into the PR description in place of the "unverified on nax" caveat.

Two of your datapoints are worth highlighting for reviewers:

• The boundary checks: qL ≤ 8 decode within noise and qL ≥ 24 returning to the existing plateau confirm the routing cut at 16 is clean on both architectures — no shape gets worse.
• The L=16 gains being larger on M5 (1.56–1.72×) than on M3 Max (1.28–1.31×) — so the fix helps most exactly where the nax bq=64 tile was wasting the most work, and holds across the whole 9–16 verify window.

Agreed that L=24–48 remains the steel split-K's job — on both machines those shapes sit back on the full-attention plateau by design, so that can be tackled independently of this PR.

[katlun-lgtm]

Bumping this one — it's been open since July 13 with no CI run against it.

The workflow run for b09124a1 is sitting in action_required, so the PR shows no checks at all. If someone can approve the run, it would at least get vetted alongside the rest of the queue. #3833 and #3863 are in the same state if that's easy to clear in the same pass.

On the substance, the main thing that changed since the description was written: @pierre427, who filed #3826, rebuilt both sides from source on an M5 Max with NAX verified active and reproduced the speedup independently — 1.95× / 2.19× at L=12 and 1.56× / 1.72× at L=16 (S=16k/32k, mask=None; causal is 1.42–1.96×), which is larger than my M3 Max numbers. His boundary checks are the part worth a reviewer's eye: qL≤8 decode is unchanged and qL≥24 snaps back to the existing full-attention plateau, so the routing cut at 16 is clean on both architectures and nothing above it regresses. Correctness holds at fp16 noise (1.4e-4 vs an fp32 reference), and a test is included.

@angeloskath, you picked up #3875 in this area recently — are you the right person for this one, or should I be pointing it at someone else?

[zcbenz]

It seems that the code does a lot more than the title "Extend SDPA vector 2-pass kernel to query sequence lengths up to 16"? We can only verify performance changes when it only improves a single thing.

[katlun-lgtm]

Fair point — the previous revision restructured the existing kernel into a dual body, so the decode path read as touched even where the codegen was meant to be unchanged. Reworked so the diff is one thing, purely additive:

• sdpa_vector_2pass_1, its dispatch, and the qL <= 8 routing are now byte-identical to main. The multi-row body lives in a new kernel, sdpa_vector_2pass_1_multirow, with its own dispatch function, used only for 8 < qL <= 16.
• The diff is +587/−2; the only removed line is the use_fallback return statement picking up the new branch. Decode cannot regress — no code it runs was touched.

Re-measured on the restructured branch (M3 Max, fp16, D=128, GQA 32/8, latency per call): L=9 goes 2.38 → 1.24 ms at S=16k and 4.51 → 2.10 ms at S=32k, tapering to 1.23–1.33× at L=16; L=8 and L≥17 measure parity, as they keep their existing routes. Full tables in the updated description. test_fast_sdpa.py passes in full; the new test is additive too.

pierre427's M5 numbers in the description are from the previous revision — the multi-row body and the routing thresholds are unchanged by the restructure, only the packaging (in-place branch → separate kernel). Also retitled since "extend" no longer describes it.

[zcbenz]

@jessegross Would you mind doing a quick check at whether this would improve MTP since you worked on qmv_wide which I assume optimize in a very similar way? This is a bit out of my scope.

[katlun-lgtm]

On the MTP question — I measured it on an M3 Max rather than guess. It helps where the kernel engages and does nothing where it doesn't, and the split turns out to be structural.

n_kv	qL	baseline ms	this PR ms	speedup
4	1	0.517	0.516	1.00× (control)
1	12	4.055	4.056	1.00×
2	12	4.353	4.355	1.00×
4	9	4.522	1.887	2.40×
8	9	4.545	2.139	2.12×
4	16	4.543	2.904	1.56×
8	16	4.545	3.630	1.25×
32	16	4.662	4.261	1.09×

The reason n_kv ≤ 2 is flat is sdpa_multirow_vector_supported's gqa_factor * ceil(qL / max_nq) <= 32 (max_nq = 4 at D = 128), which is a threadgroup budget — a low KV-head count means a high gqa factor, so those shapes are excluded by construction rather than by tuning. So typical GQA gets the win in the verify window (Llama-3-8B is n_kv=8, most Qwen are 4–8), MLA-absorbed models at n_kv=1 get nothing, and this is complementary to #3837 rather than overlapping it: the shapes where that floor is worst are exactly the ones this kernel can't take.

Method matched to #3837's so the numbers compare: bf16, D=128, n_q=32, L=32768, mask=None, warm 1 s then median of five 0.4 s wall-clock windows, eval per call, three repeats (spread under 5%, mostly under 1%). qL=1 is the control since that path is byte-identical here, and it lands at 1.00×. Baseline pip 0.32.0 against 3c543b48. For calibration, n_kv=1 qL=1 decode measures 10.8% of peak here versus the 10.3% reported on M5 Max in #3837.

[jessegross]

@zcbenz In my experience MTP is rarely effective when drafting 8 or more tokens at a time, especially at the longer context lengths where this is targeted (drafting 2-4 tokens is more typical). At long draft lengths the acceptance rate is very low for the last tokens and you spend more time verifying them against the target model than you gain from the cases where you do accept them.

If you can optimize things really well so there is a very small incremental cost to verifying additional tokens in a batch then it would make sense to draft longer. So in that sense this could help but I don't think it is enough to make those draft lengths profitable by itself.

[zcbenz]

Thanks @jessegross for sharing the production experience. Adding a new 2-pass SDPA vector kernel has non-trivial costs on both reviewing and maintenance, since it is not likely going to be used in production, I'm closing this PR, but still thanks a lot for experimenting the optimization.