Multi-head Latent Attention: the KV cache, compressed

DeepSeek-AI, 2024

The sequel to the KV-cache page. Same formula, one term attacked from a completely different direction.

Download mla_from_scratch.py — everything below runs from it. NumPy only, under a second.

The KV-cache page ends on the number that decides what you can serve: bytes per token. I'm not going to re-derive the cache here — go read that first if you haven't, because everything below assumes it. This page is about one line in that formula.

Multi-query and grouped-query attention attack kv_heads by sharing keys and values across heads, and the GQA paper's own ablations show them losing quality doing it — that's the trade Llama-2 70B took with GQA-8. MLA doesn't touch kv_heads. It changes what a cache entry is.

Store the archive, not the files

You've made this call before. A directory of files is too big to keep around, so you tar and gzip it: 

store(tar_gz(files))            # small on disk
files = untar(gunzip(archive))  # pay CPU to get them back

The usual objection is that decompression costs you on every read. Now the twist. Suppose the only thing you ever do with the decompressed files is feed them into one specific function — and that function is a matrix multiply, and so is the decompression: 

scores = (x @ W_q) @ (c @ W_uk).T     # decompress c into keys, then dot
scores = x @ (W_q @ W_uk.T) @ c.T     # same thing. Fold the two matrices
W_abs  = W_q @ W_uk.T                 # once, at load. Now:
scores = (x @ W_abs) @ c.T            # the keys are NEVER built

That's the whole paper. The decompression step folds into the operation you were going to do anyway, so at inference it costs nothing. Matrix multiplication being associative is not a deep fact. Noticing that it makes a 57× cache reduction free is.

TermMeansProgrammer's analogy
Latent cOne small vector per token per layer, x @ W_dkvThe gzip archive. The only thing on disk
dcHow many numbers the latent has. 512 in DeepSeek-V2Compression level
Wuk, WuvUp-projections: latent → keys, latent → valuesgunzip, twice, for two different outputs
AbsorptionFolding Wuk into Wq once at loadConstant folding two adjacent multiplies
Decoupled RoPE dimsA small extra block that carries position, cached separatelyThe one field you can't compress because it's indexed

The one formula

Not an attention formula. The cache formula from the previous page, and what MLA does to it:

MHA: 2 × layers × kvheads × headdim × seq × dtype
MLA:     layers × (dc + drope)     × seq × dtype

Two things changed. Read them as terms:

TermValueWhat it does
2goneK and V are compressed jointly. One latent serves both, so there is no separate V to store
kv_heads × head_dimgoneReplaced by dc. The latent is shared across every head, so adding heads no longer adds cache
dc512The latent width. This is the dial. DeepSeek-V2 runs 512 against 128 × 128 = 16384 expanded dimensions
drope64The decoupled position block. Also shared across heads. Exists only because RoPE breaks the trick — see below
layers, seq, dtypeunchangedStill the model, still your users, still fp16

The table that is the whole argument

Same arithmetic as the KV-cache page — I reuse its kv_cache_bytes function unchanged — plus one new line for MLA. Two shapes: a Llama-2-7B one you already know, and a DeepSeek-V2 one.

Don't take my word for it. Here's the actual output:

Llama-2-7B shape: layers=32 heads=32 head_dim=128 d_c=512 d_rope=64, fp16
    multi-head     :    524,288 bytes/token   1.0x
    grouped-query 8:    131,072 bytes/token   4.0x smaller
    MLA            :     36,864 bytes/token  14.2x smaller
    MLA vs GQA-8   :   3.6x smaller

DeepSeek-V2 shape: layers=60 heads=128 head_dim=128 d_c=512 d_rope=64, fp16
    multi-head     :  3,932,160 bytes/token   1.0x
    grouped-query 8:    245,760 bytes/token  16.0x smaller
    MLA            :     69,120 bytes/token  56.9x smaller ✓
    MLA vs GQA-8   :   3.6x smaller

