MLA: Compressing the KV Cache Without Decompressing It

Storing a compressed latent instead of keys and values is easy. Doing it without paying the saving straight back at read time is the hard part — and RoPE is what makes it hard. This post derives absorption from the associativity of a matmul, shows RoPE destroying it with measured errors of 42–100%, and arrives at decoupled RoPE as the repair.
llm-systems
inference
attention
Published

August 10, 2026

The KV cache post ended with a bill. Every token you generate leaves behind a key and a value in every head of every layer, and the total grows without bound as the context does.

\[\text{elements cached per token} = 2 \cdot L \cdot n_{kv} \cdot d_h\]

where \(n_{kv}\) is the number of cached key/value representations — which need not equal \(n_q\), the number of query heads.

Read as a menu rather than a fact, that formula has four levers. The layers are fixed by the architecture. That leaves \(n_{kv}\), how many independently cached key/value representations there are, and \(d_h\), how many numbers each one holds — and the leading 2, which looks structural and is not.

That 2 says: store one key and one value per token, separately. Attention certainly needs both roles. It does not follow that the cache needs two independent vectors, and challenging that premise is what separates MLA from everything else here. Multi-query and grouped-query attention leave the explicit K/V representation intact and shrink \(n_{kv}\). MLA caches one joint latent from which both are produced.

Sharing key/value representations attacks \(n_{kv}\). MLA attacks the 2 and the width together — preserve head-specific key and value projections, but make them functions of one much smaller cached latent. That is where the trap lives: compressing what you store only helps if attention can consume the compressed form directly. Otherwise you have merely converted a memory problem into an arithmetic one.

TL;DR — Cache a low-dimensional latent c instead of K and V, and the storage collapses. But a naive implementation reconstructs every historical key and value at every decode step, which hands the saving straight back. Absorption fixes that: qᵀW_UK c = (W_UKᵀq)ᵀc, so the up-projection folds into the query side once and historical keys are never built — and on the value side it folds all the way into the output projection. Then RoPE breaks it, because the matrix you wanted to absorb becomes \(R_t^\top R_j W_{UK}\), which depends on the cached token’s position. Using absorption anyway gives 42–100% relative score error while the scores still correlate at 0.69–0.92 — wrong, and plausible-looking. Decoupled RoPE is an architecture that keeps a small position-carrying branch outside the compressed path, and because its rotary key is shared across heads it adds \(d_R\) per token rather than \(n_q \cdot d_R\).

Everything runs on CPU in a couple of seconds.

import torch

torch.set_num_threads(1)
torch.manual_seed(0)
print("torch", torch.__version__)
torch 2.8.0

Act I — what has to be stored

One axis: how many independent K/V representations

Before compressing anything, the cheapest lever is the one already visible in the formula: \(n_{kv}\) counts cached key/value representations, and nothing forces it to equal \(n_q\).

Set \(n_{kv} = n_q\) and you have multi-head attention. Set \(n_{kv} = 1\) and every query head reads one shared key/value — multi-query, a factor of \(n_q\) smaller. Set \(n_{kv} = n_g\) for some group count and you land anywhere in between.

L, n_h, d_h = 60, 128, 128

print(f"{'scheme':<28}{'cached K/V heads':>18}{'elements/token':>17}{'vs MHA':>10}")
mha = 2 * L * n_h * d_h
for name, n_kv in [("multi-head (MHA)", n_h), ("grouped, 8 groups", 8),
                   ("grouped, 2 groups", 2), ("multi-query (MQA)", 1)]:
    bill = 2 * L * n_kv * d_h
    print(f"{name:<28}{n_kv:>18}{bill:>17,}{mha / bill:>9.0f}x")
scheme                        cached K/V heads   elements/token    vs MHA
multi-head (MHA)                           128        1,966,080        1x
grouped, 8 groups                            8          122,880       16x
grouped, 2 groups                            2           30,720       64x
multi-query (MQA)                            1           15,360      128x

That is one axis, and it is a good one — the cache arithmetic is exact and simple, whatever the kernel work costs. What it costs in quality is not something a post with no trained models can measure honestly, and the answer depends on the model, so take it from the literature rather than from me.

