Core ML

Multi-Head Latent Attention

Instead of caching full key and value tensors for every head, compress them into a tiny shared latent vector and reconstruct the per-head K and V on the fly — the cache stores the latent, not the expanded tensors.

Why Does This Exist?

MQA and GQA both trim the KV cache by sharing projection weight across heads. The cache shrinks by 1/H or 1/G. But you're still caching tensors whose size scales with the number of tokens times the head dimension — at long contexts and high concurrency that remains a hard ceiling.

Multi-head latent attention (MLA), introduced with DeepSeek-V2 in 2024, attacks the problem differently. Instead of sharing projections across heads, compress the K and V tensors before storing them using a low-rank down-projection, cache the compressed latent, and reconstruct per-head keys and values at attention time via an up-projection. The cache stores one $d_c$ -dimensional vector per token rather than $2H \cdot d_{head}$ per token, where $d_c$ is the latent rank — typically much smaller. The arithmetic reconstruction at query time adds a bit of compute, but you can absorb that into the weight matrix with a carefully arranged matrix product, so the effective compute overhead is small.

Think of It Like This

A compressed zip archive you unpack on demand

Imagine storing a document's key highlights not as separate copies for 32 different readers, but as a single compact summary. When a reader arrives and asks a question, you unpack the summary into the specific excerpt that reader needs. The storage cost is the compact summary — much smaller than 32 separate copies.

MLA does exactly that for the KV cache. Every token compresses its K and V information into a small latent vector. At inference time, when a new query arrives, the latent is unpacked into the per-head keys and values that query needs. What persists in memory is the small latent, not the expanded tensors.

How It Actually Works

The two projections

For each token, MLA runs a down-projection $W^{DKV}$ that maps the token's hidden state $x \in \mathbb{R}^{d_{model}}$ to a low-rank latent $c \in \mathbb{R}^{d_c}$ , where $d_c \ll H \cdot d_{head}$ :

c = W^{DKV} x

This $c$ is what gets stored in the KV cache — one vector per token, per layer, with dimension $d_c$ .

At attention time, up-projections $W^{UK}$ and $W^{UV}$ reconstruct the per-head keys and values:

K_h = W^{UK}_h \, c, \quad V_h = W^{UV}_h \, c

The full multi-head attention then proceeds as normal, using the reconstructed $K_h$ and $V_h$ .

The memory trade-off

Standard MHA caches $2H \cdot d_{head}$ floats per token. MLA caches $d_c$ floats. In DeepSeek-V2 with H=128, $d_{head}=128$ , and $d_c=512$ , that's a reduction from 32,768 floats to 512 — a 64× reduction in cache size over MHA. GQA with 8 groups on the same model would give a 16× reduction, so MLA is about 4× better than GQA here.

Absorbing the up-projection

The extra matmul at query time ( $W^{UK}_h \, c$ ) would cost $O(d_c \cdot d_{head})$ flops per head per query. But for the score computation $Q_h K_h^\top$ , you can absorb $W^{UK}$ into the query projection: compute $Q_h' = Q_h W^{UK\top}_h$ once per head during prefill, then only the inner product with $c$ is needed at decode time. This keeps the per-step compute close to GQA's.

Show Me the Code

Verifying the cache size reduction and a single round-trip through the compression-and-reconstruction.

1import numpy as np2
3rng = np.random.default_rng(0)4
5H, d_head, d_model, d_c = 8, 64, 512, 128  # small illustration6
7# Down-projection: x -> c8W_down = rng.standard_normal((d_model, d_c)) * 0.029x = rng.standard_normal((d_model,))10c = W_down.T @ x                           # shape: (d_c,)11
12# Up-projections: c -> K_h, V_h per head13W_UK = rng.standard_normal((H, d_c, d_head)) * 0.0214W_UV = rng.standard_normal((H, d_c, d_head)) * 0.0215
16K = np.einsum("hcd,c->hd", W_UK, c)       # shape: (H, d_head)17V = np.einsum("hcd,c->hd", W_UV, c)       # shape: (H, d_head)18
19# Cache sizes20standard_cache = 2 * H * d_head           # floats per token21mla_cache = d_c                            # floats per token22
23print(f"Standard KV cache (per token): {standard_cache} floats")24print(f"MLA latent cache  (per token): {mla_cache} floats")25print(f"Reduction: {standard_cache // mla_cache}×")26# Standard KV cache (per token): 1024 floats27# MLA latent cache  (per token):  128 floats28# Reduction: 8×29print(f"K shape: {K.shape}, V shape: {V.shape}")   # (8, 64) each

The reconstruction is exact given the weight matrices — no approximation, unlike sparse or linear attention.

Watch Out For

Thinking MLA is an approximation

MLA does not approximate the attention scores the way sparse attention or linear attention does. Given the same weight matrices, the reconstructed K and V are exact. The compression happens in the cache — what you store between prefill and decode — not in the attention computation itself. The attention scores are exact; you're just computing them from a compressed representation.

Underestimating the training constraint

The down-projection and up-projections have to be learned together with the rest of the model. The latent rank $d_c$ is a design choice before training, and the model has to learn to encode useful K and V information into that rank budget. You can't retrofit MLA onto a trained MHA checkpoint the way you can partially uptrain GQA — the weight structure is fundamentally different.

The Quick Version

MLA down-projects K and V into a shared low-rank latent $c$ before storing them; up-projections reconstruct per-head K and V at attention time.
Cache size drops from $2H \cdot d_{head}$ to $d_c$ floats per token — a much larger reduction than GQA.
The attention computation is exact, not an approximation — only the storage is compressed.
The up-projection can be absorbed into the query projection to keep per-step compute low.
MLA is a design choice made before training; it cannot be retrofitted to an existing MHA checkpoint.