Look at what happens when the head count goes from 32 to 128. Multi-head cache goes up 7.5×, to 3.75 MiB per token per sequence — unservable, which is why nobody builds that. GQA-8 only doubles, and that doubling comes from the layer count (32 to 60), not the heads — it pinned kv_heads at 8. MLA doesn't move at all: 576 floats per layer regardless of how many heads read them. That invariance is the property, and it's what lets DeepSeek run 128 heads in the first place.

The MLA-vs-GQA number is 3.6× at both shapes, and that's the honest comparison. The eye-catching 57× is against a multi-head baseline nobody would have shipped at 128 heads. DeepSeek report their own reduction against DeepSeek-V1's actual architecture, and separately report MLA outperforming GQA in their ablations at comparable cache. Quote both carefully; they are different claims.

Is it still attention? At full rank, exactly

A cache that changes the answer isn't a cache. So the script builds one toy model, runs textbook multi-head attention through it, then runs the same inputs through the down-projection and up-projection, and diffs.

One honesty note: to make the rank sweep mean anything I factor the joint [W_k | W_v] with an SVD and shape its singular values to decay, because trained weight matrices have a decaying spectrum and freshly randomised ones don't. 

d_model=32 n_heads=4 head_dim=8
full-rank latent d_c = 32
  max |standard - through the latent| = 7.772e-15   <- float noise ✓
  identical to 1e-10?                 True

now shrink the latent and watch what you are actually trading.
errors are relative: max abs error / max abs value.
   d_c | K reconstruction | attention output | cache bytes/token
    32 |           0.0000 |           0.0000 |         64 B
    24 |           0.0130 |           0.0113 |         48 B
    16 |           0.0429 |           0.0540 |         32 B
    12 |           0.1549 |           0.1228 |         24 B
     8 |           0.1952 |           0.1889 |         16 B
     4 |           0.5808 |           0.6061 |          8 B
     2 |           0.8738 |           0.8946 |          4 B

At full rank it's not an approximation at all — it's a lossless re-encoding, and 7.8 × 10−15 is float64 addition order. Below it you can watch the price. Halving the cache costs about 5% on the output; quartering it costs 60%. The output error tracks the key error closely, sometimes under and sometimes over, because softmax renormalises whatever it's given.

That curve is the entire design decision. Everything about picking dc = 512 is finding where it turns, on real learned weights instead of a shaped-spectrum toy. The mechanism is trivial; the calibration is the work.

Absorption: the part that makes it free

Compressing the cache is easy and obvious. The reason it isn't a disaster is the fold. Written out per head, with Wq at (dmodel, headdim) and Wuk at (dc, headdim):

(x Wq) · (c Wuk)T  =  x Wq WukT cT  =  (x (Wq WukT)) · cT

Nothing in that expression is new attention — it is still the softmax(QKT/√d) V lookup, with the K factored. The middle term is the point: nothing in W_q @ W_uk.T depends on the data, so you compute it once when the model loads and never again. The query gets projected straight into latent space and dotted against the cache. Keys are never materialised. 

cached tokens n = 64, d_c = 16, head_dim = 8
  W_q       (32, 8)   W_uk (16, 8)
  W_absorbed = W_q @ W_uk.T -> (32, 16)  (built ONCE at load)
  expanded scores[0][:4] = [ 1.4168  1.7375  0.3522 -5.6463]
  absorbed scores[0][:4] = [ 1.4168  1.7375  0.3522 -5.6463]
  max abs difference     = 8.438e-15   <- same thing ✓

operation counts, per head per decode step:
   cached n |  expanded MACs |  absorbed MACs | saving
         64 |          8,960 |          1,536 |  5.83x
        512 |         69,888 |          8,704 |  8.03x
      4,096 |        557,312 |         66,048 |  8.44x
     32,768 |      4,456,704 |        524,800 |  8.49x

Be careful with that second table in an interview. At this toy's ratio of dc to headdim the absorbed path is genuinely fewer multiply-accumulates; at DeepSeek's real ratios it isn't always. The reliable win is that the expansion tensor never exists, so it never crosses the memory bus — and decode is memory-bound, which the previous page measured at 0.48 MACs per byte.

RoPE, and why the explanations stop here

