DEV Community

jidonglab
jidonglab

Posted on

Multi-head Latent Attention: The KV Cache Trick Beyond GQA

Grouped-query attention buys you a smaller KV cache by making heads share keys and values. It works, but it's a blunt instrument: you throw away head diversity to save memory, and at some grouping ratio quality starts to slide. Multi-head Latent Attention (MLA), the scheme DeepSeek-V2 and DeepSeek-V3 are built on, does something categorically different. It keeps full per-head expressiveness and caches a compressed representation — then computes attention without ever decompressing it. That last clause is the part almost every explainer gets wrong.

TL;DR

  • Multi-head Latent Attention (MLA) caches a single low-rank latent vector per token instead of full per-head keys and values, then reconstructs K and V on the fly via up-projection matrices.
  • The real win isn't the projection — it's matrix absorption: the up-projections fold into the query and output weights, so you compute attention scores directly in latent space and never materialize the full K/V. Less memory and less compute per decode step.
  • RoPE breaks the absorption trick because rotation is position-dependent and won't commute past the up-projection. MLA fixes this with decoupled RoPE: a small separate set of dimensions carries positional information, concatenated onto the compressed content path.
  • DeepSeek reported roughly a 93% KV cache reduction versus its dense MHA baseline, with large generation-throughput gains — while matching or beating MHA quality.
  • Use MLA reasoning when you're memory-bound on long-context decode and GQA's grouping is already hurting quality. It's an architecture choice, not a runtime flag.

What is Multi-head Latent Attention (MLA)?

MLA is an attention variant that replaces the per-token KV cache with a compressed latent. In standard multi-head attention (MHA), every token you generate appends 2 * n_heads * head_dim values to the cache — one key vector and one value vector per head. That cache is what makes autoregressive decode memory-bound: every new token reads the entire cache back from HBM.

MLA changes what you store. Instead of caching keys and values, it caches one low-rank vector c_KV per token (plus a small positional piece, more on that below). Keys and values for every head are functions of that shared latent:

c_KV = W_DKV @ h        # down-project hidden state h -> latent (dim d_c, e.g. 512)
k    = W_UK  @ c_KV     # up-project latent -> per-head keys
v    = W_UV  @ c_KV     # up-project latent -> per-head values
Enter fullscreen mode Exit fullscreen mode

W_DKV is a down-projection into a bottleneck of dimension d_c that is far smaller than n_heads * head_dim. W_UK and W_UV are up-projections back to full width. Only c_KV lives in the cache.

Why can't GQA just shrink the cache further?

GQA shrinks the cache by reducing the number of distinct KV heads. Sixteen query heads might share four KV heads (a 4:1 group), cutting KV memory 4x. Push the ratio harder — 8:1, 16:1 — and you're forcing more and more query heads to attend against the same key/value projections. At the limit of one KV head you have multi-query attention, which is known to cost quality on harder tasks.

The constraint is structural: GQA's only knob is head count, and head count is coupled to representational capacity. You can't decouple "how much I cache" from "how many independent projections I have." MLA does exactly that decoupling. The cache size is governed by d_c, a bottleneck dimension you set independently, while every head keeps its own W_UK/W_UV slice. You get MHA-grade head diversity at a cache footprint smaller than aggressive GQA.

How does MLA compress the KV cache without losing heads?

The trick is low-rank factorization applied to the cache, not the weights. Keys and values across heads are highly correlated — they're all linear functions of the same hidden state. So instead of storing the redundant expanded form, you store the shared latent c_KV once and keep the per-head structure in the (uncached, static) up-projection weights.

Concretely, DeepSeek-V2 uses a latent dimension d_c = 512 for a model whose full KV width is many times larger. The per-head keys and values are never persisted; they're regenerated from the 512-dim latent whenever attention runs. Because W_UK and W_UV are model parameters, not per-token state, they cost you nothing in the cache — they're already resident.

If MLA stopped here, it would trade memory for compute: you'd save cache but pay to decompress K and V every step. The next part is why it doesn't.

Why doesn't MLA decompress K and V at inference?

Because the up-projection matrices can be absorbed into the surrounding weights, so attention runs directly on the compressed latent. This is the mechanism that makes MLA a genuine win rather than a memory-for-compute swap.

Look at the attention score between query q and key k for a single head, ignoring positional encoding for a moment. The query comes from its own down/up path, q = W_UQ @ c_Q, and the key is k = W_UK @ c_KV. The score is:

score = qᵀ k
      = (W_UQ c_Q)ᵀ (W_UK c_KV)
      = c_Qᵀ (W_UQᵀ W_UK) c_KV
Enter fullscreen mode Exit fullscreen mode

W_UQᵀ W_UK is a product of two static weight matrices. You precompute it — call it W_absorbed — and now the score is computed straight from the cached latent c_KV and the query latent c_Q. The full-width key k is never materialized. The same absorption works on the value side: W_UV folds into the output projection W_O, so the attention output is produced directly from latents too.

# Naive: reconstruct then attend (what MLA avoids)
k = W_UK @ c_KV                 # full per-head key, big
scores = q.T @ k

