Grouped-query attention: the middle ground between MHA and MQA
The KV cache note ended with the size formula and the observation that the number of key-value heads is one of its terms. Grouped-query attention is the architectural change that shrinks that term without giving up much of what multi-head attention can do. This note lays out the three options side by side, with the head counts of real models, and explains why nearly every model since 2023 picked the one in the middle.
Multi-head attention gives each of the N query heads its own key head and value head. Multi-query attention gives all N of them a single shared pair, which cuts the cache by N× but costs quality and can destabilise training. Grouped-query attention splits the query heads into G groups and gives each group one key-value pair: the cache shrinks by N/G, decode speed is close to multi-query, and quality is close to multi-head. A trained multi-head model can be converted by averaging the key and value heads within each group and training briefly.
Contents
01Three ways to assign key-value heads
In every attention layer, each query head needs a key head to match against and a value head to read from. The question the three variants answer differently is how many distinct key and value heads exist, and which query heads share them.
Multi-head (MHA) Grouped-query (GQA) Multi-query (MQA) one KV pair per Q head one KV pair per group one KV pair for all Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q │ │ │ │ │ │ │ │ └──┬──┘ └──┬──┘ └───────┬───────┘ K K K K K K K K K K K V V V V V V V V V V V 8 KV heads 2 KV heads 1 KV head
The query heads are the same in all three. The query projection has the same size, each head still computes its own attention pattern, and the output is still the concatenation of N head outputs. What changes is the size of the key and value projections, the number of distinct key and value tensors in the cache, and how much has to be read from memory at every decode step.
02Multi-head attention
The original Transformer design, and the one GPT-3 and most models before 2023 used. With N query heads there are N key heads and N value heads. Each head has a fully independent view of which tokens to attend to and what to read from them, which is the maximum of modelling capacity for a given width.
The cost is in the cache. Every token stores N key vectors and N value vectors per layer, and all of them are streamed from memory at every step of generation. For a large model at long context, that cache is bigger than the weights and the decode phase is bound by memory bandwidth, as the KV cache note works through.
03Multi-query attention
The opposite extreme, proposed in 2019: keep N query heads but a single key head and a single value head, shared by all of them. Each query head still computes its own attention weights, because its query is different, but they all score against the same keys and all read from the same values.
- What it buys. The cache is N× smaller, and so is the per-step memory traffic for attention. Generation throughput and achievable batch size go up by a large factor. PaLM and Falcon used it.
- What it costs. All heads now have to find what they need through one key space and one value space, so the model loses representational capacity, and benchmark quality drops measurably. The GQA paper also reports training instability with multi-query at large scale.
04Grouped-query attention
Grouped-query attention, from Ainslie and colleagues at Google in 2023, interpolates. The N query heads are divided into G groups, with 1 < G < N, and each group shares one key head and one value head. G = N is multi-head; G = 1 is multi-query.
The cache shrinks by N/G. The real head counts show the ratios in use:
| Model | Query heads N | KV heads G | Queries per KV head | Cache versus MHA |
|---|---|---|---|---|
| Llama 3 8B, Mistral 7B | 32 | 8 | 4 | 1/4 |
| Llama 2 70B, Llama 3 70B | 64 | 8 | 8 | 1/8 |
| Llama 3.1 405B | 128 | 8 | 16 | 1/16 |
| Gemma 2 9B | 16 | 8 | 2 | 1/2 |
| Gemma 2 27B | 32 | 16 | 2 | 1/2 |
Eight key-value heads is a recurring choice. One reason is tensor parallelism: when a model is split across eight GPUs, each device holds one key-value head and its share of query heads, with no duplication of the cache. The query count then grows with the model while the KV count stays put, so the bigger the model, the larger the saving.
05Side by side
| Multi-head | Grouped-query | Multi-query | |
|---|---|---|---|
| KV heads | N, one per query head | G, one per group | 1, shared by all |
| Cache size | Base | 1/2 to 1/16 of base; 1/8 is typical for large models | 1/N of base |
| Decode speed | Slowest, memory-bound | Close to multi-query | Fastest |
| Quality | Highest | Within noise of multi-head in most reported evaluations | Measurable drop; unstable at scale |
| Used in | Original Transformer, GPT-3, Llama 1 | Llama 2 70B, Llama 3, Mistral, Gemma 2, Qwen 2, most models since | PaLM, Falcon |
06Why GQA became the default
At the context lengths models now serve, 8k to 128k and beyond, decode is bound by memory bandwidth, and the key-value cache is a large part of what is streamed each step. Take Llama 3 70B's ratio of eight query heads to one KV head.
- Memory. The cache is 87.5 percent smaller than multi-head. At 128k context in 16-bit that is 43 GB instead of 344 GB per request. The freed memory goes to more concurrent requests or longer contexts on the same hardware.
- Speed. Attention at each decode step reads one eighth of the key and value bytes from high-bandwidth memory, so the attention part of the step is close to eight times faster, and the step as a whole is faster in proportion to how much of it attention was.
- Quality. Each group still has its own key and value space, so the model keeps most of the distinct routing that multi-head allows. The GQA paper's uptrained T5-XXL variants land within a fraction of a point of multi-head on summarisation and translation benchmarks while matching multi-query speed.
So the trade is a large, predictable gain on the bottleneck resource for a small, measured loss on quality. That is a trade almost every model builder takes, and since 2023 almost all have.
07Converting a multi-head model: uptraining
GQA did not require retraining from scratch, which is part of why it spread quickly. The paper's recipe converts an existing multi-head checkpoint in two steps.
- Mean-pool the projections. For each group, take the N/G key heads that will be merged and average their projection matrices into one; do the same for the value heads. Averaging preserved more of the original model than picking one head from the group or initialising randomly, in the paper's ablation.
- Uptrain briefly. Continue pretraining the converted model on the original objective for a small fraction of the original budget, about 5 percent of the pretraining steps in the paper. This lets the query heads in each group adapt to their now-shared keys and values.
After uptraining, the converted model recovered nearly all of the multi-head baseline's quality. The same procedure works for multi-query, with every head pooled into one, but the quality gap it leaves is larger. For a model trained with GQA from the start, as Llama 2 70B and everything after were, none of this is needed.
Grouped-query attention reduces the number of key-value heads. Multi-head latent attention, used in DeepSeek's models, goes a step further: instead of storing keys and values, it stores a low-rank latent vector per token and reconstructs them on the fly, so the cache is smaller than even an aggressive GQA ratio while keeping per-head keys and values at compute time. It is the same trade, pushed harder.
08Takeaways
- The three variants differ only in how many key-value heads exist. Query heads, attention patterns and outputs are the same shape in all of them.
- Multi-query is the extreme: one shared pair, N× smaller cache, a real quality cost.
- Grouped-query shares one pair per group. The cache shrinks by N/G; eight KV heads is the common choice, giving 4× for 7B to 8B models and 8× for 70B.
- The saving lands on the bottleneck. Decode is bandwidth-bound and the cache is what is streamed, so a smaller cache is faster decode and bigger batches.
- Quality stays close to multi-head. Each group keeps its own key and value space.
- Existing models convert cheaply. Mean-pool the heads in each group, then uptrain for a few percent of the original budget.