Every blog post I read about MLA describes the compression, describes the absorption, and then quietly moves on. Here's what they skipped.

RoPE rotates q and k by an angle proportional to their position. So the score isn't (x W_q) · (c W_uk), it's (R(m) x W_q) · (R(n) c W_uk) — and R(n) is a different matrix for every cached token. There is no way to pull it out and fold it into a fixed Wq at load time. The absorption and the rotation do not commute. Not "approximately"; at all. 

scores over 8 cached tokens, head 0, with RoPE applied:
  correct (expand, then rotate) : [ 2.5354 -2.1796  0.0678 -3.9345]
  absorbed then rotated         : [ 6.6998  2.1142  2.3629 -2.7651]
  max abs difference            = 7.4962   <- wrong
  ranking preserved?              False

Not a rounding difference. The ranking of which tokens to attend to is different, which means it isn't a slightly worse model, it's a different one. DeepSeek's fix is to stop trying: split each head in two. Most dimensions carry no position and stay absorbable. A small extra block of drope = 64 dimensions carries the rotation, is computed the ordinary way, and is cached separately — one copy, shared by every head. The score is the sum of the two parts. 

DeepSeek's fix -- split the head in two:
  content part (absorbed, no position) : [ 2.2504 -0.2611 -0.1965 -0.022 ]
  decoupled RoPE part (d_rope=8)   : [ -3.7696   8.6941 -19.5117   6.6326]
  score = sum                          : [ -1.5192   8.4329 -19.7082   6.6107]
  matches the explicit two-part sum?     True, max diff 0.000e+00 ✓

what the decoupled block costs, per token, DeepSeek-V2 shape:
  latent only        :  61,440 bytes/token
  + decoupled RoPE   :  69,120 bytes/token (+7,680, +12.5%)
  still vs multi-head: 56.9x smaller
Twelve percent more cache, and the whole architecture survives. That is what the decoupled dimensions are for. Not a tuning knob, not an ablation — a structural repair to a hole that rotary embeddings punched in the absorption. If someone explains MLA to you without mentioning it, they have explained a model that doesn't work.

Where the saving actually lands

Cache bytes per token is an abstraction. The number your finance team feels is how many concurrent sequences fit on one card. Fix a budget of 80 GiB for cache alone, on the DeepSeek-V2 shape: 

   context | multi-head |    GQA-8 |      MLA | MLA seqs in 80 GiB
     4,096 |     15.00G |    0.94G |    0.26G |    303
    16,384 |     60.00G |    3.75G |    1.05G |     75
    32,768 |    120.00G |    7.50G |    2.11G |     37
   131,072 |    480.00G |   30.00G |    8.44G |      9

  sequences that fit in 80 GiB at 32k context:
    multi-head  :      0 concurrent sequences
    GQA-8       :     10 concurrent sequences
    MLA         :     37 concurrent sequences ✓

At 32k context this shape cannot serve a single multi-head sequence in 80 GiB. GQA-8 fits 10. MLA fits 37. Per-token latency is identical in all three — the weights are in flight anyway — so that's roughly 3.7× the tokens per second per GPU, for the same silicon. This is the argument, and a lab that serves its own model at scale is exactly who would make it.

What it costs, said plainly

MLA is not free, and I'd rather you had the honest version before an interviewer hands it to you. Per attention layer, at DeepSeek-V2's shape: 

per attention layer, d_model=5120, 128 heads x 128:
  multi-head    params:   335,544,320
  GQA-8         params:   178,257,920
  MLA           params:   229,441,536  (0.68x multi-head, 1.29x GQA-8)

  cache at 32,768 tokens, batch 1:
  multi-head :  120.00 GiB
  GQA-8      :    7.50 GiB
  MLA        :    2.11 GiB ✓

Against 128-head multi-head, MLA is actually cheaper in parameters, because the latent is a bottleneck. Against GQA-8 it costs about 1.29×, plus arithmetic: an extra projection on every token and a wider per-head dimension once the RoPE block is bolted on. The trade is a fixed cost you pay once and share across the whole batch, in exchange for a per-sequence per-token cost. That's a good deal at serving scale and a bad one at batch 1 on your laptop.

