Chapter 3.5 β€” Attention Variants for Inference

Contents

  1. Where Chapter 3.4 stopped
  2. What the KV-cache actually costs
  3. Multi-Query Attention: one key-value head for all of them
  4. Grouped-Query Attention: the interpolation that won
  5. Multi-head Latent Attention: cache a compression instead
  6. Choosing between them
  7. Interview angle
  8. Self-check questions
  9. Sources

1. Where Chapter 3.4 stopped

Chapter 3.4 introduced the KV-cache as the technique that makes autoregressive generation tractable at all, and then closed the section with a warning: the cache is memory, that memory grows linearly with sequence length and batch size, and for long-context serving at high concurrency it becomes the binding constraint on how many requests a system can hold at once. Everything that chapter offered in response was a deployment-time fix β€” quantize the cache, page it, bound it with a sliding window β€” applied to a model whose architecture was already fixed.

This chapter is about the other half of the answer, the half that has to be decided before training starts. Between 2019 and 2024, the field made three successive changes to the attention mechanism itself, each one aimed squarely at the size of the KV-cache and each one accepting a different tradeoff to get it: Multi-Query Attention, Grouped-Query Attention, and Multi-head Latent Attention. These are not small implementation details. Grouped-Query Attention is used in essentially every open-weight model released since 2023, and knowing which one a model uses tells you more about its serving economics than its parameter count does. This chapter is also the first place in this book where a design decision is driven entirely by inference cost rather than by training quality or training cost β€” which is itself the point worth internalizing, because it's a pattern that repeats.

2. What the KV-cache actually costs

Start by writing the cost down precisely. For a single sequence, at each layer, the cache holds one key vector and one value vector per token per head. With $L$ layers, $h$ heads of dimension $d_h$, sequence length $n$, batch size $b$, and $p$ bytes per element, the total is $$\text{cache bytes} = 2 \cdot L \cdot h \cdot d_h \cdot n \cdot b \cdot p$$ where the leading $2$ is for keys and values. Note what's not in this formula: nothing about the number of parameters, and nothing quadratic. The cache grows linearly in context and in batch, and it does so with a constant that's set entirely by the model's shape.

Put real numbers in. LLaMA-2 70B has $L = 80$ layers, $h = 64$ query heads, and $d_h = 128$. Serving a single 4096-token sequence in fp16, with a conventional multi-head cache, would need $2 \cdot 80 \cdot 64 \cdot 128 \cdot 4096 \cdot 2$ bytes β€” about $10.7$ GB. For one sequence, at a context length that is small by 2024 standards, on a model whose fp16 weights are $140$ GB. Batch sixteen of those requests together, or extend to 32K context, and the cache alone exceeds the weights. This is the number that makes the rest of this chapter necessary: a serving system's throughput is roughly "how many concurrent sequences fit in the memory left over after the weights," and the cache is what determines that.

The memory footprint is only half the problem, and arguably the less important half. Chapter 3.4's section 2 established that autoregressive decoding is memory-bandwidth-bound rather than compute-bound, and the KV-cache is the clearest illustration of why. At each decoding step, the attention computation reads the entire cache β€” every key and value for every previous token β€” and does roughly one multiply-accumulate per element read. That's an arithmetic intensity of about one FLOP per byte, against modern accelerators that need something in the range of a few hundred FLOPs per byte to keep their compute units busy. The GPU is idle, waiting on memory, for the overwhelming majority of every decoding step. Shazeer's framing of this in the paper that started this whole line of work is worth stating in his terms: the bottleneck isn't the arithmetic, it's the ratio of memory access to arithmetic, and the way to speed up decoding is to load fewer bytes per step. Every technique in this chapter shrinks the same number, and gets both the memory saving and the bandwidth saving from the same change.

3. Multi-Query Attention: one key-value head for all of them

Shazeer's 2019 proposal is the most aggressive move available and the easiest to state: keep all $h$ query heads, but give the entire layer a single key head and a single value head, shared by every query head. The queries stay fully multi-headed β€” each one still projects into its own subspace and computes its own attention distribution β€” but they all attend against the same keys and the same values. The cache term $h \cdot d_h$ collapses to $d_h$, shrinking the cache by a factor of $h$: for the LLaMA-2 70B shape above, $10.7$ GB becomes $167$ MB.

It's worth being precise about what this does and doesn't cost, because the intuition "you've removed most of the model's attention capacity" is wrong. The number of attention patterns is unchanged β€” there are still $h$ distinct query projections producing $h$ distinct distributions over the sequence. What's lost is the ability for different heads to look at different projections of the same tokens: every head now scores against one shared view of the context and retrieves from one shared value space. Chapter 2.1's argument for multi-head attention was that heads specialize, one tracking syntactic dependency and another coreference; MQA lets them keep specializing in what they ask for while forcing them to share what they can see.

