MQA and GQA: Sharing Keys and Values
Lesson 3 of 3 in Multi-Head, Multi-Query, Grouped-Query Attention.
So far heads were about quality. This lesson is about what they cost you at serving time. When a model generates, each new token must attend over every previous token’s keys and values. Recomputing those for the whole context on every step would be ruinous, so every serious serving stack caches them — that store is the KV cache.
Here is the problem: with classic multi-head attention, every head in every layer keeps its own K and V for every token in the context. The cache grows linearly with context length, and it lives in the same accelerator memory as the model Weights. Longer conversations and bigger batches eat that memory fast — in practice, the KV cache, not arithmetic, is often what limits how many users one accelerator can serve.
Multi-query attention (MQA) attacks the cache directly. Shazeer (2019, arXiv:1911.02150) noticed that at decode time the expensive part is loading cached keys and values, while queries are computed fresh for the current token anyway. So: keep all h query heads, but share one key head and one value head across all of them. The cache shrinks by roughly the head count. The price is a measurable quality drop on some tasks, and — as later work noted — training can be less stable in some setups.
Grouped-query attention (GQA) (Ainslie et al. 2023, arXiv:2305.13245) softens the deal: instead of one shared K/V pair, use g groups, each shared by h / g query heads. Set g = h and you have MHA; set g = 1 and you have MQA; anything between interpolates. The paper’s evals found quality close to full MHA at a fraction of the cache — and, just as usefully, an existing MHA checkpoint can be converted (mean-pool each group’s K/V heads) and then ‘uptrained’ briefly rather than retrained from scratch. Grouped-query attention has since become a common choice in openly documented model families precisely because of this trade.
| Multi-head (MHA) | Multi-query (MQA) | Grouped-query (GQA) | |
|---|---|---|---|
Query heads |
|
|
|
Key/value heads |
|
|
|
KV-cache size | baseline | ~ | ~ |
Attention quality | the reference point | measurable drop on some tasks | close to MHA in published evals |
How you get there | the default when training from scratch | train from scratch, or convert and uptrain | convert an MHA checkpoint (mean-pool K/V per group), then uptrain briefly |
The KV-cache arithmetic
Per sequence, the cache size is a straight product:
KV bytes ≈ 2 × n_layers × n_kv_heads × head_dim × seq_len × bytes_per_value
The 2 is keys and values. n_kv_heads is h for MHA, 1 for MQA, g for GQA. bytes_per_value comes from precision — 2 bytes for 16-bit floats, less if the cache itself is quantized. Multiply the whole thing by the number of concurrent sequences to size a server. Notice which factors you cannot easily touch: layers, head_dim, and precision move model quality directly, and sequence length is the product requirement. n_kv_heads is the one lever that shrinks the cache without narrowing the model — exactly the knob MQA and GQA turn.
The cache is a bandwidth problem too, not just a capacity one: every generated token must read the entire cache back through the memory system. Halving the cache roughly halves that per-token memory traffic, which is why fewer KV heads shows up as faster decode and bigger feasible batches — not merely as more sessions fitting in memory.
In production
GQA exists because serving memory, not arithmetic, is the bottleneck of generation. The KV cache is the mechanism behind a pattern you will meet on every cloud: long context costs disproportionately more than short.
AWS
Self-hosting on AWS GPU instances, weights and KV cache share one fixed pool of accelerator memory. Fewer KV heads means more room for concurrent sequences on the same hardware — that is throughput, and therefore cost per token. Capacity planning for an endpoint starts from tokens-in-flight (sum of active contexts), not requests per second, because every live conversation holds its cache for its whole lifetime.
Azure
The same mechanism sits underneath reserved-capacity offerings on Azure: what a capacity unit fundamentally rations is memory and bandwidth for tokens in flight. Two workloads with identical request counts but different context lengths consume very different capacity. Prompt trimming, streaming, and session limits are all KV-cache management wearing product names.
Google Cloud
On Google Cloud accelerators, decode speed is bounded by how fast the chip can re-read the KV cache for every generated token, so cache size drives latency as well as capacity. Smaller caches (GQA/MQA) mean less memory traffic per token — faster decode and larger feasible batches. This is also the mechanism behind long-context pricing tiers: past a point, every extra thousand tokens of context imposes real memory and bandwidth cost on each subsequent token.
Interactive checkpoint quiz (1 questions) — open this page in a browser to take it.