Fast Attention and Long Context
Multi-head latent attention
Multi-head latent attention stores one small compressed vector per token instead of full keys and values, then expands it when needed, cutting the cache by around seventy times.
- 12 min read
- 3 reading levels
- Updated
Read these first
On this page 10
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
The short answer
Multi-head latent attention stores a tiny summary of each word, and rebuilds the full notes on demand.
The squash bottle in your kitchen
You do not keep fifty glasses of sharbat in the fridge. You keep one small bottle of concentrate.
When someone wants a drink, you pour a little and add water. One small bottle on the shelf, many glasses on demand.
Multi-head latent attention does that with memory. It stores the concentrate, and mixes the full drink at the moment of use.
What was in the fridge before
For every word it has read, a model stores notes for every head. A head is one independent reader inside the model, and its notes are called keys and values.
The previous two lessons shrank those notes by sharing them. Multi-query attention shared one set. Grouped-query attention shared one set per small group.
Sharing has a floor. You cannot go below one set, and one set costs quality.
The different idea
Do not share the notes. Compress them.
For each word, squeeze everything the heads would need into one short vector. Store only that. When a head needs its keys and values, expand the short vector back out.
grouped-query attention latent attention
----------------------- ----------------
store 8 sets of full notes store 1 short summary
read them back directly expand it when reading
per word: thousands of per word: a few hundred
numbers numbersThe compression is learned during training, not designed by hand. The model works out for itself what is worth keeping.
Does it lose information?
Yes, some. This is squeezing, and squeezing throws things away.
Here is the surprise that made this idea famous. The team that built it found the compressed version scored better than the uncompressed one on their tests, not worse.
The most likely reason is that the squeeze acts like a discipline. It stops the model storing noise, in the same way a word limit often improves an essay. That is a plausible explanation, not a proven one.
The clever bit that costs nothing
If you have to expand the summary before every use, have you saved any work? It sounds like you have moved the cost rather than removed it.
You have not, because of a piece of arithmetic. The expanding step can be folded into the question-asking step ahead of time, once, when the model is loaded. At run time the questions are compared against the summaries directly, and the full notes are never built.
This part is genuinely fiddly and almost everyone needs two passes at it. That is normal.
The wrinkle worth knowing
Models track word order using a rotation applied to keys. That rotation does not survive the folding trick.
So a small slice of each key is kept outside the compression, uncompressed, carrying the position information. The design is a compromise, not a clean idea, and the paper says as much.
Where you have already seen this
- DeepSeek's models, which is where this was invented and first shipped at scale.
- Any service offering a very long context at a low price per word.
- Open models you can run at long context on fewer cards than you expected.
Remember this
- Store a small learned summary per word, not full notes for every head.
- Expand it only when needed, and fold the expansion into the question step so it is free.
- Around seventy times less stored per word than plain attention, at similar or better quality.
What to learn next
- Sliding-window attention — bounding the cache instead of compressing it.
- Sparse and block attention patterns — reading fewer tokens rather than storing them smaller.
- LoRA — the other well-known use of low-rank matrices, for a different purpose.
Developer — Code and libraries.
Setup
pip install torchWritten against PyTorch 2.5.1, Python 3.10. Runs on CPU in under a second.
The mechanism, and the identity that makes it fast
import torch
torch.manual_seed(0)
d_model, n_heads, d_head, d_c = 128, 4, 32, 16 # d_c is the latent size
W_DKV = torch.randn(d_model, d_c) / d_model**0.5 # down-project: what gets cached
W_UK = torch.randn(d_c, n_heads * d_head) / d_c**0.5 # up-project to keys
W_UV = torch.randn(d_c, n_heads * d_head) / d_c**0.5 # up-project to values
W_Q = torch.randn(d_model, n_heads * d_head) / d_model**0.5
x = torch.randn(6, d_model) # 6 tokens
c = x @ W_DKV # THE CACHE: 6 x 16
K = (c @ W_UK).view(6, n_heads, d_head) # reconstructed on the fly
V = (c @ W_UV).view(6, n_heads, d_head)
Q = (x @ W_Q).view(6, n_heads, d_head)
print("cached latent c:", tuple(c.shape), "=", c.numel(), "numbers")
print("keys + values if cached directly:", K.numel() + V.numel(), "numbers")
print("ratio:", (K.numel() + V.numel()) / c.numel(), "x smaller")
# The absorption identity: scores can be computed against the latent directly.
q_h = Q[:, 0, :] # head 0 queries, 6 x d_head
k_h = K[:, 0, :]
scores_normal = q_h @ k_h.T # reconstruct K, then score
W_UK_h = W_UK[:, 0:d_head] # this head's slice, d_c x d_head
scores_absorbed = (q_h @ W_UK_h.T) @ c.T # fold W_UK into the query instead
print()
print("absorbed scores match:", torch.allclose(scores_normal, scores_absorbed, atol=1e-4))
print("max difference:", (scores_normal - scores_absorbed).abs().max().item())
print()
print("DeepSeek-V3, from its published config.json")
L, heads = 61, 128
kv_lora_rank, qk_rope, qk_nope, v_dim = 512, 64, 128, 128
mla = 2 * L * (kv_lora_rank + qk_rope) # 2 bytes per bf16 number
mha = 2 * L * heads * ((qk_nope + qk_rope) + v_dim)
print(f" MLA cache : {mla:>10,} bytes/token ({mla/1024:.1f} KiB)")
print(f" same shape, MHA: {mha:>10,} bytes/token ({mha/1024/1024:.2f} MiB)")
print(f" ratio : {mha/mla:.1f}x")
for ctx in (32_768, 131_072):
print(f" at {ctx:>7,} tokens: MLA {mla*ctx/2**30:6.2f} GB vs MHA {mha*ctx/2**30:7.2f} GB")cached latent c: (6, 16) = 96 numbers keys + values if cached directly: 1536 numbers ratio: 16.0 x smaller absorbed scores match: True max difference: 1.9073486328125e-06 DeepSeek-V3, from its published config.json MLA cache : 70,272 bytes/token (68.6 KiB) same shape, MHA: 4,997,120 bytes/token (4.77 MiB) ratio : 71.1x at 32,768 tokens: MLA 2.14 GB vs MHA 152.50 GB at 131,072 tokens: MLA 8.58 GB vs MHA 610.00 GB
Reading the output
scores_absorbed equals scores_normal. This is the line that matters most in the file. (q W_UK^T) c^T and q (c W_UK)^T are the same product, associated differently.
The consequence: at decode time you never build K. You pre-multiply the query projection by the key up-projection once at load time, and every step scores queries directly against cached latents. The "expansion" cost is zero per token.
16x on the toy, 71x on the real thing. The toy compresses 4 heads x 32 dims x 2 down to 16. DeepSeek-V3 compresses 128 heads x 320 dims down to 512 + 64.
8.58 GB against 610 GB at 128k context. The right-hand figure is not a real design anyone tried; it is what DeepSeek-V3's attention shape would cost without compression. It is the reason MLA exists.
The real hyperparameters
From DeepSeek-V3's config on the Hub, which you can fetch yourself:
| Field | Value | Meaning |
|---|---|---|
kv_lora_rank | 512 | width of the cached latent |
q_lora_rank | 1536 | queries are compressed too, for parameters not cache |
qk_nope_head_dim | 128 | per-head key dims that come from the latent |
qk_rope_head_dim | 64 | per-head key dims carrying position, cached separately |
v_head_dim | 128 | per-head value dims |
num_attention_heads | 128 | |
num_hidden_layers | 61 |
Cache per token per layer is kv_lora_rank + qk_rope_head_dim = 576 numbers. Note that the second term is shared across heads, not per head. That is the whole trick in one line.
Why the position part is separate
Rotary position embedding rotates each key by an angle that depends on its position. Position-dependent means it cannot be folded into a fixed matrix at load time, so the absorption identity breaks.
DeepSeek's answer is decoupled RoPE. Split the key: 128 dims per head reconstructed from the latent and never rotated, plus 64 dims that are rotated and cached directly, shared across all heads. Scores are the sum of a compressed part and a positional part.
It is an engineering compromise. The paper presents it as one.
Common mistakes
Reading num_key_value_heads from a DeepSeek config and concluding it uses MHA. The field says 128, equal to the query heads. The cache is 576 numbers per layer regardless, because the runtime caches latents, not keys. Config fields describe compute shape, not cache shape.
Implementing MLA without absorption and benchmarking it. Reconstructing keys and values every step is correct and slow. You will measure MLA as worse than GQA and conclude the wrong thing.
Assuming any inference server supports it. MLA needs bespoke kernels. Support in vLLM, SGLang and llama.cpp arrived model by model. Check your server's release notes for the specific model, not for "MLA".
Trying to convert a GQA checkpoint to MLA. There is no cheap uptraining recipe of the kind GQA has. MLA is a pre-training decision today.
Confusing this with LoRA. Both use low-rank matrices. LoRA adds a low-rank update to frozen weights for fine-tuning; MLA makes the key/value path low-rank in the architecture. Unrelated goals. See LoRA.
Try it yourself
Set d_c = 4 and print the difference between K and its best rank-4 reconstruction. Then raise d_c to 128 and confirm the ratio drops below 1, which is the point where compression stops paying.
What to learn next
- Sliding-window attention — bounding the cache instead of compressing it.
- Sparse and block attention patterns — reading fewer tokens rather than storing them smaller.
- LoRA — the other well-known use of low-rank matrices, for a different purpose.
Researcher — Mathematics and papers.
Formulation
From DeepSeek-AI (2024), DeepSeek-V2 (arXiv:2405.04434), refined in DeepSeek-V3 (arXiv:2412.19437).
For token $t$ with hidden state $h_t \in \mathbb{R}^{d}$, compress keys and values jointly:
$$ c^{KV}_t = W^{DKV} h_t \in \mathbb{R}^{d_c}, \qquad d_c \ll h \cdot d_h $$
$W^{DKV} \in \mathbb{R}^{d_c \times d}$ is the down-projection, $h$ the head count and $d_h$ the per-head dimension. Reconstruct per head:
$$ k^{C}_{t,i} = W^{UK}_i c^{KV}t, \qquad v{t,i} = W^{UV}_i c^{KV}_t $$
Only $c^{KV}_t$ is cached. Queries are compressed too, through $c^{Q}_t = W^{DQ} h_t$ with rank $d_c'$, which reduces activation memory during training but not the cache.
Decoupled rotary embeddings
RoPE applies a position-dependent rotation $R_t$ to keys. Then $q_{t,i}^{\top} R_{s} W^{UK}_i c^{KV}_s$ contains $R_s$ between the two weight matrices, so $W^{UK}_i$ cannot be absorbed into $W^{Q}_i$ at load time.
The fix splits the key into a content part and a positional part:
$$ k_{t,i} = \left[\, \underbrace{W^{UK}_i c^{KV}t}{d_h^{\text{nope}} = 128} \;;\; \underbrace{R_t W^{KR} h_t}_{d_h^{R} = 64,\ \text{shared across heads}} \,\right] $$
$W^{KR} \in \mathbb{R}^{d_h^R \times d}$ produces one rotated vector per token, shared by all heads. Cache per token per layer is therefore
$$ d_c + d_h^{R} = 512 + 64 = 576 \text{ elements} $$
against $h(d_h^{\text{nope}} + d_h^{R} + d_h^{v}) = 128 \times 320 = 40{,}960$ for the equivalent MHA. At 61 layers and bf16 this is 70,272 bytes per token, close to 68.6 KiB, against 4.77 MiB.
Matrix absorption
The identity exploited at inference:
$$ (q_{t,i})^{\top} \left(W^{UK}_i c^{KV}_s\right) = \left((W^{UK}i)^{\top} q{t,i}\right)^{\top} c^{KV}_s $$
Precompute $\tilde{W}^{Q}_i = (W^{UK}_i)^{\top} W^{UQ}_i$ once. Then scoring reads only $c^{KV}$ from the cache. Similarly $W^{UV}_i$ absorbs into the output projection $W^{O}_i$. Neither $K$ nor $V$ is ever materialised during decode.
The cost is a wider effective query: scores are computed in $\mathbb{R}^{d_c}$ rather than $\mathbb{R}^{d_h}$, so attention FLOPs per token rise. Since decode is memory-bound (see memory-bound vs compute-bound), trading FLOPs for bytes is the correct direction. During compute-bound prefill the calculus reverses, and implementations often reconstruct $K$ and $V$ explicitly for prefill while using absorption for decode.
The quality result
DeepSeek-V2 reports MLA outperforming MHA on their evaluation suite while using a fraction of the cache. This is the claim that separates MLA from every other cache-reduction method, all of which trade quality for memory.
Treat the mechanism as unsettled. Candidate explanations include a regularisation effect from the bottleneck, the higher effective head dimension during scoring, and confounds from the surrounding architecture and data. The result has been reproduced by others at smaller scale, but a clean ablation isolating the bottleneck from the rest of DeepSeek's recipe is not, to my reading, in the public literature.
Position in the design space
| Method | Cache per token per layer | Quality vs MHA |
|---|---|---|
| MHA | $2 h d_h$ | baseline |
| GQA ($g$ groups) | $2 g d_h$ | close |
| MQA | $2 d_h$ | measurably worse |
| MLA | $d_c + d_h^{R}$ | reported equal or better |
GQA reduces the head count multiplier. MLA removes the head count from the expression entirely, replacing it with a rank chosen independently. That is a different axis, not a better point on the same one.
What came next
DeepSeek-V3.2 layers DeepSeek Sparse Attention on top of MLA: a lightweight "lightning indexer" scores query-key relevance cheaply, selects a top-$k$ subset of tokens (2048 in the released code), and runs fine-grained attention only over those. Cache compression and attention sparsity are orthogonal, and combining them is the current frontier for long-context serving. See sparse and block attention patterns.
Adoption outside DeepSeek has been slower than the numbers suggest it should be. The reason is not the mathematics but the kernels: GQA is a two-line change to any serving stack, and MLA is a rewrite.
What to learn next
- Sliding-window attention — bounding the cache instead of compressing it.
- Sparse and block attention patterns — reading fewer tokens rather than storing them smaller.
- LoRA — the other well-known use of low-rank matrices, for a different purpose.