One more thing to keep straight, because it's the easiest place to be confidently wrong. DeepSeek also apply a low-rank compression to the queries. That one has nothing to do with the KV cache — queries are never cached, as the previous page's break-it table showed. It reduces activation memory during training. Same technique, different problem, and the papers are clear about it even though summaries usually aren't.

A caveat on quality. That MLA matches or beats grouped-query attention on quality at comparable cache size is DeepSeek's reported claim from their own ablations. Nothing on this page measures it. What I measured here is bytes, MACs and reconstruction error — the parts that fit on a laptop.

The code

Pure NumPy. No torch, no training — a forward pass and some byte counting, which is all a linear-algebra paper needs. Runs in about a seventh of a second. Every function under 20 lines.

Download the full script — 567 lines plus 7 runnable demos. Then: 

python3 mla_from_scratch.py
def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, dtype_bytes=2,
                   batch=1):
    """Standard KV cache. Same formula as the KV-cache page, unchanged."""
    return 2 * layers * kv_heads * head_dim * seq_len * dtype_bytes * batch


def mla_cache_bytes(layers, d_c, d_rope, seq_len, dtype_bytes=2, batch=1):
    """MLA cache: one latent per token per layer, plus the decoupled RoPE key.
    Both are shared across ALL heads, which is where the saving comes from."""
    return layers * (d_c + d_rope) * seq_len * dtype_bytes * batch


def low_rank_factors(W, rank):
    """W (m, n) -> (W_down (m, rank), W_up (rank, n)) with W ~ W_down @ W_up."""
    U, S, Vt = np.linalg.svd(W, full_matrices=False)
    return U[:, :rank] * S[:rank], Vt[:rank]


class ToyMLA:
    def __init__(self, d_model, n_heads, head_dim, d_c, seed=0):
        rng = np.random.default_rng(seed)
        self.d_model, self.n_heads, self.head_dim = d_model, n_heads, head_dim
        self.d_c = d_c
        d_all = n_heads * head_dim

        self.W_q = rng.normal(0, 0.4, (d_model, d_all))

        # The joint K/V projection. I shape its singular values to DECAY,
        # because a trained weight matrix has a decaying spectrum and a
        # freshly randomised one does not. Without this the rank sweep below
        # is meaningless -- every direction matters equally, so truncation is
        # uniformly catastrophic and you learn nothing from it.
        raw = rng.normal(0, 0.4, (d_model, 2 * d_all))
        U, _, Vt = np.linalg.svd(raw, full_matrices=False)
        spectrum = 6.0 * (0.82 ** np.arange(min(U.shape[1], Vt.shape[0])))
        W_kv = (U * spectrum) @ Vt                            # (d, 2*d_all)
        self.W_k, self.W_v = W_kv[:, :d_all], W_kv[:, d_all:]

        # Compress K and V JOINTLY -- one latent serves both. That joint-ness
        # is the paper's word and it is why there is no factor of 2 above.
        self.W_dkv, W_ukv = low_rank_factors(W_kv, d_c)
        self.W_uk = W_ukv[:, :d_all]      # (d_c, d_all)
        self.W_uv = W_ukv[:, d_all:]      # (d_c, d_all)

    def latent(self, X):
        """The ONLY thing that gets cached: (n, d_c)."""
        return X @ self.W_dkv

    def attention_via_latent(self, X):
        """MLA the naive way: compress, then re-expand K and V and carry on.
        Correct, but it materialises the full K and V -- no saving at all
        except in what you STORE. Demo 4 removes the expansion."""
        C = self.latent(X)                           # (n, d_c)  <- cached
        Q = self.heads(X @ self.W_q)
        K = self.heads(C @ self.W_uk)                # re-expanded
        V = self.heads(C @ self.W_uv)                # re-expanded
        w = softmax(Q @ K.transpose(0, 2, 1) / np.sqrt(self.head_dim))
        out = w @ V
        return out.transpose(1, 0, 2).reshape(X.shape[0], -1)