The rest of this post is about the other route: one shared cached latent, read by head-specific projections.

Act II — cache what produced them, not them

Store a latent

Keys and values are linear functions of the hidden state. Rather than caching their outputs, cache a low-dimensional thing they can both be computed from:

\[c_t = W_D h_t, \qquad k_t = W_{UK}\,c_t, \qquad v_t = W_{UV}\,c_t\]

with c of width \(d_c\) much smaller than \(n_q \cdot d_h\). Cache only c.

d_model, d_c, d_head, n_heads, T = 96, 32, 24, 4, 12

h = torch.randn(T, d_model)                                     # hidden states so far
W_D  = torch.randn(d_c, d_model) / d_model ** 0.5               # down-projection
W_UK = torch.randn(n_heads, d_head, d_c) / d_c ** 0.5           # per-head key up-proj
W_UV = torch.randn(n_heads, d_head, d_c) / d_c ** 0.5           # per-head value up-proj
W_Q  = torch.randn(n_heads, d_head, d_model) / d_model ** 0.5

C = h @ W_D.T                                                   # the cache: (T, d_c)
print(f"cached per token — MHA style : {2 * n_heads * d_head} elements")
print(f"cached per token — latent    : {d_c} elements")
cached per token — MHA style : 192 elements
cached per token — latent    : 32 elements

Ignoring positional state for a moment, the content bill has fallen from \(2 L n_q d_h\) to \(d_c L\). That is the first satisfying result, and it is also where most explanations stop.

The trap: you moved the cost, you did not remove it

Decoding step t needs to attend over all previous keys and values. If all you kept is c_1 … c_t, the obvious implementation reconstructs them:

K = C @ W_UK.T          # rebuild every historical key
V = C @ W_UV.T          # rebuild every historical value

and then does it again at step t+1, and again at t+2. You are recomputing the entire history’s projections on every single token.

That is the structural dilemma, and it is worth stating before any timing:

cache C only                    cache reconstructed K, V
      ↓                                   ↓
rebuild all history, every step    memory saving is gone
      ↓
arithmetic that grows with context

Neither branch preserves both the memory saving and the inference efficiency you wanted. The escape is not a better implementation of either one.

Absorb the key up-projection

Look at what an attention score actually is, for one head i and one cached token j:

\[q_{t,i}^\top k_{j,i} = q_{t,i}^\top W^{UK}_i c_j\]

Three matrices in a row, and matrix products associate. Group them the other way:

\[q_{t,i}^\top W^{UK}_i c_j = \left(\big(W^{UK}_i\big)^\top q_{t,i}\right)^\top c_j\]

The left factor depends only on the current query. Compute it once per step, and then score directly against the cached latents. No historical key is ever built.

q = torch.stack([h[-1] @ W_Q[i].T for i in range(n_heads)])     # (n_heads, d_head)

scores_naive, scores_absorbed = [], []
for i in range(n_heads):
    K_i = C @ W_UK[i].T                                          # rebuild keys  (T, d_head)
    scores_naive.append(K_i @ q[i])
    scores_absorbed.append(C @ (W_UK[i].T @ q[i]))               # fold into the query

diff = max((a - b).abs().max().item()
           for a, b in zip(scores_naive, scores_absorbed))
print(f"bit-identical : {all(torch.equal(a, b) for a, b in zip(scores_naive, scores_absorbed))}")
print(f"max abs diff  : {diff:.2e}")
bit-identical : False
max abs diff  : 2.86e-06

Not bit-identical, and the distinction matters. The two expressions are the same over the reals; they are different orders of floating-point operations, so they differ in the last few digits. Regrouping a matmul is exactly the kind of reassociation that changes rounding — the same phenomenon as the SmoothQuant identity in the quantization post. Say “algebraically identical, differing by a few times 1e-06 in fp32,” never “the same.”

Absorb the value up-projection too

Key absorption is only half of the algebra; values admit the same reassociation, and skipping it leaves half the win on the table. The output of one head is

\[o_i = \sum_j a_{ij} v_{j,i} = \sum_j a_{ij} W^{UV}_i c_j = W^{UV}_i \left(\sum_j a_{ij} c_j\right)\]

