Attention Variants · MLA Part 1

Understanding MLA (Multi-head Latent Attention)

DeepSeek-style models store one small latent per token instead of full keys and values. This post follows one token through the forward, then explains why the head must split into a content part and a position part.

MLA is a low-rank version of multi-head attention. It keeps every query head, but stores keys and values as one shared latent per token — like Multi-Query Attention's tiny cache, yet each head still gets its own learned view. This post is the architecture; a follow-up covers the two ways to run it at inference.

Where MLA is used

MLA is the attention behind DeepSeek-V2 / V3, Kimi K2.5, and GLM-5.2 (MLA-256). Every walkthrough number here uses DeepSeek-V3; the decode post lets you switch models in the charts.

Assumed background

Multi-head attention (Q, K, V, one cache per head) and RoPE. If RoPE is rusty, the RoPE post covers it. Everything else is defined here.

1.Where MLA sits: MHA → GQA → MQA → MLA

All four schemes keep many query heads. They differ only in what the KV cache holds per token.

SchemeWhat is cached per tokenPer-head views?Cache size
MHAone full key and value per headyes, independentlargest
GQAkey/value for a few shared groupsshared in groupsmedium
MQAone key and one value, shared by all headsno, identicalsmallest
MLAone shared latent (plus a small position sidecar — Section 3)yes — reconstructed per headnear-MQA

MQA is cheap because it shares one key/value, but every head sees the same thing. MLA also caches one shared object, but it is a richer latent from which each head reconstructs its own key and value with its own learned matrix. So MLA gets MQA-like storage with MHA-like head diversity.

One-line placement: MLA = the small cache of MQA, the per-head expressiveness of MHA.

2.What is actually "MLA" (the common trap)

Reading the code, it is easy to think MLA is "LoRA on Q and K." It is not. Two independent things live in the attention block, and only one of them is MLA.

PieceWhat it isIs it MLA?
Query low-rankoptional parameter compression of the query projectionNo — just saves weights; the query is not cached
KV latentcompress K and V for all heads into one latent, decompress per headYes — this is MLA

The query low-rank path is optional. The KV latent is the whole point and is always present. Formally, for a token's hidden vector \(x\):

\[ c \;=\; x\,W_{\text{down}}, \qquad c\in\mathbb{R}^{512} \]

and each head \(h\) reconstructs its own key and value from that same latent:

\[ K^{(h)} \;=\; c\,W_K^{(h)}, \qquad V^{(h)} \;=\; c\,W_V^{(h)}. \]

The per-head matrices are weights, not cache — which is why each head can differ while the storage stays small. For now, think of \(c\) as “the thing that replaces full K and V.”

That is the core idea. The real forward has one more piece — position — which forces a split we have not named yet. Section 3 walks through it; Section 5 states exactly what the cache holds.

3.One token through MLA

Follow a single token's hidden state through one attention layer. DeepSeek-V3 numbers; each width is introduced when it first appears.

hidden 7168 heads 128 latent 512

Four phases. Each does one job.

Phase ACompresshidden → what gets cached
Phase BExpandlatent → per-head K, V
Phase CQuerybuild Q the usual way
Phase DAttendconcat, score, output

Phase ACompress: make the thing that will be cached

Start with the token's hidden state, width 7168. One small linear layer projects it — but not to a single latent. The output is 576-wide, and it immediately splits:

hidden state  (7168)
↓ one linear:  7168 → 576
compressed  (576)
↓ split 576 = 512 + 64
c latent
(512)
K_rope
(64)
same projection, two roles — we are about to see why
A1Down-project and split(1, 576) → (1, 512) + (1, 64)

The 512-wide piece is the KV latent from Section 2 — compressed content, the source of every head's key and value. The 64-wide piece is new: a shared rope key that carries position. It is produced here, beside the latent, not inside it.

Why two pieces? The short answer: content can live in a fixed low-rank map; position cannot. Section 4 proves that. For now, just notice the architecture peels them apart at the first step.
A2RMSNorm the latent only(1, 512) → (1, 512)

Normalize the 512-wide latent. Leave the 64-wide rope key alone.

This is not QK-norm. QK-norm (as in Qwen3) normalizes each head's full query and key after projection. Here the norm sits on the shared content bottleneck before any head is expanded. Rope is left out so its magnitude does not change the content scale (RMSNorm divides by one shared RMS over the whole vector). Also, because RMSNorm is nonlinear: \(W\cdot\operatorname{RMSNorm}(c) \neq \operatorname{RMSNorm}(W\,c)\).
Notation — what we just named

