Story of Attention — Part 3: KV-Cache Bottleneck: MQA and GQA

Part 3 of 6 in the Story of Attention series.

1. Before the Transformer — how n-gram counting and RNN memory hit their ceilings, and how additive/multiplicative attention first let a decoder look back.
2. Attention All You Need — scaled dot-product attention, multi-head attention, and the four-year search for a way to encode position.
3. KV-Cache Bottleneck: MQA and GQA (this post) — sharing and grouping key/value heads to shrink inference memory.
4. Sparse Attention, Then and Now — fixed, windowed, hashed, and learned sparsity patterns, from 2019 to 2025.
5. Linear Attention: Kernels, Decay, and Gating — replacing softmax with decomposable kernels, decay, and gated recurrent state.
6. Where Attention Stands Today — latent compression, differential attention, and what frontier LLMs actually ship.

---

Each layer of a Transformer stores one key and one value per head, per past token, across the entire sequence generated so far. At a 128,000-token context, that cache alone takes up gigabytes of GPU memory. It just sits there, reread in full for every generated token, without doing any computation itself. Multi-Query and Grouped-Query Attention answer one very simple question: do all those heads really need their own keys and values?

Two problems carried over from the last part: attention computation grows quadratically with sequence length, and the KV cache grows with the number of heads, layers, and context length. This part tackles the cache, because at inference time that's what actually hurts.

---

1. Multi-Query Attention — MQA

Do all the heads really need their own keys and values? Shazeer asked this question in 2019, in a paper titled "Fast Transformer Decoding: One Write-Head is All You Need" [1]. The answer was to keep all $h$ query heads, but share a single key head and a single value head across all of them. So each head still asks its own question through its own $W_i^Q$, but all the heads look up the same shared keys and values.

$Q_i = XW_i^Q \in \mathbb{R}^{n \times d_q} \quad \text{for } i = 1, \ldots, h$

$K = XW^K \in \mathbb{R}^{n \times d_k} \quad \text{(single shared K)}$

$V = XW^V \in \mathbb{R}^{n \times d_v} \quad \text{(single shared V)}$

$\text{head}_i = \text{softmax}\!\left(\frac{Q_i K^T}{\sqrt{d_k}}\right) V$

With this change, the KV cache drops from $2h$ matrices to just 2, a reduction by a factor of $h$. For a model with 64 heads, that is a 64× reduction in cache size.

But for the next three years, not much happened with this idea, because in 2019 the context windows were around 512 or 1,024 tokens, and a cache that small costs almost nothing to store. So the problem MQA solves had not shown up yet. The paper got its recognition in 2022, when Google's PaLM [2] shipped MQA at scale and showed what the savings looked like in production.

Memory Bandwidth and Incremental Decoding

In the previous part we saw that memory bandwidth is what limits decoding. A useful way to measure this is the ratio of memory access to arithmetic. When this ratio is high, the workload spends its time waiting on memory rather than computing. For MHA during incremental decoding:

$\text{Ratio}_\text{MHA} = \frac{n}{d} + \frac{1}{b}$

For MQA, the KV cache term shrinks by a factor of $h$:

$\text{Ratio}_\text{MQA} = \Theta\!\left(\frac{1}{d} + \frac{n}{d \cdot h} + \frac{1}{b}\right)$

The $n/d$ term is the one that grows with context length, and MQA cuts it by a factor of $h$. Where does this factor come from? Across $n$ decoding steps, MHA's cumulative KV-cache memory traffic is $\Theta(b \cdot n^2 \cdot d)$, while MQA's is $\Theta(b \cdot n^2 \cdot k)$ where $k = d/h$, and substituting $k$ turns the middle term into $n/(d \cdot h)$. In simple words, decoding was spending most of its time reloading keys and values from memory, and MQA removes most of that traffic.

The numbers reported in the paper show the effect clearly [1]:

The quality is essentially unchanged, the small BLEU difference is just noise, while decoding became an order of magnitude faster. The model still computes the same matrix multiplication as before, each query head runs its own full attention pass, exactly as in MHA. The speedup comes from every head using the same $K$ and $V$ instead of its own copy, so there is far less memory to load at each step.

PaLM [2], StarCoder [3], and Falcon-7B [4] shipped MQA. But as models and evaluations grew more demanding, a real cost showed up. When every query head reads the same keys and values, the heads lose some of their ability to specialize. On aggregate benchmarks this degradation is hard to see, but it becomes measurable on harder tasks, the ones that need the heads to look at the context from several different angles. So one shared KV head turned out to be slightly too few. The obivious question becomes, is there a middle ground between $h$ KV heads and one?

---

