Part III · Inference Internals & Production Serving
Chapter 33 Deep dive ~22 min read

The attention compression family: MHA, MQA, GQA, MLA

"The KV cache is the new memory bottleneck. Attention head architecture is the new compression target"

Welcome to Stage 3 of Part III: the research frontier. The previous chapters (Stage 2) covered the practitioner internals you need to be effective in production. The next nine chapters cover the cutting-edge research that’s actively shaping how the next generation of models is built. By the end of Stage 3 you should be able to read any new architecture paper, place it in context, and predict whether it will matter for production.

We open Stage 3 with the attention compression family: a sequence of architectural choices about how to organize the K and V projections in attention. Each choice trades off model quality, KV cache size, and training stability. The progression from MHA → MQA → GQA → MLA is one of the cleanest examples of “research finds an elegant solution to a production problem” in the LLM era.

This chapter goes deep on why these variants exist, how they differ, and which is right for which use case.

Outline:

  1. The KV cache problem revisited.
  2. MHA — multi-head attention, the original.
  3. MQA — multi-query attention.
  4. GQA — grouped-query attention.
  5. MLA — multi-head latent attention (DeepSeek).
  6. How each affects KV cache size.
  7. How each affects model quality.
  8. The training-quality-vs-inference-cost frontier.
  9. The choices in current models.

33.1 The KV cache problem revisited

Recall from Chapter 22: the per-token KV cache size is

2 × n_layers × n_kv_heads × d_h × bytes

For Llama 3 70B with 80 layers, 8 KV heads, 128 head dim, in bf16, that’s 320 KB per token. For a context of 32k tokens, that’s 10.5 GB of KV cache per request. For 100 concurrent users, 1 TB.

The KV cache is the dominant memory cost of serving any non-trivial LLM. The model weights are static and shared across requests; the KV cache scales as users × context_length. Anything that reduces the per-token KV cache size translates directly to more concurrent users and longer contexts at the same hardware cost.

The most direct way to reduce per-token cache size is to reduce n_kv_heads. The original transformer (and BERT, GPT-2, and most pre-2023 models) has n_kv_heads = n_attention_heads. So if a model has 64 attention heads, it has 64 KV heads, and the KV cache is at full size. Modern models reduce n_kv_heads aggressively, often to 8 or fewer, with techniques that preserve model quality despite the reduction.

This is the attention compression family. Each variant is a different way to use fewer KV heads than query heads.

KV cache memory per token scales linearly with n_kv_heads; MHA at 64 heads is 64x larger than MQA at 1 head. n_kv_heads (log scale) MQA 1 GQA-8 8 GQA-16 16 GQA-32 32 MHA 64 MLA ~latent
KV cache per token scales with n_kv_heads — MHA at 64 heads costs 64× more than MQA, while GQA-8 (the modern default) sits at 8× savings with negligible quality loss.

33.2 MHA — multi-head attention, the original

Multi-head attention (MHA) is the original attention from Attention Is All You Need (2017). Each attention head has its own Q, K, and V projections. For a model with H heads of dimension d_h:

  • W_Q^(i) of shape (d_model, d_h) for each head i
  • W_K^(i) of shape (d_model, d_h) for each head i
  • W_V^(i) of shape (d_model, d_h) for each head i

The total parameter count for the QKV projections per layer is 3 × d_model × (H × d_h) = 3 × d_model².

The KV cache stores K and V for every head: n_kv_heads = H. For Llama 1 7B with H = 32, d_h = 128, the per-token KV cache is 2 × 32 × 32 × 128 × 2 = 524,288 bytes = 512 KB. For a 32k context, that’s 16 GB per request.

MHA is the maximum-quality, maximum-cost option. Every head is fully independent and can learn its own attention pattern. The cost is the unaffordably large KV cache.

