Why grouped-query attention exists

Multi-head attention is a memory-bandwidth disaster at decode time. GQA keeps most of the quality and throws away most of the bandwidth bill.

The picture version

Five pictures for a reader who has never looked inside a chat model. The prose below fills in the seams the pictures skip.

Why it exists

You’ve had the long chat — a few hundred messages deep, a pasted document or two in there somewhere. You ask one more question and watch the reply stream out, word by word, at that same steady drip. Here’s the part that isn’t obvious: on a standard full-attention model, every single one of those words costs the machine a re-read of the whole conversation’s stored state. Not a re-think — a re-read, off the memory chips, once per word, once per layer. That re-read is a large part of what you’re waiting on. Hold onto that chat; it’s the running example for the whole post.

If you stare at where time actually goes during LLM autoregressive decoding — generating tokens one at a time after the prompt is processed — the answer is unintuitive: the arithmetic usually isn’t what you’re waiting on. Moving bytes out of VRAM is. Decode is memory-bandwidth-bound, and the bytes come from two places: the model weights (fixed) and the KV cache, re-read from HBM for every layer and every head. Which of the two dominates depends on your context length and batch size — and it’s the KV half that grows without bound as your chat gets longer.

Standard multi-head attention makes that bill enormous. With 64 heads, your 300-message chat leaves behind 64 separate K tensors and 64 separate V tensors per layer per token, and all of them get streamed again for every new word. That’s the memory-bandwidth wall the industry kept hitting once context windows grew.

Grouped-query attention is the flinch. It says: keep all 64 query heads — those are the cheap, expressive part — but share keys and values across groups of them. A small tweak to the attention block, the same overall transformer recipe, almost the same quality, a fraction of the KV traffic.

Why it matters now

Open the config of most modern open-weight models — Llama 2 70B, Llama 3, Mistral 7B, Qwen 2.5, Gemma 2 — and you’ll see num_attention_heads and num_key_value_heads as separate fields, with the second one smaller. That’s GQA. (The pattern isn’t universal: Gemma 7B, for instance, ships equal counts, i.e. plain multi-head.) Those two fields are the config-file fingerprint of your 300-message chat costing several times less cache than it otherwise would — how many times over depends on the ratio that model picked.

GQA isn’t the only thing that made long-context serving practical — quantization, paged KV caches, fused kernels, and bigger HBM all contributed — but it’s the one that attacks the cache size at the architecture level, before any of the others get a turn. Without it, you’re choosing between small context, more expensive inference, or the most aggressive option (full multi-query attention), whose quality cost is real enough to worry about and small enough to argue over.

The short answer

GQA = multi-head queries + shared key/value heads in groups

Picture to keep: 64 people each asking their own question of the same 8 filing cabinets — instead of 64 people each hauling around a private, nearly-identical copy of the archive. The picture breaks in one place: the 64 people don’t all get the same answer, because each one still asks a different question. Only the cabinets are shared, not what comes out of them.

You keep the full set of query heads, but you tell groups of them to share a single K head and a single V head. With 64 query heads and 8 KV heads, every group of 8 queries reads the same K and V. For your 300-message chat, the KV cache shrinks by 8×, and the KV-cache portion of decode traffic shrinks by the same factor. (Total bytes per token also include the model weights, which don’t change — so the wall-clock speedup is smaller than 8×.)

How it works

Follow the bill down, one flinch at a time.

Naive version. In multi-head attention with H heads, each head h has its own Q_h, K_h, V_h projections. The KV cache stores H copies of K and V per token per layer. Your 300-message chat, at 64 heads, leaves 64 notes per token per layer — and generating one more word means streaming all of them out of HBM. Double the chat, double the wait.

Why it breaks. Nothing here is compute-bound; it’s bandwidth-bound. The GPU spends its time reading, not multiplying. So the fix has to remove bytes, not flops.

