When the K/V copy costs more than the attention

At 64 query heads and 8 KV heads over a 128K KV cache, cuDNN's SM80 attention path spent 7.28 ms per call copying K and V. The attention kernel itself took 6.0-6.7 ms. Removing the copy was a small change that made batched steps 2.6-3.3× faster and let 10 more batched rows run at all.

None of that code merged. NVIDIA shipped native grouped-query attention (GQA) through its own PR, so I closed mine. What merged from me is tests, and they turned out to be the more interesting result.

What I was looking at

I started on decode, not GQA. In the SM80 row of cuDNN frontend's support tracker, decode=False, so single-token steps ran on the cuDNN backend. I asked to take that row, which meant serving them from the open-source kernels on an A100.

Flipping that flag in memory gave a clear answer. The unmodified prefill kernels matched an fp32 reference on 125 of 127 decode cases. They also took 3.9× the backend's device time and were faster on none of the 76 timed shapes. Correct, but slow, so lifting the gate alone would make decode worse.

The code and a profile pointed to three causes:

1. The adapter copied K and V up to the query head count on every GQA call. At 32/8 heads, those copies took 34-69% of GPU time.
2. The KV walk wasn't split.
3. The 128-row Q tile left 1 of 128 rows live at decode.

I proposed tackling them in that order, one PR each. GQA has several query heads sharing one KV head, so the copy writes the same data out many times. The kernels already knew how to avoid it: they compute kv_head_idx = head_idx // (H_q // H_kv). A maintainer had also written that the expansion would disappear once the native path was qualified on A100.

The change, and how I measured it

The prototype changed the adapter and no kernel file. compile() got the real h_kv, and _kv_operand stopped repeating heads. At 64/8 heads, d=128, batch 1 and a 128K cache, the expansion it removed is 2 x 131072 x 64 x 128 x 2 bytes, or 4 GiB per call.

I had one rented A100 with clocks not locked, so a speedup could easily be noise. The setup was built around comparison:

• Harness: the repo's own benchmark/attention_inference, one subprocess per case, with the median of 20 samples after 5 warm-ups and a 256 MiB L2 flush before each.
• Metric: torch.profiler device time of the attention kernels, which on develop includes the two K/V copy kernels.
• Same pod: develop and the PR alternated in one pod.
• Controls: rows that run identical code on both sides, 88 backend rows and 18 FROST rows with no GQA, had a median PR/develop of 1.000. Those rows set a control band, and I only claimed rows outside it.

For correctness, 66 fp16 cases compared native GQA against the same inputs pre-expanded and run as multi-head attention, using torch.equal on O and LSE across all three schedulers.

The tooling could lie too. Before the main run, a review of it found defects that were mostly false greens. The diff tool reported OK on empty, unhashed or crashed runs, one wrapper dropped pytest's return code, and the benchmark always exited 0. Any of these could have produced a clean comparison of nothing, so they were fixed first.

What it did, and what it didn't

The time saved depends on how much the copy writes. At 64/8 it's large, and at 8/1 it's small, because the expansion is 8x smaller while the kernel takes about the same time.

A100, bf16, batch 1, S_q=2, S_kv=128K.

The other results:

• Batched steps: 6 rows ran 2.6-3.3× faster at batch 32 and 128. Those rows are really 3 configs, because S_q 2 and 8 give nearly the same times.
• Out of memory: develop ran out on 10 of the 16 batched rows (5 configs), needing a 64-512 GiB expanded workspace. The PR ran all 10.
• Prefill: 1.26-1.64× at 1K and 1.10-1.66× at 4K. I didn't claim anything at 16K, because those rows sit inside a control band of up to 22%.

The last column is the part that matters most. The backend stayed faster on 27 of 31 multi-token rows, and at 64/8 d128 it takes 3.30 ms against the PR's 6.66 ms. Removing the copy removed waste. It didn't make the kernel fast, and I never measured the other two causes on their own.

Going back over the write-up before posting, several claims turned out too strong:

• "Bit-identical" became "same checksum". The harness compares an XOR fold of torch.hash_tensor, so only the 66-case test is truly bitwise.
• 941 cases were 922 distinct ones. 19 ids repeat.
• The first description overstated "zero scratch" and "bitwise", and got the history of the expansion wrong.

Why it closed, and the tests that merged

Getting reviewed was its own surprise. In a look through the repo, 0 of 38 draft PRs had received a maintainer review while still a draft. Mine sat for about a day and a half, then was approved the same day I marked it ready. CI then hung on an SM100 lane, hitting a 25-minute watchdog. My change only touched SM80 code, and the same test had hung on #1183 earlier that day, so I asked for a rerun.

The next day, #1239 opened and merged within the same day, as part of NVIDIA's prepared-launch work. It served dense GQA natively on SM80. By my read it covered most of what mine set out to do, though the legacy path still expanded K/V for some head widths, for RoPE and for strides not divisible by 8. I closed mine and asked on the issue where an outside contribution would help most. One option I offered was a tests-only PR, and that's what was suggested back.

The tests were the part still worth keeping. By my read, develop had no bitwise GQA check at that point.

The tests carry a sharper finding than the speedup. Develop's GQA tests used 4/2 and 16/4 heads, and at those ratios two different quantities happen to be equal: query heads per KV head, and KV heads. A kernel that confuses them still passes.

So the tests add 8/4 and 8/2. I then planted bugs one at a time on an A100, restoring the kernel after each:

The control shows the old suite does test head mapping, just not at a ratio that separates these two quantities. A maintainer then ran the tests on an A100, with 68 and 114 passing, and reproduced the main claim: with the second planted bug in the kernel, the existing tests still passed while 31 of 68 and 6 of 12 new tests failed. Before that, the holdup had been a milestone and project that only maintainers can set. The tests merged on Oct 6.

Limits, and what I'd do next

• "Longer than the kernel" is one shape. The 7.28 ms is 64/8 d128 at S_q=4, and develop's run there was bimodal, at 11.96-14.57 ms in one run. At 8/1 the same kind of step gains only 1.15×.
• Clocks weren't locked. Pass-to-pass spread reached 20% on some GQA rows, more than their control bands.
• Bitwise equality is fp16 only. bf16 and the d64 and d192x128 flavors went through tolerance checks and an order-blind checksum.
• Some of these shapes aren't servable. A 64/8 d128 cache at 128K is 0.5 GiB per layer, 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. I can't say how it compares with my prototype.

Next, I'd add a small-Q tile and split-KV one at a time and time each alone, because the gap to the backend is still unattributed. I'd also rerun this harness on develop after #1239, which would answer the comparison directly.

Sources

PR #1223, PR #1287, #1239 and issue #377.