Extend SDPA vector 2-pass kernel to query sequence lengths up to 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 PR:

• Lets each simdgroup in sdpa_vector_2pass_1 process up to 4 consecutive query rows serially (2 for head dims > 128 to bound per-lane register state), reusing the K/V rows already in registers. The block-parallel reduction over the key sequence is retained, K/V is still read once per (block, threadgroup), and pass 2 is unchanged.
• Keeps the original single-row body verbatim under an n_per_simd == 1 function constant, so decode (qL ≤ 8) codegen is unchanged.
• Routes 8 < qL <= 16 to the vector path when kL >= 1024 (the speculative-decoding verify window from the issue); qL ≤ 8 and qL > 16 keep their existing routes. A first cut also moved qL 17–32, but that was slower than the steel path at L=32 on M3 Max, so v1 stops at 16.
• Adds a dispatch-side fallback that increases rows-per-simdgroup if register pressure ever lowers the pipeline's threadgroup limit.

Benchmarks (M3 Max, fp16, D=128, mask=None, 40-iter timed, baseline/patched interleaved A/B/A/B — numbers stable across both rounds):

Shape	S	L=8	L=12	L=16	L=24…64
GQA 32/8	16384	1.13 → 1.14 ms	2.49 → 1.44 ms (1.73×)	2.50 → 1.96 ms (1.28×)	unchanged
GQA 32/8	32768	2.03 → 2.03 ms	4.63 → 2.53 ms (1.83×)	4.65 → 3.56 ms (1.31×)	unchanged
GQA 32/8 causal	16384	1.18 → 1.19 ms	2.58 → 1.52 ms (1.70×)	2.55 → 1.99 ms (1.28×)	unchanged
MHA 32/32	16384	1.28 → 1.32 ms	2.52 → 1.75 ms (1.44×)	2.53 → 2.23 ms (1.13×)	unchanged

Decode shapes (qL ≤ 8, S 4096/32768, D 128 and 256) measure parity within noise, as expected from the unchanged single-row body.

Correctness, beyond the updated unit tests: swept qL {1,7,8,9,12,16,24,31,32} × gqa {1,4,8} × D {64,128,256} × masks {none, causal, bool, float} × S {512,4096,16384}, plus bf16, the transposed-query layout, and sinks (gqa {4,8} × qL {1,8,9,12,16}) against an unfused fp32 reference — all within fp16 tolerance (2.5e-3).

Independent validation on M5 Max (nax path) — confirmed by pierre427, who built the merge-base and this branch from source with the NAX kernels verified active, so the baseline genuinely takes sdpa_full_self_attention_nax:

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 matches the design: qL ≤ 8 decode unchanged (within noise), qL ≥ 24 returns to the existing full-attention route with no regression, correctness at fp16 noise (max abs err 1.4e-4). Notably the L=16 gains on M5 (1.56–1.72×) exceed the M3 Max ones (1.28–1.31×) — the nax path benefits across the whole 9–16 window.

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_few_query now sweeps Lq {8,12,16} × Nkv {1,2} at kL=4096 through the new path)
• 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.