Empirically that costs something. Shazeer reported a small but measurable degradation in quality on translation, and the later GQA paper found a larger effect at scale, plus a second problem that matters more in practice: MQA models can be unstable to train, particularly with long sequences, and the quality gap doesn't close by simply training longer. The combination β€” real quality loss and real training difficulty, in exchange for a very large cache reduction β€” is why MQA saw genuine but limited adoption (PaLM and Falcon among the notable users) rather than becoming the default. It defined the axis, though, and once you see attention as having a tunable number of key-value heads, the obvious question is whether the two endpoints are really the only options.

4. Grouped-Query Attention: the interpolation that won

Ainslie et al.'s answer, published in 2023, is exactly the interpolation you'd guess, and its success is a good lesson in how often the useful contribution is the boring middle of an existing spectrum. Partition the $h$ query heads into $g$ groups; give each group its own key head and value head, shared by the $h/g$ query heads inside it. Setting $g = h$ recovers standard multi-head attention; setting $g = 1$ recovers MQA; anything in between is Grouped-Query Attention, and the cache shrinks by a factor of $h/g$. The typical choice in production models is $g = 8$: LLaMA-2 70B uses eight key-value heads against sixty-four query heads, cutting its cache from $10.7$ GB to $1.34$ GB at 4K context while keeping quality within noise of the full multi-head model.

The empirical finding that makes GQA worth using is that the quality-versus-cache curve is sharply non-linear near the MQA end. Going from $h$ key-value heads down to $8$ costs almost nothing measurable; going from $8$ down to $1$ costs noticeably more. Nearly all of the memory reduction is in the first move (going from 64 heads to 8 recovers 87.5% of the possible saving) and nearly all of the quality loss is in the second, so the interior of the range is not a compromise between two good options β€” it's strictly better than one of the endpoints for any realistic serving workload. GQA also fixes MQA's training instability, which is arguably as important to its adoption as the quality number.

The paper's second contribution is the practical one, and it's the part interviewers ask about. If GQA has to be trained in from scratch, then every existing multi-head checkpoint is stuck with its cache forever β€” and in 2023 that meant every good open model. Ainslie et al. showed you can convert an existing checkpoint instead: mean-pool the key and value projection matrices of the heads within each group into a single projection per group, which gives a model that is structurally GQA and immediately produces sensible (if degraded) outputs, then continue pretraining briefly to let the rest of the network adapt. They call this uptraining, and the striking result is how little of it is needed β€” roughly 5% of the original pretraining compute is enough to recover essentially full quality. That number is what turned GQA from an architecture proposal into something anyone with an existing model could adopt in an afternoon of cluster time, and it's why the technique propagated through the open-weight ecosystem as fast as it did.

5. Multi-head Latent Attention: cache a compression instead

MQA and GQA both work by reducing the number of key-value heads. DeepSeek-V2's Multi-head Latent Attention attacks the same term differently: keep all the heads, and compress what you store. For each token, project the input down into a single low-rank latent vector $c_t \in \mathbb{R}^{d_c}$ with $d_c$ much smaller than $h \cdot d_h$, cache only that latent, and reconstruct each head's full keys and values on demand with up-projection matrices $W^{UK}$ and $W^{UV}$: $$c_t = x_t W^{DKV}, \qquad k_t = c_t W^{UK}, \qquad v_t = c_t W^{UV}$$ DeepSeek-V2 uses $d_c = 512$ against 128 heads of dimension 128, so the cached object per token per layer is $512$ numbers rather than $2 \cdot 128 \cdot 128 = 32768$.

Stated that way it sounds like it just moves the cost: you've saved memory but now you have to decompress the cache at every step, which is exactly the arithmetic you were trying to avoid. The elegant part is that you don't. Attention scores need $q_t^\top k_s = (x_t W^Q)(c_s W^{UK})^\top = x_t \left(W^Q {W^{UK}}^\top\right) c_s^\top$ β€” and $W^Q {W^{UK}}^\top$ is a product of two fixed weight matrices, so it can be computed once and folded into the query projection. The same absorption works on the output side, where $W^{UV}$ merges into $W^O$. At inference the up-projections never run: queries are projected straight into the latent space and attend against the cached latents directly. You get the small cache and the reduced bandwidth, with no decompression step at all.

