Multi-Head and Multi-Latent Attention

Llama-3.2-3B keeps 112 kibibytes of keys and values for every token of every sequence. Had it been built with a separate key/value head for each of its 24 query heads, instead of 8 shared ones, that would be 336. On the 12 GB card of this book, with the weights loaded, the first design leaves room for 11 sequences of 4,096 tokens and the second for 3.

That saving is a decision made when the model was designed, and a serving system can lose it without noticing. The way the book’s own starter code wrote grouped-query attention (repeat each shared key/value head for its queries, then run ordinary attention) reads the same bytes as the multi-head version. In the measurement below, going from 32 key/value heads to 1 changed its decode time by almost nothing (29 to 31 milliseconds). Written so that each cached byte is read once, the same step at Llama-3.2’s 8 key/value heads took 1.85 milliseconds instead of 30.

This chapter does the accounting for the four ways a model can decide how much to keep (multi-head, multi-query, grouped-query and multi-head latent attention), reads the real numbers from the models cached on this machine, measures what each costs to read on the GPU, converts a real multi-head model into grouped ones to see what that does to quality, and builds a latent-attention layer from the DeepSeek-V2 paper’s equations to check that its cache saving survives contact with the arithmetic.

Where this sits

Chapter 10 derived the cache’s size from layers, heads and head width, and Chapters 14 to 16 spent their effort managing it in blocks. Chapter 9 and 11 showed that decode is limited by the bytes it reads, and the cache is a growing share of those bytes. This chapter asks where the cache’s size comes from and what a runtime can do to honour it. Chapter 22 adds the other major lever, restricting which earlier tokens attention reads at all.

The Inference Lab piece added here is the attention accounting and decode kernels: cache-size arithmetic from a config, four ways to write a grouped decode step with their timings, a latent-attention layer in expanded and absorbed form, and a GPT-2 converter that pools key/value heads. For the Pi series, the point is that a number the memory chapters treated as a constant (bytes per token) is a property of the model artefact, which Chapter 2’s manifest should record, and which a conversion tool can change only by changing the model.

What the cache holds

Attention at a layer has \(h_q\) query heads and, in the cache, \(h_{kv}\) key/value heads of width \(d\). Each token adds a key and a value for each of them, in each of \(L\) layers:

Multi-head attention (MHA) has \(h_{kv} = h_q\). Multi-query attention (MQA, Shazeer 2019) shares a single key/value head across all query heads, which the paper argues reduces the memory bandwidth of incremental decoding at a minor quality cost. Grouped-query attention (GQA, Ainslie et al. 2023) picks an intermediate number: query heads are divided into groups and each group shares one key/value head. Query heads are never cut: the compute of the queries and the output projection is unchanged, only the size of the cache and the keys and values read from it.

I read the head counts from the config.json of every model cached here and computed the bytes per token in fp16:

Model	Layers	Query heads	KV heads	Head width	KV bytes per token	Against MHA of the same shape
GPT-2	12	12	12	64	36,864	1×
Phi-3-mini	32	32	32	96	393,216	1×
TinyLlama-1.1B	22	32	4	64	22,528	8× smaller
Llama-3.2-1B	16	32	8	64	32,768	4×
Llama-3.2-3B	28	24	8	128	114,688	3×
Mistral-7B-v0.3	32	32	8	128	131,072	4×
Llama-3.1-8B	32	32	8	128	131,072	4×
Gemma-2-2B	26	8	4	256	106,496	2×

(Pythia-160m has GPT-2’s shape; Qwen3-1.7B, whose weights are not cached, has 28 layers, 16 query and 8 key/value heads of width 128, 114,688 bytes. Gemma-2 alternates full and windowed layers, which Chapter 22 takes up; the row counts every layer at full length.)

The formula is not an approximation. For GPT-2, TinyLlama and Llama-3.2-1B I ran a 256-token forward pass and summed the bytes of the keys and values the library left in its cache: 9,437,184, 5,767,168 and 8,388,608 bytes, exactly 256 times the per-token figure from the config in each case.

Two things in the table are worth reading carefully. A design choice dominates: Phi-3-mini, a multi-head model considerably smaller than Mistral-7B, needs 393 KB per token, three times Mistral-7B’s 131 KB. And for the models with grouped attention, the saving is large and does not look dramatic in the model’s name: a “3B” Llama keeps 3 times less than it would with multi-head attention.

In capacity terms, using Chapter 14’s accounting: with Llama-3.2-3B in fp16 (6.4 GB of weights) and 90% of the 12 GiB card, 4.8 GiB remains for the cache. As built, that holds 11 sequences of 4,096 tokens, or one of 32,768. With 24 key/value heads the same budget would hold 3, and none at 32,768. This is arithmetic from the config, with activations and fragmentation ignored.