latent \(c\) (512) = compressed content, source of every head's K and V.
K_rope (64) = position key, produced beside the latent, not inside it.
cache entry = latent + K_rope = 576 — two objects side by side, not one fused vector.

Phase BExpand: give every head its own K and V

The latent is shared. The heads are not. One up-projection turns the single latent into 128 heads' worth of content keys and values. The rope key does not go through this step.

B1Up-project the latent → content key and value(1, 512) → (1, 128 heads, 256)

A linear layer (kv_b_proj) maps the latent to num_heads × (nope_dim + value_dim) = 128 × (128 + 128). Then split:

\[ K_{\text{nope}}^{(h)} \in \mathbb{R}^{128}, \qquad V^{(h)} \in \mathbb{R}^{128} \quad\text{for each head }h. \]

“Nope” means no positional encoding — the content half of the key. That is all this matrix produces. Rope never enters it.

One latent, two content roles. The same cached vector is the source of both the content key and the value. Each head has its own slice \(W_K^{(h)}\) and \(W_V^{(h)}\). The superscript \((h)\) means head index — head 0 uses matrix 0, head 1 uses matrix 1, and so on.
Why “latent” can be wider than a head

latent_dim = 512 looks big next to nope_dim = 128. That is fine. The latent is shared across all heads, so one 512-vector replaces 128 × (128 + 128) = 32,768 content values per token. “Latent” means low-rank relative to the full multi-head KV, not smaller than one head.

B2RoPE the rope key, then broadcast(1, 1, 64) → (1, 128, 64)

Apply rotary position embedding to the 64-wide rope key. Then copy that one vector to every head — no matmul, just a broadcast. Width stays 64. Still no up-projection.

Why shared, not per-head? Position is a fact about the token, not about the head. Giving each head its own rope key would multiply storage by 128 for no new information. And there is nothing to up-project: rope was never compressed into the latent.

After Phase B, every head has:

The full key is not assembled yet. That happens in Phase D, after the query is ready.

Phase CQuery: build Q the ordinary way, then split

The query path does not go through the KV latent. Queries are not cached, so there is nothing to compress for the cache. (DeepSeek-V3 does use an optional low-rank factorization on the query weights to save parameters — that is the trap from Section 2, and it is unrelated to the KV latent.)

C1Project to per-head queries(1, 7168) → (1, 128 heads, 192)

Each head gets a query of width qk_head_dim = 192. That width is not free — it is the sum of the two parts we are about to split:

\[ \underbrace{192}_{\textit{query/key width}} \;=\; \underbrace{128}_{\textit{nope}} \;+\; \underbrace{64}_{\textit{rope}}. \]
C2Split and RoPE the rope part(1, 128, 192) → [128 | 64]

Split each head's query into Q_nope (128) and Q_rope (64). Apply RoPE only to Q_rope, using the same position as the key. Leave Q_nope untouched.

Same split as the key, for the same reason. The content half will later dot with K_nope. The position half will later dot with K_rope. RoPE on both sides is what makes the score depend on the relative offset (see the RoPE post).

Phase DAttend: concat, score, project out

D1Concatenate nope and ropeQ, K: (1, 128, 192)

Per head:

\[ Q \;=\; [\,Q_{\text{nope}} \;\Vert\; Q_{\text{rope}}\,], \qquad K \;=\; [\,K_{\text{nope}} \;\Vert\; K_{\text{rope}}\,]. \]

Value stays width 128 — no rope split on V. Value width is independent of query/key width; the score matmul and the value matmul are two different operations, so they can have two different widths.

D2Scaled attention(1, 128, 128)

Ordinary multi-head attention, with the full query/key width in the scale:

\[ \mathrm{Attention}(Q,K,V) \;=\; \mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{192}}\right)V. \]

Output per head is width 128 (the value width). Flatten the 128 heads and project back to hidden 7168.

Putting the two paths side by side, the asymmetry is the whole design:

Query path

Q

hidden (7168)
↓ project (optionally low-rank)
per-head query (192)
↓ split 192 = 128 + 64
Q_nope (128)  |  Q_rope (64)
↓ RoPE on the 64 only
Q ready
Key / Value path

K, V

