Author(s): Naimish Baranwal

Originally published on Towards AI.

From Vanilla Transformers to Modern LLMs: The Evolution of Transformer Architecture — Part 3: Attention Optimization (MHA, MQA, GQA, MLA)

TL;DR

The problem. When an LLM generates text, each new token has to look back at every earlier token. To avoid recomputing the Keys (K) and Values (V) of those earlier tokens, the model stores them in memory. This is the KV cache. It grows with every token, every layer and every attention head, so for long conversations it becomes one of the biggest memory costs during inference.

The question. Can we shrink the KV cache without hurting quality?

Four answers, in the order they were invented:

MHA (Multi-Head Attention) is the original design. Every Query head gets its own Key and Value. It gives the model the most independent “perspectives” on the text, but it stores the most data.

MQA (Multi-Query Attention) makes all Query heads share a single Key and Value. The cache becomes tiny, but every head has to work from the same K/V, so quality can drop.

GQA (Grouped-Query Attention) is the middle ground. Query heads are split into groups, and each group shares one Key and Value. You choose how many groups, which gives you a dial between MHA’s quality and MQA’s memory savings.

MLA (Multi-Head Latent Attention) takes a different route. Instead of sharing K/V between heads, it compresses them into one small vector per token (a latent) and rebuilds each head’s own Key and Value from that vector when needed. Every head keeps its own perspective, but the cache shrinks a lot.

Cache size (smaller is better) MHA >> GQA > MQA MHA >> MLA (about GQA with 2.25 groups in DeepSeek-V2)

Head diversity (higher is better) MHA = MLA > GQA > MQA

MHA: Intuition

Before understanding GQA and MQA, we first need to understand exactly how Multi-Head Attention (MHA) works.

In the previous part, we saw that KV Cache can become a major memory bottleneck during inference, especially when every attention head maintains its own Key and Value cache.

So the natural question is:

Why do we need multiple attention heads in the first place?

The main idea behind MHA is simple:

Instead of using a single Query, Key, and Value representation for the entire model, we divide the attention representation into multiple heads.

Each head operates on a different portion of the representation and can learn a different type of relationship between tokens.

You can think of it as asking several people to read the same sentence, but each person focuses on a different aspect.

For example:

Head 1 → focuses more on syntactic relationships
Head 2 → focuses more on semantic relationships
Head 3 → focuses more on nearby-token relationships
Head 4 → focuses more on long-range relationships
...

The exact behavior is learned by the model; we do not manually assign a specific purpose to each head.

1. How Does MHA Create These Different Heads?

Suppose our model has:

d_model = 4096 #embedding dimension
n_head = 32 #number of heads

Then each attention head operates on a smaller dimension:
d_head = d_model / n_head #embedding dimension per head
d_head = 4096 / 32= 128
 

Conceptually, the model first projects the input into Query, Key, and Value representations:

Q = XWq
K = XWk
V = XWv
The resulting representations are then divided into n_head separate heads:
Q → Q₁, Q₂, Q₃, ..., Q₃₂
K → K₁, K₂, K₃, ..., K₃₂
V → V₁, V₂, V₃, ..., V₃₂

Each head has:
Qᵢ : 128 dimensions
Kᵢ : 128 dimensions
Vᵢ : 128 dimensions
and performs its own attention:
Head i:
Attentionᵢ =softmax(QᵢKᵢᵀ / √d_head)Vᵢ

Finally, the outputs from all heads are combined:

Head 1 ──┐
Head 2 ──┤
Head 3 ──┤
 ... ├──→ Concatenate → Output Projection
Head 32 ─┘

This is the basic idea of Multi-Head Attention.

2. Why Multiple Heads?

The advantage is that the model does not have to represent every relationship inside a single attention space.

Instead, it gets multiple learned representation spaces:

Input
 ↓
 Q, K, V projections
 ↓
 ┌─────────┼─────────┐
 ↓ ↓ ↓
 Head 1 Head 2 Head 3 ... Head N
 ↓ ↓ ↓
 Attention Attention Attention
 └─────────┼─────────┘
 ↓
 Concatenate
 ↓
 Output Projection

This gives the model multiple learned perspectives over the same sequence.

But there is a cost.

Because every head has its own Key and Value representations, MHA also requires a separate K/V cache for every head during autoregressive inference.

For example, with 32 heads:

Head 1 → K₁, V₁
Head 2 → K₂, V₂
Head 3 → K₃, V₃
...
Head 32 → K₃₂, V₃₂

All of these K/V representations must be stored for the previous tokens.

This is exactly the issue we discussed in Part 2.

More attention heads can provide greater representational flexibility, but they also mean:

More KV heads
 ↓
Larger KV Cache
 ↓
Higher memory usage during inference

3. How Many Attention Heads Do We Need?

A natural question is:

If multiple attention heads are useful, how many heads should a Transformer have?

The number of attention heads is an architectural hyperparameter. It is decided when designing the model architecture and remains fixed during training and inference.

For example, a model might be configured as:

d_model = 4096
n_head = 32
This gives:
d_head = d_model / n_head = 4096/32 =128

So each attention head works with a 128-dimensional Query, Key, and Value representation.

Another model could use:
d_head = 4096 / 64 = 64

So increasing the number of heads generally means that each individual head operates on a smaller subspace.

Also, there is no single universal formula that says:

“A 7B model must have exactly 32 heads.”

Instead, model designers choose n_head based on several considerations:

1. Model quality
2. Training cost
3. Inference speed
4. Memory usage
5. KV Cache size

For many Transformer architectures, designers prefer a head dimension that is large enough to represent useful relationships, while using enough heads to provide multiple learned attention patterns.

For example:
d_model = 4096

Option 1:
n_head = 32
d_head = 128
Option 2:
n_head = 64
d_head = 64

Both are mathematically possible, but they create different architectures and computational trade-offs. The final configuration is then selected as part of the model architecture.

There is also an important practical constraint:

d_model must usually be divisible by n_head so that the hidden representation can be evenly divided among the attention heads.

For example:

4096 / 32 = 128 ✓
4096 / 64 = 64 ✓
4096 / 30 = 136.53 ✗

The number of attention heads is an architectural hyperparameter. In principle, we can use hyperparameter optimization frameworks such as Optuna to search over different candidate values — for example, 16, 32, or 64 heads — by training and evaluating multiple model configurations.

I used a similar approach while training my Small Language Model (SLM) to experiment with different numbers of attention heads and identify a configuration that provided a good balance between model performance and computational cost.

However, for large-scale LLMs, fully training every candidate configuration can be prohibitively expensive. Therefore, the number of attention heads is typically chosen based on architectural design, smaller-scale experiments, and compute/memory trade-offs, rather than exhaustive hyperparameter search.

4. More Heads Are Not Always Better

It might seem that increasing the number of heads should always improve the model because it gives the Transformer more “perspectives.”

But there is a trade-off.

More heads
 ↓
More attention subspaces
 ↓
Potentially greater representational flexibility
 ↓
But also more KV heads in standard MHA
 ↓
Larger KV Cache
 ↓
Higher inference memory

This is exactly where the story connects back to KV Cache from Part 2. If a model uses standard MHA with 32 attention heads, it stores K and V for all 32 heads. This leads to the next important question:

Do all Query heads really need their own Key and Value heads?

Instead of:
32 Query heads
32 Key heads
32 Value heads

what if we keep:
32 Query heads
8 Key-Value heads
or even:
32 Query heads
1 Key-Value head

This is the fundamental idea behind GQA and MQA.

MQA (Multi-Query Attention) and GQA (Grouped-Query Attention)

Now we have a trade-off:

Can we preserve the benefits of multiple Query heads while reducing the number of Key-Value heads that need to be stored?

This question leads directly to the next evolution of attention:

MHA → MQA → GQA

MHA
↓
Every Query head has its own K and V
↓
Large KV Cache
↓
MQA
↓
All Query heads share K and V
↓
Very small KV Cache
↓
GQA
↓
Groups of Query heads share K and V
↓
Balance between quality and memory

1. MQA

In Multi-Query Attention (MQA), the model still produces multiple query heads, but all of them share a single key head and a single value head.

