passdrill

KV cache size: MHA vs MQA vs GQA vs MLA, with a worked example

Plenty of pages explain multi-head attention (MHA), multi-query attention (MQA), grouped-query attention (GQA), and multi-head latent attention (MLA) one at a time, with their own separate toy numbers. Few put all four side by side on one consistent model shape so you can see exactly how many gigabytes of KV cache each one costs. This page does that arithmetic in full, using one real model's real dimensions as the baseline, so you can check every step by hand.

What the KV cache stores, and the formula behind its size

During autoregressive generation, a transformer decoder would have to re-run attention over every previous token for every new token if it started from scratch each step. Instead it caches each layer's key (K) and value (V) vectors the first time they're computed, then reuses them. That cache grows with every generated token, and at long context lengths or large batch sizes it's often the main limit on how many requests a GPU can serve at once — not the model weights.

For a standard attention layer (MHA, GQA, or MQA), the cache holds one K vector and one V vector per key-value head, per layer, per token. In bytes:

bytes = 2 (K and V) × num_layers × num_kv_heads × head_dim × bytes_per_element × context_tokens

The only thing that changes between MHA, GQA, and MQA is num_kv_heads. Everything else about the model — layers, head dimension, context length — stays fixed. MLA replaces this formula entirely, which is the point of the last section below.

The reference model: Llama 3.1 8B's real shape

To keep the comparison honest, every variant below is applied to one real architecture: Llama 3.1 8B, which has 32 transformer layers, a head dimension of 128, and (as shipped) 8 key-value heads shared across 32 query heads — a 4:1 ratio. Using FP16 (2 bytes per element) and an 8,192-token context:

2 × 32 layers × 8 kv heads × 128 × 2 bytes × 8,192 tokens = 1,073,741,824 bytes = 1.0 GiB

That 1.0 GiB-per-sequence figure is Llama 3.1 8B's real, shipped GQA cache cost at that context length. If the same model had been built with full MHA instead — one key-value head per query head, so 32 kv heads instead of 8 — the identical formula gives exactly 4.0 GiB: four times larger, because the cache scales linearly with num_kv_heads. This 4x-vs-real-1x comparison is the standard way GQA's savings get reported, and it checks out from the formula alone.

Sliding to the extreme: MQA

Shazeer's 2019 multi-query attention paper proposed the opposite extreme: collapse every query head down to a single shared key-value head (num_kv_heads = 1), maximizing cache savings at some cost to quality on tasks sensitive to attention diversity. Applying the same formula to Llama 3.1 8B's shape with 1 kv head instead of 8:

2 × 32 × 1 × 128 × 2 × 8,192 = 134,217,728 bytes = 0.125 GiB (128 MiB)

That's an 8x reduction versus the real GQA setting, and 32x versus full MHA — but GQA exists specifically because going all the way to one shared head measurably hurts quality on some benchmarks, which is why Llama 3.1, Mistral, and most current production models land on GQA with 4–8 kv heads rather than MQA's single head.

MLA: compression instead of fewer heads

DeepSeek-V2 (DeepSeek-AI, 2024) takes a different approach entirely. Instead of sharing K/V vectors across query heads, multi-head latent attention compresses each token's key and value information into one shared low-rank latent vector (dimension d_c), plus one small shared vector carrying rotary position information (dimension d_h^R), and reconstructs per-head keys and values from that latent on the fly during attention. The cache only ever stores the compressed latent and the shared RoPE vector — not separate K/V vectors per head — so num_kv_heads drops out of the formula entirely:

bytes = num_layers × (d_c + d_h^R) × bytes_per_element × context_tokens

DeepSeek-V2's published configuration sets d_c = 512 (its kv_lora_rank) and d_h^R = 64 (its qk_rope_head_dim) — not coincidentally, 512 = 4 × head_dim and 64 = head_dim / 2 for a head dimension of 128. Applying those same two numbers to the Llama 3.1 8B reference shape (32 layers, 8,192 tokens, FP16), purely to compare on equal footing:

32 × (512 + 64) × 2 × 8,192 = 301,989,888 bytes ≈ 0.281 GiB

Note this is not DeepSeek-V2's own real cache size — the actual released model has more layers and a different head count than this 32-layer reference shape, so its real total differs. It's the same compression dimensions applied to one common architecture purely so all four numbers below are comparable.

All four, side by side

Attention typeKV heads cachedCache per sequence (32 layers, 8,192 tokens, FP16)vs. MHA
MHA (32 kv heads)324.0 GiB1x (baseline)
GQA (Llama 3.1 8B's real setting, 8 kv heads)81.0 GiB4x smaller
MLA (DeepSeek-V2's published d_c/d_h^R, applied here)n/a (compressed latent)≈0.281 GiB≈14.2x smaller
MQA (1 kv head)10.125 GiB32x smaller

The ordering is worth pausing on: MLA's cache here lands between GQA and MQA, not below MQA. That might look unimpressive against MQA's raw number, but it's missing the real point — because MLA's cache size doesn't depend on num_kv_heads or query-head count at all, its relative advantage grows on models with far more attention heads than this 8B reference. DeepSeek-V2's own paper reports a 93.3% smaller KV cache and up to 5.76x higher maximum generation throughput compared to its own prior dense model, DeepSeek 67B — a much larger, many-headed MHA model, not the 32-head shape used above, so that figure isn't directly comparable to the 93.0%-smaller-than-MHA result the table above happens to produce (4.0 GiB down to 0.281 GiB); treat the similarity as coincidental rather than a confirmation of the same comparison. The paper also reports MLA matching or exceeding standard MHA's output quality, unlike MQA, which is the actual reason DeepSeek chose compression over simply sharing fewer heads.

Common mistakes

For more on how these architecture choices interact with context-window extension, positional encodings, and inference-time tricks like speculative decoding, see the How LLMs Work practice quiz.

Source: Vaswani et al., "Attention Is All You Need" (arXiv:1706.03762); Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need" (arXiv:1911.02150); Ainslie et al., "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" (arXiv:2305.13245); DeepSeek-AI, "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model" (arXiv:2405.04434); Meta Llama 3.1 model architecture configuration (32 layers, 8 key-value heads, head dimension 128).

Drill How LLMs Work: Transformers & Training practice questions →