Same associativity, and the systems consequence is the interesting part:

BAD                                  GOOD
rebuild V_j for every cached token   weighted-sum the small latents first
then take the weighted sum           then up-project once, for this token only

You never touch a historical value either. And there is one more step available, because every head’s output passes through the output projection W_O anyway. Since each W^{UV}_i is a fixed matrix, it can be folded into W_O ahead of time — the absorbed path doesn’t need a separate value up-projection stage at all.

outs_naive, outs_absorbed, latent_sums = [], [], []
for i in range(n_heads):
    a_i = torch.softmax(scores_naive[i] / d_head ** 0.5, dim=0)
    outs_naive.append((C @ W_UV[i].T).T @ a_i)                   # rebuild V, then sum
    z_i = C.T @ a_i                                              # sum latents first
    latent_sums.append(z_i)
    outs_absorbed.append(W_UV[i] @ z_i)                          # one small projection

print(f"value absorption      max diff : "
      f"{max((a - b).abs().max().item() for a, b in zip(outs_naive, outs_absorbed)):.2e}")

W_O = torch.randn(d_model, n_heads * d_head) / (n_heads * d_head) ** 0.5
W_O_eff = torch.cat([W_O[:, i * d_head:(i + 1) * d_head] @ W_UV[i]
                     for i in range(n_heads)], dim=1)            # fold W_UV into W_O

u_naive = W_O @ torch.cat(outs_naive)
u_folded = W_O_eff @ torch.cat(latent_sums)
print(f"W_UV folded into W_O  max diff : {(u_naive - u_folded).abs().max():.2e}")
value absorption      max diff : 2.38e-07
W_UV folded into W_O  max diff : 2.38e-07

So the absorbed path is: project the query into the latent space once, score against c, weight the latents, and apply one effective output projection. The cache stays compressed from end to end and nothing in the history is ever reconstructed.

Act III — position breaks the algebra

What RoPE does to the expression

Everything above assumed the score is qᵀ W_UK c. Rotary embeddings change it. RoPE rotates the query by its position and each key by its position:

\[(R_t q_t)^\top R_j \big(W^{UK} c_j\big) = q_t^\top\, R_t^\top R_j W^{UK}\, c_j\]

Stare at the middle:

\[\boxed{\;R_t^\top R_j W^{UK}\;}\]

Absorption required a matrix that could be folded into the query once and reused against every cached token. This one has a j in it. It is a different matrix for every position in the cache, so there is no single position-independent matrix that can be absorbed once and reused against every cached token. The trick is gone — not degraded, gone.

The uncomfortable part is what happens if you use it anyway:

def rope(x, positions, dim):
    """Rotate pairs of channels by an angle proportional to position."""
    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
    ang = positions[:, None] * inv_freq[None, :]
    cos, sin = torch.cos(ang), torch.sin(ang)
    out = torch.empty_like(x)
    even, odd = x[..., 0::2], x[..., 1::2]
    out[..., 0::2] = even * cos - odd * sin
    out[..., 1::2] = even * sin + odd * cos
    return out


print(f"{'head':>6}{'relative error':>17}{'correlation':>14}")
for i in range(n_heads):
    K_roped = rope(C @ W_UK[i].T, torch.arange(T).float(), d_head)   # key by its own pos
    q_roped = rope(q[i][None], torch.tensor([float(T - 1)]), d_head)[0]

    true_scores = K_roped @ q_roped                        # what RoPE attention actually is
    pretend = C @ (W_UK[i].T @ q_roped)                    # absorption, used anyway

    rel = ((true_scores - pretend).norm() / true_scores.norm()).item()
    corr = torch.corrcoef(torch.stack([true_scores, pretend]))[0, 1].item()
    print(f"{i:>6}{rel:>16.1%}{corr:>14.3f}")
  head   relative error   correlation
     0           41.5%         0.915
     1           71.2%         0.702
     2          100.0%         0.691
     3           78.5%         0.703

Relative error in Frobenius norm from 42% to 100% across the four heads, with correlations between 0.69 and 0.92. The wrong scores stay substantially correlated with the right ones, so gross structure alone would not necessarily reveal the failure. That combination is the dangerous one — a different function, output that does not obviously look broken, and nothing to catch it. The exact magnitudes belong to this controlled test and will vary with dimensions and initialization; the durable result is binary. The fixed absorption identity used above no longer holds.