There is one genuine complication, and it's the detail that separates having read about MLA from having understood it. The absorption trick requires that nothing position-dependent sit between $W^Q$ and $W^{UK}$ β€” but RoPE (Chapter 2.2) works precisely by inserting a position-dependent rotation there, and $R_{t-s}$ depends on the specific query-key pair, so it can't be folded into a fixed matrix. RoPE and weight absorption are directly incompatible. DeepSeek's fix, which they call decoupled RoPE, is to split the head dimension in two: most of it carries no positional rotation and stays absorbable, while a small extra slice (64 dimensions in DeepSeek-V2) carries RoPE and is handled separately, as a single shared key head in the MQA style, cached alongside the latent. The final attention score is just the sum of the two pieces, $q_t^\top k_s = x_t\left(W^Q {W^{UK}}^\top\right)c_s^\top + (q_t^R)^\top R_{t-s}\, k_s^R$, the absorbed term plus the RoPE-carrying term. The per-token cache becomes $d_c + d_h^R = 512 + 64 = 576$ numbers per layer, which is a ~98% reduction against the $2 \cdot 128 \cdot 128 = 32768$ a multi-head layer of the same shape would store; DeepSeek-V2 reports it as a 93.3% cut relative to their previous dense DeepSeek 67B. Their more interesting claim is on the other axis: they measure MLA as matching or slightly exceeding full multi-head attention on quality, which β€” if it holds up β€” would make it the first point on this spectrum that isn't a tradeoff at all. That result is newer and less independently replicated than GQA's, and worth quoting with that caveat attached.

6. Choosing between them

Writing the four schemes' per-token, per-layer cache in the same units makes the progression clear. Multi-head attention stores $2 \cdot h \cdot d_h$ numbers; MQA stores $2 \cdot d_h$; GQA stores $2 \cdot g \cdot d_h$; MLA stores $d_c + d_h^R$. For a model shaped like LLaMA-2 70B ($h = 64$, $d_h = 128$), that's $16384$, $256$, $2048$ at $g=8$, and β€” using DeepSeek's ratios β€” something in the neighborhood of $576$. The quality ordering is roughly the reverse for the first three, with MLA claiming to break the pattern.

In practice the decision is usually made for you by what you're doing. If you're training a new model from scratch, GQA at $g = 8$ is the default that essentially nobody gets criticized for, and MLA is the more ambitious option worth considering if you're willing to take on the decoupled-RoPE complexity and are targeting long contexts where the cache dominates. If you have an existing multi-head checkpoint and a serving cost problem, uptraining to GQA is the answer, and it's cheap. If you can't retrain at all, then you're back in Chapter 3.4's world β€” quantize the cache, page it, window it β€” and none of this chapter helps you, which is precisely why it matters that these are architectural decisions rather than deployment ones.

