Grouped-query attention and the KV cache

Training processes a whole sequence at once. Generation doesn't. It produces one token, appends it, and runs again. And if you're not careful, every step recomputes the keys and values of every earlier token, even though they never change. The fix is to just keep them, and what it costs to keep them is what this post is about.

Keys and values don't change, so keep them

In causal attention, the key and value of position $j$ only depend on tokens up to $j$. When token $t + 1$ comes in, positions 1 to $t$ have exactly the same keys and values as before. So a decoder can store them in a cache and only compute the new token's query, key and value each step. The new query attends to the cached keys, and the new key and value get appended.

The cache costs memory for every layer, every key and value head, every head dimension and every cached position:

The leading 2 is one key and one value. For a real model at a long context this gets big, and generation reads all of it for every new token. Decoding one token does very little arithmetic and a lot of reading, so its speed is set by memory bandwidth. A smaller cache means less memory and also faster decoding.

And the only term a design choice can change without losing information is $H_{kv}$.

Let query heads share key and value heads

Standard multi-head attention (MHA) has one key head and one value head per query head. Multi-query attention (MQA) keeps all the query heads and shares one single key and value head across all of them. Grouped-query attention (GQA) is in between, the query heads split into groups and each group shares one key and value head.

So what happened ?

• The cache shrinks by exactly the ratio of query heads to KV heads. Nothing else in the formula moves, so you know the saving before any training.
• The paper-exercise shape, 12 layers at width 512 and 4K positions in bfloat16, needs 96 MiB with MHA and 24 MiB with two KV heads. That's the Day 2 exercise, done by hand in the note.
• Each query head still has its own query projection. Heads in a group can look for different things, they just search the same keys.

Why MQA is fastest and still loses quality

The GQA paper measured this on T5-XXL. Inference time went from 1.51 s with MHA to 0.24 s with MQA, because decoding reads a cache one eighth the size or less. But quality dropped too, from 47.2 to 46.6 average ROUGE. With one key and value head, every query head retrieves from the same subspace, so heads can't specialize in what they look up anymore. GQA with 8 groups took 0.28 s and scored 47.1, so most of the speed and almost none of the loss. The paper picked 8 groups because the slowdown from MQA stays small at first and grows as the group count gets close to MHA.

The note flags one thing not to copy. Eight groups was a choice for a 64-head model. octlm has 8 query heads, where 8 KV heads is just plain MHA. The ratio and the shape of the curve transfer, the number doesn't. So EXP-013 swept 1, 2, 4 and 8.

Fewer key and value heads, broadcast at attention time

The key and value projections output kv_heads × head_size columns instead of d_model. The fused kernel broadcasts each KV head across its group with enable_gqa=True. The handwritten Day 1 path does the same thing explicitly with repeat_interleave, which is how a test checks the two agree.

def kv_cache_bytes(config: DecoderConfig, length: int, element_bytes: int = 2) -> int:
    """Bytes held by the key and value cache for one sequence, all layers."""
    per_token = config.cache_dims * element_bytes
    return per_token * config.n_layers * config.cached_positions(length)

The cache, measured at the Day 2 shape

At four layers, width 256, eight query heads and bfloat16, from runs/day2-exp013-cache.jsonl:

Every time you halve the KV heads, the cache halves at every length.

Quality, measured in the grid

So what happened ?

• All three grouped variants land inside the baseline's seed band on code. gqa-4 is 0.034 better, gqa-2 0.011 worse, mqa-1 0.005 better. None of that is a signal.
• They have fewer parameters, because the key and value projections shrink, 3,346,944 for gqa-2 against 3,740,160.
• All three are 20 to 24 percent slower per training step. Training never reads a cache, so it gets none of the benefit, and the broadcast inside the kernel costs time.

Decision, keep two KV heads. It cuts the 4K cache from 16 MiB to 4 MiB, and its quality difference is inside the noise. One KV head would halve the cache again at the same measured quality, but the note doesn't take it, because the GQA paper's quality penalty for MQA grows with model size, and a 400-step run at width 256 just can't see it. Two heads is the hedge. The Day 4 model has two.

What we didn't build, and why

• The cache itself, on Day 2. Day 2 measured cache sizes from the formula and didn't actually cache tensors during generation. Day 5 built and measured it, exact greedy parity, 256 KiB to 2 MiB of cache, and 10.6 to 15.8 percent faster generation in single T4 timings.
• Converting an MHA checkpoint to GQA. The paper mean-pools the key and value projections in each group, then trains for 5 percent of the original steps. octlm trains GQA from scratch, so it never needed the conversion.