Suppose the model has H Query heads:

Q₁, Q₂, Q₃, ..., Qₕ

In Multi-Query Attention (MQA), every Query head shares the same Key and Value representations:

K₁ = K₂ = K₃ = ... = Kₕ = K
V₁ = V₂ = V₃ = ... = Vₕ = V

So although every Query head has its own representation, all of them use the same Key and Value:

Q₁ ─┐
Q₂ ─┤
Q₃ ─┼────→ K, V
... ┤
Qₕ ─┘


Ignoring batch size and data type, the KV Cache is proportional to:
MHA KV Cache ∝ 2 × H × L × dₕ

In MQA, there is only one Key head and one Value head shared by all Query heads:
MQA KV Cache ∝ 2 × 1 × L × dₕ

H = number of Query heads
L = sequence length
dₕ = dimension of each attention head
2 = one Key + one Value

Therefore, compared with MHA, MQA reduces the KV-cache requirement by approximately a factor of H.

The main trade-off is that all query heads must use the same key-value information. This greatly improves memory efficiency and decoding speed, but it reduces the diversity of attention patterns compared with MHA.

2. GQA

MQA reduces the KV Cache dramatically by making all Query heads share a single Key head and a single Value head.

But this creates a new trade-off:

Can we keep multiple Key-Value heads, while still reducing the KV Cache compared with MHA?

This leads to Grouped-Query Attention (GQA).The basic idea is simple:

Instead of every Query head having its own K/V, or all Query heads sharing one K/V, we divide the Query heads into groups. Each group shares one Key head and one Value head.

Q₁ ─┐
Q₂ ─┤
Q₃ ─┤──→ K₁, V₁
Q₄ ─┘

Q₅ ─┐
Q₆ ─┤
Q₇ ─┤──→ K₂, V₂
Q₈ ─┘

8 Query heads
 ↓
2 Key-Value heads

GQA KV Cache ∝ 2 × G × L × dₕ

H = 32 Query heads
G = 8 KV heads
MHA:
32 KV heads

GQA:
8 KV heads

MQA:
1 KV head

So GQA is a middle ground between MHA and MQA. MHA maximizes independent Key-Value representations, MQA maximizes KV-cache efficiency, and GQA provides a configurable trade-off between the two.

This naturally leads to the next question:

Can we compress the information stored in the KV Cache even further without simply reducing the number of KV heads?

That takes us to MLA — Multi-Head Latent Attention.

MLA (Multi-Head Latent Attention)

It is introduced with DeepSeek-V2. Its main goal is to reduce the memory required by the KV cache during autoregressive inference without reducing the model to a single key-value head.

Instead of caching a separate full-size key and value vector for every attention head, MLA compresses the key-value information for each token into a smaller shared latent vector.

So rather than caching:

K₁, V₁
K₂, V₂
K₃, V₃
...
Kₙ, Vₙ

in their full-dimensional form, MLA stores a compressed latent representation for each token:

z₁
z₂
z₃
...
zₙ

where each latent vector is smaller than the original Key-Value representation.

I explained MLA in two ways: first with equations, then with a worked example, so you can check every calculation yourself. The numbers may feel heavy at first, so I tried to keep each step as simple and detailed as possible. Once you get a hold of it, you will feel satisfied. It is just simple matrix multiplication.

1. A Simple Intuition: The 32 Folders

Imagine you have 32 researchers (the Query heads) who need to repeatedly look up historical information to write the next word of a document.

• MHA is like keeping 32 completely separate, massive filing cabinets — one dedicated to each researcher. Every time a new fact is written, it gets copied 32 times. (Maximum flexibility, but a massive waste of physical space).
• MQA (Multi-Query Attention) is like forcing all 32 researchers to share just one single filing cabinet in the middle of the room. (Incredibly space-efficient, but the researchers lose the ability to organize information for their specific needs, causing a drop in quality).
• GQA is like grouping the researchers into teams of 4, and giving each team one shared filing cabinet. We only keep 8 copies of the information. (A practical compromise between space and speed).
• MLA (DeepSeek’s Multi-Head Latent Attention) takes a completely different approach. Instead of keeping physical copies of the folders, it compresses the core contents into an ultra-dense .zip or .rar file (a low-rank latent vector). When the researchers need the information, they mathematically "unzip" just what they need on the fly. You keep all 32 researchers, but your storage footprint becomes a small fraction.

2. MLA: Mathematical Deep Dive

The core idea of MLA is low-rank joint compression of the Key and Value information. Instead of caching a full K and V for every head, we cache one small latent vector per token and reconstruct what we need from it. DeepSeek-V2 introduced this mechanism specifically to reduce KV-cache memory during autoregressive inference. Unlike ZIP, it is lossy, and the model learns what is worth keeping.

We’ll build this in five steps: the MHA baseline, KV compression, query compression, the absorption trick, and the decoupled RoPE fix

2a. Notation and the MHA baseline

In the previous article, we used:

Q = XWq
K = XWk
V = XWv

Let’s continue with the same notation. Suppose our sequence contains four tokens and assume each token has a 5-dimensional representation.

X =
[
 [1, 0, 1, 0, 1], → my
 [0, 1, 0, 1, 0], → name
 [1, 1, 0, 0, 1], → is
 [0, 1, 1, 1, 0] → naimish
]

Here:
X : 4 × 5

The four rows represent the four token representations.When we focus on one token, say "naimish", we can call its row:

h₄ = [0, 1, 1, 1, 0]
The row for token t is called the hidden state of token t:
hₜ = row t of X shape: (1 × d)

So hₜ is just "one row of X".

Whole sequence: Q = X Wq

One token t: qₜ = hₜ Wq
 kₜ = hₜ Wk
 vₜ = hₜ Wv

We switch to the per-token view because the KV cache grows one token at a time during generation.

Notation and MHA :

** Notation **
d = model (embedding) dimension
nₕ = number of attention heads
dₕ = dimension per head
D = Down (project to a smaller size)
U = Up (project back to a bigger size)
cₜᴷⱽ = compressed Key-Value latent representation
Wᴰᴷⱽ= KV down-projection matrix (Compresses hₜ into the small latent cₜᴷⱽ)
Wᵁᴷ = Key up-projection matrix (Expands the latent back into the content key kₜᶜ)
Wᵁⱽ = Value up-projection matrix(Expands the latent back into the content value vₜᶜ)
W^DQ : Query down-projection matrix hₜ → cₜ^Q
W^UQ : Query up-projection matrix cₜ^Q→ qₜᶜ
W^QR : Query RoPE projection cₜ^Q → qₜᴿ (R = RoPE)
Wᴷᴿ : Key RoPE projection hₜ → kₜᴿ (R = RoPE)
Wᴼ : Output projection concatenated heads → uₜ (O = Output)

Using the zip-file analogy
Wᴰᴷⱽ = "zip" (compress)
cₜᴷⱽ = the zipped file ← only this is cached
Wᵁᴷ = "unzip, give me the Keys"
Wᵁⱽ = "unzip, give me the Values"

** MHA**
Shapes of the weight matrices in standard MHA:
Wq : (d × nₕ·dₕ)
Wk : (d × nₕ·dₕ)
Wv : (d × nₕ·dₕ)

So each token produces full-size vectors:
qₜ = hₜ Wq shape: (1 × nₕ·dₕ)
kₜ = hₜ Wk shape: (1 × nₕ·dₕ)
vₜ = hₜ Wv shape: (1 × nₕ·dₕ)

These are split into nₕ heads:
qₜ = [ qₜ,₁ ; qₜ,₂ ; ... ; qₜ,ₙₕ ]
kₜ = [ kₜ,₁ ; kₜ,₂ ; ... ; kₜ,ₙₕ ]
vₜ = [ vₜ,₁ ; vₜ,₂ ; ... ; vₜ,ₙₕ ]

Each head i attends over all tokens j up to the current token t:
i is the attention-head index
t = position index of the token making the query (the newest token while generating)
j = position index of the token being attended to (j = 0, 1, ..., t)

For head (i), the raw attention score between the current token (t) and token (j) is:
Attention scoreᵢ(t,j) = qₜ,ᵢ · kⱼ,ᵢᵀ


