MHA, MQA, GQA, MLA: The KV-Cache Trade-offs Behind Every Modern Transformer

Ask someone to sketch a transformer on a whiteboard and they’ll draw the 2017 version: softmax attention over every pair of tokens, one set of projection matrices, heads stacked in parallel. That picture is still right for the math. But it’s been quietly wrong as a description of production models for years. The attention mechanisms inside Llama, Mistral, DeepSeek, and most current open-weight releases differ from vanilla multi-head attention in one specific dimension — the shape of the key/value cache — and that single difference determines how much VRAM a model needs, how long a context it can serve, and how fast each generated token appears.

This post is a tour of that one dimension: multi-head attention (MHA), multi-query attention (MQA), grouped-query attention (GQA), and multi-head latent attention (MLA). Each is a different answer to the same question — how many copies of keys and values do you actually need to keep around? — and the trade-offs between them explain a surprising amount of recent model design. Understanding this progression also makes recent architecture papers legible: a good fraction of what they optimize is KV-cache economics.

The Problem: The KV Cache Grows With Context

Autoregressive generation has an awkward asymmetry: generating token N requires attending to all N−1 previous tokens, but nothing needs recomputing if you cache each layer’s keys and values. That cache is the KV cache, and its size is easy to compute:

cache bytes = 2 (K and V) × layers × kv_heads × head_dim × seq_len × batch × bytes_per_element

For a 7B-class model with 32 layers, 32 heads of dimension 128, in FP16, at a 32K context: 2 × 32 × 32 × 128 × 32,768 × 2 bytes ≈ 16 GB per sequence. That’s on the order of the weights themselves. Serve a batch of 8 and the KV cache alone wants 128 GB. Everything that follows is an attempt to shrink one specific factor of that product — the number of heads whose keys and values you must retain.

MHA: The Baseline With a Cost Problem

Vanilla multi-head attention gives every query head its own key and value head. Thirty-two query heads means thirty-two key/value pairs per layer, each fully materialized in the cache. It’s the highest-quality configuration and the most expensive one — which is why almost nothing ships as pure MHA anymore. The cost isn’t just memory; it’s bandwidth. While decoding, every generated token streams the entire KV cache from GPU memory to the compute units, so cache size directly caps generation speed for long contexts.

MQA: Sharing One K/V Head Across Everything

The multi-query attention paper proposed a brutal simplification: keep all 32 query heads, but share a single key and value head across all of them. The KV cache shrinks by a factor of 32 — our 16 GB example becomes half a gigabyte. Decoding gets dramatically faster, especially in the memory-bound regime of large-batch serving.

The cost is quality. Keys and values are now a bottleneck: every query head must extract what it needs from one shared representation. The original paper reported modest degradation, but follow-up work found that a single shared K/V head can make training unstable and measurably hurts quality at scale. MQA ships in a few production systems, but the industry mostly treated it as a proof that K/V sharing works at all — and went looking for a middle setting.

GQA: The Interpolation That Won

Grouped-query attention is exactly that middle setting: use G key/value heads, each shared by a group of query heads. G=1 recovers MQA; G equal to the number of query heads recovers MHA; the points in between trade a slice of quality for most of the memory savings. The GQA paper showed that an intermediate grouping recovers most of MQA’s speed while retaining quality close to full MHA — and described an uptraining recipe: take an existing MHA checkpoint, mean-pool its K/V projections into groups, and continue training for a small fraction of the original compute. You get GQA’s inference economics without pretraining from scratch.

Adoption was fast and near-universal: GQA is the attention in Llama 2 70B and the Llama 3 family, Mistral 7B, and most open-weight releases since. If a model card lists 32 query heads and 8 KV heads, that’s GQA with 4 query heads per group, and you can price its KV cache at exactly one quarter of the MHA equivalent. It became the default because it sits at the knee of the trade-off curve: pushing K/V reduction past GQA’s typical 4–8x starts costing quality faster than it saves memory.