GPT-3, the original Llama 1, and most pre-2023 transformers used MHA. The reason: at the time, KV cache size wasn’t the dominant concern (people weren’t yet thinking about long-context production serving). Once long-context production serving became the primary use case, MHA’s cost became unsustainable.

Head layout comparison: MHA has one K/V per query head; MQA shares a single K/V; GQA shares one K/V per group of query heads. MHA H Q-heads, H KV-heads Q₁ K₁ V₁ Q₂ K₂ V₂ Q₃ K₃ V₃ Q₄ K₄ V₄ 4 independent K/V pairs MQA H Q-heads, 1 shared KV Q₁ Q₂ Q₃ Q₄ K V 1 shared K/V pair (64× smaller cache)
MHA gives every query head its own independent K and V; MQA forces all query heads to share a single K and V, collapsing the KV cache by a factor of H but limiting expressiveness.

33.3 MQA — multi-query attention

Multi-query attention (MQA) (Shazeer, 2019) is the extreme opposite of MHA: share a single K and V across all heads. There’s only one K projection and one V projection for the whole layer, but still H separate Q projections.

The K and V are computed once per token and broadcast to all heads. Each head has its own Q but shares the K and V with every other head.

The KV cache is now n_kv_heads = 1. For the same model dimensions, the per-token cache shrinks by a factor of H. For Llama-style models with 32 heads, MQA gives a 32× reduction in KV cache size.

That’s the win. The cost is quality loss. Sharing K and V means every head has to attend to the same projected representation of the past. Heads can no longer specialize their attention patterns based on different views of the input. The model is less expressive.

In practice, MQA at the same parameter count costs about 1-2 percentage points on standard benchmarks compared to MHA. Some models can be retrained from scratch with MQA and recover most of the gap; others can’t.

MQA was used in PaLM and a few other early models. It was a clever idea but the quality loss was real, and the field moved to a compromise: GQA.

33.4 GQA — grouped-query attention

Grouped-query attention (GQA) (Ainslie et al., 2023) is the middle ground between MHA and MQA. Instead of having all heads share one K/V (MQA) or each head having its own K/V (MHA), groups of heads share a K/V.

For a model with H = 32 query heads and n_kv_heads = 8, each KV head is shared by 4 query heads. The 32 query heads are partitioned into 8 groups of 4; each group shares one K projection and one V projection.

The KV cache size is now n_kv_heads = 8, which is 4× smaller than MHA but 8× larger than MQA. The quality is much closer to MHA than to MQA — typically within 0.5 points on benchmarks. GQA gets 75% of the cache savings of MQA with 75% less quality loss. That’s the kind of trade-off engineers love.

GQA can be applied as an “uptrain” — take an MHA-trained model and continue training it with GQA, where the K and V projections of the head groups are initialized as averages of the original heads. After a small amount of additional training, the model recovers its original quality and now has a much smaller KV cache. Llama 2 70B was uptrained from MHA to GQA-8 this way.

GQA has become the default attention architecture for modern LLMs:

  • Llama 2 / 3 / 3.1 / 3.2 — GQA with 8 KV heads.
  • Mistral / Mixtral — GQA with 8 KV heads.
  • Qwen 2 / 2.5 — GQA with 4-8 KV heads depending on model size.
  • Gemma 2 / 3 — GQA.
  • Phi-3 / 4 — GQA.

The choice of n_kv_heads = 8 is not magical; it was the empirically discovered sweet spot in the GQA paper. Some models go lower (4 KV heads) for tighter memory budgets at small additional quality cost.

GQA groups query heads so that each group of 4 shares one K/V pair, balancing cache savings against expressiveness. GQA — Grouped-Query Attention (8 KV heads, 32 Q heads → 4 per group) Group 1 Q₁ Q₂ Q₃ Q₄ K₁ V₁ Group 2 Q₅ Q₆ Q₇ Q₈ K₂ V₂ · · · Group 8 Q₂₉ Q₃₀ Q₃₁ Q₃₂ K₈ V₈ 32 Q-heads / 8 KV-heads = 4 heads per group. KV cache is 4× smaller than MHA. Quality: <0.5 points below MHA at GQA-8. The production sweet spot.
GQA partitions H query heads into G groups, each sharing one K/V pair — the ratio H/G is the cache compression factor, and GQA-8 delivers 8× compression with under 0.5 points of quality cost.