Reading less at decode

A decode step reads the whole cache for its sequence once per layer, for one query token. Chapter 11’s picture says the time of a long-context decode attention is the cache’s size over the bandwidth. If so, time should fall in proportion to \(h_{kv}\), which is the benefit MQA and GQA are for. Whether it does depends on how the kernel is written, and the measurement shows that the obvious way to write it does not.

I timed one decode step of attention for 32 query heads of width 128, in float16, with 32, 8, 4, 2 and 1 key/value heads, on random tensors, in four formulations:

• Repeat: copy each shared head for its group of queries (repeat_interleave), then ordinary attention. This is how the starter’s GQA module was written.
• Grouped: view the queries as \([B, h_{kv}, G, d]\), so each shared head attends for its whole group in one batched matrix multiply, without copying.
• enable_gqa: PyTorch’s scaled-dot-product attention with its grouped option.
• Grouped queries into the fused kernel: the same grouping, but passed to the library’s fused attention as if the group were the query positions of a single head, with no mask needed for a single decode step.

Best of three rounds, milliseconds, batch 16 and 8,192 tokens of context (the cache is 2.1 GB at 32 heads):

KV heads	Cache	Bandwidth floor	Repeat	Grouped	enable_gqa	Grouped queries, fused
32 (MHA)	2,147 MB	7.06	29.1	12.7	7.2	7.19
8	537 MB	1.77	30.0	3.48	67.4	1.85
4	268 MB	0.88	30.0	2.50	64.3	0.97
2	134 MB	0.44	30.0	0.82	62.8	0.54
1 (MQA)	67 MB	0.22	31.0	0.64	63.9	0.49

(The floor is the cache over the machine’s measured 304 GB/s. The enable_gqa figure at 32 heads is the ordinary kernel, since no grouping was needed.) Results at 4,096 tokens have the same shape. Every formulation matched the repeat version’s output to 5e-4 or better.

Three findings.

Repeating the shared heads removes the entire saving. It takes about 30 milliseconds whatever the number of key/value heads, because the copy rebuilds 32 heads’ worth of keys and values in memory and the attention then reads them. The model’s design said 4 times less cache; the code read as much as the multi-head model and wrote it too. The saving appears only when the implementation reads each cached byte once.

The fused grouped-queries form follows the floor. It is within about 5% of the bandwidth floor at 32 and 8 heads, 10% at 4, and slower relative to the floor as the cache shrinks and fixed costs begin to show (0.49 against 0.22 ms at one head). Time falls by 3.9 times from 32 heads to 8, 7.4 times to 4, and 14.7 times to 1, against cache ratios of 4, 8 and 16.

The library’s grouped option was 36 to 130 times slower than the fused grouped form, and 19 to 100 times slower than even the hand-grouped one. On this build (PyTorch 2.6.0, a GPU of compute capability 7.5) enable_gqa=True took 62 to 67 milliseconds where the fused grouped form took between 0.5 and 1.9. I did not determine which kernel it selected; the timing is what a path that materialises the repeated heads would show, and it is a fact about this build and card, not about the option. It is also a reminder that “the framework supports it” is a claim to time. With a single sequence (batch 1, 4,096 tokens) the fused grouped-queries time sat at 0.22 to 0.29 ms whatever the head count, because there the step is launch-bound (Chapters 9 and 18) and the cache is too small to matter.

A model choice, not a switch

Nothing in a multi-head model’s weights can be swapped for fewer key/value heads at serving time. The GQA paper converts an existing multi-head checkpoint into a grouped one by mean-pooling the key and value projection heads of each group (which it found to work better than keeping a single head or initialising new ones) and then uptraining, in its recipe with about 5% of the original pre-training compute. It reports, for T5-XXL, that GQA with 8 groups was already reasonable right after conversion while MQA needed the uptraining. I did the conversion on a real model to see what the first step does alone, and what a brief version of the second recovers.

GPT-2 has 12 heads, so it can be converted to 6, 4, 3, 2 or 1 key/value heads. I rewrote the model with a configurable number of key/value heads, initialised from the Hugging Face weights; with 12 it reproduces the library’s logits to 1.4e-4. The pooling is verified by a test that makes the heads within each group identical and checks that the pooled model then equals the original exactly.

The text is this book (GPT-2’s tokenizer; chapters 1 to 14 and their code for training, 263,365 tokens; chapters 15 to 16 for calibration, 40,012; chapters 17 to 19 as the test split, 83,621). Bits per token on the test split:

KV heads	Cache per token	Converted, no training	After 400 steps of uptraining	Gap to the uptrained 12-head model
12 (original)	36,864 B	4.762	3.538 (control)	0
6	18,432 B	8.337
4	12,288 B	9.087	4.169	+0.63
3	9,216 B	10.201
2	6,144 B	10.753	4.614	+1.08
1	3,072 B	11.336	4.916	+1.38

Without training, pooling destroys this model, unlike the paper’s report for T5-XXL (a much larger encoder-decoder model, scored on summarisation and question answering, so the two are not comparable): even halving the key/value heads takes the test score from 4.76 to 8.34 bits per token (perplexity from 27 to 324), and one head gives 11.3, against 15.6 for a uniform guess over GPT-2’s vocabulary. A checkpoint cannot be converted by a serving flag.

Uptraining recovers most of the damage quickly (at 4 key/value heads the calibration score fell from 8.84 to 5.45 within 25 steps and to 3.88 by step 400), but not all of it, and the control matters. A model trained for 400 steps on the book’s text improves by 1.2 bits just from adapting to the domain, so the converted models’ scores must be compared with the uptrained 12-head model, not with the original. Against it, 4 heads were 0.63 bits per token worse, 2 heads 1.08 worse and 1 head 1.38 worse, still improving at the end of training. The test split chose nothing here: each model is scored at the step that was best on the calibration split.

What this does and does not show. It shows the shape of the cost: it grows as heads are removed, it is nearly all recoverable in principle (the curves were still falling), and none of it is free at serving time. It does not show how much a model trained for GQA from the start loses, which is what the models in the table above are; for them the papers report small losses, and I have not measured any. And it is a 124M-parameter model, 400 steps (about 0.8 million tokens, three passes over a small corpus) at one seed, with a training split small enough that the 12-head model overfits it (training cross-entropy 0.85 nats against about 2.5 on the test split). The paper’s uptraining budget is a fraction of pre-training compute that I cannot compute for GPT-2 and that is certainly far larger than this.

Latent attention

Multi-head latent attention (MLA), introduced with DeepSeek-V2, takes a different route to a small cache: instead of keeping fewer heads, it keeps one compressed vector per token from which all heads’ keys and values can be rebuilt. I wrote the layer from the paper’s equations (section 2.1 and its Table 1, which I read in the paper’s text) with random weights, at a reduced size. The computation is:

where \(c_t\) has dimension \(d_c\) much smaller than \(n_h d_h\). Only \(c_t\) needs to be cached. There is a catch with rotary position embeddings: applying rotation to \(k^C\) would put a position-dependent matrix between \(W_Q\) and \(W_{UK}\), and \(W_{UK}\) could no longer be folded into the query. The paper's remedy is **decoupled rotary**: extra query components \(q^R\) per head and a single shared key \(k^R_t = \text{RoPE}(W_{KR} h_t)\) of dimension \(d_R\) carry the position, are concatenated to the content parts, and \(k^R\) is cached too. The attention score is scaled by \(1/\sqrt{d_h + d_R}\). The cache is therefore \(d_c + d_R\) elements per token per layer. For DeepSeek-V2 the paper sets \(d_c = 4 d_h\) and \(d_R = d_h/2\), which is \(4.5\, d_h\), and states that this equals GQA with 2.25 groups.

My reduced layer has 16 heads of width 128, \(d_c = 512\) and \(d_R = 64\) (the same ratios), so the cache per token per layer is:

	Elements	Against MHA
MHA (16 key/value heads)	4,096	1×
GQA, 4 heads	1,024	4× smaller
GQA, 2 heads	512	8×
MLA	576	7.1×
MQA	256	16×

576 elements is exactly 2.25 times the 256 of one grouped head, as the paper says.

Absorption, checked

The point of the construction is that keys and values need never be rebuilt. For the scores, \(q^{C\top}_{t} k^C_j = (W_{UK}^\top q^C_t)^\top c_j\): fold \(W_{UK}\) into the query once per step and attend directly against the cached latents. For the output, \(\sum_j a_j\, W_{UV} c_j = W_{UV} \sum_j a_j c_j\): sum the latents with the attention weights, and apply \(W_{UV}\) once afterwards (the paper folds it into the output projection). I implemented the expanded step, which rebuilds every key and value from the cache, and the absorbed step, and checked both against the full-sequence causal computation at several positions, in float32: the largest differences were \(5.2 \times 10^{-8}\) and \(8.2 \times 10^{-8}\) on outputs of typical magnitude 0.015. A test in float64 holds them to \(10^{-10}\), another that a later token cannot change an earlier output, and another that rotary scores depend only on the offset between positions. The absorbed form is exact, not an approximation.

