Multi-Head Latent Attention (MLA)

Multi-Head Latent Attention (MLA) is the mechanism DeepSeek-V2 introduced to solve the same problem Grouped Query Attention (GQA) solves, an oversized KV cache throttling decode throughput, but by a fundamentally different route. Where GQA shrinks the cache by making groups of query heads share a smaller number of whole key/value heads, MLA shrinks it by compressing every token’s keys and values, across all heads, into a single low-rank latent vector before caching, then reconstructing full-size, per-head keys and values from that latent vector on the fly at attention time. DeepSeek-V2 reported that this compresses the KV cache far below what an equivalent-quality GQA setup achieves, while matching or exceeding standard multi-head attention’s benchmark scores, a combination that made MLA the default attention layer across the DeepSeek-V2, V3, and R1 model family, and pushed other labs (GLM, Kimi, MiniMax) to adopt it or a close variant.

Why the KV Cache Needed a Different Kind of Compression

During autoregressive decoding, a transformer must keep every previous token’s key and value vectors resident in GPU memory, one slot per layer, per head, because each new token’s query attends back over all of them. That cache grows linearly with sequence length, batch size, and the number of KV heads, and reading it back off high-bandwidth memory, not the matrix multiplications themselves, is what actually limits how many tokens per second a deployment can decode at long context. GQA attacks this by cutting the number of KV heads: instead of one KV head per query head, several query heads share one. It is an effective, easy-to-implement fix, and it’s why virtually every serving-optimized open-weight model since Llama 2 ships with some GQA group count g rather than full multi-head attention.

But GQA’s dial is coarse. g has to be a positive integer that evenly divides the number of query heads, so the achievable cache reductions are quantized: 64 heads can only be grouped into 32, 16, 8, 4, 2, or 1 KV heads, not, say, an “11.3x” reduction. And because whole heads are being merged, at aggressive group counts some of what different heads used to specialize in gets averaged away. DeepSeek’s MLA paper asked a different question: instead of picking which whole heads to share, what if the cache stored the minimum information needed to reconstruct all heads, at a size that’s a free-standing hyperparameter rather than a divisor of the head count?

The Mechanism: Compress First, Reconstruct at Attention Time

MLA replaces the standard “project to keys and values, then attend” pipeline with a compress/decompress step. For each token, its hidden state is first projected down into a low-rank latent vector of dimension r (DeepSeek-V2 uses r = 512, against a would-be uncompressed key+value width of 2 × 128 heads × 128 head_dim = 32,768). That latent vector, not the full keys and values, is what gets cached. When attention is actually computed, each head’s key and value are reconstructed from the shared latent vector through a per-head up-projection matrix, so every head still gets its own distinct key and value, they’re just derived from a shared compressed representation instead of stored independently.

There’s one wrinkle: rotary positional embeddings (RoPE) are position-dependent and don’t compress cleanly through a shared low-rank bottleneck, since rotating a compressed vector doesn’t commute with decompressing it head-by-head. DeepSeek-V2’s fix, called “decoupled RoPE,” splits the query and key into two parts: a larger portion that goes through the compression/reconstruction pipeline uncompressed by position, and a small, fixed-size slice (64 dimensions in DeepSeek-V2) that carries RoPE directly and gets cached alongside the latent vector. That’s why the actual per-token cache size is r + d_rope (512 + 64 = 576 elements) rather than just r.

graph LR
    classDef default fill:#ffffff,stroke:#4338CA,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef data    fill:#EEF0F7,stroke:#0D9488,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef process fill:#F7F8FC,stroke:#6366F1,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef output  fill:#4338CA,stroke:#4338CA,stroke-width:2px,color:#ffffff,rx:8px,ry:8px;

    H([Token hidden state]):::data --> C[Compress: down-project to latent r + RoPE slice]:::process
    C --> Cache[(Cached: r + d_rope elements)]:::output
    Cache --> D1[Up-project: reconstruct K, V for head 1]:::process
    Cache --> D2[Up-project: reconstruct K, V for head 2]:::process
    Cache --> Dn[Up-project: reconstruct K, V for head h]:::process
    D1 --> A[Scaled dot-product attention, per head]:::process
    D2 --> A
    Dn --> A

# scenario: the decode-time attention forward pass for one MLA layer,
# reconstructing per-head K/V from a cached latent vector instead of
# reading h separate full-size K/V heads off HBM
import torch

def mla_attention(hidden, cached_latent, w_up_k, w_up_v, w_q, num_heads, head_dim):
    # cached_latent: (batch, seq_len, r + d_rope), the ONLY thing read from HBM
    # w_up_k, w_up_v: (num_heads, head_dim, r) per-head up-projection matrices,
    # small enough to sit in fast on-chip memory, not re-read per token
    q = w_q(hidden).view(hidden.shape[0], -1, num_heads, head_dim)
    k = torch.einsum('bsr,hdr->bshd', cached_latent, w_up_k)   # reconstruct K
    v = torch.einsum('bsr,hdr->bshd', cached_latent, w_up_v)   # reconstruct V
    return torch.nn.functional.scaled_dot_product_attention(
        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True
    )