2. Grouped-Query Attention — GQA

In 2023, Ainslie et al. [5] answered this with Grouped-Query Attention. The idea is to partition the $h$ query heads into $G$ groups, and give each group one shared key head and one shared value head. This makes MHA and MQA the two extremes of the same design. With $G = h$ every head gets its own keys and values, which is full MHA, and with $G = 1$ all the heads share one set, which is MQA. The useful settings sit in between.

_{Multi-head, grouped-query, and multi-query attention. All three keep every query head; what changes is how many key/value heads the queries share: one per query head (MHA), one per group (GQA), one for all (MQA). Redrawn from Ainslie et al., 2023 [5].}

$K_g = XW_g^K, \quad V_g = XW_g^V \quad \text{for } g = 1, \ldots, G$

For query head $i$ belonging to group $g(i) = \lceil iG/h \rceil$:

$\text{head}_i = \text{softmax}\!\left(\frac{Q_i K_{g(i)}^T}{\sqrt{d_k}}\right) V_{g(i)}$

$\text{MultiHead}_\text{GQA} = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W^O$

Key result: For $h = 32$ heads, $d_k = 128$, $n = 2{,}048$, in float16 (2 bytes/element):

Method	KV Entries	Memory per Layer
MHA ($G = h = 32$)	$2 \times 32 \times 2048 \times 128 \times 2$	33.6 MB
GQA ($G = 8$)	$2 \times 8 \times 2048 \times 128 \times 2$	8.4 MB
MQA ($G = 1$)	$2 \times 1 \times 2048 \times 128 \times 2$	1.0 MB

GQA with 8 groups delivers a 4× cache reduction over MHA while retaining most of MHA's quality.

The GQA paper also comes with a very practical benefit. An existing MHA checkpoint can be converted into a GQA model, by mean-pooling the KV heads within each group and then uptraining for a further 5% of the original pretraining compute, so there is no need to pretrain from scratch. The paper measured the tradeoff as follows [5]:

So GQA-8 gives up almost nothing in quality for a 5.4× speedup, while MQA, faster still at 6.3×, gives up a bit more. Between these two options, the field picked GQA. It became the architectural standard for open-weight LLMs, and LLaMA 2's larger variants (34B/70B) [6], all of LLaMA 3 [7], Mistral [8], Mixtral [9], Gemma 2 [10], and Qwen2 [11] all use it. In fact, when a model from 2023 onward does not mention its attention variant, it is usually GQA.

But it is equally important to see what GQA does not do. It compresses the cache by a constant factor. The attention computation inside every head is still quadratic in sequence length. At $n = 128{,}000$, each head's $n \times n$ matrix still holds 16.4 billion entries, and it gets recomputed for every head in every layer. Sharing key/value heads does nothing to make this matrix smaller. To shrink the matrix itself, we need a different kind of idea, which is to simply not compute most of it.

---

What comes next

In this part, we saw how sharing and grouping key/value heads shrank the KV cache. The change was cheap, and the quality cost was small enough that it almost became the default mechanism in this field. But GQA leaves the $n \times n$ attention matrix untouched, and at long context lengths that matrix is what actually dominates the cost.

The next part goes after the matrix directly, by not computing most of it. It follows the sparse-attention line of work from the first hand-designed patterns (Sparse Transformer, Longformer, BigBird), through content-based hashing (Reformer) and the attention-sink observation, up to the learned sparsity that ships in DeepSeek's 2025 models.

Continue to Part 4: Sparse Attention, Then and Now →

---

References

[1] Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150.

[2] Chowdhery, A., Narang, S., Devlin, J., et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311.

[3] Li, R., Allal, L. B., Zi, Y., et al. (2023). StarCoder: May the Source Be with You! arXiv:2305.06161.

[4] Almazrouei, E., Alobeidli, H., Alshamsi, A., et al. (2023). The Falcon Series of Open Language Models. arXiv:2311.16867.

[5] Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245.

[6] Touvron, H., Martin, L., Stone, K., et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288.

[7] Grattafiori, A., et al. (2024). The Llama 3 Herd of Models. arXiv:2407.21783.

[8] Jiang, A. Q., Sablayrolles, A., Mensch, A., et al. (2023). Mistral 7B. arXiv:2310.06825.

[9] Jiang, A. Q., Sablayrolles, A., Roux, A., et al. (2024). Mixtral of Experts. arXiv:2401.04088.

[10] Gemma Team, Riviere, M., et al. (2024). Gemma 2: Improving Open Language Models at a Practical Size. arXiv:2408.00118.

[11] Yang, A., Yang, B., Hui, B., et al. (2024). Qwen2 Technical Report. arXiv:2407.10671.