First fix: stop duplicating. In multi-query attention (Shazeer, 2019, Fast Transformer Decoding: One Write-Head is All You Need), there’s just one K and one V per layer, shared across all query heads. The KV cache is H times smaller. Shazeer’s own paper reports only minor quality loss; the GQA paper motivates its own existence partly by appealing to MQA’s quality-and-stability trade-offs. I’d treat “MQA always degrades quality noticeably” as folklore rather than an established result.

Why that breaks (or at least worries people). One shared K/V per layer is the maximally aggressive setting, and it’s the one where the quality-and-stability concerns show up. You’ve gone from 64 knobs to 1, with nothing in between.

The fix: make it a dial. GQA (Ainslie et al., 2023, EMNLP) splits the difference. Pick G groups, where 1 ≤ G ≤ H. Each group has its own K and V; all the query heads in that group share them. G = H recovers multi-head attention; G = 1 recovers multi-query. Llama 2 70B uses H = 64, G = 8 (the paper says the 34B and 70B variants use GQA; the exact 64/8 split is from the released config). Llama 3 ships GQA across all sizes, but the head ratios vary by size — Llama 3 8B is 32/8, for instance.

One more break, one more fix. If GQA only worked when you trained from scratch, nobody with an existing 70B checkpoint could use it. So the paper’s other contribution is the recipe for converting an existing multi-head checkpoint to GQA — they call it “uptraining.” Mean-pool the K heads inside each group, mean-pool the V heads inside each group, then continue pretraining for a small fraction (around 5%) of the original pretraining compute, on the same data recipe. The uptrained model lands close to the original multi-head quality with multi-query-class inference speed.

Why does sharing K and V work but sharing Q wouldn’t? Intuition, not proof: the queries are what each token uses to ask its own question of the past, so they need to stay diverse. Keys and values are the answers the past offers — and apparently those are redundant enough across heads that you can compress them hard. This is the standard hand-wave; there is no clean theoretical account of why the asymmetry holds, only empirical evidence that it does.

The seam worth noticing: GQA helps decode throughput much more than it helps prefill (processing the prompt). Decode is usually memory-bandwidth-bound; prefill is usually compute-bound, with K and V computed in parallel across the prompt. So for your long chat, GQA barely touches the pause before the first word — it shortens everything after.

You started with GQA = multi-head queries + shared K/V heads in groups. What did this post add that the compression line hides? — + a dial, not a switch. MQA already shrank the cache; GQA’s real contribution is that G is tunable between the two extremes, plus a cheap uptraining path so you don’t have to pretrain a new model to move along that dial.

Check yourself

Before you go — someone tells you they switched a model from H=32, G=32 to H=32, G=4 and got a 3× end-to-end speedup on their workload. Should you believe it?

And one more — would GQA help a workload that sends huge prompts but only ever asks for a one-word answer?

Famous related terms

• Multi-query attention (MQA) — MQA = multi-head queries + 1 shared K/V — the 2019 extreme that GQA generalizes; faster than GQA, with a quality trade-off whose size is debated.
• Multi-head attention (MHA) — MHA = H independent (Q, K, V) heads — the original “Attention Is All You Need” recipe; high quality, expensive at inference.
• KV cache — KV cache = past keys and values + reused per token — the data structure GQA is shrinking.
• Memory bandwidth — bandwidth = bytes/sec from HBM to compute — the actual quantity GQA is conserving.
• Multi-head latent attention (MLA) — MLA = MHA + low-rank latent K/V cache + decoupled RoPE channel. DeepSeek’s later variant that compresses K/V into a low-rank latent rather than sharing heads. Different idea, same goal: shrink the cache.

Going deeper

• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (EMNLP 2023) — arXiv:2305.13245. The primary source: read it for the exact definition of a group and for the uptraining recipe that converts an existing MHA checkpoint.
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (2019) — arXiv:1911.02150. Read this if you want the original argument for why decode is bandwidth-bound in the first place — it’s where the whole line of work starts.
• The Llama 2 paper (arXiv:2307.09288) — for what the trade-off looked like when a production team actually had to pick G for a model people would use.

A gap worth naming rather than papering over: there isn’t a third-party explainer for GQA specifically that beats the paper. The Ainslie paper is short and readable enough that it doubles as one.