The KV cache savings are real and substantial. For Llama 3 70B (80 layers, 64 attention heads, GQA-8, 128 head dim), the per-token cache is 2 × 80 × 8 × 128 × 2 = 320 KB. With MHA (n_kv_heads = 64), it would be 2 × 80 × 64 × 128 × 2 = 2.6 MB. 8× smaller cache with negligible quality cost.

This is one of the most consequential architectural changes in the modern LLM era. It’s a single hyperparameter — n_kv_heads — and choosing it well is what makes long-context serving affordable.

33.5 MLA — multi-head latent attention

Multi-head latent attention (MLA) is DeepSeek’s contribution (DeepSeek-V2, 2024). It’s the most ambitious compression scheme to date and goes further than GQA.

The trick: instead of reducing the number of KV heads, compress the K and V into a low-rank latent vector and decompress them on the fly during attention.

The architecture:

MLA compresses the token representation into a small latent vector and expands it per-head at attention time; only the latent is cached. token d_model=5120 compress W_DKV d_model→d_kv latent d_kv=512 ← cached (tiny) expand W_UK, W_UV per-head d_kv→H×d_h K₁ K₂ … Kₕ V₁ V₂ … Vₕ full H heads at attn time (not cached)
MLA caches only the low-rank latent vector (512 dims) per token instead of full per-head K and V (H × d_h dims), giving 32–64× cache reduction with no measurable quality loss.
  1. The model has a shared latent dimension d_kv that’s much smaller than d_model. For DeepSeek-V2 with d_model = 5120, d_kv = 512.
  2. Each token’s K and V information is compressed into a single vector of dimension d_kv. This is the latent: one vector per token, not one K vector and one V vector per head.
  3. At attention time, the latent is expanded through learned matrices to produce the per-head K and V. The expansion is per-head (different heads decompress to different K/V) but the underlying latent is shared.
  4. The KV cache stores the latent vectors, not the per-head K and V. So the per-token cache size is 2 × d_kv × bytes, dramatically smaller than even GQA.

Wait, why “2 × d_kv” if we’re storing one latent per token? Because there are two pieces: one for K-related info and one for V-related info. They’re each d_kv / 2 typically. The point is the cache stores compressed latents, not full K and V.

For DeepSeek-V2 with d_kv = 512, the per-token latent KV cache is roughly ~1 KB per layer (compared to ~20 KB per layer for a similar-sized GQA model). Across 60 layers, the total per-token KV cache is ~70 KB, vs ~200 KB+ for GQA-8 in a similar-sized model.

The quality cost is essentially zero. DeepSeek-V2 and V3 demonstrate that MLA matches or exceeds GQA quality at the same parameter count. This is one of the most surprising results of the 2024 architecture wave.

The catches:

  • More complex inference. The K and V have to be decompressed from the latent at each attention step, which adds a small amount of extra compute. The extra compute is small relative to the overall attention cost, so the trade-off is favorable.
  • More complex training. MLA requires careful initialization and training tuning. Most labs haven’t replicated it yet.
  • Special hardware-software integration. Efficient MLA inference requires the serving stack to know about the latent decompression. vLLM and SGLang both have MLA-specific kernels for DeepSeek models.

MLA is the future direction for KV cache compression, but adoption is slow because:

  • It’s tied to DeepSeek’s specific implementation.
  • Most labs are still optimizing GQA before reaching for MLA.
  • The training quality story is less proven outside DeepSeek’s models.

I expect more models to adopt MLA over the next year, especially as DeepSeek-style training recipes become more standard.

