KV cache and attention variants

When a model writes a reply, it reuses the keys and values of every earlier token instead of recomputing them. That store is the KV cache. Its size grows with the number of layers, the number of key-value heads, the head width and the sequence length. Attention variants such as MQA, GQA and MLA exist mainly to shrink it.

The big idea

Full multi-head attention stores separate keys and values for each head. The variants change how many key-value sets are kept:

GQA is the practical default for open dense models. Quality stays close to full multi-head attention, and the cache is much smaller. MLA compresses further, at the cost of more complicated code.

Walking through the cache arithmetic

For one sequence, the cache per token is roughly:

bytes per token ≈ 2 · layers · kv_heads · head_dim · bytes per element

The factor 2 covers keys and values. Take a model with 80 layers, head width 128, and 16-bit numbers (2 bytes each).

MHA with 64 KV heads: 2 · 80 · 64 · 128 · 2 = 2,621,440 bytes ≈ 2.6 MB per token
GQA with 8 KV heads:  2 · 80 · 8 · 128 · 2  =   327,680 bytes ≈ 0.33 MB per token

Now multiply by 128,000 tokens, for one sequence:

MHA: 2.6 MB × 128,000 ≈ 335 GB
GQA: 0.33 MB × 128,000 ≈ 42 GB

The GQA cache is eight times smaller, and that difference decides whether a long context fits on one GPU.

Common mistakes

• Counting query heads. The cache scales with key-value heads. Query heads do not add cache.
• Over-shrinking the key-value heads. Too few heads can hurt tasks that need exact recall. Measure on the tasks you care about.
• Comparing cache sizes at different precisions. A 16-bit cache and an 8-bit cache differ by 2×. Check the precision before comparing.

Check yourself

1. A model has 32 key-value heads where another has 8, with everything else the same. How much bigger is the cache? Answer: Four times, because the cache scales with key-value heads.
2. Why does GQA usually keep quality close to full attention? Answer: Each query head still has its own queries. Only the keys and values are shared within a group.

Key takeaways

• The KV cache stores keys and values for every past token, and its size grows with the sequence.
• Cache bytes per token scale with layers times key-value heads times head width.
• GQA is the usual trade: a much smaller cache, close to full-attention quality.
• MLA compresses the cache with a latent, at higher implementation cost.

Further learning

The GQA paper introduces grouped-query attention and its training recipe:

GQA: Training Generalized Multi-Query Transformer ModelsMulti-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation, and moreover it may not be desirable to train a separate mod…arXiv.org

Next, see how a model can hold many more parameters than it runs per token: mixture of experts.