Every query head used to carry its own Key and Value — that's what fills the KV cache. Let several query heads share one K/V head and the cache shrinks by exactly n_heads / n_kv_heads, with the query side untouched.
In the attention you built back in Part 1 · lesson 6, every query head had its own Key and Value head — that's MHA, multi-head attention. The KV cache (Part 2 · lesson 3) stores one K and one V per token per head, so its size scales with the number of KV heads.
GQA keeps all 32 query heads but lets groups of them share a smaller set of K/V heads. 32 query heads over 8 K/V heads → 4 queries share each K/V → the cache is 4× smaller. The queries still differ; only the K/V are pooled.
MQA is the limit: one K/V head for all 32 query heads → a 32× smaller cache, at a little more quality cost. Modern models (Llama-3) sit at GQA-8: most of the memory win, almost none of the quality loss.
TRY → Toggle MHA → GQA-8 → GQA-4 → MQA and watch 32 query dots re-bracket into fewer K/V heads — the cyan bar collapses against the MHA baseline (32× at MQA). Now push sequence length to 131,072 and raise batch to 2+: the MHA baseline shoots past the amber 80 GB line while the GQA config stays well under it. Every number recomputes from the formula — nothing is stored.
Per token, all layers (fp16, 32 layers, head_dim 128): 2 (K&V) × 32 layers × n_kv_heads × 128 × 2 B. Only n_kv_heads changes between the four configs — everything else is fixed by the model:
| config | KV heads | KV / token | shrink | queries / K/V |
|---|---|---|---|---|
| MHA | 32 | 512 KiB | 1× | 1 |
| GQA-8 | 8 | 128 KiB | 4× | 4 |
| GQA-4 | 4 | 64 KiB | 8× | 8 |
| MQA | 1 | 16 KiB | 32× | 32 |
Full cache at seq 8,192, batch 1 (fp16): MHA 512 KiB × 8192 = 4.0 GiB (4.29 GB); GQA-8 1.0 GiB (1.07 GB); GQA-4 0.5 GiB; MQA 0.13 GiB (134 MB). The shrink is exactly n_heads / n_kv_heads = 32 / n_kv.
Only the K and V are shared. Each query head still computes its own Query and its own attention weights over the shared Keys/Values, so the model keeps most of its expressiveness. The queries in the diagram above never merge — only the cyan K/V chips do. This is why GQA is nearly-free memory: the thing that grows the cache (K/V per head, Part 2 · lesson 3) is exactly the thing GQA pools.
Tie-back: kv-memory.html named this lever — "cut n_kv_heads 32→8, the cache shrinks 4× with no change to the query side." This page is that lever with the dial in your hand.
The KV cache scales with the number of KV heads, not query heads. Sharing 32 query heads over n_kv K/V heads shrinks the cache by exactly 32 / n_kv — 4× at GQA-8, 32× at MQA — while every query head still runs its own attention.
n_heads / n_kv_heads = 32/8 = 4×. The queries stay independent — they just point at pooled Keys/Values.