On an A100 at 64/8 heads and a 128K KV cache, K/V head copies in NVIDIA's cuDNN frontend took 7.28 ms per call against a 6.0-6.7 ms attention kernel. That is one shape. Removing the copies made 6 batched rows 2.6-3.3× faster, and 10 more stopped running out of memory. Those rows are really 3 configs and 5 configs. Multi-token steps at batch 1 ran 1.15-2.19× faster at that cache size. Clocks weren't locked, and the cuDNN backend stayed faster on 27 of 31 multi-token rows.
The merged code is tests, not a kernel. Those tests catch 3 injected bugs that NVIDIA's SM80 tests missed. I closed my prototype after NVIDIA shipped native GQA in #1239, and I didn't measure #1239.
What I built
In grouped-query attention (GQA), several query heads share one KV head. The SM80 adapter copied K and V up to the query head count on every call. The kernels already map each query head to its KV head, with h // (H_q // H_kv).
The prototype changed the adapter and no kernel file. That removed an expansion of 4 GiB per call at 64/8 heads. After #1239, I closed it. On issue #377, one option I offered was a tests-only PR, and that's what was suggested back.
That PR adds 80 tests and changes no library code. Of those, 66 cases check that native GQA gives bit-for-bit the same O and LSE as the same inputs with K/V copied up. NVIDIA's GQA tests used 4/2 and 16/4 heads. At those ratios, query heads per KV head equals KV heads, so a kernel that mixes up the two still passes. The new tests add 8/4 and 8/2.
What I measured
Timing
I used one rented A100-SXM4-80GB and bf16. Each timing is the median of 20 samples of kernel device time, which on the develop branch includes the copies. Develop and the prototype alternated in one pod. Rows with identical code on both sides had a median ratio of 1.000, and I only claimed rows outside their control band.
| Heads and d | Develop (ms) | Prototype (ms) | Faster by (×) | cuDNN backend (ms) |
|---|---|---|---|---|
| 64/8 d128 | 14.57 | 6.66 | 3.30 | |
| 64/8 d64 | 8.23 | 4.23 | 4.09 | |
| 32/8 d128 | 9.91 | 6.66 | 3.30 | |
| 32/2 d256 | 17.15 | 12.55 | 9.96 | |
| 8/1 d128 | 7.59 | 6.63 | 3.30 |
I read the time saved as following the copy's size. At d128, the expansion is 8x smaller at 8/1 heads than at 64/8 heads, and the kernel takes about the same time.
Develop ran out of memory on 10 of the 16 batched rows, because it needed a 64-512 GiB expanded K/V workspace. The prototype ran all 10.
| Heads, d and batch | Develop (ms) | Prototype (ms) | Faster by | cuDNN backend (ms) |
|---|---|---|---|---|
| 64/8 d64, batch 32 | 155.9 | 47.1 | 3.31× | 53.5 |
| 8/1 d128, batch 128 | 160.3 | 54.3 | 2.95× | 21.7 |
| 8/1 d128, batch 32 | 42.4 | 16.3 | 2.60× | 8.5 |
| 64/8 d64, batch 128 | Out of memory | 179.0 | n/a | 212.7 |
| 64/8 d128, batch 32 | Out of memory | 103.8 | n/a | 43.0 |
| 32/8 d128, batch 32 | Out of memory | 54.5 | n/a | 22.1 |
| 32/2 d256, batch 32 | Out of memory | 102.8 | n/a | 93.7 |
| 32/2 d256, batch 128 | Out of memory | 391.2 | n/a | 367.8 |
The backend stayed faster on 27 of 31 multi-token rows. The other 4 are 64/8 d64 at batch 32 and 128.
Planted bugs
I planted bugs one at a time on an A100, restoring the kernel after each.
| Planted bug | New tests failing | NVIDIA's SM80 tests |
|---|---|---|
| LPT_L2 short last block split by the full group count | 4 of 80 | All 411 pass |
| KV head (h // H_kv) % H_kv | 70 of 80 | All 411 pass |
| KV head h // H_kv, d256 kernel only | 37 of 80 | All 60 d256 cases pass |
| KV head h % H_kv (control) | 70 of 80 | Caught, 30 fail |
I read the control as showing that NVIDIA's suite does test head mapping, just not at a ratio that separates the two.
What it doesn't show
- Correct output at a 128K KV cache. The timings don't check it.
- It doesn't show a fast kernel. At 64/8 d128 the prototype takes 6.66 ms against the backend's 3.30 ms.
- It doesn't explain that gap. I named a non-split KV walk and a mostly empty Q tile as other causes, but never timed either alone.
- The 7.28 ms is 64/8 d128 at S_q=4, and develop's run there was bimodal. At 8/1 d128 the gain is only 1.15×.
- Pass-to-pass spread ran up to 20% on some GQA rows, more than their control bands. The batched rows ran once per side.
- Bitwise equality is fp16 only. The timings are bf16, which went through tolerance checks and an order-blind checksum.
- Some shapes aren't servable. K and V take 0.5 GiB per layer at 64/8 d128, so batch 32 wouldn't fit an A100 across a whole model. The rows measure the kernel, not a deployment.
- I didn't measure #1239, so I can't compare it with my prototype.
Links
- The longer worklog: the decode background and next steps.
- #1223: the prototype, closed.
- #1287: the tests, merged.
- #1239: NVIDIA's native GQA.
- Issue #377: where the tests-only PR was suggested.