The score is then scaled:
scaled_scoreᵢ(t,j) = scoreᵢ(t,j) / √dₕ
and converted into an attention weight using Softmax over all tokens j:

aᵢ(t,j)=Softmaxⱼ(scoreᵢ(t,j) / √dₕ)


Finally, the head output is:
oₜ,ᵢ=Σⱼ aᵢ(t,j) · vⱼ,ᵢ
 or
oₜ,ᵢ = Σⱼ softmax( qₜ,ᵢ · kⱼ,ᵢᵀ / √dₕ ) · vⱼ,ᵢ (j = 1 to t)

All head outputs are concatenated and passed through the output projection:
uₜ = [oₜ,₁ ; oₜ,₂ ; ... ; oₜ,ₙₕ] Wᴼ

This is the standard MHA computation.

Now the memory problem. To generate token t, we need kⱼ and vⱼ of every previous token, so we must cache them:

MHA cache per token, per layer = 2 · nₕ · dₕ
 ↑
 one K and one V

MLA’s goal is to shrink this number without losing multi-head behaviour.

2.b. Low-rank KV compression

Standard MHA goes straight from hₜ to K and V. MLA adds a bottleneck in the middle.

Step 1: compress. Project hₜ down into a small latent vector:

cₜᴷⱽ = hₜ Wᴰᴷⱽ

cᴷⱽ is compressed by one ordinary matrix multiplication: cₜᴷⱽ = hₜ · Wᴰᴷⱽ. There is no special compression algorithm.

Step 2: expand. Project the latent up to get keys and values:

kₜᶜ = cₜᴷⱽ Wᵁᴷ
vₜᶜ = cₜᴷⱽ Wᵁⱽ

(The small ᶜ on kₜᶜ and vₜᶜ means "content", to separate them from the positional part we add later.)

Here are the shapes:

hₜ : (1 × d)
Wᴰᴷⱽ : (d × d꜀) ← down-projection
cₜᴷⱽ : (1 × d꜀) ← the latent ★ THIS is what we cache
d꜀ is the size of the compressed latent vector cₜᴷⱽ.
Wᵁᴷ : (d꜀ × nₕ·dₕ) ← key up-projection
Wᵁⱽ : (d꜀ × nₕ·dₕ) ← value up-projection
kₜᶜ : (1 × nₕ·dₕ)
vₜᶜ : (1 × nₕ·dₕ)
Key constraint: d꜀ << nₕ · dₕ

Compare the two routes:

MHA: hₜ ──────── Wk ────────→ kₜ (one big matrix-Wk)
MLA: hₜ ── Wᴰᴷⱽ ──→ cₜᴷⱽ ── Wᵁᴷ ──→ kₜᶜ (two thin matrices-Wᴰᴷⱽ and Wᵁᴷ)

The big matrix Wk is replaced by the product of two thin matrices:
Wk ≈ Wᴰᴷⱽ · Wᵁᴷ
(d × nₕ·dₕ) ≈ (d × d꜀) · (d꜀ × nₕ·dₕ)

This is called a low-rank factorization. The product can have rank at most d꜀, so all the key and value information of a token has to pass through a d꜀-sized bottleneck.

Notice that K and V share one latent cₜᴷⱽ. That is why it is called joint compression: one small vector replaces both K and V.

2.c. Query compression

DeepSeek-V2 compresses queries the same way:

cₜ^Q= hₜ W^{DQ} (1 × d′) with d′ << nₕ·dₕ
qₜᶜ = cₜ^Q W^{UQ} (1 × nₕ·dₕ)

This does not reduce the KV cache, because queries are never cached. It only reduces memory used during training, so it is an optimization rather than the core idea.

2.d. The absorption trick (never rebuild K and V)

There is an obvious catch. If we cache cⱼᴷⱽ but must expand it into full K and V for every past token at every step, we save memory but waste compute. MLA avoids this with matrix associativity: (AB)C = A(BC).

Let Wᵢᵁᴷ mean the slice of Wᵁᴷ belonging to head i, and W^UQᵢ the matching slice of W^UQ.

For keys. The attention score of head i between current token t and past token j is:

Attention scoreᵢ(t, j) = qₜ,ᵢᶜ · (kⱼ,ᵢᶜ)ᵀ

Substitute the definitions:
qₜ,ᵢᶜ = cₜ^Q · Wᵢ^UQ
kⱼ,ᵢᶜ = cⱼᴷⱽ · Wᵢᵁᴷ

So the score becomes:
Attention scoreᵢ(t, j) = ( cₜ^Q ·Wᵢ^UQ) · ( cⱼᴷⱽ · Wᵢᵁᴷ )ᵀ

 = cₜ^Q · [ Wᵢ^UQ· (Wᵢᵁᴷ)ᵀ ] · (cⱼᴷⱽ)ᵀ
 └───────────────┘
 only learned weights, no token

The bracketed part depends only on trained weights, not on any token. So we can precompute it once and merge it into a single matrix. At inference, the model multiplies the query by this merged matrix and dots it directly with the cached latent cⱼᴷⱽ. The full key kⱼ is never built.

For values. Let aᵢⱼ be the softmax attention weight of head i on token j (for the current token t):

The Attention score is then scaled:
scaled_scoreᵢ(t,j) = scoreᵢ(t,j) / √dₕ
and converted into an attention weight using Softmax over all tokens j:

attention weight(aᵢⱼ)=Softmaxⱼ(scoreᵢ(t,j) / √dₕ)

from step-2 Project the latent up to get values:
vⱼ,ᵢᶜ = cⱼᴷⱽ · Wᵢᵁⱽ

oₜ,ᵢ = Σⱼ aᵢⱼ · vⱼ,ᵢᶜ
 = Σⱼ aᵢⱼ · ( cⱼᴷⱽ · Wᵢᵁⱽ )
 = ( Σⱼ aᵢⱼ · cⱼᴷⱽ ) · Wᵢᵁⱽ

We take the weighted sum over the small latents first, then apply Wᵢᵁⱽ once. That matrix can then merge with the output projection. Let Wᵢᴼ be the rows of Wᴼ for head i:

uₜ = Σᵢ oₜ,ᵢ · Wᵢᴼ
 = Σᵢ ( Σⱼ aᵢⱼ · cⱼᴷⱽ ) · [ Wᵢᵁⱽ · Wᵢᴼ ]
 └──────────┘
 precomputed, merged

After absorption, attention runs directly on compressed latents, and the up-projection matrices disappear into other weights.

2.e. The decoupled RoPE fix

Problem: The trick above has a catch. RoPE rotates queries and keys by a position-dependent rotation matrix: Rₜ for the query at position t and Rⱼ for the key at position j.

If we applied RoPE to the up-projected key:

kⱼ,ᵢ = cⱼᴷⱽ · Wᵢᵁᴷ · Rⱼ
qₜ,ᵢ = cₜ^Q· Wᵢ^UQ · Rₜ

then the score becomes:

scoreᵢ(t, j) = cₜ^Q· Wᵢ^UQ · [ Rₜ · Rⱼᵀ ] · (Wᵢᵁᴷ)ᵀ · (cⱼᴷⱽ)ᵀ
 └────────┘
 Rₜ₋ⱼ : depends on the distance (t − j)

The rotation Rₜ₋ⱼ now sits between Wᵢ^UQ and (Wᵢᵁᴷ)ᵀ. It changes for every (t, j) pair, so these two matrices can no longer be merged into one fixed matrix. RoPE breaks the absorption trick.

The fix: Decoupled RoPE

DeepSeek splits each query and key into two parts:

Content part → compressed, NO RoPE (keeps the absorption trick)
Position part → small, WITH RoPE (carries position information)

Position part:

qₜᴿ = RoPE( cₜ^Q· W^QR ) one per head, size dₕᴿ each
kₜᴿ = RoPE( hₜ · Wᴷᴿ ) ONE vector shared by all heads, size dₕᴿ

Then concatenate content and position for each head:

qₜ,ᵢ = [ qₜ,ᵢᶜ ; qₜ,ᵢᴿ ] size: dₕ + dₕᴿ
kₜ,ᵢ = [ kₜ,ᵢᶜ ; kₜᴿ ] size: dₕ + dₕᴿ