hidden (7168)
↓ compress to 576, split
latent (512)  |  K_rope (64)
↓ RMSNorm latent; up-project
K_nope & V, per head
↓ RoPE K_rope; broadcast to heads
K, V ready
PartComes fromPer head?Cached?
K_nope, Vlatent, via per-head up-projectionyes — own matrixonly the latent
K_ropedirect slice of the compress projectionno — sharedthe tiny 64-wide key
Q_nope, Q_ropequery projection of the hidden stateyesnever — queries are not cached

4.Why position cannot live in the latent

Section 3 showed that position is peeled off before the latent. Here is why.

The content part works because its reconstruction matrix is fixed: the same \(W_K^{(h)}\) applies to every cached token. Written out for one head, one query position \(m\), one key position \(n\):

\[ Q_{\text{nope}}\cdot K_{\text{nope}} \;=\; Q_{\text{nope}}\,\big(W_K^{(h)}\,c\big) \;=\; \big(Q_{\text{nope}}\,W_K^{(h)}\big)\cdot c. \]

The matrix \(W_K^{(h)}\) does not depend on \(n\), so it can be moved onto the query side once. That regrouping is what makes efficient decoding possible — the subject of Part 2.

Position breaks that. RoPE inserts a relative rotation \(R_{\,n-m}\) that depends on both positions, sitting between query and key:

\[ \big(R_m\,q_{\text{rope}}\big)\cdot\big(R_n\,k_{\text{rope}}\big) \;=\; q_{\text{rope}}^{\top}\,R_{\,n-m}\,k_{\text{rope}}. \]

There is no single fixed matrix here — \(R_{\,n-m}\) is different for every cached position \(n\):

key at position 0  → relative rotation  R(0 − m)
key at position 1  → relative rotation  R(1 − m)
key at position 2  → relative rotation  R(2 − m)
...                                (a different rotation per cached token)

So position cannot be folded into the fixed latent reconstruction. DeepSeek's fix is to carry position separately: a small 64-wide rope slice, shared across heads, rotated per position, concatenated back on. Content lives in the compressible latent; position rides alongside in a tiny explicit vector.

Small cache?Fixed reconstruction?
nope (content)yes — in the latentyes — fixed \(W_K^{(h)}\)
rope (position)yes — but tiny & shared, not in the latentno — rotation changes per position

5.What gets cached

Per token, per layer, the MLA design stores two objects — do not call the whole thing “the latent”:

ObjectWidthWhat it isNormed?Up-projected?
latent \(c\)512compressed content (source of K_nope + V)yes — RMSNormyes — per head
K_rope64position keynonever
cache entry576latent + K_rope, side by side

Compare to a full multi-head cache that would store, per head, a 192-wide key and a 128-wide value:

MHA-style full cache
40,960 values
MLA cache (c + rope)
576 values

128 × (192 + 128) = 40,960 versus 512 + 64 = 576 — about 71× smaller.

Phases B–D can always be recomputed from those two objects plus the weights. That is the whole storage win.

Reference code vs serving engines

The Hugging Face path walked above is the training / reference form: it expands the latent into full K and V, then (if caching) stores those expanded tensors — so you will not see a 576-wide cache in that file. Serving engines such as SGLang and vLLM keep the latent itself and never materialize full K/V for every past token. Same weights, different factorization. How you recompute at decode — rebuild full K/V, or fold the up-projections into the query — is the subject of Part 2.

6.Training vs decode (preview)

Same math, two factorizations: training expands the latent and runs ordinary attention; decode keeps the latent and folds reconstruction into the query/output. The crossover between those two algorithms is Part 2.

7.Summary

QuestionAnswer
What is MLA?compress content K/V into one latent; reconstruct per head; keep rope separate
Is the query low-rank part MLA?no — optional weight compression; MLA is the KV latent
Latent vs rope?latent = content (512); K_rope = position (64); cache = both side by side (576)
Does W_UK include rope?no — only latent → nope (+ W_UV for V)
Is rope ever up-projected?no — not in training, materialize, or absorb
What does the norm touch?latent only — never the rope key
How big is the win?~71× smaller KV cache for DeepSeek-V3

MLA in one line: cache one latent per token, let every head learn its own key/value from it, and keep a tiny shared rope key because position refuses to compress.

References

How to cite this post

Dong, S. (2026). Understanding MLA (Multi-head Latent Attention).
https://simondong1.github.io/mla.html