Separate content from position

If the problem is that position and content are entangled inside one dot product, the repair is to stop entangling them. Split each head’s query and key into two parts:

\[q = [\,q^C ; q^R\,], \qquad k = [\,k^C ; k^R\,], \qquad q^\top k = (q^C)^\top k^C + (q^R)^\top k^R\]

The content part q^C, k^C carries no positional information and stays on the compressed, absorbable path. The positional part q^R, k^R carries RoPE and is cached uncompressed — but it is small.

Be precise about what this achieves. It does not repair full-RoPE MLA into the same function; that function is simply not absorbable. It is a different attention parameterization, designed so that compression survives. What can be demonstrated exactly is narrower: for this architecture, the naive and absorbed implementations compute the same thing.

d_rope = 16
W_QR = torch.randn(n_heads, d_rope, d_model) / d_model ** 0.5   # per-head rotary query
W_KR = torch.randn(d_rope, d_model) / d_model ** 0.5            # SHARED rotary key

k_rope = rope(h @ W_KR.T, torch.arange(T).float(), d_rope)      # cached, one per token
pos_t = torch.tensor([float(T - 1)])
scale = (d_head + d_rope) ** 0.5

S_n, S_a, O_n, O_a = [], [], [], []
for i in range(n_heads):
    qC = h[-1] @ W_Q[i].T
    qR = rope((h[-1] @ W_QR[i].T)[None], pos_t, d_rope)[0]

    s_naive = ((C @ W_UK[i].T) @ qC + k_rope @ qR) / scale       # rebuild keys
    s_absorbed = (C @ (W_UK[i].T @ qC) + k_rope @ qR) / scale    # absorbed content term
    S_n.append(s_naive); S_a.append(s_absorbed)

    a_n, a_a = torch.softmax(s_naive, 0), torch.softmax(s_absorbed, 0)
    O_n.append((C @ W_UV[i].T).T @ a_n)                          # rebuild values
    O_a.append(W_UV[i] @ (C.T @ a_a))                            # absorbed value term

worst = lambda A, B: max((a - b).abs().max().item() for a, b in zip(A, B))
print(f"scores            max diff {worst(S_n, S_a):.2e}")
print(f"attention weights max diff "
      f"{worst([torch.softmax(s, 0) for s in S_n], [torch.softmax(s, 0) for s in S_a]):.2e}")
print(f"outputs           max diff {worst(O_n, O_a):.2e}")
scores            max diff 7.15e-07
attention weights max diff 1.49e-07
outputs           max diff 4.17e-07

All three stages agree to floating-point tolerance, so both key-side and value-side absorption survive the positional branch.

Why the rotary key is shared

There is one asymmetry in that code that is easy to skim past and explains the entire cache formula. W_QR has a per-head dimension. W_KR does not — one rotary key vector serves every head.

The reason is that queries are transient and keys are historical. A query is computed fresh each step and thrown away, so making it per-head adds parameters but no per-token cache state. A key is cached forever. Making the rotary key per-head would multiply the positional part of the cache by \(n_q\):

d_c_big, d_rope_big = 512, 64                      # illustrative latent dimensions

shared = L * (d_c_big + d_rope_big)
per_head_rope = L * (d_c_big + n_h * d_rope_big)
print(f"shared rotary key   : L·(d_c + d_R)     = {shared:,} elements/token")
print(f"per-head rotary key : L·(d_c + n_h·d_R) = {per_head_rope:,}"
      f"   ({per_head_rope / shared:.0f}x larger)")
shared rotary key   : L·(d_c + d_R)     = 34,560 elements/token
per-head rotary key : L·(d_c + n_h·d_R) = 522,240   (15x larger)

With these dimensions, a per-head rotary key would make the positional component \(n_q = 128\) times larger and the total cached state about 15 times larger. Sharing the rotary key is the architectural choice that keeps the positional cache from scaling with the head count. Whether a per-head rotary key would represent anything usefully different is not something this post can measure; the arithmetic above is only about what it would cost.