# Absorbed: attend in latent space (what MLA does)
W_absorbed = W_UQ.T @ W_UK      # precomputed once, static
scores = c_Q.T @ W_absorbed @ c_KV   # touches only the 512-dim latent
Enter fullscreen mode Exit fullscreen mode

The consequence: per-head attention math effectively happens against a 512-dim latent instead of the full KV width. You cache less and the dot products read less. This is why MLA improves decode throughput rather than just memory.

Why does RoPE break MLA, and what is decoupled RoPE?

RoPE breaks the absorption because rotary position encoding applies a position-dependent rotation to keys, and that rotation refuses to commute past the up-projection. Here's the collision. With RoPE, the score between a query at position m and a key at position n is:

score = (R_m q)ᵀ (R_n k) = qᵀ R_{n-m} k
Enter fullscreen mode Exit fullscreen mode

R_{n-m} is a rotation matrix that depends on the relative position n - m. Now substitute k = W_UK c_KV:

score = qᵀ R_{n-m} W_UK c_KV
Enter fullscreen mode Exit fullscreen mode

You wanted to precompute something like W_UQᵀ W_UK. But R_{n-m} sits between the query weights and W_UK, and it changes with every query/key position pair. There is no single static matrix to absorb; you'd have to materialize a rotated key per position, which reintroduces the full-width K you were trying to avoid.

MLA's fix is decoupled RoPE. Split the key into two parts:

  • A content part derived from the compressed latent c_KV, carrying no position — this stays absorbable.
  • A small positional part k^R that carries RoPE, generated directly from the hidden state and shared across all heads.
k^R    = RoPE(W_KR @ h)         # small, e.g. 64 dims, shared across heads
k_head = [ W_UK @ c_KV ; k^R ]  # concat: content (no RoPE) + position (RoPE)
Enter fullscreen mode Exit fullscreen mode

The query gets a matching split: a latent-derived content part and a small RoPE-carrying part. Attention scores are the sum of a content term (computed via absorption from c_KV) and a positional term (computed from the tiny k^R). The positional dimensions are few — 64 in DeepSeek-V2 — and shared, so the cache addition is small. You cache c_KV plus one shared k^R per token, and absorption still works on the content half.

How much memory does MLA actually save?

Per token, MLA caches d_c + d_R values — the content latent plus the decoupled RoPE key. In DeepSeek-V2 that's 512 + 64 = 576 values per token per layer. Compare that against what MHA would store: 2 * n_heads * head_dim, which for a model of that class is many thousands of values per token. A back-of-envelope comparison:

def kv_bytes_per_token(layers, values_per_token, dtype_bytes=2):
    return layers * values_per_token * dtype_bytes

# Illustrative shapes for a large model
layers = 60

mha  = kv_bytes_per_token(layers, 2 * 128 * 128)   # 2 * n_heads * head_dim
mla  = kv_bytes_per_token(layers, 512 + 64)        # d_c + d_R

print(mha / mla)   # ~57x smaller cache in this illustrative setup
Enter fullscreen mode Exit fullscreen mode

The exact ratio depends on architecture, but the shape of the result is real: DeepSeek reported roughly a 93% KV cache reduction relative to its dense MHA baseline and a large boost to maximum generation throughput. Crucially, they also reported MLA matching or beating MHA on quality — the low-rank bottleneck acts as a mild regularizer rather than a lossy compromise, unlike aggressive GQA grouping.

At long context this compounds. KV cache scales linearly with sequence length; shrinking the per-token footprint ~50x is the difference between a 128K-token conversation fitting on one accelerator or spilling. Smaller cache also means fewer bytes read per decode step, and since decode is memory-bandwidth-bound, that reads through directly to tokens per second.

When should you reach for MLA in production?

MLA is an architecture decision, not a serving flag you toggle — you get it by choosing a model built on it (DeepSeek-V2/V3 and derivatives) or by designing one. If you're serving open-weights models and your bottleneck is long-context decode memory, MLA-based models let you hold far more concurrent sequences or far longer contexts per GPU than a comparable GQA model at equal quality. If you're already on a GQA model and pushing the grouping ratio to claw back memory at a measurable quality cost, that's the exact regime MLA was designed to escape. If your workloads are short-context and prefill-dominated, the KV cache isn't your constraint and MLA's advantage is muted.

One caveat: matrix absorption changes the compute shape, and not every inference stack implements the absorbed path. Naive implementations that reconstruct K and V per step get MLA's memory savings but pay decompression compute, erasing the throughput win. Check that your runtime actually attends in latent space.

The direct answer

Multi-head Latent Attention (MLA) is a KV cache compression scheme that stores one low-rank latent vector per token instead of full per-head keys and values, reconstructing K and V through up-projection matrices that fold into the query and output weights so attention runs directly on the compressed latent — less memory and less compute per step. Because rotary position encoding is position-dependent and won't absorb past those projections, MLA carries positional information in a small separate decoupled-RoPE path concatenated onto the content latent. The result, in DeepSeek-V2/V3, is roughly a 93% smaller KV cache than dense MHA with equal-or-better quality — a cleaner trade than GQA, which can only shrink the cache by sacrificing head diversity.

Top comments (0)