What it costs to read

The layer step (projections included, cache write not timed, float16), best of three rounds, milliseconds:

Batch, context	MLA absorbed	MLA expanded	MHA	GQA-4	MQA	Cache: MLA / MHA / GQA-4
16, 1,024	1.08	5.35	0.72	0.49	0.41	19 / 134 / 34 MB
16, 4,096	1.53	19.5	3.50	1.15	0.53	76 / 537 / 134 MB
16, 8,192	2.17	not run	6.75	1.98	0.76	151 / 1,074 / 268 MB
1, 8,192	0.60	2.36	0.65	0.68	0.40	9 / 67 / 17 MB

Absorption is what makes the cache saving usable: the expanded step is 13 times slower than the absorbed one at 4,096 tokens, rebuilding keys and values for the whole context at every step. The absorbed MLA step is three times faster than MHA at long context and with 7 times less cache. But it did not beat grouped-query attention with 4 heads in time (2.17 against 1.98 ms at 8,192 tokens), despite a cache 44% smaller. My hypothesis is that the absorbed form does more arithmetic per cached byte, since every one of 16 heads works against the same latent of 576 elements instead of a head’s own 128, and that my version is a sequence of unfused matrix multiplies and a softmax, not a single kernel. That is a hypothesis about my implementation; I did not profile it. The batch-1 timings are in the noise (the median for MLA absorbed was 1.38 ms against a best of 0.60), because they are launch-bound.

The quality of MLA is not measured here at all. The layer has random weights, so it shows that the cache is small and the decode step can be fast and exact. The paper reports that MLA with this cache matches or beats multi-head attention in its experiments (its Appendix D, which I did not read); I have no model to test that on. The paper’s 93.3% reduction in KV cache is relative to DeepSeek 67B, in a model that changes other things at once, and is not what I measured: relative to a 16-head layer of the same width, my layer’s cache is 86% smaller.

What the numbers do not show

• Random tensors for the decode timings. They show bytes read and kernel behaviour, not what real activations do; the attention weights’ values do not change the work.
• One GPU, one library build. The enable_gqa result in particular belongs to PyTorch 2.6.0 on a compute-capability-7.5 card.
• My MLA is a reference implementation, unfused, with random weights, query compression omitted (the paper says it saves training memory, not cache). It checks exactness and cache size; its speed is not a production kernel’s.
• The conversion experiment is small: 124M parameters, one seed, a small corpus the models overfit, 400 steps, and a domain (this book) unlike GPT-2’s training text. The pooling follows the paper’s description, but I did not compare it with the paper’s other ways of initialising.
• Cache size is not everything. Reducing key/value heads also reduces the keys and values computed per token, and some systems make other choices (quantised or windowed caches, Chapter 22) that combine with these.
• No model trained with GQA or MLA was evaluated for quality. The papers’ claims are the authors'.

Inference Lab increment

The Lab now has:

• kv_elements_per_token, mla_elements_per_token, equivalent_gqa_groups, sequences_that_fit, kv_spec_from_config: the accounting, verified against the bytes a real cache holds.
• attn_decode_repeat, attn_decode_grouped, attn_decode_sdpa, attn_decode_sdpa_grouped: four decode-step formulations that agree and differ by up to 60 times in time.
• build_mla: DeepSeek-V2’s attention with decoupled rotary, in full-sequence, expanded-step and absorbed-step forms.
• build_gpt2_grouped, uptrain, bits_per_token: GPT-2 with a configurable number of key/value heads, initialised by mean-pooling.
• kvspec, decode, mla and quality commands.

The rule this adds: the architecture sets the cache; the runtime’s job is not to undo it. A model’s key/value head count is part of its identity and cannot be changed without retraining. An implementation honours it only if every cached byte is read once per step, and each claim of a saving (a framework option, a fused kernel, a latent form) is something to time on the hardware in use.

What you should now be able to do

• Compute a model’s KV bytes per token from its config and the number of sequences a memory budget holds.
• Explain why GQA and MQA reduce decode time in proportion to the cache and why repeating the shared heads defeats it.
• Write a grouped decode step that reads each key/value head once, and check it against a reference.
• Say what converting a multi-head checkpoint to fewer key/value heads costs without and with training, and why a control is needed.
• Derive the MLA cache size, explain why rotary embeddings force a decoupled key, and show that the absorbed form equals the expanded one.
• Say which claims about MLA this chapter checked (cache size, exactness, decode speed of one implementation) and which it did not (quality).

Key terms

