MHA to MLA: the KV cache math that decides whether you can serve the model at all

Attention variants aren’t an architecture curiosity. They exist because of one brutal constraint: the KV cache is what actually limits how many users you can serve and how long a context you can afford, and every attention head in a transformer caches a Key and Value tensor per token it’s seen. Every variant covered here is a different answer to the same question, how do you shrink that cache without destroying quality. Skip the arithmetic and you can nod along in an interview right up until someone asks why Llama 2 70B needed GQA and Llama 2 7B didn’t.

The number that explains everything else in this post#

For standard multi-head attention, every layer caches K and V per token, sized by heads times head dimension:

$\text{KV cache} = 2 \times \text{layers} \times \text{heads} \times \text{head\_dim} \times \text{seq\_len} \times \text{bytes}$

Take a 70B-class model: 80 layers, 64 heads, head dimension 128, BF16, 128K context.

$2 \times 80 \times 64 \times 128 \times 131072 \times 2 \approx 320 \text{ GB per single request}$

Sit with that. An H100 has 80GB of HBM. One long-context request from one user needs four entire GPUs of memory just for cache, before a single model weight is loaded. This isn’t a theoretical concern, it’s why MHA at 70B-plus scale with long context is economically dead on arrival for serving, and it’s why “attention variants” is inseparable from “how do we serve this to more than one person at a time.”

GQA: the industry default, and why it won#

Instead of every one of the $H$ query heads getting its own K/V head, GQA shares K/V heads across groups: $G$ KV heads, $H$ query heads, $G < H$, each KV head serving $H/G$ query heads. The KV cache shrinks by exactly that factor of $H/G$. Llama 3 70B uses 64 query heads and 8 KV heads, an 8x reduction, dropping 320GB toward 40GB. Still enormous. The difference between impossible and expensive.

The real implementation detail: $W_K$ and $W_V$ are shaped $[G \times d_{\text{head}}, d_{\text{model}}]$ instead of $[H \times d_{\text{head}}, d_{\text{model}}]$, and at attention time you expand the $G$ KV heads back out to $H$ using repeat_interleave before computing scores, the same attention math as MHA, just against fewer distinct K/V tensors.

Here’s the failure mode nobody warns you about, and it’s genuinely one of the worst kinds. This is a one-line implementation with a completely silent failure mode. Call repeat_interleave on the wrong dimension and the model still trains, loss still goes down, gradients still flow, nothing crashes. What happens is every query head attends to the wrong KV head’s content. The model trains on a self-consistent but incorrect attention pattern. You don’t find this in a stack trace. You find it three weeks later when eval numbers come in worse than a smaller ablation model, and someone git-bisects through tensor reshape code to find a dim=1 that should have been dim=2. This is the single most common silent bug first-year engineers introduce in a from-scratch attention rewrite. If you’re implementing GQA for the first time, write a numerical unit test against a known-good MHA output before you trust it, not after.

Why the industry actually converged here matters more than it sounds: Llama 2 70B adopted GQA specifically because MHA at 70B made serving long contexts economically impossible. Not a marginal quality benefit. The alternative wasn’t shippable.

MQA: the extreme case, mostly abandoned#

Multi-Query Attention is GQA with $G=1$, one shared KV head for every query head, maximum possible cache reduction, used in PaLM specifically for serving efficiency. The industry mostly moved past it: collapsing every head down to a single shared KV representation costs more quality than it’s worth relative to GQA at $G=8$‘s incremental cache savings. $G=1$ is a knob, not a destination, and $G=8$ empirically sits closer to the Pareto frontier. Worth knowing it exists and knowing exactly why it lost, because “why not just use MQA for maximum compression” is a real interview question with a real answer.

MLA: a genuinely different move, not a bigger GQA#

GQA shares heads. MLA compresses the representation itself, projecting K and V into a shared low-rank latent vector and only reconstructing the full K/V at attention time:

$c = x \cdot W_{DKV} \quad (\text{down-project into a latent, dim} \sim 512), \qquad K = c \cdot W_{UK}, \qquad V = c \cdot W_{UV}$

You only cache $c$, the roughly 512-dimensional latent, not the full per-head K and V. Structurally different from GQA: compressing the representation, not reducing the count of representations.

MHA at 320GB, GQA at roughly 40GB, MLA at roughly 10GB per request, a 32x reduction from the naive baseline and roughly 4x better than GQA alone. The official framing matters as much as the number: this significantly reduces the KV cache during inference, supporting a larger batch size for better throughput. Cache size isn’t an abstract efficiency metric, it directly multiplies into how many concurrent users one GPU cluster can serve, which directly multiplies into unit economics per token served.

The failure mode here is worse than GQA’s, not better. With GQA, a broadcast bug produces attention against the wrong-but-real KV head. With MLA, a wrong low-rank projection shape produces an attention matrix that’s subtly numerically wrong in a way that doesn’t show up as a shape error, the matmuls still go through, dimensions still line up, and you get a plausible-looking but incorrect output. This is why MLA implementations get validated numerically against a reference implementation before anyone trusts a training run built on top of them. If asked how you’d verify a from-scratch MLA implementation, the answer isn’t “check the loss curve,” the loss curve can look fine while the attention math is wrong. You need a direct tensor-level comparison against a known-correct reference on a fixed input.

Sliding window: an orthogonal axis, not a competitor#

Everything above compresses the KV cache per token. Sliding window attention instead limits how many tokens get cached at all, each token attends only to the last $W$ tokens, giving $O(S \times W)$ memory instead of $O(S^2)$. Mistral 7B ships with $W=4096$.

This is a genuinely different tradeoff, not a strictly better one. It works when the task’s relevant context is local, recent conversation turns, nearby code, a recent paragraph. It fails hard on anything needing true long-range dependency, recalling a fact from 50K tokens back, connecting a conclusion at the end of a document to a premise at the start. Don’t confuse “cheaper long context” (GQA, MLA) with “no long context” (SWA), they solve different problems, and conflating them is one of the more common mistakes people make after reading only the summary paragraph on each.

The debate that isn’t settled, and why that’s the actual signal#

Every course treats this as a clean progression: MHA was expensive, GQA fixed it, MLA fixed it better, done. That’s not the current state of frontier engineering opinion, and knowing that is exactly what separates someone who read a paper from someone tracking what labs are doing right now.

From Moonshot AI’s Kimi founder AMA, December 2025, an active architecture decision at a frontier lab, not a retrospective: KDA hybrids with NOPE MLA perform better than full MLA with RoPE in their apples-to-apples comparison across pretraining and RL, achieving higher benchmark scores while being faster and more economical. Unpack what that’s saying: full MLA, the DeepSeek-V2/V3 approach above, with rotary position encoding, is not the final answer. A hybrid combining a different attention mechanism (KDA) with a variant of MLA that drops rotary encoding (NOPE) won head-to-head, on real pretraining and RL runs, at a lab shipping trillion-parameter production models. This happened after MLA was already public and already famous. The field didn’t stop iterating the moment DeepSeek published, and most explainers freeze at exactly the wrong point in time.

If you present MLA as “the solved problem, better than GQA, end of story,” you sound like you read the DeepSeek-V2 paper and stopped there. If you can say MLA was the right answer for a specific KV-compression tradeoff, but recent lab-reported results suggest hybrid approaches with modified positional encoding are outperforming it on pretraining and RL efficiency, you sound like someone tracking the space instead of reciting last year’s conclusion.

Where this actually gets tested#