def absorb_q_into_latent(W_q_head, W_uk_head):
    """(d_model, head_dim) x (d_c, head_dim) -> (d_model, d_c). Done ONCE."""
    return W_q_head @ W_uk_head.T


def scores_absorbed(x, W_absorbed, C):
    """The MLA way: project into latent space, dot against the cache."""
    q_c = x @ W_absorbed                 # (1, d_c)
    return q_c @ C.T                     # (1, n)


def decoupled_score(x_q, pos_q, C, positions, W_absorbed, W_qr, K_rope):
    """DeepSeek's score: absorbable content part + decoupled RoPE part.

    K_rope is (n, d_rope), already rotated and cached -- ONE copy for all
    heads. That sharing is why the extra cache is cheap."""
    content = (x_q @ W_absorbed) @ C.T                       # no position
    q_r = apply_rope(x_q @ W_qr, np.array([pos_q]))          # (1, d_rope)
    return content + q_r @ K_rope.T

Now break it

This is the part that actually makes it stick. Reading won't do it. You debug for a living — and two of the rows below are bugs that raise nothing, pass every shape assertion you'd think to write, and just quietly give you a worse model.

Break thisWatch for
Set the latent too small: ToyMLA(..., d_c=2) instead of full rank Key reconstruction error climbs to 0.8738 relative, attention output to 0.8946, and the absolute output error is 3.0672. The sweep in demo 2 shows where it turns: output is still within about 5% at d_c=16, past 60% wrong at d_c=4.
Cache the expanded keys and values instead of the latent — swap mla_cache_bytes for kv_cache_bytes(60, 128, 128, 1) 69,120 bytes per token becomes 3,932,160. That is 56.9× worse, and it is the entire point of the paper gone. You have paid for the projections and kept none of the benefit.
Apply RoPE with no decoupled dimensions: apply_rope(x @ W_abs, pos) @ C.T Scores go from [-4.5353 18.2531 -10.3902] to [-8.3733 4.9212 2.3044] — max difference 28.9771, and ranking preserved? False. Different attention, not noisier attention.
Absorb the wrong matrix: W_q @ W_uv.T instead of W_q @ W_uk.T Identical shape (32, 16). No exception. Scores off by 22.4897. Meanwhile the genuinely reversed order, W_uk @ W_q.T, is just the transpose and is harmless — 0.0000. Shape checks are not correctness checks.
Rebuild the absorbed matrix inside the decode loop instead of once at load Still correct, and pointless: 536,870,912 extra multiply-accumulates over 32,768 steps on a four-head toy. On a real model this is the difference between absorption being free and absorption being a tax.
Then read the papers — narrowly. In DeepSeek-V2, the architecture section on MLA: the compression, then the decoupled RoPE subsection, then the KV-cache comparison table. Read the query compression paragraph too, and note where it says training. In DeepSeek-V3, just enough to confirm they kept MLA unchanged and to see it alongside their MoE work. Skip the MoE, the routing, the training infrastructure and every benchmark table — all real, all a different paper.

On the maths anxiety

This one is genuinely just linear algebra, and less of it than the attention page needed. You need three things: that a matrix product is associative, that a matrix can be written as a thin product of two smaller ones, and that throwing away the smallest singular values is the best approximation of a given rank. The first is the whole trick. The second is what "low-rank" means. The third is only how I built the toy — a real model learns its factors by gradient descent and never runs an SVD.

What you do need is the byte arithmetic, and you need it fast, because the interview question is not "explain MLA". It's "you have an 80 GB card and a 32k context — what's your maximum batch size, and which architecture gets you there". That's multiplication and division on numbers you can read off a config file, and demo 5 is exactly that computation.

Where to go next

This page gave you one thing: the cache arithmetic and the absorption identity, in code you can break. It says nothing about how MLA trains, nothing about the MoE it ships inside, and nothing about how the kernels are written.

Order I'd take them in: Raschka for where MLA sits among the alternatives, then Banatt for the algebra, then the DeepSeek-V2 paper, which reads easily once you've seen both. V3 last, and only if you care that it stuck.