MLA: Compressing the Cache Instead of Shrinking It

Multi-head latent attention, introduced with DeepSeek-V2, attacks the problem from a different angle. Instead of sharing or dropping K/V heads, it compresses them: each token’s keys and values are jointly down-projected into a small latent vector, and the cache stores that compressed form. Projections back up to per-head K/V happen on the fly during attention. The cache holds only the latent — roughly one-tenth the size of an equivalent MHA cache for DeepSeek-V2’s configuration — while attention still operates with per-head representations at compute time.

The subtle engineering cost is matrix absorption: because K and V only exist as projections of the latent, naively materializing them would undo the memory savings. The DeepSeek-V2 paper describes how the latent projections get algebraically folded into the attention computation — and how RoPE interacts awkwardly with low-rank caching, requiring a decoupled RoPE structure where a small slice of each key and query bypasses the compression. The result is a cache dramatically smaller than GQA’s at comparable quality, at the price of a more complex attention kernel. This is the architecture behind the DeepSeek-V2/V3 lineage and a large part of why those models have an unusual serving envelope: enormous effective context at costs that GQA-only designs of the same scale can’t match.

Side Quests Worth Knowing

Two more ideas appear often enough in papers and inference engines to warrant a mental bookmark.

YOCO (you only cache once) restructures the stack itself: a single self-decoder block at the bottom produces one K/V cache that all later layers reuse, with the upper layers forming a non-attention “cross-decoder” that reads it. The cache stops scaling with depth — instead of 32 layers each keeping their own cache, 32 layers share one. The YOCO paper reports competitive quality with large savings in both memory and prefill compute. No frontier open-weight release has adopted it yet, but it’s the logical endpoint of the trend: the KV cache is a cost center, and YOCO minimizes the number of places that cost is paid.

Attention sinks came out of a failure mode in streaming language models: models with windowed attention collapsed when sequences exceeded the window. The cause was that softmax attention must sum to one, so when all “real” tokens fall outside the window, attention piles onto whatever tokens remain — and the first tokens had absorbed a huge share of attention mass during training. Keeping those first tokens permanently in the window as “sinks” fixes the collapse. The insight matters beyond streaming: it explains why serving stacks for very long generations care about which tokens stay pinned in the cache regardless of the head-count scheme above.

Choosing: A Practical Comparison

From the serving side, the choice is mostly already made — but the trade-offs matter when you pick checkpoints or estimate serving capacity:

  • MHA: highest quality ceiling, highest cache cost. Rare in production; mostly a training-time reference point now.
  • MQA: extreme cache savings, measurable quality risk. Occasional use in latency-critical small models.
  • GQA: the industry default. KV heads at one quarter to one eighth of query heads keeps quality near MHA with 4–8x smaller caches.
  • MLA: the smallest caches at scale, with the most complex kernels; the DeepSeek lineage’s defining choice.

The reliable back-of-envelope: KV-cache bytes ≈ 2 × layers × kv_heads × head_dim × seq_len × batch × bytes_per_element. When a model card lists kv_heads, you can immediately compare serving footprints across checkpoints — and when a new architecture paper claims better efficiency, check which factor of that product it’s attacking. Quantized caches, sliding-window eviction, and sink-aware caching layer on top of whatever head-count scheme a model ships with, so the same checkpoint can have very different memory profiles depending on inference-engine settings.

Wrapping Up

The arc from MHA to MQA to GQA to MLA is really one idea applied with increasing sophistication: the key/value cache is the dominant cost of autoregressive inference, so every mechanism that shrinks it without proportionally hurting quality wins. GQA won the mainstream by interpolating between sharing and quality; MLA went further by compressing rather than sharing; YOCO and attention-sink techniques attack the remaining dimensions of depth and sequence position. Next time a model card lists 32 query heads and 8 KV heads, you’ll know exactly what that number is buying — and what it costs.

Leave a Reply

Your email address will not be published. Required fields are marked *