It's also worth naming what all of these are and aren't orthogonal to. They compose freely with FlashAttention (which changes how attention is computed, not what's stored), with cache quantization (which changes the bytes per element, not the element count), with PagedAttention (which changes allocation, not size), and with sliding-window attention (which bounds $n$ rather than the per-token constant). They do not compose with each other β€” a layer has one key-value-head arrangement β€” so this is a genuine choice rather than another item to stack. Part III has now covered making attention cheaper in sequence length (Chapter 3.1), making position generalize further (Chapter 3.2), buying parameters without compute (Chapter 3.3), serving the result (Chapter 3.4), and shrinking what serving has to remember (this chapter). All of it has assumed that softmax attention is the mechanism and the only question is how to afford it. Chapter 3.6 closes Part III by taking that assumption apart.

7. Interview angle

A senior LLM interview will very often ask about the KV-cache, and this material is where a candidate either recites a definition or demonstrates they've had to pay for one:

  • "Your KV-cache is dominating GPU memory at high concurrency. Walk me through your options." A strong answer sorts the options by whether they require retraining. Deployment-time: quantize the cache, use paged non-contiguous allocation to eliminate fragmentation waste, bound the window. Architectural: reduce key-value heads via GQA, or move to a latent cache via MLA. The strongest answers add that if the model is an existing multi-head checkpoint, GQA is reachable through uptraining rather than a full retrain, which makes it a live option rather than a "next model" item.
  • "What exactly does Multi-Query Attention give up relative to multi-head, and why did GQA replace it?" The answer should be precise that MQA keeps all query heads and all attention patterns, and shares only the keys and values β€” so heads can still specialize in what they ask for, but no longer see different projections of the context. GQA won because the quality-versus-cache curve is sharply non-linear: eight key-value heads recover most of the memory saving at almost no quality cost, and going to one costs disproportionately more, plus MQA has training-stability problems that GQA doesn't.
  • "How would you convert a trained multi-head model to GQA?" Mean-pool the key and value projections within each group to initialize the grouped projections, then uptrain for a small fraction β€” Ainslie et al. used about 5% β€” of the original pretraining compute. A candidate who knows the mean-pooling initialization specifically, rather than just "fine-tune it," has read the paper.
  • "Explain how MLA can cache a compressed latent without paying to decompress it, and what breaks when you add RoPE." This is the sharpest version of the question. The expected answer is that $W^{UK}$ can be absorbed into $W^Q$ (and $W^{UV}$ into $W^O$) because both are fixed matrices, so queries attend directly in latent space β€” and that RoPE inserts a position-dependent rotation between them that can't be folded into a fixed product, which is why DeepSeek splits off a small separate RoPE-carrying dimension handled MQA-style.

The through-line is whether you understand that the cost of an LLM has two largely independent halves. Almost everything about training cost is a function of parameters and FLOPs; almost everything about serving cost, at the concurrency levels real products run at, is a function of a number that doesn't appear in the parameter count at all. Candidates who have only trained models tend to reach for quantization for everything; candidates who have served them reach for the cache first.

8. Self-check questions

  1. Write the KV-cache size formula and identify which terms are set by the architecture, which by the workload, and which by the serving configuration. Which of them can you change without retraining?
  2. Explain why autoregressive decoding is memory-bandwidth-bound, using the arithmetic intensity of the attention step against the cache, and why that means shrinking the cache buys you latency as well as memory.
  3. In MQA, what precisely is shared and what stays per-head? Use that to explain why "MQA removes most of the model's attention capacity" is the wrong description.
  4. Why is $g = 8$ a common choice for GQA rather than $g = 2$ or $g = 32$? Frame your answer in terms of the shape of the quality-versus-memory curve rather than as an arbitrary convention.
  5. Describe the uptraining procedure for converting a multi-head checkpoint to GQA, including how the grouped projections are initialized and roughly what it costs.
  6. Explain the weight-absorption argument that lets MLA avoid decompressing its cache, and then explain precisely why adding RoPE breaks it and what decoupled RoPE does about it.
  7. For each of FlashAttention, cache quantization, PagedAttention, and sliding-window attention, say whether it composes with GQA and why β€” and explain why GQA and MLA don't compose with each other.

9. Sources

  • Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150. https://arxiv.org/abs/1911.02150
  • Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., LebrΓ³n, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245. https://arxiv.org/abs/2305.13245
  • DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434. https://arxiv.org/abs/2405.04434
  • DeepSeek-AI (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437. https://arxiv.org/abs/2412.19437
  • Touvron, H., Martin, L., Stone, K., et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288. https://arxiv.org/abs/2307.09288
  • Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., & Dean, J. (2022). Efficiently Scaling Transformer Inference. arXiv:2211.05102. https://arxiv.org/abs/2211.05102

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

A model has $L = 30$ layers, $h = 40$ heads, $d_h = 128$. Serving one 2048-token sequence in fp16 with a full multi-head KV-cache, how many gibibytes does the cache occupy?

Explanation: $2 \cdot 30 \cdot 40 \cdot 128 \cdot 2048 \cdot 2 = 1{,}258{,}291{,}200$ bytes $= 1.17$ GiB β€” for a single sequence.

Why is shrinking the KV-cache a latency win and not only a memory win?

Explanation: Each decoding step streams the entire cache from memory and performs about one multiply-accumulate per element read. At ~1 FLOP/byte against accelerators needing hundreds, the GPU sits idle waiting on memory, so loading fewer bytes directly shortens the step.

In Multi-Query Attention, what is shared across heads and what stays per-head?

Explanation: MQA keeps all $h$ query projections, so there are still $h$ distinct attention distributions. What is lost is that heads can no longer see different projections of the context β€” they all score against one shared view.

A model has 64 query heads. Converting it to GQA with 8 key-value groups shrinks the KV-cache by what factor?

Explanation: The cache scales with the number of key-value heads, so the factor is $h/g = 64/8 = 8$ β€” recovering 87.5% of the saving MQA would have given.

Why is $g = 8$ a common GQA setting rather than $g = 1$?

Explanation: Going from 64 key-value heads to 8 costs almost nothing measurable; going from 8 to 1 costs noticeably more, and MQA additionally has training-stability problems GQA avoids.

How does Ainslie et al.'s uptraining initialize the grouped key/value projections when converting a multi-head checkpoint?

Explanation: Mean-pooling within each group produces a structurally-GQA model that already outputs something sensible; roughly 5% of the original pretraining compute then recovers essentially full quality.

What makes MLA's compressed cache free of a decompression cost at inference time?

Explanation: $q_t^\top k_s = x_t (W^Q {W^{UK}}^\top) c_s^\top$, and $W^Q {W^{UK}}^\top$ is a product of two fixed weight matrices computed once. The up-projections never run at inference.

Why does RoPE break MLA's weight-absorption trick, and what is DeepSeek's fix?

Explanation: $R_{t-s}$ depends on the specific query-key pair, so it can't be premultiplied into fixed weights. Decoupled RoPE keeps most of the head dimension absorbable and routes RoPE through a small extra slice (64 dims in DeepSeek-V2) cached as one shared key head.

Which of these does NOT compose with GQA?

Explanation: MLA and GQA are alternative answers to the same question β€” a layer has exactly one key-value-head arrangement. The other three change how attention is computed, how many bytes each element takes, and how memory is allocated, none of which conflict with GQA.