33.6 How each affects KV cache size

Putting it all together with concrete numbers for a hypothetical 70B model with H = 64, d_h = 128, n_layers = 80:

Variantn_kv_headsPer-token KV cachePer-token cache vs MHA
MHA642.6 MB1× (baseline)
GQA-32321.3 MB2× smaller
GQA-1616660 KB4× smaller
GQA-88330 KB8× smaller
GQA-44165 KB16× smaller
GQA-2282 KB32× smaller
MQA141 KB64× smaller
MLA (d_kv=512)n/a~80 KB32× smaller (approx)

The 8× savings of GQA-8 over MHA is the modern default. MLA pushes this further to 32-64×. MQA pushes it even further but at quality cost.

For a serving fleet, the choice of attention variant directly determines:

  • How many concurrent users fit in a given GPU memory budget.
  • How long contexts can be without running out of memory.
  • The cost-per-token for any non-trivial workload.

A 70B model with GQA-8 vs MHA has roughly 8× the serving capacity at the same hardware cost — and that’s just from the architectural choice, before any other optimization.

33.7 How each affects quality

The quality story, summarized:

  • MHA: maximum quality, maximum cost. The baseline.
  • GQA-32: essentially no quality loss vs MHA. But the cache savings are small (2×).
  • GQA-16: very small quality loss (<0.1 points). Cache 4× smaller.
  • GQA-8: small quality loss (<0.5 points typically). Cache 8× smaller. Modern default.
  • GQA-4: noticeable but small loss (~0.5-1 points). Cache 16× smaller.
  • GQA-2: larger loss (~1-2 points). Cache 32× smaller.
  • MQA: largest loss (~1-3 points). Cache H× smaller.
  • MLA: essentially no quality loss vs MHA. Cache 32-64× smaller. The frontier.

The “uptrain” trick — taking an MHA-trained model and converting to GQA via continued training — works for going from MHA to GQA-8 with a few hundred billion training tokens. The converted model recovers most of its quality. This is how Llama 2 70B went from MHA (in the original 65B) to GQA-8.

Going further (uptraining to GQA-4 or MQA) is harder. Quality recovery is less complete; the model may end up worse than if it had been trained with the smaller n_kv_heads from scratch.

For new training runs, GQA-8 is the default. For models specifically optimized for long-context inference, MLA (or even more aggressive compression) is increasingly common.

33.8 The training-quality-vs-inference-cost frontier

The deeper observation: the choice of n_kv_heads is a Pareto front between training quality and inference cost. There’s no free lunch.

Quality vs KV cache size Pareto front: GQA-8 is near-optimal; MLA extends the frontier further. KV cache size per token (log scale, smaller = better →) Model quality (higher = better ↑) MHA (H=64) GQA-32 GQA-8 ← default GQA-4 MQA MLA ← frontier
The Pareto frontier of quality vs KV cache size — GQA-8 is the current production sweet spot, and MLA extends the frontier by achieving MHA-level quality at near-MQA cache cost.
  • Lower n_kv_heads → cheaper inference, harder training, slightly lower quality at convergence.
  • Higher n_kv_heads → more expensive inference, easier training, slightly higher quality at convergence.

Different models pick different points on this frontier based on their priorities:

  • Frontier labs serving billions of requests want the cheapest inference, so they lean toward GQA-4 or MLA. They can afford to pay for the harder training to recover quality.
  • Open-source labs releasing once and moving on care about training simplicity, so they default to GQA-8 (the empirical sweet spot).
  • Specialized models (e.g., reasoning-focused) might prefer slightly higher quality and accept more expensive serving — Reasoning models often use GQA-16 or even MHA in some experimental variants.

The frontier is moving. As MLA becomes more battle-tested, the “optimal” point shifts toward more aggressive compression. As context windows get longer (because production demands it), KV cache size becomes more dominant relative to other costs, which pushes toward more compression.

33.9 The choices in current models