• MHA: multi-head attention; one key/value head per query head.
• MQA: multi-query attention; one key/value head for all query heads.
• GQA: grouped-query attention; an intermediate number of key/value heads.
• MLA: multi-head latent attention; a compressed per-token latent from which keys and values are rebuilt.
• Uptraining: further training of a converted checkpoint to recover quality.
• Decoupled rotary key: the shared, position-carrying key that lets the content keys be absorbed.
• Absorption: folding the key and value up-projections into the query and output projections.
• Bandwidth floor: the time to read a byte count at the measured memory bandwidth.

Reader exercises

1. Add the keys and values of a sliding window to the cache accounting for Gemma-2, using the config’s window and layer pattern. How far does 106,496 bytes per token fall at 32,768 tokens?
2. Time attn_decode_sdpa_grouped on a different GPU or PyTorch version. Does enable_gqa behave the same?
3. Quantise the cache to 8 bits (Chapter 12) for the grouped kernel. How close to the new floor does it get?
4. Uptrain the 1-head GPT-2 for 4,000 steps on a larger corpus. How much of the 1.38-bit gap closes?
5. Convert by a different initialisation (take the first head of each group instead of the mean). How does the unconverted score compare?
6. Fuse the absorbed MLA step into fewer kernels. How close to the 0.5 ms floor of its 151 MB cache can it get at 16 × 8,192?
7. Add query compression to build_mla and check that it changes nothing about the cache or the exactness.
8. Compute the cache per token for a DeepSeek-V2-sized layer (128 heads, \(d_h = 128\)) under MHA, GQA-8 and MLA from the paper’s settings.

References

• Fast Transformer Decoding: One Write-Head is All You Need — Shazeer, 2019; multi-query attention; read on 21 October 2026.
• GQA — Ainslie et al., EMNLP 2023; grouped-query attention and uptraining from multi-head checkpoints; read on 21 October 2026.
• DeepSeek-V2 — DeepSeek-AI, 2024; multi-head latent attention, decoupled rotary embedding, cache comparison (section 2.1, Table 1); read on 21 October 2026.

Evidence in this chapter

Claim	Category
MQA shares one key/value head; GQA uses an intermediate number; both reduce decode bandwidth	Established background (the papers’ abstracts)
MLA caches \(d_c + d_R\) per layer, needs decoupled rotary for absorption, equals GQA with 2.25 groups at DeepSeek-V2’s settings	Established background (the paper’s text); the 576 against 256 arithmetic reproduces it
KV bytes per token for ten cached models, and that a real cache holds exactly that many bytes	Book result (three models measured, the rest read from config)
Llama-3.2-3B fits 11 sequences of 4,096 tokens against 3 if it were multi-head	Arithmetic from config; activations ignored
Repeating shared heads gives no saving (about 30 ms at every head count); a fused grouped-query form follows the bandwidth floor	Book result, one GPU, random tensors
enable_gqa was 36 to 130 times slower than the fused grouped form on this build	Book result for PyTorch 2.6.0 and a compute-capability-7.5 GPU only
Mean-pooling heads and uptraining as the recipe of the GQA paper (checked in its text)	Established background; the paper’s T5-XXL results are its authors'
Pooling GPT-2’s key/value heads without training destroys quality (4.76 to 8.34 bits at 6 heads); brief uptraining leaves a gap of 0.63 to 1.38 bits to an uptrained control	Book result, one seed, a small corpus the model overfits
The absorbed MLA step equals the full-sequence computation (8.2e-8) and is 13 times faster than the expanded step at 4,096 tokens	Book result, random weights
Absorbed MLA was not faster than GQA-4 despite a smaller cache	Book result for this implementation; the explanation is a hypothesis
MLA quality; the 93.3% cache reduction; quality of models trained for GQA	Not measured / authors’ claims

Implementation status

• Manuscript: content/books/inference/21-chapter.md
• Runnable examples: examples/ch21/attention_variants.py (original toy modules; its “MLA” is a sketch only), examples/ch21/attention_lab.py (accounting, decode kernels, MLA and GPT-2 conversion)
• Tests: tests/test_ch21_attention.py, tests/test_ch21_attention_lab.py (cache arithmetic and config parsing, the four decode formulations agreeing, shared heads really shared, rotary invariants, MLA absorbed and expanded steps equal to the full sequence, causality, GPT-2 conversion equal to the original at full heads and exact under identical groups, uptraining reducing loss; all on CPU)
• Evidence: research/ch21/sources.md, research/ch21/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about the quality of any model trained with grouped or latent attention, or about any production attention kernel. Chapter 22 turns to the other way of keeping less: sliding-window attention, where the cache stops growing with the context.