MLA: Low-Rank Joint Compression — Full Derivation

The Real Problem

Sections 4.1 and 4.2 left us in an uncomfortable place. The KV cache, not the parameter count, sets the memory wall of long-context inference. Standard multi-head attention (MHA) stores two full-width tensors per token per layer — keys and values — and the total memory grows linearly in sequence length, number of heads, head dimension, and number of layers. For a 128K-token request through a 128-head, 61-layer model that is several hundred gigabytes per sequence. No commercially viable serving stack can pay that.

Multi-query attention (MQA) and grouped-query attention (GQA) chip away at the problem by sharing K and V across heads. MQA collapses all heads to one shared KV; GQA picks an intermediate number of groups. Both give large constant-factor savings, but they impose a hard floor on the cache: you still spend $(n_{heads}^{kv})\cdot d_h$ floats per token, and pushing that floor lower means giving away expressivity.

The question that defines this section. Can we drop below the GQA/MQA floor — store far fewer numbers per token — without losing the per-head specialisation that gives multi-head attention its modelling power?

Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2 and reused in DeepSeek-V3 and DeepSeek-R1, answers yes — by reframing the cache as a low-rank joint code, not a tensor of keys and values at all.

The Intuition: An Autoencoder for the KV Cache

Look at what a single layer of MHA stores per token: the keys $K_t \in \mathbb{R}^{n_{heads}\cdot d_h}$ and the values $V_t \in \mathbb{R}^{n_{heads}\cdot d_h}$. Both are deterministic linear projections of the same hidden state $h_t \in \mathbb{R}^{d_{model}}$. So we are caching two views of the same source — and we are caching them at full rank, even though in practice trained attention weights and SVD analyses of $W_K$ and $W_V$ show that the effective rank is much lower than $n_{heads}\cdot d_h$.

Here is the MLA idea in one sentence: compress $h_t$ once into a small latent vector $c_{KV,t}$, store only that, and reconstruct the keys and values on the fly when attention needs them.

Mathematically this is a learned low-rank factorisation of the would-be key and value projections, with the factor shared between K and V. Two ideas you have already met merge here:

1. Linear bottleneck — the latent dim $d_c \ll n_{heads}\cdot d_h$ forces information to compress, the way the middle layer of an autoencoder does.
2. Joint coding — keys and values are decoded from the same code, exploiting their statistical dependence. Two separate low-rank decompositions would not see that dependence.

From this picture two questions immediately follow: how small can $d_c$ get before performance breaks? and does the up-projection at every attention call destroy the savings? The math below makes both answerable.

Standard MHA: What We Want to Cache Less Of

To keep the derivation honest, start from the standard MHA cache. Given a token's hidden state $h_t \in \mathbb{R}^{d_{model}}$, MHA computes per head $i$:

$q_{t,i} = h_t W_i^Q$, $k_{t,i} = h_t W_i^K$, $v_{t,i} = h_t W_i^V$, with $W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_{model}\times d_h}$.

At decode step $t$ we must hold the keys and values for all previous tokens, across all heads: $\text{cache}_t^{MHA} = \{(k_{s,i}, v_{s,i})\}_{s\le t,\ i\le n_{heads}}$. That is $2 \cdot t \cdot n_{heads} \cdot d_h$ floats per layer.

Notice the last row: MLA breaks out of the MQA floor entirely. Where MQA must keep at least $2 d_h$ floats, MLA caches one latent of size $d_c$, and nothing forces $d_c$ to be $\ge 2 d_h$. In DeepSeek-V3 $d_c = 512$ and $d_h = 128$, so MLA actually caches more per token than MQA would — but it does so once for all heads instead of once per head, and the quality difference is the whole reason MLA exists.

Step 1: Down-Projection into a Latent

Introduce one new learned matrix:

$W^{DKV} \in \mathbb{R}^{d_{model}\times d_c}$, with $d_c \ll n_{heads}\cdot d_h$.

For every token, compute the latent code once:

$c_{KV,t} = h_t\, W^{DKV} \in \mathbb{R}^{d_c}$.

Three properties of $W^{DKV}$ matter:

1. Shared across heads. Unlike $W_i^K$, there is no per-head copy. One encoder, one cache row.
2. Shared between K and V. The same latent feeds both up-projections below. That is the "joint" in joint compression — it forces the model to find statistics common to keys and values.
3. Bottleneck rank. $\text{rank}(W^{DKV}) \le d_c$. Every downstream projection that flows through it inherits that rank ceiling.

Step 2: Up-Projection at Forward Time

Attention still needs full-width keys and values to do its job. Introduce two up-projection matrices, one for each:

$W^{UK} \in \mathbb{R}^{d_c \times n_{heads}\cdot d_h}$, $W^{UV} \in \mathbb{R}^{d_c \times n_{heads}\cdot d_h}$.

Reconstruct keys and values from the cached latent:

$K_t = c_{KV,t}\, W^{UK}$, $V_t = c_{KV,t}\, W^{UV}$.

Then split each across heads — column block $i$ of $K_t$ is head $i$'s key, and similarly for values. From here, attention is exactly the standard scaled dot-product:

$\text{Attn}_i(Q_t, K_{\le t}, V_{\le t}) = \text{softmax}\!\left(\frac{q_{t,i}\, K_{\le t,i}^\top}{\sqrt{d_h}}\right) V_{\le t,i}$

Reading this honestly, you should worry. We saved memory on the cache, but now every attention call does an extra $O(t\cdot d_c\cdot n_{heads}\cdot d_h)$ matmul to recover the keys and values. Has MLA only moved cost from memory to compute? The next subsection answers no.

Step 3: The Absorption Trick (Why MLA Is Fast)

Write out a single attention score in terms of the latent. For the new query at step $t$ attending to past position $s$, head $i$:

$q_{t,i} \cdot k_{s,i}^\top = (h_t W_i^Q) \cdot (c_{KV,s} W_i^{UK})^\top = h_t \big( W_i^Q\, (W_i^{UK})^\top \big)\, c_{KV,s}^\top$

The two learned matrices $W_i^Q$ and $W_i^{UK}$ are constant once training is done. Define

$\tilde W_i^{Q} \;\triangleq\; W_i^Q\,(W_i^{UK})^\top \in \mathbb{R}^{d_{model}\times d_c}$.

Then at inference the score becomes

$q_{t,i} \cdot k_{s,i}^\top = (h_t\, \tilde W_i^{Q}) \cdot c_{KV,s}^\top$.

Read this slowly. The dot product against $k_{s,i}$ has become a dot product against $c_{KV,s}$. The up-projection $W_i^{UK}$ has been absorbed into a modified query matrix $\tilde W_i^{Q}$. At inference we never materialise the full per-head key — we project the query into latent space and dot it against the cached latent directly.

The same trick applies on the value side. The output of head $i$ mixes value vectors weighted by attention probabilities, then passes through the output projection $W_i^O$. The chain $W_i^{UV}\, W_i^O$ can likewise be pre-multiplied into a single matrix from the latent to the residual stream, so the up-projection of $V$ never runs at decode time either.

At training the absorption is not done. We keep $W^{DKV}, W^{UK}, W^{UV}, W^Q, W^O$ as separate learnable matrices so each gradient has a clear signal. At inference we collapse adjacent linear maps into single matrices because nothing in between them is non-linear or per-input.

The Value Stream and the Output Projection

It is worth spelling out the value side because the asymmetry with the query side is sometimes confusing. The attention output for one token, one head, is

$o_{t,i} = \sum_{s\le t} a_{t,s,i}\, v_{s,i} = \sum_{s\le t} a_{t,s,i}\, c_{KV,s} W_i^{UV}$

and the layer output mixes heads through $W^O \in \mathbb{R}^{n_{heads}\cdot d_h \times d_{model}}$:

$y_t = \text{concat}_i(o_{t,i})\, W^O = \sum_i o_{t,i}\, W_i^O$, with $W_i^O$ the block of $W^O$ corresponding to head $i$.

Substituting and grouping:

$y_t = \sum_i \sum_{s\le t} a_{t,s,i}\, c_{KV,s}\, \big( W_i^{UV} W_i^O \big)$

The factor $W_i^{UV} W_i^O$ is constant — precompute it as $\tilde W_i^{OV}$. At inference, no value matmul ever expands the latent; the latent flows straight into the residual stream through one absorbed projection per head.

Manual Numerical Walkthrough

Click to expand a fully calculated three-token MLA pass with $T=3,\ d_{model}=4,\ n_{heads}=2,\ d_h=2,\ d_c=2$. The same numbers reappear in the NumPy and PyTorch implementations below — and you can multiply every step by hand.

1H = [[1, 0, 1, 0],   # token 0
2     [0, 1, 0, 1],   # token 1
3     [1, 1, 0, 0]]   # token 2     shape (3, 4)

1W_DKV = [[1, 0],         W_UK = [[1, 0, 0, 1],     W_UV = [[0, 1, 1, 0],
2         [0, 1],                 [0, 1, 1, 0]]            [1, 0, 0, 1]]
3         [1, 0],          shape (2, 4)              shape (2, 4)
4         [0, 1]]
5   shape (4, 2)
6
7W_Q = [[1, 0, 0, 1],
8       [0, 1, 1, 0],
9       [0, 0, 0, 0],
10       [0, 0, 0, 0]]    shape (4, 4)

1token 0: [1,0,1,0] · W_DKV = (1+1, 0+0) = [2, 0]
2token 1: [0,1,0,1] · W_DKV = (0+0, 1+1) = [0, 2]
3token 2: [1,1,0,0] · W_DKV = (1+0, 0+1) = [1, 1]
4
5c_KV = [[2, 0],
6        [0, 2],
7        [1, 1]]   shape (3, 2)   ← THE CACHE (6 floats for 3 tokens)

1K = [[2, 0, 0, 2],     V = [[0, 2, 2, 0],
2     [0, 2, 2, 0],          [2, 0, 0, 2],
3     [1, 1, 1, 1]]          [1, 1, 1, 1]]   shape (3, 4) each
4
5split into n_heads = 2 heads of d_h = 2:
6
7K_h0 = [[2,0],[0,2],[1,1]]   K_h1 = [[0,2],[2,0],[1,1]]
8V_h0 = [[0,2],[2,0],[1,1]]   V_h1 = [[2,0],[0,2],[1,1]]

1Q_h0 = [[1,0],[0,1],[1,1]]
2Q_h1 = [[0,1],[1,0],[1,1]]

1raw = [[2, 0, 1],
2       [0, 2, 1],
3       [2, 2, 2]]
4
5scaled = raw / sqrt(d_h) = raw / sqrt(2)
6       ≈ [[1.41, 0.00, 0.71],
7          [0.00, 1.41, 0.71],
8          [1.41, 1.41, 1.41]]
9
10masked (causal):
11       = [[ 1.41, -inf, -inf],
12          [ 0.00, 1.41, -inf],
13          [ 1.41, 1.41, 1.41]]
14
15softmax row-wise:
16row 0: [1.000, 0.000, 0.000]
17row 1: e^0=1, e^1.41≈4.114, sum=5.114 → [0.196, 0.804, 0.000]
18row 2: uniform → [0.333, 0.333, 0.333]

1token 0: 1.000·[0,2] + 0.000·[2,0] + 0.000·[1,1] = [0.000, 2.000]
2token 1: 0.196·[0,2] + 0.804·[2,0] + 0.000·[1,1] = [1.608, 0.392]
3token 2: 0.333·([0,2]+[2,0]+[1,1])             = [1.000, 1.000]

1q_{1, h0} · k_{0, h0}^T = [0, 1] · [2, 0] = 0   ✓

1# W_Q_h0 = [[1,0],[0,1],[0,0],[0,0]]
2# W_UK_h0 (cols 0,1 of W_UK) = [[1,0],[0,1]]
3# absorbed: W_Q_h0 @ W_UK_h0^T = [[1,0],[0,1],[0,0],[0,0]]
4#   (here happens to equal W_Q_h0 because W_UK_h0 is identity-shaped)
5
6h_1 · absorbed = [0,1,0,1] · [[1,0],[0,1],[0,0],[0,0]] = [0, 1]
7[0, 1] · c_KV_0^T = [0, 1] · [2, 0] = 0    ✓ same answer