MLA vs. GQA: A Continuous Knob vs. a Discrete One

The two techniques attack the same bandwidth problem with a different shape of trade-off. GQA’s group count g must divide the head count evenly and directly removes whole heads’ worth of distinct information; MLA’s latent rank r is a free integer that can be tuned independently of the head count, and every head still gets a reconstructed, head-specific key and value rather than a literally shared one.

TransMLA’s key result is that this relationship is directional: because MLA’s reconstruction is strictly more expressive than head-sharing (a wide enough latent can represent any GQA configuration, but not vice versa), any GQA model can be converted into an MLA model post-hoc and further compressed, whereas going the other way loses information. That reframes MLA not as a competing, incompatible design but as a generalization of GQA’s underlying idea, don’t store more key/value information per token than attention needs to read back, taken to its more expressive limit.

Try It: Latent Rank vs. KV Cache Size

The chart below uses DeepSeek-V2’s real shape (128 heads, head_dim 128, 60 layers, d_rope fixed at 64) and the exact formula for per-token, per-layer cache size: MLA caches (r + 64) elements, while an equivalent GQA setup with g KV heads caches 2 × g × 128 elements. Drag the latent rank slider down and watch the MLA cache size fall below even aggressive GQA settings, since r isn’t constrained to be a divisor of 128.

Drag r down toward 64 (the RoPE floor) to see MLA's cache undercut even MQA-tight GQA settings; drag it up toward 1024 to see where MLA stops paying off versus a small GQA group count.

The Cost: A More Involved Reconstruction Step

MLA’s compression isn’t free. Reconstructing per-head keys and values from the latent vector at every attention step requires the per-head up-projection matrices (w_up_k, w_up_v above) to be applied on every forward pass, adding matrix multiplications that GQA’s direct head-repetition doesn’t need. In practice this cost is small relative to the bandwidth saved, since the up-projection weights are tiny compared to the KV cache traffic they replace and stay resident in fast on-chip memory rather than round-tripping through HBM, but it does mean MLA kernels are more complex to write and optimize than GQA’s straightforward repeat_kv expansion. This is part of why GQA remained the default for so long even after MLA’s results were published: DeepSeek had to build custom decoding kernels to realize MLA’s theoretical bandwidth savings in practice, and it took the broader inference-serving ecosystem (vLLM, SGLang, TensorRT-LLM) time to ship equally optimized MLA kernels.

There’s also a subtler quality question. MLA forces all heads to reconstruct from one shared latent, so if attention heads truly need to specialize in mutually incompatible ways, that specialization has to be recoverable through the up-projection matrices rather than stored directly. DeepSeek-V2’s ablations found this cost negligible at their chosen rank, and the model’s benchmark scores support that, but it means the “how low can r go” question is empirical per model family rather than a closed-form guarantee the way a GQA group count’s KV cache savings are.

Where MLA Is Used Today

What’s New (2025-2026)

• TransMLA and post-hoc conversion. Rather than training MLA models from scratch, TransMLA and related work show an existing GQA checkpoint can be algebraically re-expressed as an equivalent MLA model and then further compressed, reporting up to a 10.6x KV cache reduction and a measured decode speedup at long context on a converted Llama-2-7B, without a full retrain.
• MLA reaching vision-language models. Follow-on work (MHA2MLA-VLM) extends the same latent-compression recipe to multimodal transformers, where the KV cache pressure from long visual token sequences makes the savings even more valuable than in text-only decoding.
• Hardware-centric re-analysis. As MLA adoption grew, dedicated hardware studies have examined how MLA’s reconstruction step maps onto GPU memory hierarchies and matrix units differently than GQA’s simpler head-repetition, informing kernel design for the next generation of inference accelerators rather than treating MLA as a drop-in replacement.
• MLA as a default rather than a DeepSeek-specific trick. What started as one lab’s architectural choice in DeepSeek-V2 is now treated by several other frontier labs as the more expressive alternative to GQA, evaluated by default for new frontier training runs rather than as a niche efficiency technique.

from transformers import AutoConfig, AutoModelForCausalLM

# DeepSeek-V2's actual MLA shape: 128 attention heads, but keys and
# values are cached as a single 512-dim latent plus a 64-dim
# decoupled RoPE key, not as 128 full-size key/value heads.
config = AutoConfig.from_pretrained(
    "deepseek-ai/DeepSeek-V2",
    num_attention_heads=128,
    kv_lora_rank=512,       # <- the latent rank r: this is what gets cached
    qk_rope_head_dim=64,    # <- small fixed RoPE slice, also cached
    qk_nope_head_dim=128,
    v_head_dim=128,
)
model = AutoModelForCausalLM.from_config(config)

# Per-token, per-layer cache: kv_lora_rank + qk_rope_head_dim = 576 elements,
# versus 2 * num_attention_heads * v_head_dim = 32,768 for standard MHA.