Extending commit for PR#4077 for GQA size 12 and 16

[dudududukim]

• ✅ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Claude was used to update the codes and I have cross-validated written code and tested.

---

Following PR for #4077 (extending group size of 12 and 16 w\ head_dim=128)

As described in #4077 reading KV Cache has redundancy wrt Query head number.
While former PR solved GQA (gqa factor=8), this PR further expand to gqa factor=12or16.

All measurement was taken with Power-connected with power mode 2(performance mode) with M5 pro 24GB macbook pro.

Reference-kernel deficit, re-measured (32K, hd128, H_kv=4, unique KV 67.1MB, measured peak 265 GB/s):

	amplification	time	unique GB/s
gqa 12 ref	12x	540.6us	124
gqa 16 ref	16x	675.3us	99

About Heads Per Threadgroup(HPT) values

HPT is bounded by threadgroup memory (hard 32KB).
The staging buffer is G * HPT * V floats:

G	HPT	usage (hd128)
8	4	16 KB	existing
12	4	24 KB	chosen ✅
16	4	32.5 KB	over by 512 B — impossible
16	2	16 KB	chosen ✅

Registers bound it from the other side: HPT=8 spills at hd128 (established in the #4077 work), so HPT <= 4. HPT=1 makes NT=G, which is mathematically identical to the reference kernel, so gqa 16 has no third option here.
Further improvement is planned to be added with follow-up PR.

Target models

As we could not run full model with 24GB memory, we just used config.json to run the same compute operation rather than loading the weight. (It would be great to be verified by other resource sufficient user! @erwinzhang7)

config.json verified, all head_dim=128:

gqa	models
12	GLM-4.5-Air (106B), GLM-4.6 (357B), Mistral-Large-2411 (123B)
16	Qwen3-235B-A22B, Llama-3.1-405B

Kernel benchmarks

Following the CONTRIBUTING.md test script was tested both on main and PR branch with everything un-changed.

(a) Two-build comparison — origin/main build vs PR build, via bench_sdpa_gqa.py :

python bench_sdpa_gqa.py > main.txt      # on origin/main
python bench_sdpa_gqa.py > pr.txt        # on the PR branch
python bench_sdpa_gqa.py --compare main.txt pr.txt

fp16, hd128, H_kv=4. Median of 9 trials x 16 iterations. main is 0.32.2.dev20260823+7789905ed, PR is 0.32.2.dev20260823+b7270a558:

gqa	ctx	main us	PR us	main GB/s	PR GB/s	speedup
4	8k	85.5	89.1	196	188	0.96x
4	16k	132.7	131.7	253	255	1.01x
4	32k	261.0	259.1	257	259	1.01x
8	8k	98.9	97.0	170	173	1.02x
8	16k	178.7	178.4	188	188	1.00x
8	32k	321.5	321.6	209	209	1.00x
12	8k	154.2	136.9	109	123	1.13x
12	16k	303.2	257.0	111	131	1.18x
12	32k	539.3	464.3	124	145	1.16x
16	8k	199.7	184.7	84	91	1.08x
16	16k	391.9	353.8	86	95	1.11x
16	32k	678.5	632.5	99	106	1.07x

gqa 4 is outside the gate on both builds. gqa 8 is inside the gate but runs byte-identical code on both builds, which makes it the stronger control: it isolates cross-build drift from the effect, and lands at 1.00-1.02x.

(b) Speedup measured : sdpa_vector_2pass_1 (reference) vs sdpa_vector_2pass_1_gqa_N (new kernel)

gqa	8K	16K	32K
4 (control)	1.01x	1.00x	1.00x
8 (existing #4077)	1.19x	1.24x	1.24x
12	1.14x	1.18x	1.17x
16	1.09x	1.10x	1.07x

Truncated end-to-end speedup

Models using gqa 12/16 cannot be fit in our 24GB, same as #4077, we tested with keeping the real per-layer model structure and truncating some last layers.
(reduced the expert pool from 128 to 16 while keeping num_experts_per_tok at 8)

GLM-4.5-Air (gqa 12), 16 of 46 layers:

ctx	speedup
4096 (below gate)	1.001x <- control
8192	1.019x
16384	1.049x
32768	1.063x
65536	1.071x
131072	1.135x

Qwen3-235B (gqa 16), 16 of 94 layers: 32K 1.031x, 64K 1.033x, 128K 1.061x

Truncation deflates the result, so these are floors: halving depth halves attention and MLP together while the per-step lm_head tail is fixed, making a shallow model less attention-dominated. Measured directly — GLM @32k: 4L 1.053x -> 8L 1.055x -> 16L 1.063x.

Extrapolation to full depth (Expected speed-up)

Regressing the depth sweep gives step(L) = L x block + tail, so the fixed lm_head/embedding tail can be post-computed expanding the expected latency to full depth.

GLM-4.5-Air — block 1.754 ms (main) / 1.642 ms (PR), tail 1.36 ms; the tail is ~2% of a 46-layer step, so truncation costs at most +0.2pp:

ctx	16L measured	46L predicted
4096	1.001x	1.001x
8192	1.019x	1.020x
16384	1.049x	1.051x
32768	1.063x	1.065x
65536	1.071x	1.072x
131072	1.135x	1.137x

Qwen3-235B — block 1.256 ms (main) / 1.216 ms (PR), same 1.36 ms tail:

ctx	16L measured	94L predicted
32768	1.031x	1.033x
65536	1.033x	1.035x
131072	1.061x	1.063x

Both models recompute the same final part. Since they have the same ~152k vocabulary size and 4096 hidden size, they use the same lm_head.

The numbers above are linearly extrapolated from measurements on this machine. They are not performance predictions for hardware that can fit the entire model in memory.

Functionality check

vs fp32 naive attention, ctx 8k:

dtype	gqa 8	gqa 12	gqa 16	tol
fp16	2.63e-05	2.35e-05	2.61e-05	3e-03
bf16	2.53e-04	2.34e-04	2.43e-04	2e-02

Gate execution was checked with Xcode GPU trace: sdpa_vector_2pass_1_gqa_12_float16_t_128_128 dispatched.

Modified Test script

One line, extending the list #4077 introduced in test_sdpa_vector_gqa_long:

-        for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64)]:
+        for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64), (48, 4, 128), (64, 4, 128)]:

(48, 4, 128) is gqa 12, (64, 4, 128) is gqa 16, both hd128, both inside the gate at the existing L in [8192, 8201]. Those two L values give an exact split and a remainder split: at L=8201 with blocks=256 the chunk is 33
tokens, so gqa 16's HPT=2 sub-chunk is 17 with the last one partial.

test_fast_sdpa 23/23 pass. Sensitivity: measured gqa 16 fp32 maxdiff against fp32 naive is 5.59e-08 versus the test's atol=1e-4, ~1800x margin.

Honest note on "confirm the new tests fail in the main branch": a correctness test for these shapes PASSES on main, because main falls back to the reference kernel and produces correct output. That half of the rule cannot be satisfied by a routing change with identical numerics. The test guards the new path; it does not reproduce a bug.

Formatting

uvx pre-commit run --all ✅ PASSED