Interactive Visualization

Use the explorer below to walk through MLA on a 5-token toy sequence. The Pipeline view shows the original key, the latent code, and the reconstruction side-by-side per token. The Cache Size view scales the latent dim $d_c$ and lets you watch the per-token memory cost move past MQA and below. The Heatmap view overlays standard attention weights on MLA weights — the bottleneck warps the similarity space, but the global pattern survives.

Watch what happens at $d_c = 1$: every token is squeezed onto a single number, the reconstruction collapses, and the heatmap loses most of its structure. That is what the bottleneck looks like when it is too tight. Practical MLA picks $d_c$ well above the rank floor at which loss starts to climb.

Plain Python (NumPy) Implementation

Before wrapping anything in nn.Module, run one MLA forward pass in NumPy. Every weight is hardcoded, every value printable, no autograd in the way. Click any line on the right to see exactly what arrives, what each operation does, and what leaves.

[[1, 0, 1, 0],   # token 0
 [0, 1, 0, 1],   # token 1
 [1, 1, 0, 0]]   # token 2
shape (3, 4) = (T, d_model)

[[1, 0],
 [0, 1],
 [1, 0],
 [0, 1]]
shape (4, 2) = (d_model, d_c)
Role: maps every hidden state h_t to a 2-D code c_KV[t] = h_t @ W_DKV.

[[1, 0, 0, 1],
 [0, 1, 1, 0]]
shape (2, 4) = (d_c, n_heads * d_h)
Role: K = c_KV @ W_UK gives the full per-head key block.

[[0, 1, 1, 0],
 [1, 0, 0, 1]]
shape (2, 4) = (d_c, n_heads * d_h)
Role: V = c_KV @ W_UV gives the full per-head value block.

[[1, 0, 0, 1],
 [0, 1, 1, 0],
 [0, 0, 0, 0],
 [0, 0, 0, 0]]
shape (4, 4) = (d_model, n_heads * d_h)
Role: Q = H @ W_Q. The last two rows are zeros so the toy query depends only on the first two H columns.

[[1,0,1,0],[0,1,0,1],[1,1,0,0]] — shape (3, 4)

Row 0: 1·[1,0] + 0·[0,1] + 1·[1,0] + 0·[0,1] = [2, 0]
Row 1: 0+[0,1]+0+[0,1] = [0, 2]
Row 2: [1,0]+[0,1]+0+0 = [1, 1]
c_KV = [[2,0],[0,2],[1,1]]   shape (3, 2)
→ THIS is what we store. 6 floats for 3 tokens, regardless of n_heads.

[[2,0],[0,2],[1,1]] — shape (3, 2)

Row 0: [2,0] · [[1,0,0,1],[0,1,1,0]] = [2, 0, 0, 2]
Row 1: [0,2] · ⋯ = [0, 2, 2, 0]
Row 2: [1,1] · ⋯ = [1, 1, 1, 1]
shape (3, 4)

Per head 0 (cols 0,1): [[2,0],[0,2],[1,1]]
Per head 1 (cols 2,3): [[0,2],[2,0],[1,1]]
shape (3, 2, 2)

Row 0: [2,0] · [[0,1,1,0],[1,0,0,1]] = [0, 2, 2, 0]
Row 1: [0,2] · ⋯ = [2, 0, 0, 2]
Row 2: [1,1] · ⋯ = [1, 1, 1, 1]
shape (3, 4)

Per head 0: [[0,2],[2,0],[1,1]]
Per head 1: [[2,0],[0,2],[1,1]]
shape (3, 2, 2)

Row 0 = [1,0,1,0]·W_Q = [1, 0, 0, 1]
Row 1 = [0,1,0,1]·W_Q = [0, 1, 1, 0]
Row 2 = [1,1,0,0]·W_Q = [1, 1, 1, 1]
shape (3, 4)

