Read each K/V byte once in gqa-8 decode attention

[dudududukim]

Closes #4076.

Proposed changes

sdpa_vector_2pass_1 has every simdgroup of a threadgroup stream the same K/V block, so loads scale with query heads instead of kv heads and gqa-8 models decode well below the machine's bandwidth (195 GB/s for 32q/4kv/hd128 and 120 GB/s for 64q/8kv/hd64 at 32K context, against 266 GB/s for a dense qmv on the same M5 Pro).

This adds sdpa_vector_2pass_1_gqa: each simdgroup owns a contiguous token sub-chunk of the block and computes HPT of its group's query heads from registers, so each K/V byte is read gqa_factor / HPT times instead of gqa_factor. The simdgroups' partials are merged once through threadgroup memory at the end; the partials layout and sdpa_vector_2pass_2 are unchanged. HPT is a template parameter because register pressure decides it: 8 heads per simdgroup wins at head_dim 64 but spills and loses at head_dim 128, where 4 is right.

Dispatch is deliberately narrow: single-token queries without mask or sinks, gqa_factor == 8, head_dim 64 or 128, N >= 8192. Everything else stays on the existing kernel. Sinks and other gqa factors can follow if this direction looks good.

All numbers below are from an M5 Pro 24 GB on AC, bf16, with K/V passed as strided cache slices and distinct buffers cycled between calls so the SLC doesn't flatter either arm. Both arms run in one process with the dispatch flipped between them (through a temporary env check that is not part of this diff), interleaved ABBA over several rounds; the tables give the median.

Kernel level:

geometry	ctx	main	this PR	speedup
32q/4kv/hd128	8192	90.8 us (185 GB/s)	84.8 us (198 GB/s)	1.07x
32q/4kv/hd128	16384	176.9 us (190 GB/s)	162.4 us (207 GB/s)	1.09x
32q/4kv/hd128	32768	343.5 us (195 GB/s)	301.5 us (223 GB/s)	1.14x
64q/8kv/hd64	8192	147.1 us (114 GB/s)	127.7 us (131 GB/s)	1.15x
64q/8kv/hd64	16384	288.5 us (116 GB/s)	247.5 us (136 GB/s)	1.17x
64q/8kv/hd64	32768	561.6 us (120 GB/s)	440.6 us (152 GB/s)	1.27x

Per-round spread is within 0.01x of the median for every row above.

End-to-end decode, ms/step, on Qwen3-30B-A3B-4bit truncated to 24 layers (so 32K of KV fits in 24 GB; 32q/4kv/hd128) and on the full 36-layer Qwen2.5-3B-Instruct-4bit (16q/2kv/hd128). Context 4096 is below the dispatch gate, so both arms run the identical kernel there and it serves as a control:

ctx	30B main	30B this PR	delta	3B main	3B this PR	delta
4096 (control)	7.289	7.292	-0.0%	9.626	9.626	-0.0%
8192	8.596	8.301	+3.5%	10.451	10.178	+2.7%
16384	10.757	9.936	+8.3%	12.152	11.479	+5.9%
32760	15.025	13.726	+9.5%	15.674	14.355	+9.2%

Individual control rounds span -0.2%..+0.3% (30B) and -0.6%..+1.6% (3B), so the gated contexts sit well outside the noise floor.

Logits stay within bf16 rounding of the current kernel with argmax stable over greedy decode, and outputs match an fp32 reference to 1e-5 (fp16) and 2e-4 (bf16) on both gated and ungated shapes. The added test covers both instantiated head dims at a context that is not a multiple of the block count.

Checklist

• 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
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[erwinzhang7]

Verified on an M5 Max, 40 GPU cores, 128 GB, macOS 26.2. Your numbers are from an
M5 Pro 24 GB, so this is a different part rather than a bigger one, and the 128 GB
removes the truncation your 30B run needed.

base throughout is 6a0dd0fd2, this branch's own parent, rather than current
main, since the branch carries a merge commit and diffing against main would
credit it with three weeks of unrelated work.

The kernel win holds on wider silicon. bf16, K/V as
strided cache slices with eight buffer sets cycled so the SLC cannot hold them,
12 rounds of 20 calls with one eval per round:

geometry	ctx	base	this PR	speedup	yours
32q/4kv/hd128	8192	53.4 us	46.2 us	1.16x	1.07x
32q/4kv/hd128	16384	87.6 us	80.9 us	1.08x	1.09x
32q/4kv/hd128	32768	156.2 us	139.3 us	1.12x	1.14x
64q/8kv/hd64	8192	73.5 us	61.4 us	1.20x	1.15x
64q/8kv/hd64	16384	130.2 us	109.4 us	1.19x	1.17x
64q/8kv/hd64	32768	244.0 us	186.6 us	1.31x	1.27x

The baseline on this machine is 1.7x to 2.3x faster than yours on the same kernel
before any change. The worry with a bandwidth argument is that it stops applying on a part with
bandwidth to spare. The argument holds: larger on four of the six
geometries, and within a couple of points on the other two.

Nothing outside the gate moves at all. Both builds scored against each other
and against an fp32 reference:

case	max abs diff vs base	vs fp32
ungated, ctx 4096 below the gate	0	3.5e-4
ungated, gqa_factor 4	0	3.3e-4
ungated, head_dim 80	0	4.6e-4
ungated, 2 query tokens	0	3.4e-4
gated 32q/4kv/hd128	2.4e-4	2.8e-4
gated 64q/8kv/hd64	2.4e-4	1.9e-4
gated, ctx 8192+137	4.9e-4	2.4e-4

Bit-identical on every ungated shape, so the dispatch condition fires exactly
where you say and nowhere else. On the gated shapes the differences are bf16
rounding and the new kernel sits as close to fp32 as the old one, so it isn't
sacrificing accuracy for speed. The 8192+137 row would expose an off-by-one in the sub-chunk
partitioning, and it is clean.

End to end, on all 48 layers rather than 24. Same model, Qwen3-30B-A3B-4bit,
untruncated, 128 GB holds 32K of KV for the whole thing. 24 decode steps per
context after a full prefill, median:

ctx	base	this PR	delta	yours, 24 layers
4096	9.303 ms	9.260 ms	+0.5%	-0.0%
8192	10.682 ms	10.264 ms	+3.9%	+3.5%
16384	12.757 ms	12.177 ms	+4.5%	+8.3%
32760	17.724 ms	16.491 ms	+7.0%	+9.5%

The win is real at full depth and clears the control comfortably. My 4096 control
comes out at +0.5% rather than zero (the two arms are separate processes with
separate model loads) so read 0.5% as this harness's noise floor.

However, the magnitude is smaller than the truncated run shows, and consistently:
+4.5% against +8.3% at 16K, +7.0% against +9.5% at 32K. That's from halving the
layer count: it removes half the non-attention work per step, so attention
occupies a larger share of what is left and any speedup measured against the step
scales up with it.

I'd quote +7.0% at 32760 on the full model rather than +9.5%. It is still a
good number for a decode path and it is one nobody can argue with. Good work!

[zcbenz]

Can you add the affected shapes to sdpa_bench.py?

Sorry just realized this is sdpa_vector, ignore this comment.

[zcbenz]

It seems that most of the host code is duplicated, I think you can just change the kname?

[dudududukim]

Right!
Fixed in d1ab9de, the gate now just appends _gqa to the kname.
Tests pass and the dispatch still switches at the boundary.

[zcbenz]

Nice improvement, thanks!