The dot product of two concatenated vectors is the sum of the dot products of their parts, so the score splits into two clean terms:

qₜ,ᵢ · kⱼ,ᵢᵀ = qₜ,ᵢᶜ · (kⱼ,ᵢᶜ)ᵀ + qₜ,ᵢᴿ · (kⱼᴿ)ᵀ
 └─ content ─────┘ └─ position ────┘
 (absorbable, no RoPE) (RoPE lives here)

The content term keeps the absorption trick, and the position term carries RoPE without getting in the way. The full attention is:

oₜ,ᵢ = Σⱼ softmax( qₜ,ᵢ · kⱼ,ᵢᵀ / √(dₕ + dₕᴿ) ) · vⱼ,ᵢᶜ
uₜ = [ oₜ,₁ ; oₜ,₂ ; ... ; oₜ,ₙₕ ] Wᴼ

Because kₜᴿ is shared across all heads, the extra cache cost is small.

For every token, in every layer, MLA stores only two things:

cₜᴷⱽ : (1 × d꜀) compressed content latent
kₜᴿ : (1 × dₕᴿ) decoupled positional key (shared by all heads)

MLA cache per token, per layer = d꜀ + dₕᴿ

Side-by-side comparison:

MHA : 2 · nₕ · dₕ
GQA : 2 · nᵍ · dₕ (nᵍ = number of KV groups)
MQA : 2 · dₕ
MLA : d꜀ + dₕᴿ

The full MLA pipeline

 hₜ
 ┌───────────────┼────────────────┐
 ↓ ↓ ↓
 hₜ·W^DQ hₜ·Wᴰᴷⱽ hₜ·Wᴷᴿ
 ↓ ↓ ↓
 cₜ^Q cₜᴷⱽ RoPE
 / \ │ ↓
 W^UQ W^QR │ ★ CACHED kₜᴿ ★ CACHED
 ↓ ↓ │
 qₜᶜ RoPE(·) ├──→ Wᵁᴷ → kₜᶜ (absorbed at inference)
 ↓ └──→ Wᵁⱽ → vₜᶜ (absorbed at inference)
 qₜᴿ

 qₜ = [ qₜᶜ ; qₜᴿ ] kₜ = [ kₜᶜ ; kₜᴿ ]
 ↓ 
 softmax( qₜ · kⱼᵀ / √(dₕ + dₕᴿ) ) · vᶜ
 ↓
 concat heads → Wᴼ → uₜ

It seems like a lot of equations and formulas to understand, so let’s understand with a worked example:

3. MLA: Mathematical Example

Let’s go deeper with a mathematical example, so it will be more clearer. Below is one full MLA step with my tokens. Also, I compared it with standard MHA so it will be more understandable. I made up small integer weights so every number can be checked by hand (real weights are learned decimals). I skip query compression to keep it readable. RoPE is added later, in 3.i. I do include the 1/√dₕ scaling and softmax.

3. a. Setup

Tokens are numbered 1 to 4. RoPE positions start at 0, so token 4 is at position 3. Similarly, j runs 0 to t in one place and 1 to t in another

Tokens: 1 = "my" 2 = "name" 3 = "is" 4 = "naimish"

d = 5 input embedding size
nₕ = 2 number of heads
dₕ = 2 size per head (nₕ·dₕ = 4)
d꜀ = 2 latent size (smaller than nₕ·dₕ = 4, this is the compression)


Hidden states (each row is one hₜ, so X is 4 × 5):
h₁ "my" = [1, 0, 1, 0, 1]
h₂ "name" = [0, 1, 0, 1, 0]
h₃ "is" = [1, 1, 0, 0, 1]
h₄ "naimish" = [0, 1, 1, 1, 0]

We are at the moment of generating with token 4 (“naimish”). Tokens 1 to 3 are already in the cache.

In standard Multi-Head Attention (MHA), each token’s hidden state is projected separately into Key and Value vectors:

d = 5 (input embedding size)
nₕ = 2 (number of attention heads)
dₕ = 2 (size per head)

Total K dimensions = nₕ × dₕ = 4
Total V dimensions = nₕ × dₕ = 4
For simplicity, let's use these illustrative weight matrices:
Wk (5 × 4) = [ [1, 0, 0, 1],
 [0, 1, 1, 0],
 [1, 0, 1, 0],
 [0, 1, 0, 1],
 [1, 1, 0, 0] ]

Wv (5 × 4) = [ [1, 0, 1, 0],
 [0, 1, 0, 1],
 [1, 1, 0, 0],
 [0, 0, 1, 1],
 [1, 0, 0, 1] ]

Each token's Key and Value vectors are calculated as:
K = XWk, V = XWv


Here, X is our (4×5) matrix containing all four token representations.

3.b. Compress each token and cache it

Wᴰᴷⱽ (5 × 2) = [ [1, 0],
 [0, 1],
 [1, 0],
 [0, 1],
 [1, 1] ]

cⱼᴷⱽ = hⱼ · Wᴰᴷⱽ. Worked out for "my":

c₁ = [1,0,1,0,1] · Wᴰᴷⱽ
 = [ 1·1 + 0·0 + 1·1 + 0·0 + 1·1 , 1·0 + 0·1 + 1·0 + 0·1 + 1·1 ]
 = [ 3 , 1 ]

The others are done the same way:
c₁ "my" = [3, 1]
c₂ "name" = [0, 2]
c₃ "is" = [2, 2]
c₄ "naimish" = [1, 2]

The KV cache is now only these four small vectors (8 numbers). Standard MHA would store a full K (4 numbers) and V (4 numbers) per token, which is 32 numbers. Here is the standard MHA calculation:

For token 1, "my":
Calculate the Key vector:
h₁ = [1, 0, 1, 0, 1]

k₁ = h₁ Wk
 = [1, 0, 1, 0, 1] Wᴷ
 = row₁ + row₃ + row₅ of Wᴷ
 = [1, 0, 0, 1] + [1, 0, 1, 0] + [1, 1, 0, 0]
 = [3, 1, 1, 1]

Applying the same multiplication to the remaining tokens gives:

k₁ ("my") = [3, 1, 1, 1]
k₂ ("name") = [0, 2, 1, 1]
k₃ ("is") = [2, 2, 1, 1]
k₄ ("naimish") = [1, 2, 2, 1]

Each 4-dimensional Key vector is split into two heads:

Token 1:
Head 1 → [3, 1]
Head 2 → [1, 1]


Calculate the Value vector:

v₁ = h₁ Wv
 = [1, 0, 1, 0, 1] Wv
 = row₁ + row₃ + row₅ of Wv
 = [1, 0, 1, 0] + [1, 1, 0, 0] + [1, 0, 0, 1]
 = [3, 1, 1, 1]

Applying the same multiplication to the remaining tokens gives:

v₁ ("my") = [3, 1, 1, 1]
v₂ ("name") = [0, 1, 1, 2]
v₃ ("is") = [2, 1, 1, 2]
v₄ ("naimish") = [1, 2, 1, 2]

Again, each Value vector is split into two heads.

Token 1:
Head 1 → [3, 1]
Head 2 → [1, 1]


At each token, MHA stores four Key numbers and four Value numbers.

Token 1 → 4 K numbers + 4 V numbers = 8
Token 2 → 4 K numbers + 4 V numbers = 8
Token 3 → 4 K numbers + 4 V numbers = 8
Token 4 → 4 K numbers + 4 V numbers = 8

Total = 4 × 8 = 32 numbers


MHA KV Cache

Token 1 → K₁ [3,1,1,1] V₁ [3,1,1,1]
Token 2 → K₂ [0,2,1,1] V₂ [0,1,1,2]
Token 3 → K₃ [2,2,1,1] V₃ [2,1,1,2]
Token 4 → K₄ [1,2,2,1] V₄ [1,2,1,2]

Total cached numbers = 32

3.c. The up-projection matrices (fixed weights, NOT cached)

Wᵁᴷ (2 × 4) = [ [1, 2, 1, 1],
 [0, 1, 1, 0] ]
 └head 1┘ └head 2┘