Per head 0: [[1,0],[0,1],[1,1]]
Per head 1: [[0,1],[1,0],[1,1]]
shape (3, 2, 2)

Any tensor of scores. We softmax along the last axis so each row sums to 1 — the standard attention-weight pattern.

Q_h0 @ K_h0.T = [[1,0],[0,1],[1,1]] @ [[2,0,1],[0,2,1]]
   = [[2, 0, 1],
      [0, 2, 1],
      [2, 2, 2]]

raw / sqrt(2) ≈ [[1.41, 0, 0.71],[0, 1.41, 0.71],[1.41, 1.41, 1.41]]

[[False, True,  True ],
 [False, False, True ],
 [False, False, False]]

[[ 1.41, -1e9, -1e9],
 [ 0.00,  1.41, -1e9],
 [ 1.41,  1.41,  1.41]]

Row 0: softmax([1.41, -∞, -∞]) = [1.0, 0.0, 0.0]
Row 1: softmax([0, 1.41, -∞]). e^0=1, e^1.41≈4.11. → [0.196, 0.804, 0.0]
Row 2: softmax([1.41, 1.41, 1.41]) = [0.333, 0.333, 0.333]

t=0: 1.0 * V_h0[0] = [0, 2]
t=1: 0.196*[0,2] + 0.804*[2,0] = [1.608, 0.392]
t=2: (1/3)([0,2]+[2,0]+[1,1]) = [1.0, 1.0]

Per-token row = [head0 | head1]
 token 0: [0, 2, 2, 0]
 token 1: [1.608, 0.392, 0.196, 1.608]
 token 2: [1.0, 1.0, 1.0, 1.0]
shape (3, 4)

1import numpy as np
2
3np.random.seed(0)
4
5# --- shared toy config (matches the worked example above) ---
6T, d_model, n_heads, d_h, d_c = 3, 4, 2, 2, 2
7
8# --- hidden states for 3 tokens (T, d_model) ---
9H = np.array([[1., 0., 1., 0.],
10              [0., 1., 0., 1.],
11              [1., 1., 0., 0.]])
12
13# --- MLA learned matrices ---
14W_DKV = np.array([[1., 0.],
15                  [0., 1.],
16                  [1., 0.],
17                  [0., 1.]])               # (d_model, d_c)
18W_UK  = np.array([[1., 0., 0., 1.],
19                  [0., 1., 1., 0.]])       # (d_c, n_heads * d_h)
20W_UV  = np.array([[0., 1., 1., 0.],
21                  [1., 0., 0., 1.]])       # (d_c, n_heads * d_h)
22W_Q   = np.array([[1., 0., 0., 1.],
23                  [0., 1., 1., 0.],
24                  [0., 0., 0., 0.],
25                  [0., 0., 0., 0.]])       # (d_model, n_heads * d_h)
26
27# --- Step 1: down-project hidden states into the latent ---
28c_KV = H @ W_DKV                            # (T, d_c) — THIS is what we cache
29
30# --- Step 2: up-project keys and values from the latent ---
31K = (c_KV @ W_UK).reshape(T, n_heads, d_h)  # (T, H, d_h)
32V = (c_KV @ W_UV).reshape(T, n_heads, d_h)  # (T, H, d_h)
33
34# --- Queries are independent of the cache ---
35Q = (H @ W_Q).reshape(T, n_heads, d_h)      # (T, H, d_h)
36
37# --- attention per head, with causal mask ---
38def softmax(x):
39    x = x - x.max(axis=-1, keepdims=True)
40    e = np.exp(x)
41    return e / e.sum(axis=-1, keepdims=True)
42
43scores = np.einsum("thd,Thd->htT", Q, K) / np.sqrt(d_h)   # (H, T, T)
44mask   = np.triu(np.ones((T, T), dtype=bool), k=1)        # block future
45scores = np.where(mask, -1e9, scores)
46attn   = softmax(scores)                                   # (H, T, T)
47
48out_per_head = np.einsum("htT,Thd->thd", attn, V)         # (T, H, d_h)
49out          = out_per_head.reshape(T, n_heads * d_h)     # (T, n_heads*d_h)
50
51print("c_KV (cached):\n", c_KV)
52print("attn head 0:\n", attn[0])
53print("out shape:", out.shape)