This isn’t asked as “define GQA.” It shows up as a systems-tradeoff question, usually framed around serving economics. “Why did Llama 2 need GQA at 70B but not at 7B” tests whether you did the KV cache arithmetic above or just memorized “GQA reduces memory.” “Walk me through the KV cache size for a 70B model at 128K context under MHA, then under GQA with 8 groups” is a mental-math check disguised as a systems question, if you can’t produce 320GB and 40GB on a whiteboard, that’s the tell. “Your GQA implementation trains fine, loss curve looks normal, but downstream eval is 3 points worse than expected, where do you look” has one correct answer, the silent repeat_interleave broadcast bug, not “check the learning rate.”

Interview language across labs consistently frames this as economic engineering, not academic architecture, “supporting larger batch size” and “serving more users,” not “improving benchmark scores.” If your answer to “why GQA or MLA” doesn’t mention serving throughput or cost, you’re missing the actual pressure the labs themselves cite.

What this takes to be frontier-job-ready#

The technical axis is exactly what this post has been building toward: Anthropic’s own Pretraining Scaling role asks engineers to “debug and resolve complex issues across the full stack, from hardware errors and networking to training dynamics and evaluation infrastructure.” A silent repeat_interleave bug that trains fine and shows up as a 3-point eval regression three weeks later is precisely the kind of cross-stack issue that quote is naming, not a hypothetical.

The operational axis shows up as diagnostic discipline under uncertainty. Poolside’s Reinforcement Learning role lists “diagnosing training instabilities in RL runs” and “high-quality, reproducible, maintainable code” as hard requirements, not aspirations, and a numerical unit test against a known-good MHA output before trusting a new GQA implementation is exactly that discipline in practice, not extra credit.

The autonomy axis is the unsettled-debate section made concrete. OpenAI’s own hiring language states that people who need a well-defined roadmap won’t succeed, and asks for engineers who “own ambiguous problems end-to-end without needing a tightly specified roadmap.” Whether MLA or a KDA-plus-NOPE hybrid is the right call for a specific model is not a question with a textbook answer right now, and forming a defensible opinion from the actual evidence, not from whichever paper you read first, is the skill that axis is naming.

Common mistakes#

Treating GQA’s group count as a knob to maximize rather than a tradeoff to balance: $G=1$ (MQA) exists, and the industry mostly rejected it, more compression isn’t automatically better once quality starts paying for it.

Trusting a passing loss curve as proof an attention implementation is correct: both the GQA broadcast bug and MLA’s projection-shape bug are silent by construction, and the loss curve is exactly the signal that fails to catch either one.

Presenting MLA as a closed question: the Kimi AMA result shows the field moved past “MLA is simply better than GQA” within the same year MLA became famous.

Conflating sliding window attention with GQA or MLA as if they’re on the same axis: SWA trades away long-range capability entirely; GQA and MLA keep it and make it cheaper. Confusing the two means recommending the wrong fix for the wrong constraint.

Try it yourself#

Beginner. Using the KV cache formula, compute the cache size for a 7B-class model (32 layers, 32 heads, head_dim 128, BF16) at a 4K context, and compare it against the 70B/128K figure above to see why 7B models never needed GQA as urgently.

Intermediate. For Llama 3 70B’s GQA configuration (64 query heads, 8 KV heads), work out exactly which dimension repeat_interleave needs to act on to correctly expand 8 KV heads back to 64, and describe what the resulting attention pattern would look like if that dimension were off by one.

Advanced. Using MLA’s projection formulas, construct a toy example with $d_{\text{model}}=8$, a latent dimension of 2, and a deliberately wrong $W_{UK}$ shape, and show that the resulting matmul still produces a shape-valid but numerically incorrect attention output, the exact failure mode described above.

---

The one-sentence version: GQA shares KV heads across groups of query heads and MLA compresses the K/V representation itself into a shared low-rank latent, two structurally different answers to the same 320-gigabyte problem, both with a failure mode that trains fine and fails silently, and the field kept moving past MLA the moment it got famous rather than stopping there. Getting the attention mechanism cheap enough to run is only half the problem; actually serving it to a hundred million people at once is the other half.