The ledger

print(f"{'scheme':<34}{'elements/token':>17}{'vs MHA':>10}")
schemes = [("MHA", 2 * L * n_h * d_h),
           ("GQA, 8 groups", 2 * L * 8 * d_h),
           ("MQA", 2 * L * 1 * d_h),
           ("MLA (d_c=512, d_R=64)", L * (512 + 64))]
base = schemes[0][1]
for name, bill in schemes:
    print(f"{name:<34}{bill:>17,}{base / bill:>9.0f}x")
scheme                               elements/token    vs MHA
MHA                                       1,966,080        1x
GQA, 8 groups                               122,880       16x
MQA                                          15,360      128x
MLA (d_c=512, d_R=64)                        34,560       57x

Two things in that table are worth not glossing over. MLA’s saving is large but it is not the largest number in the column — multi-query compresses further in raw elements, because it keeps one K/V head for the entire layer. But MQA shares one K/V representation across all query heads. The queries stay head-specific, so the attention patterns can still differ head to head; what disappears is head-specific key and value projection space. MLA keeps head-specific W_UK and W_UV while caching only the shared latent they read from. The comparison is not element count alone.

And the MLA row depends entirely on d_c and d_R, which are architectural choices, not constants. Move them and the row moves.

The costs, stated honestly:

  • Additional projection stages. A down-projection on the way in, up-projections that must exist even if they get folded away at inference.
  • A different attention parameterization, not simply a different stored tensor shape. Whether total parameters go up or down depends on the dimensions.
  • A separate positional path with its own cache buffer, its own projections, and its own concatenation in the score.
  • Specialised execution. The absorbed path is a genuinely different computation from the naive one — different matmul shapes, different memory access, two cache buffers instead of one. It is the kind of thing that lives behind an implementation switch, not a config flag.

A question this post deliberately skipped

Real MLA also applies low-rank compression to the queries: c^Q = W_DQ h, then up-project. Why did that never appear above?

Because it solves a different problem. Queries are not historical state — each one is used at the step that produced it and then discarded. Compressing them reduces activation memory during training and changes nothing about the KV-cache bill, which is what this post is about. Leaving it out of the derivation keeps the causal chain intact.

What to carry away

  • The conventional K/V bill is a menu — and MLA changes the menu. \(2 L n_{kv} d_h\) has levers for how many key/value representations you cache and how wide each is; MQA and GQA reduce \(n_{kv}\), while MLA replaces explicit K/V storage entirely with one learned joint latent plus a small positional state.
  • Caching a latent is the easy half. A naive implementation rebuilds every historical key and value at every step, converting a memory problem into an arithmetic one.
  • Absorption is just associativity. qᵀW_UK c = (W_UKᵀq)ᵀc on the key side, and on the value side the up-projection folds all the way into W_O. Algebraically identical — but not bit-identical, because reassociation changes the rounding; the key-side expressions differ by a few times 1e-06 in fp32.
  • RoPE destroys it. The matrix to absorb becomes R_tᵀR_j W_UK, which depends on the cached token’s position. Using absorption anyway gave 42–100% relative error across heads here, at correlations of 0.69–0.92 — wrong, but plausible-looking.
  • Decoupled RoPE is a different architecture, not a patch. It keeps a small positional branch outside the compressed path so that content compression stays absorbable.
  • The rotary key is shared across heads because it is the part that gets cached. Per-head would multiply the positional component by \(n_q\) and the total cache by about 15× at the dimensions used above.

If one sentence survives: the useful representation is not the smallest one, it is the smallest one the next operator can consume without undoing the compression.

That generalizes well past attention. It is the same argument as the quantization post’s reduction-axis section — a scale that has to be applied inside a sum is not free, and a cache that has to be decompressed before use is not compressed.

Where this goes next

Other families attack the cache differently: quantizing K and V in place, evicting tokens judged unimportant, or bounding attention to a window. Their error sources and constraints are different enough to deserve separate treatment. Converting an already-trained MHA model into an MLA one is also a distinct problem — that is where the question of whether a pretrained cache is empirically low-rank actually belongs, and it does not arise here, because MLA learns the bottleneck as part of the architecture.

References