A summary of what current open models use:

ModelYearVariantn_kv_heads (or equivalent)
GPT-32020MHA96
Llama 12023MHA32
Llama 2 70B2023GQA8
Llama 3 70B2024GQA8
Llama 3.1 405B2024GQA8
Mistral 7B2023GQA8
Mixtral 8x7B2023GQA8
Qwen 2.5 72B2024GQA8
Gemma 2 27B2024GQA16
Phi-3 medium2024GQA10
DeepSeek-V22024MLAn/a (latent dim 512)
DeepSeek-V32024MLAn/a (latent dim 512)

The pattern: everyone is on GQA-8 as the safe default, with DeepSeek pushing into MLA. New models in 2025 are increasingly experimenting with MLA or smaller GQA variants.

The Llama 4 architecture (when released) is widely expected to push further on KV cache compression, possibly adopting MLA or another novel scheme. This is one of the most actively iterated parts of the architecture.

33.10 The mental model

Eight points to take into Chapter 34:

  1. The KV cache is the dominant memory cost of serving. Reducing it directly increases serving capacity.
  2. MHA is the original. Each head has its own K and V. Maximum quality, maximum cost.
  3. MQA shares one K and V across all heads. Maximum compression, real quality cost.
  4. GQA groups heads to share K and V in groups. The compromise. GQA-8 is the modern default.
  5. MLA compresses K and V into a low-rank latent vector. The frontier. Lossless quality, big compression.
  6. GQA-8 vs MHA gives 8× cache savings at <0.5 points quality cost. The biggest single architectural lever.
  7. MLA can give 32-64× cache savings at essentially no quality cost. Used by DeepSeek-V2/V3.
  8. The choice is a Pareto front between training quality and inference cost. No free lunch, but the frontier is moving toward more compression.

In Chapter 34 we cover the other major architectural compression: Mixture of Experts.


Read it yourself

  • Vaswani et al., Attention Is All You Need (2017). The MHA original.
  • Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (2019). The MQA paper.
  • Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023).
  • The DeepSeek-V2 technical report. The MLA section is the introduction of the technique.
  • The DeepSeek-V3 technical report. Continued use and refinement of MLA.
  • The Llama 2 paper (Touvron et al., 2023), section on attention. Justifies the GQA choice.

Practice

  1. Compute the KV cache size per token for a 7B model with H=32, d_h=128, n_layers=32 under MHA, GQA-8, GQA-4, and MQA.
  2. Why does GQA-8 give 8× cache savings over MHA without 8× quality loss? Argue at the level of “what are the heads doing.”
  3. Read the DeepSeek-V2 paper’s MLA section. Walk through the latent compression and decompression steps.
  4. Why was MQA quickly displaced by GQA? Identify the specific quality issue with MQA.
  5. For a 70B model serving 100 concurrent users at 4k context each, compute the total KV cache memory under MHA, GQA-8, and MLA. What’s the difference in serving capacity?
  6. Why hasn’t every lab adopted MLA? Argue both technical and adoption-curve reasons.
  7. Stretch: Read the GQA paper’s uptrain experiment (section 4). Why does uptraining from MHA to GQA work? Reproduce the quality recovery curve mentally.

Concept check

4 questions. Click a choice to check. Your score is saved locally.

Score
0 / 4
  1. 1. Llama 3 70B uses GQA with 8 KV heads versus MHA with 64 query heads. By what factor does this reduce per-token KV cache size compared to MHA?
  2. 2. Multi-Query Attention (MQA) uses a single KV head shared by all query heads. What is the main quality downside compared to GQA?
  3. 3. DeepSeek's MLA compresses KV heads by projecting K and V into a low-rank latent space. What does this allow that GQA does not?
  4. 4. A model was trained with MHA and you want to continue pretraining it with GQA to reduce serving costs. What is the standard approach for initializing the GQA KV heads from the MHA checkpoint?
Related chapters