Wᵁⱽ (2 × 4) = [ [1, 1, 0, 1],
 [0, 1, 1, 0] ]

Each head owns 2 columns. The head slices of Wᵁᴷ:

Wᵁᴷ head 1 = [ [1, 2], Wᵁᴷ head 2 = [ [1, 1],
 [0, 1] ] [1, 0] ]

Similarly for standard MHA:

 Head 1 Head 2
Token K V K V

"my" [3,1] [3,1] [1,1] [1,1]
"name" [0,2] [0,1] [1,1] [1,2]
"is" [2,2] [2,1] [1,1] [1,2]
"naimish" [1,2] [1,2] [2,1] [1,2]

3.d. The current query

Wq (5 × 4) = [ [1,0,0,1],
 [1,1,0,0],
 [0,1,1,0],
 [0,0,1,1],
 [1,0,1,0] ]

q₄ = h₄ · Wq = [0,1,1,1,0] · Wq
 = row2 + row3 + row4
 = [1,1,0,0] + [0,1,1,0] + [0,0,1,1]
 = [1, 2, 2, 1]
head 1 query: q₄,₁ = [1, 2]
head 2 query: q₄,₂ = [2, 1]

3.e.1. The naive MLA approach — reconstruct every Key

The straightforward approach reconstructs a full Key for every cached latent:

kⱼ = cⱼ · Wᵁᴷ
For token 1, "my":
k₁ = [3,1] · Wᵁᴷ
= [3·1 + 1·0,
 3·2 + 1·1,
 3·1 + 1·1,
 3·1 + 1·0]
 = [3,7,4,3]
Repeat this for the other tokens:
Token Head 1 Key Head 2 Key
-------------------------------------
"my" [3,7] [4,3]
"name" [0,2] [2,0]
"is" [2,6] [4,2]
"naimish" [1,4] [3,1]

MHA and MLA here are two different illustrative models. MHA usesWᴷ, while MLA’s effective key matrix is Wᴰᴷⱽ · Wᵁᴷ, which is a different matrix. So MHA’s k₄ = [1,2,2,1] and MLA’s k₄ = [1,4,3,1] don’t match.

Now calculate the attention scores using the dot product between each query and its corresponding Key q · kⱼ:

Head 1 query (q = [1,2]):
"my": 1·3 + 2·7 = 17
"name": 1·0 + 2·2 = 4
"is": 1·2 + 2·6 = 14
"naimish": 1·1 + 2·4 = 9
Scores = [17, 4, 14, 9]

Head 2 query (q = [2,1]):
"my": 2·4 + 1·3 = 11
"name": 2·2 + 1·0 = 4
"is": 2·4 + 1·2 = 10
"naimish": 2·3 + 1·1 = 7
Scores = [11, 4, 10, 7]

Cost: 4 expansions of c · Wᵁᴷ (8 multiplications each = 32), and when token 5 arrives, all 5 latents get expanded again.

3.e.2. The absorbed approach — calculate scores without building Keys

The key insight is that we can rearrange the matrix multiplication.

The score is q · kⱼᵀ with kⱼ = cⱼ · Wᵁᴷ. Using (AB)ᵀ = BᵀAᵀ:

We know:
kⱼ = cⱼ · Wᵁᴷ
The attention score is:
scoreⱼ = q · kⱼᵀ
Substitute the expression for the Key:
scoreⱼ = q · (cⱼ · Wᵁᴷ)ᵀ
 = q · (Wᵁᴷ)ᵀ · cⱼᵀ
 = [q · (Wᵁᴷ)ᵀ] · cⱼᵀ
 └────────────┘
 depends ONLY on the query, not on token j

Notice what changed: the query and projection weights can be combined before interacting with the cached latent.

So compute this once for the current query:

qᴬ = q · (Wᵁᴷ)ᵀ (per head, using that head's slice)

We calculate this once per head for the current token, then compare it directly with every cached latent.

Head 1 - query = [1,2]
(Wᵁᴷ head 1)ᵀ = [[1,0],
 [2,1]]
qᴬ₁ = [1,2] · [[1,0],
 [2,1]]
 = [1·1 + 2·2, 1·0 + 2·1]
 = [5,2]



Head 2 - query = [2,1]
(Wᵁᴷ head 2)ᵀ = [[1,1],
 [1,0]]
qᴬ₂ = [2,1] · [[1,1],
 [1,0]]
 = [2·1 + 1·1, 2·1 + 1·0]
 = [3,2]

Now calculate scores directly against the cached latents. No full Keys need to be reconstructed.

Head 1 - transformed query = [5,2]:

"my": [5,2] · [3,1] = 15 + 2 = 17
"name": [5,2] · [0,2] = 0 + 4 = 4
"is": [5,2] · [2,2] = 10 + 4 = 14
"naimish": [5,2] · [1,2] = 5 + 4 = 9
Scores = [17, 4, 14, 9]

Head 2 - transformed query = [3,2]
"my": [3,2] · [3,1] = 9 + 2 = 11
"name": [3,2] · [0,2] = 0 + 4 = 4
"is": [3,2] · [2,2] = 6 + 4 = 10
"naimish": [3,2] · [1,2] = 3 + 4 = 7
Scores = [11, 4, 10, 7]

The scores are identical to the naive route (17, 4, 14, 9 and 11, 4, 10, 7), but k₁ ... k₄ were never created. The difference is that the absorbed approach never constructs the full historical Keys.

The difference is in what has to be redone at each step:

NAIVE every step: expand ALL cached latents (grows with context length)
ABSORBED every step: transform ONLY the new query (always 1 token)

When token 5 arrives:

NAIVE: re-expand c₁ ... c₅ → 5 × 8 = 40 multiplications
ABSORBED: compute qᴬ for token 5 only → 8 multiplications,
 then dot with the cached c₁ ... c₅ (cheap, only d꜀ = 2 numbers each)

The old latents are never touched by Wᵁᴷ. The multiplication by Wᵁᴷ moved from the many (all past tokens) to the one (the current query). For our example, counting only the operations shown above:

NAIVE MLA:
 Expand 4 latents: 4 × 8 = 32 multiplications
 Score calculations: 4 × 2 heads × 2 dimensions = 16
 Total = 48 multiplications

ABSORBED MLA:
 Transform 2 queries: 2 heads × 2 × 2 = 8 multiplications
 Score calculations: 4 × 2 heads × 2 dimensions = 16
 Total = 24 multiplications

That is a 2× reduction in this score calculation. The potential advantage grows with head dimension: reconstructing full Keys for all past tokens repeats the up-projection work across the context, whereas absorption moves that projection to the current query. The per-token expansion costs d꜀ · nₕ · dₕ and the absorbed dot costs nₕ · d꜀, so the ratio is about dₕ (128× for DeepSeek’s dₕ = 128), and it applies to every past token.

MLA reduces what must be stored, and absorption reduces the need to reconstruct full historical Keys when computing attention scores.

Standard MHA caches full keys and values. Each token produces a 4-dimensional Key and a 4-dimensional Value:

Kⱼ = hⱼ · Wᴷ # 5 × 4 projection
Vⱼ = hⱼ · Wⱽ # 5 × 4 projection
For each token, the projection requires:
Key projection: 5 × 4 = 20 multiplications
Value projection: 5 × 4 = 20 multiplications
Total per token = 40 multiplications
For four tokens, creating the cache requires:
4 × 40 = 160 multiplications

These projections happen as the tokens enter the model. Once cached, the full Keys do not need to be reconstructed at each subsequent decoding step.

Calculating attention scores for token 4:

Each query head has dimension 2. It calculates a dot product against the corresponding Key for each of the four tokens.

Head 1: 4 keys × 2 multiplications = 8
Head 2: 4 keys × 2 multiplications = 8
Total score calculation = 16 multiplications

So, at token 4, standard MHA requires 16 multiplications for the attention-score dot products, excluding query projection, softmax, Value aggregation, and other operations.

Compare the three approaches to the calculation operation

3.f. Softmax (with 1/√dₕ scaling)

Head 1: scores [17, 4, 14, 9] / √2 → [12.02, 2.83, 9.90, 6.36]
 softmax → a₁ = [0.890, 0.000, 0.107, 0.003]

Head 2: scores [11, 4, 10, 7] / √2 → [7.78, 2.83, 7.07, 4.95]
 softmax → a₂ = [0.641, 0.005, 0.316, 0.038]

This is the weight each head gives to my / name / is / naimish.

3.g. Values, using the same trick as Keys

Naive: expand values vⱼ = cⱼ · Wᵁⱽ, then take the weighted sum.

Head 1 values: my [3,4] name [0,2] is [2,4] naimish [1,3]
o₁ = 0.890·[3,4] + 0.000·[0,2] + 0.107·[2,4] + 0.003·[1,3]
 = [2.887, 3.997]

Absorbed: weighted sum of the latents first, then apply Wᵁⱽ once.

Head 1:
 c̄₁ = 0.890·[3,1] + 0.000·[0,2] + 0.107·[2,2] + 0.003·[1,2]
 = [2.887, 1.110]

 Wᵁⱽ head 1 = [ [1,1],
 [0,1] ]

 o₁ = c̄₁ · Wᵁⱽ head 1
 = [ 2.887·1 + 1.110·0 , 2.887·1 + 1.110·1 ]
 = [ 2.887, 3.997 ] ✓ same as naive


Head 2:
 c̄₂ = 0.641·[3,1] + 0.005·[0,2] + 0.316·[2,2] + 0.038·[1,2]
 = [2.594, 1.359]

 Wᵁⱽ head 2 = [ [0,1],
 [1,0] ]

 o₂ = c̄₂ · Wᵁⱽ head 2
 = [ 2.594·0 + 1.359·1 , 2.594·1 + 1.359·0 ]
 = [ 1.359, 2.594 ] ✓ same as naive

The weighted sum runs over 2-number latents, and Wᵁⱽ is applied once per head instead of once per past token.

3.h. Output projection

Wᴼ (4 × 5) = [ [1,0,0,0,1],
 [0,1,0,0,0],
 [0,0,1,0,0],
 [0,0,0,1,1] ]

[o₁ ; o₂] = [2.887, 3.997, 1.359, 2.594]

u₄ = [o₁ ; o₂] · Wᴼ
 = [2.887, 3.997, 1.359, 2.594, 2.887 + 2.594]
 = [2.887, 3.997, 1.359, 2.594, 5.481]

This is a 5-number vector, the same size as h₄, ready to go into the next part of the Transformer layer. In the real model, Wᵁⱽ is merged with Wᴼ ahead of time, so even Step 3.g’s “apply Wᵁⱽ once” disappears into one precomputed matrix.

3.i. The decoupled RoPE: Why RoPE breaks the absorption trick

Recap: the trick without RoPE

The efficiency in Step 3.e.2 came from regrouping the score so that the query part does not depend on the cached token j:

score = q · (cⱼ · Wᵁᴷ)ᵀ
 = [ q · (Wᵁᴷ)ᵀ ] · cⱼᵀ
 └────────────┘
 same matrix for every token j

qᴬ = [1,2] · [[1,0],[2,1]] = [5, 2] ← computed ONCE, reused for my, name, is, naimish

What RoPE does

RoPE gives each token a position and rotates the query and key by an angle that depends on that position. Positions start at 0:

"my" = position 0 "name" = 1 "is" = 2 "naimish" = 3

To keep the arithmetic clean, I use a toy rotation of 90° per position. A 2-number vector [x, y] rotated by m steps becomes:

position 0 (0°) : [ x, y]
position 1 (90°) : [−y, x]
position 2 (180°) : [−x, −y]
position 3 (270°) : [ y, −x]

As matrices, multiplying a row vector [x, y] by Rₘ rotates it by m steps:

R₀ = [[ 1, 0], R₁ = [[ 0, 1], R₂ = [[−1, 0], R₃ = [[ 0,−1],
 [ 0, 1]] [−1, 0]] [ 0,−1]] [ 1, 0]]

Real RoPE uses different angles for different pairs of dimensions, but the idea is the same: a rotation whose angle depends on the token’s position.

What happens if we put RoPE on the content key

Suppose we rotate the query by its position t and the content key by its position j:

score = (q · Rₜ) · (cⱼ · Wᵁᴷ · Rⱼ)ᵀ
 = q · Rₜ · Rⱼᵀ · (Wᵁᴷ)ᵀ · cⱼᵀ
 └────┘
 Rₜ₋ⱼ (rotate by t steps, back by j steps = rotate by the distance t − j)
 
 = [ q · Rₜ₋ⱼ · (Wᵁᴷ)ᵀ ] · cⱼᵀ
 └────────────────────┘
 the "fixed matrix" is now Rₜ₋ⱼ · (Wᵁᴷ)ᵀ, which changes with j

So the transformed query for cached token j becomes:
qᴬ for token j = q · Rₜ₋ⱼ · (Wᵁᴷ)ᵀ

The numbers

Head 1: q = [1, 2], (Wᵁᴷ head 1)ᵀ = [[1,0],[2,1]], current token “naimish” at position 3.

past token position j distance t−j q rotated by t−j qᴬ for THIS token
my 0 3 [ 2, −1] [ 0, −1]
name 1 2 [−1, −2] [−5, −2]
is 2 1 [−2, 1] [ 0, 1]
naimish 3 0 [ 1, 2] [ 5, 2]

One row worked out in full ("my"):
q rotated by 3 steps: [1, 2] → [2, −1]
qᴬ = [2, −1] · [[1,0],[2,1]] = [ 2·1 + (−1)·2 , 2·0 + (−1)·1 ] = [0, −1]

Each qᴬ is dotted with that token's cached latent cⱼᴷⱽ:
my: [ 0,−1]·[3,1] = −1
name: [−5,−2]·[0,2] = −4
is: [ 0, 1]·[2,2] = 2
naimish: [ 5, 2]·[1,2] = 9

To check that these are the true RoPE scores, take “is” and do it the long way:

k₃ (head 1) = [2,2] · Wᵁᴷ head 1 = [2, 6]
rotate by its position 2 → [−2, −6]
rotate q by position 3 → [ 2, −1]
score = 2·(−2) + (−1)·(−6) = −4 + 6 = 2 ✓ matches

So [−1, −4, 2, 9] is a correct set of position-aware scores. The scores are not the problem. The cost is.

Why this breaks absorption

Compare the transformed queries with and without RoPE:

Without RoPE: my [5,2] name [5,2] is [5,2] naimish [5,2] ← one vector
With RoPE: my [0,−1] name [−5,−2] is [0,1] naimish [5,2] ← four different vectors

The rotation depends on the distance t − j, which is different for every cached token. So the transformed query must also be different for every cached token. That is the same kind of work we were trying to avoid. In the naive approach we did one transformation per cached token (expand cⱼ into kⱼ). Here we also do one transformation per cached token (build a separate qᴬ for each j). The cost grows with context length again, so absorption saved nothing.

That is the same kind of work we were trying to avoid. In the naive approach we did one transformation per cached token (expand cⱼ into kⱼ). Here we also do one transformation per cached token (build a separate qᴬ for each j). The cost grows with context length again, so absorption saved nothing.

Counting only the transformation work (the dot-product work is the same in every route):

Naive (expand each latent into a key): 4 tokens × 8 = 32
Absorbed, no RoPE (one qᴬ per head): 2 heads × (2×2) = 8
Absorbed, RoPE on content (qᴬ per token): 4 tokens × 2 heads × (2×2) = 32

It also gets worse over time: when token 5 arrives, every distance grows by 1, so all earlier qᴬ vectors are invalid and must be recomputed.

The same thing blocks precomputing a merged weight matrix. Absorption relies on Wᵁᴷ and the query-side weights being fused into one fixed matrix ahead of time. With the rotation in between, you would need a different fused matrix for every possible distance. The next step, 3.j, shows how DeepSeek avoids this.

3.j. The decoupled RoPE fix: split each vector into content and position

Content part (dₕ = 2 numbers): no RoPE, uses c, Wᵁᴷ, absorption as in Step 3.e.2
Position part (dₕᴿ = 2 numbers): small, WITH RoPE, carries position

(dₕᴿ = 2 here because a rotation needs a pair of numbers. DeepSeek-V2 uses dₕᴿ = 64.)

The content part is unchanged. We still have:

qᴬ₁ = [5, 2] qᴬ₂ = [3, 2]
content scores head 1 = [17, 4, 14, 9]
content scores head 2 = [11, 4, 10, 7]

Now we build the position part.

3.k. Positional key kᴿ (one vector, shared by both heads, CACHED)

It is computed straight from hⱼ, not from the latent:

Wᴷᴿ (5 × 2) = [ [1, 0],
 [0, 1],
 [1, 1],
 [0, 0],
 [1, 0] ]

Before RoPE (kⱼᴿ′ = hⱼ · Wᴷᴿ). Worked out for “my” = [1,0,1,0,1] (rows 1, 3, 5 of Wᴷᴿ):

kᴿ′ for "my" = [1,0] + [1,1] + [1,0] = [3, 1]

All four tokens:

my [3, 1]
name [0, 1]
is [2, 1]
naimish [1, 2]

Now rotate each by its own position:

my (pos 0): [3, 1] → [ 3, 1]
name (pos 1): [0, 1] → [−1, 0] ([−y, x])
is (pos 2): [2, 1] → [−2, −1] ([−x, −y])
naimish (pos 3): [1, 2] → [ 2, −1] ([ y, −x])

These rotated vectors are what goes into the cache. A token’s rotation depends only on its own position, so it never needs to be redone later.

3.l. Positional query qᴿ (one per head, current token only)

The paper computes this from the query latent cQ. Since we skipped query compression, we compute it from h₄ directly:

Wᵠᴿ (5 × 4) = [ [1,0,0,1],
 [1,0,1,0],
 [0,1,0,2],
 [1,2,0,0],
 [0,0,1,1] ]

qᴿ′ = h₄ · Wᵠᴿ = [0,1,1,1,0] · Wᵠᴿ = row2 + row3 + row4
 = [1,0,1,0] + [0,1,0,2] + [1,2,0,0]
 = [2, 3, 1, 2]
head 1: [2, 3] head 2: [1, 2]

(Wᵠᴿ stands for the W^QR, since there is no superscript Q in Unicode.)

“naimish” is at position 3, which is 270° in our toy rotation, so [x, y] → [y, −x] :

qᴿ head 1 = [ 3, −2]
qᴿ head 2 = [ 2, −1]

3.m. Positional scores qᴿ · kⱼᴿ

Head 1 (qᴿ = [3, −2]):
 my [3,−2]·[ 3, 1] = 9 − 2 = 7
 name [3,−2]·[−1, 0] = −3 + 0 = −3
 is [3,−2]·[−2,−1] = −6 + 2 = −4
 naimish [3,−2]·[ 2,−1] = 6 + 2 = 8 → [7, −3, −4, 8]

Head 2 (qᴿ = [2, −1]):
 my [2,−1]·[ 3, 1] = 6 − 1 = 5
 name [2,−1]·[−1, 0] = −2 + 0 = −2
 is [2,−1]·[−2,−1] = −4 + 1 = −3
 naimish [2,−1]·[ 2,−1] = 4 + 1 = 5 → [5, −2, −3, 5]

3.n. Why rotating makes this relative-position aware

Check head 1, token “is” (position 2) against the current token (position 3). Take the vectors before rotation:

qᴿ′ = [2, 3] positional query of "naimish"
kᴿ′ = [2, 1] positional key of "is"
Each is rotated by its own position:
q: position 3 (270°): [x, y] → [ y, −x] [2, 3] → [ 3, −2]
k: position 2 (180°): [x, y] → [−x, −y] [2, 1] → [−2, −1]

score = [3, −2] · [−2, −1] = 3·(−2) + (−2)·(−1) = −6 + 2 = −4

Instead of rotating both, rotate only qᴿ′ by the distance t − j = 1 step (90°: [x, y] → [−y, x]) and dot with the unrotated key:

rotate [2,3] by 1 step → [−3, 2]
[−3, 2] · [2, 1] = −6 + 2 = −4

That equals the −4 from Step 3.m. The score depends only on the distance between the tokens, which is what RoPE is for. A token pair at positions (2, 3) scores the same as that same pair at positions (100, 101). The rotation is applied to small, separate vectors, so it never touches the compressed content path.

3.o. Combine content and position

The concatenated dot product equals the sum of the two parts:

[qᶜ ; qᴿ] · [kᶜ ; kᴿ]ᵀ = qᶜ · (kᶜ)ᵀ + qᴿ · (kᴿ)ᵀ
 └ absorbed ┘ └ RoPE ┘

Head 1: content [17, 4, 14, 9] + position [ 7, −3, −4, 8] = [24, 1, 10, 17]
Head 2: content [11, 4, 10, 7] + position [ 5, −2, −3, 5] = [16, 2, 7, 12]

The scale factor is now 1/√(dₕ + dₕᴿ) = 1/√4 = 1/2:
Head 1: [24, 1, 10, 17] / 2 = [12.0, 0.5, 5.0, 8.5]
Head 2: [16, 2, 7, 12] / 2 = [ 8.0, 1.0, 3.5, 6.0]

Softmax:
Head 1: a₁ = [0.970, 0.000, 0.001, 0.029]
Head 2: a₂ = [0.872, 0.001, 0.010, 0.118]

Compared with Step 3.f, “naimish” gets more attention (0.029 and 0.118). Two things changed: the position scores added +8 and +5 for “naimish”, and the divisor became 2 instead of √2, which flattens the softmax.

Option 1: no RoPE on content → one qᴬ = [5,2] works, scores [17, 4, 14, 9]
 efficient, but NO position information

Option 2: RoPE on content key → scores [−1, −4, 2, 9]
 position-aware, but a different qᴬ per token,
 so absorption is lost

Option 3: decoupled RoPE → content [17, 4, 14, 9] (absorbed, one qᴬ)
 + position [7, −3, −4, 8] (small separate vectors)
 = [24, 1, 10, 17]
 efficient AND position-aware

3.p. Values and output (same trick as Step 3.g)

Position entered the model only through the scores (qᴿ · kᴿ). Values are built from the content latent alone (vⱼᶜ = cⱼᴷⱽ · Wᵁⱽ), so there is no RoPE on them. Position still influences the output indirectly, because it changed the weights a. The value step itself is identical to 3.g.

Head 1:
Absorbed values: blend the cached latents with those weights first
 (c̄ = Σ aⱼ · cⱼ), then apply Wᵁⱽ once. 
 c̄₁ = 0.970·[3,1] + 0.000·[0,2] + 0.001·[2,2] + 0.029·[1,2] = [2.941, 1.030]
 c̄₁ (c-bar) is “the average latent this head listens to”.
apply Wᵁⱽ once. Head 1 owns columns 1 and 2 of Wᵁⱽ = [[1,1,0,1],[0,1,1,0]], 
so its slice is [[1,1],[0,1]]:
 o₁ = c̄₁ · [[1,1],[0,1]] = [2.941, 3.971]

Similarily Head 2:
 c̄₂ = 0.872·[3,1] + 0.001·[0,2] + 0.010·[2,2] + 0.118·[1,2] = [2.752, 1.128]
 o₂ = c̄₂ · [[0,1],[1,0]] = [1.128, 2.752]

Output projection with the same Wᴼ as Step 3.h:
[o₁ ; o₂] = [2.941, 3.971, 1.128, 2.752]
u₄ = [o₁ ; o₂] · Wᴼ
 = [2.941, 3.971, 1.128, 2.752, 2.941 + 2.752]
 = [2.941, 3.971, 1.128, 2.752, 5.693]

3.q: What the cache holds now

For each token, per layer:

What is stored per token:
When a token arrives, two small vectors are computed and kept:
cⱼᴷⱽ : the compressed content latent (2 numbers)
kⱼᴿ : the positional key, ALREADY ROTATED (2 numbers)
token cⱼᴷⱽ kⱼᴿ
my [3, 1] [ 3, 1]
name [0, 2] [−1, 0]
is [2, 2] [−2, −1]
naimish [1, 2] [ 2, −1]
kⱼᴿ is stored after rotation because a token’s rotation depends only on its own position, which never changes. So it never needs to be redone.

What is not stored:
Wᵁᴷ, Wᵁⱽ, Wᴼ, ... fixed weights, the same for every token (stored once in the model)
qₜ, qᴬ, qᴿ queries, used once and thrown away
full K and V never built (that was the point of absorption)

The counting:
MLA cache per token = d꜀ + dₕᴿ = 2 + 2 = 4 numbers
MHA cache per token = 2 · nₕ · dₕ = 2 · 2 · 2 = 8 numbers

For 4 tokens, that is 16 numbers versus 32. The saving is only 2× because the toy sizes are tiny. With DeepSeek-V2’s sizes it is 576 versus 32,768 (about 57×).

Summary of what each part does

CONTENT PATH hⱼ → cⱼᴷⱽ (cached) → absorbed via qᴬ · cⱼᵀ no RoPE, one qᴬ for all tokens
POSITION PATH hⱼ → kⱼᴿ, rotated at its own position (cached) RoPE, tiny, shared by all heads
score = content score + position score

Why RoPE needs its own path

RoPE rotates keys by an angle that depends on position. If we rotate the content key, the rotation lands between W^UQ and (Wᵁᴷ)ᵀ and changes for every token pair, so the absorbed query would differ for every cached token. Decoupled RoPE gives position its own small vectors (qᴿ, kᴿ). The content score stays rotation-free and absorbable. The position score carries RoPE. The final score is their sum.

What Actually Goes Into the KV Cache?

MHA
For every token, MHA needs to cache:
K₁, V₁
K₂, V₂
...
Kₙ, Vₙ
The cache per token is proportional to:2 × nₕ × dₕ × L
where:
nₕ = number of attention heads
dₕ = dimension of each head
L = number of Transformer layers

GQA
GQA stores fewer K/V heads:
2 × nᵍ × dₕ × L

MQA
MQA stores just one K/V head:
2 × dₕ × L

MLA
MLA stores the compressed latent representation plus its decoupled RoPE Key:
(d꜀ + dₕᴿ) × L

where:
d꜀ = KV compression dimension
dₕᴿ = dimension of the decoupled RoPE Key

DeepSeek-V2 specifically uses:
d꜀ = 4dₕ
dₕᴿ = dₕ / 2
which gives:
MLA cache
= (4dₕ + dₕ/2)L
= 4.5dₕL

For comparison, MHA requires:
2nₕdₕL

So when the number of attention heads is large, MLA can reduce the cache dramatically while retaining multi-head attention behavior. DeepSeek-V2 reports that its MLA cache is equivalent to GQA with roughly 2.25 groups, while achieving stronger benchmark results than its MHA comparison.

MHA vs MQA vs GQA vs MLA

Now we can compare the four approaches from the perspective of the problem we started with:

Ultimately, MLA offers the best of both worlds: MLA’s cache (4.5·dₕ per token per layer in DeepSeek-V2) is far smaller than MHA’s and smaller than common GQA settings (for example 8 KV heads = 16·dₕ). It is comparable to GQA with about 2.25 groups, but still larger than MQA’s 2·dₕ. In DeepSeek’s reported comparisons, its quality matches or beats MHA.

What About Speed?

There is an important distinction here:

Lower KV-cache memory does not automatically mean a fixed percentage increase in tokens/second.

During autoregressive decoding, smaller KV caches can reduce memory traffic and make inference more efficient, especially when decoding is memory-bandwidth limited. But actual speed depends on GPU architecture, batch size, context length, kernels, parallelism, and implementation.

Conceptually:

MHA
↓
Large KV Cache
↓
More memory traffic
↓
Potentially slower decoding

MQA / GQA
↓
Smaller KV Cache
↓
Less memory traffic
↓
Potentially faster decoding

MLA
↓
Much smaller cached representation
↓
Lower KV-memory bandwidth requirement
↓
Potentially higher decoding efficiency

A hardware-focused study of MLA similarly finds that its smaller KV cache reduces memory-bandwidth demands and can shift parts of the workload toward being more compute-bound; the actual benefit depends on the hardware and execution strategy.

DeepSeek-V2 also reported a 5.76× higher maximum generation throughput and a 93.3% reduction in KV-cache size compared with its DeepSeek 67B baseline. That comparison is between complete model systems, however, so it should not be interpreted as “MLA itself is 5.76× faster than MHA.”

The Evolution at a Glance

The entire journey can now be summarized as:

 MHA
 ↓
 Every Query has its own K/V
 ↓
 Large KV Cache
 ↓
 ┌─────────┴─────────┐
 ↓ ↓
 MQA GQA
 ↓ ↓
 One shared K/V Groups share K/V
 ↓ ↓
Very small cache Smaller cache
 └─────────┬─────────┘
 ↓
 MLA
 ↓
 Compress K/V information
 ↓
 Cache a small latent vector
 ↓
 Very small KV representation

The key architectural progression is therefore:

MHA reduces neither the number nor the size of KV representations. MQA reduces the number to one. GQA reduces the number to several groups. MLA takes a different route by compressing the information itself into a low-dimensional latent representation.

What we covered in Part 3

Problem: Every new token must look back at all previous tokens.
 Their Keys and Values are cached, and that cache grows with
 context length, layers and heads.
MHA → every Query head has its own K/V largest cache
MQA → all Query heads share one K/V smallest cache, quality can drop
GQA → groups of Query heads share a K/V middle ground
MLA → cache one small latent per token small cache, all heads kept
 (+ a small decoupled RoPE key)

Key takeaways

• MHA, MQA and GQA reduce the cache by sharing. They store fewer K/V heads, and the cost is less diversity between heads.
• MLA reduces the cache by compressing. One latent vector cₜᴷⱽ per token replaces the full K and V, and every head can still rebuild its own keys and values from it.
• The absorption trick avoids rebuilding K and V at every step. Because matrix multiplication is associative, Wᵁᴷ moves to the query side and the current query is transformed once.
• RoPE breaks absorption if it is applied to the content key. The rotation depends on the distance t − j, so each cached token would need its own transformed query. Decoupled RoPE keeps content (compressed, no RoPE) and position (small, with RoPE) on separate paths.
• The toy example showed every step as plain matrix multiplication. Nothing in MLA is more exotic than that.

References

[1] Vaswani et al., 2017. Attention Is All You Need. arXiv:1706.03762
[2] Su et al., 2021. RoFormer: Enhanced Transformer with Rotary
Position Embedding. arXiv:2104.09864
[3] Shazeer, 2019. Fast Transformer Decoding: One Write-Head is All
You Need. (MQA) arXiv:1911.02150
[4] Ainslie et al., 2023. GQA: Training Generalized Multi-Query
Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245
[5] DeepSeek-AI, 2024. DeepSeek-V2: A Strong, Economical, and Efficient
Mixture-of-Experts Language Model. (MLA) arXiv:2405.04434

Next: Part 4 — Model Capacity and Parameter Efficiency (Mixture of Experts)

Where Part 3 leaves us

Part 3 was about the attention side of the Transformer: how to make the KV cache smaller. But a Transformer block has a second major component, the feed-forward network (FFN), and in the classic design the FFN holds roughly two-thirds of a block’s parameters.

Transformer block
├── Attention ← Part 3: MHA → MQA → GQA → MLA (cache and memory)
└── FFN ← Part 4: where most of the parameters live

The question for Part 4

Bigger models tend to know and do more, which usually means more parameters. In a standard (dense) Transformer, however, every token passes through every parameter:

Dense model: more parameters → more capacity
 → more compute for EVERY token


So a natural question follows:
Can we build a model with a very large number of parameters,
but let each token use only a small part of it?

That is the idea behind Mixture of Experts (MoE).

The idea in one picture

Dense FFN: MoE FFN:
token ──→ one big FFN token ──→ Router ──→ picks top-k experts
 ┌─ Expert 1
 ├─ Expert 2 ← only these run
 ├─ Expert 3
 └─ Expert N
 combine selected outputs (weighted)

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI