Attention variants and the KV-cache tradeoff (MHA, MQA, GQA, MLA)

The KV-cache size formula has a term called kv_heads, and it is not a fixed fact of nature — it is a design knob. Turning it down shrinks the cache and speeds up decode, at some cost to quality. The choice of attention variant — MHA, MQA, GQA, or MLA — is exactly the choice of where to set that knob, and it is one of the highest- leverage architecture decisions for serving. This chapter walks the spectrum.

What you will understand by the end

• Why the number of key/value heads directly sets KV-cache size and decode bandwidth.
• What MHA, MQA and GQA are, and the memory-vs-quality tradeoff between them.
• Why GQA became the production default in modern open models.
• What MLA does differently — and why it is not just another point on the same line.

The lever: KV memory scales with key/value heads

Recall the KV-cache size: 2 × layers × kv_heads × head_dim × … × seq × batch. In classic attention, every query head has its own key/value head, so kv_heads equals the number of attention heads — and the cache (and the bytes streamed from HBM every decode step) scales straight with it. That is the pressure point. Every variant below is a different answer to one question: how many key/value heads do we actually need to cache?

The spectrum: MHA → GQA → MQA

  MHA — every query head has its own K/V            cache = 8 K/V heads  (full)
    Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
    K1 K2 K3 K4 K5 K6 K7 K8
    V1 V2 V3 V4 V5 V6 V7 V8

  GQA — query heads grouped; K/V shared per group    cache = 2 K/V heads  (¼)
    [ Q1 Q2 Q3 Q4 ]      [ Q5 Q6 Q7 Q8 ]
          K1                    K2
          V1                    V2

  MQA — all query heads share one K/V               cache = 1 K/V head   (1/8)
    Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
              K  (shared)
              V  (shared)

• MHA — multi-head attention (the baseline). N query heads, N key/value heads. Every head attends its own way; maximum expressiveness, maximum cache. This is "classic" attention from the original Transformer.
• MQA — multi-query attention. All query heads share a single K/V head. The cache shrinks by a factor of N, and decode reads far fewer K/V bytes per token, so it is faster. The cost is quality: the heads can still ask different questions (distinct queries) but must look at the same keys and values, which measurably weakens the model on harder tasks.
• GQA — grouped-query attention. The middle ground: split the query heads into G groups, and share one K/V head per group. The cache is G/N of MHA — a big cut with only a small quality loss, because each group still gets its own keys and values. With, say, 64 query heads and 8 groups, the KV cache is 1/8 the size while quality stays close to full MHA.

MLA: compress instead of share

MHA, MQA and GQA all sit on one line: they trade fewer K/V heads for less quality. MLA — multi-head latent attention (introduced by DeepSeek) takes a different route. Instead of dropping heads, it compresses the keys and values into a small shared low-rank latent vector per token, caches that, and reconstructs the per-head K/V on the fly during attention.

The payoff is that the cached thing is tiny — comparable to or smaller than GQA — while the model keeps something close to full multi-head expressiveness, because no head was actually thrown away. In the language of compression, it exploits the fact that the K/V vectors are low-rank: most of their information survives a squeeze through a narrow latent. It is the current frontier answer to "small cache and high quality."

Why this is a double win for inference

Shrinking kv_heads helps inference in two independent ways at once:

• Capacity — a smaller cache means the KV budget holds more concurrent requests and longer contexts. Bigger batches, same GPU.
• Bandwidth — decode re-reads the cache from HBM every token; fewer K/V heads means fewer bytes streamed per step, so decode itself gets faster. This is a direct hit on the memory-bound bottleneck.

Mental model

Query heads are free; key/value heads are what you cache and stream. MHA keeps one K/V per query head; MQA collapses to one; GQA shares one per group (the sweet spot); MLA compresses K/V to a latent instead of dropping heads. Fewer/smaller K/V → bigger batches and faster decode, paid for in quality you must measure.

Common mistakes

• Assuming all attention is MHA. Most models you serve today are GQA; the KV cache is already smaller than the naive heads × head_dim estimate. Use the model's real kv_heads.
• Treating a smaller cache as quality-free. MQA and aggressive GQA can hurt hard tasks; the savings are real but so is the tradeoff — re-measure.
• Confusing query heads with K/V heads. Cutting K/V heads is what saves memory; the model can still have many query heads. They are separate counts.
• Expecting to "turn on GQA" at serve time. It is trained into the model. Your serving lever is picking a model with the attention you want.

Practical guidance

• Read num_key_value_heads (or equivalent) off the model config, not the query-head count, when you budget the KV cache — for a GQA model it is often a small fraction of the heads.
• On decode-heavy or long-context workloads, prefer models with fewer K/V heads (GQA) or MLA — the cache and the decode bandwidth both drop.
• When weighing two models with different attention, fix the eval set and compare quality and tokens/sec together; do not let the smaller cache decide alone.
• If the KV cache is not your bottleneck (measure it — like this project's compute-bound L4), do not over-optimise the attention variant; spend your effort where the knee actually is.

Summary

• The number of key/value heads sets both KV-cache size and per-token decode bandwidth; it is a design knob, not a constant.
• MHA keeps one K/V per query head; MQA collapses to one (biggest cut, biggest quality hit); GQA shares one per group (small loss, big saving) and is the modern default.
• MLA compresses K/V into a low-rank latent instead of dropping heads — small cache with near-full quality.
• Fewer/smaller K/V heads is a double win (capacity + bandwidth), but quality is a measured tradeoff, and the win only matters when the KV cache is actually your bottleneck.

Knowledge check

A model card says "32 query heads, 8 key/value heads." Which attention variant is this, and roughly how much smaller is its KV cache than plain MHA?

It is GQA (grouped-query attention): 32 query heads sharing 8 K/V heads means 4 query heads per group. The KV cache scales with K/V heads, so it is about 8/32 = ¼ the size of full MHA (which would cache 32 K/V heads) — while keeping quality close to MHA. Pure MQA would be 1 K/V head (1/32); pure MHA would be 32.

Two models score the same on your eval, but one is MQA and one is MHA. On a long-context, high-concurrency decode workload, which do you expect to serve more cheaply, and why — on two counts?

The MQA model, on both counts. Its KV cache is far smaller, so (1) the KV budget holds more concurrent requests and longer contexts (more batch on the same GPU — capacity), and (2) decode streams fewer K/V bytes from HBM per token, so tokens/sec is higher (bandwidth). Since quality already matched on your eval, the smaller-cache model wins the serving economics.

Related chapters

• The KV cache and context growth — the kv_heads term this chapter turns into a knob
• Transformer and attention intuition — where query, key and value heads come from
• Prefill and decode — why fewer K/V bytes make decode faster
• Quantization — the other "represent it in less" lever, and the low-rank idea behind MLA