Shape table

The eight tensors a plain-MLA forward actually touches:

PyTorch Implementation

The PyTorch version of the same forward pass is a direct transcription of the NumPy trace into nn.Module form, with one important addition: it accepts an optional past latent cache so it can be called repeatedly during autoregressive decoding. The cache it returns is the only thing the next step needs.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class MultiHeadLatentAttention(nn.Module):
6    """
7    Plain MLA (no RoPE — section 4.4 adds it).
8    Caches only c_KV ∈ R^{d_c} per token instead of K and V.
9    """
10
11    def __init__(self, d_model: int, n_heads: int, d_h: int, d_c: int):
12        super().__init__()
13        self.n_heads, self.d_h, self.d_c = n_heads, d_h, d_c
14
15        # Queries are produced fresh from the current hidden state.
16        self.W_Q = nn.Linear(d_model, n_heads * d_h, bias=False)
17
18        # Down-projection — produces the cached latent.
19        self.W_DKV = nn.Linear(d_model, d_c, bias=False)
20
21        # Up-projections — decompress the latent into K and V at forward time.
22        self.W_UK = nn.Linear(d_c, n_heads * d_h, bias=False)
23        self.W_UV = nn.Linear(d_c, n_heads * d_h, bias=False)
24
25        # Output projection — same role as in standard MHA.
26        self.W_O = nn.Linear(n_heads * d_h, d_model, bias=False)
27
28    def forward(self, h: torch.Tensor, c_KV_past: torch.Tensor | None = None):
29        # h: (B, T_new, d_model). c_KV_past: (B, T_past, d_c) or None.
30        B, T_new, _ = h.shape
31
32        # 1. Compress the new tokens once. THIS is what we cache.
33        c_KV_new = self.W_DKV(h)                                # (B, T_new, d_c)
34        c_KV     = c_KV_new if c_KV_past is None \
35                   else torch.cat([c_KV_past, c_KV_new], dim=1) # (B, T, d_c)
36        T = c_KV.shape[1]
37
38        # 2. Up-project the FULL latent stream to keys and values.
39        K = self.W_UK(c_KV).view(B, T, self.n_heads, self.d_h)
40        V = self.W_UV(c_KV).view(B, T, self.n_heads, self.d_h)
41
42        # 3. Queries from the new hidden states only.
43        Q = self.W_Q(h).view(B, T_new, self.n_heads, self.d_h)
44
45        # 4. Standard scaled dot-product attention with a causal mask.
46        Q = Q.transpose(1, 2)                                   # (B, H, T_new, d_h)
47        K = K.transpose(1, 2)                                   # (B, H, T,     d_h)
48        V = V.transpose(1, 2)                                   # (B, H, T,     d_h)
49
50        scores = (Q @ K.transpose(-2, -1)) / (self.d_h ** 0.5)  # (B, H, T_new, T)
51        causal = torch.triu(torch.ones(T_new, T, dtype=torch.bool,
52                                       device=h.device),
53                            diagonal=T - T_new + 1)
54        scores = scores.masked_fill(causal, float("-inf"))
55        attn   = F.softmax(scores, dim=-1)
56
57        out = attn @ V                                          # (B, H, T_new, d_h)
58        out = out.transpose(1, 2).reshape(B, T_new, -1)         # concat heads
59        return self.W_O(out), c_KV                              # return new cache

Note on the absorption trick. The class above does not implement the absorbed forward path. At training and prefill, materialising K and V is fine because we process the whole sequence and there is no per-step compression payoff. The absorbed version of forward — where queries are projected straight into $d_c$ space and dotted against $c_{KV}$ — is what production decoders run. Section 4.6 ("Implementing MLA in PyTorch") builds the absorbed kernel explicitly.

At Massive Scale: DeepSeek-V3 Numbers

Plug DeepSeek-V3's actual hyperparameters into the MLA cache formula and the engineering story becomes vivid:

Standard MHA cache. Per token, per layer: $2\cdot n_{heads}\cdot d_h = 32{,}768$ floats = 64 KB at BF16. Across 61 layers: 3.91 MB per token. At 128 K context: ~500 GB per sequence. Not servable on a single node.

MLA cache. Per token, per layer: $d_c + d_h^R = 576$ floats = 1.15 KB at BF16. Across 61 layers: 70.3 KB per token. At 128 K context: ~9.2 GB per sequence. Servable on a single H100.

The implication for serving is direct. KV cache is the dominant memory cost of long-context inference, more than parameter weights when sequences run long. Cutting it by 56× changes which models can run on which hardware, how many concurrent users a node can hold, and how aggressively a batch can be packed. In DeepSeek-V3's reported serving setup, MLA is the reason the model fits within commodity multi-GPU nodes at 128 K context lengths.

The arithmetic of attention cost itself

Memory is half the story. Counting the FLOPs of one attention step at decode (one new token, $T$ in the cache):

• Standard MHA: $2 n_{heads} d_h T$ for the query-key dot products plus the same for the value mix.
• MLA (absorbed): $2 n_{heads} d_c T$ for the query-latent dots plus a one-time $d_{model}\to d_c$ projection of the new query.

At V3 numbers the FLOP ratio for the dominant attention term is $d_h / d_c = 128/512 = 1/4$ — MLA is cheaper per step than the MHA baseline it replaces, not just smaller in memory. This is the rare case in transformer engineering where the memory-saving move is also the compute-saving move; usually you trade one for the other.

What Can Go Wrong in Practice

MLA looks like a free lunch on paper. Three things bite in practice.

1. The latent dimension is a real architectural knob

Too small $d_c$ and the model loses per-head resolution — heatmaps blur, downstream perplexity climbs. The DeepSeek-V2 ablations sweep $d_c$ over a wide range and identify a sharp elbow below which loss degrades quickly. The chosen value ($d_c = 512$) sits well above the elbow with a safety margin.

2. RoPE does not pass through the latent cleanly

Rotary positional embedding rotates K (and Q) by position-dependent block-diagonal matrices before the dot product. If you cache $c_{KV}$ and apply RoPE only at decompression time, you need RoPE to commute with the up-projection $W^{UK}$ — which it does not in general. This is not a small detail; naive MLA + RoPE breaks long-range position sensitivity. The fix is the decoupled-RoPE channel covered in §4.4: an additional small per-head positional key $k_t^R \in \mathbb{R}^{d_h^R}$ is cached alongside the latent and concatenated into the score computation. It is the $d_h^R = 64$ in the V3 table above.

3. Absorption inflates the weight matrices

Multiplying $W^Q$ by $(W^{UK})^\top$ gives a denser matrix than either factor. For V3 numbers the absorbed query matrix is $d_{model}\cdot d_c$ per head — larger than the original per-head $d_{model}\cdot d_h$ because $d_c > d_h$. The weight overhead is paid once at load time; the inference savings are paid back per token. Worth it for any non-trivial context length, but a real memory tradeoff that the serving stack must budget for.

Summary

• MLA replaces the standard $(K_t, V_t)$ KV cache with a single per-token latent $c_{KV,t} = h_t W^{DKV}$ of dimension $d_c \ll n_{heads}\cdot d_h$.
• Keys and values are reconstructed at training/prefill time via $W^{UK}, W^{UV}$. The two up-projections share the same latent — "joint compression".
• At inference, the up-projections absorb into the query and output projections. The decode-time forward never materialises full K or V. Score computation runs in latent space: query→$d_c$, then dot against $c_{KV}$.
• On DeepSeek-V3 dimensions, MLA cuts the KV cache by ~57× per layer and is also cheaper in attention FLOPs at decode than the MHA baseline it replaces.
• The two real complications are an extra positional channel (§4.4's decoupled RoPE) and a one-time weight inflation when the absorbed matrices are stored.

The next section adds the missing piece. Section 4.4 derives the decoupled-RoPE channel from the failure mode it fixes, and completes MLA's connection to long-context positional reasoning.