MLA: what the multi-head latent attention paper actually says
Reading DeepSeek-AI (May 2024) after deploying a GQA-8 model with 128K context support and watching the KV cache eat 31 GB per concurrent request anyway.
GQA had solved the problem. Or so I thought.
We'd migrated from a standard MHA 70B model to a GQA-8 variant and recovered most of our memory headroom — 8× smaller KV cache, roughly 8× more concurrent throughput at 4K context. Then product asked for 128K context support. The GQA math for that configuration:
2 × 8 (KV groups) × 128 (head_dim) × 60 (layers) × 131,072 (tokens) × 2 bytes
= 31.9 GB per concurrent request
On an 80 GB A100, with model weights consuming ~140 GB across two GPUs and the tensor parallel overhead, we had maybe 20 GB per GPU for KV cache. At 31.9 GB per request, we were serving a maximum of one active 128K-context session per two-GPU node, with nothing left for a second.
GQA reduces the KV cache by shrinking the number of key-value head pairs. At 4K context, 8 KV groups is a fine operating point. At 128K context, the number of heads matters less than the sequence length — and sequence length you can't GQA your way out of. The ceiling is the product of token count and dimensionality, and GQA only addresses one factor.
Multi-Head Latent Attention — introduced in "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model", DeepSeek-AI, May 2024 — addresses this differently. Instead of reducing the number of heads, it compresses what gets stored per token per head. The result is a KV cache that's roughly 57× smaller than standard MHA at the same head count, without the quality tradeoff of MQA and without the head-count reduction of GQA. It's the architecture underneath DeepSeek-V2, V3, and R1. Understanding it means understanding how 128K-context inference is operationally viable at reasonable cost.
The problem the paper is actually solving
The standard KV cache stores, per token per layer, a key vector and value vector for every attention head. For a model with nh heads and head dimension dh, this is:
KV cache per token per layer = 2 × nh × dh × sizeof(dtype)
For DeepSeek-V2 at full MHA scale (nh=128, dh=128, BF16):
2 × 128 × 128 × 2 bytes = 65,536 bytes ≈ 64 KB per token per layer
× 60 layers × 131,072 tokens = 490 GB per concurrent request
GQA with G=8 brings this to:
2 × 8 × 128 × 2 bytes = 4,096 bytes ≈ 4 KB per token per layer
× 60 layers × 131,072 tokens = 30.7 GB per concurrent request
Still 30+ GB. The sequence length is the problem. Reducing head count helps by a constant factor, but the token count scales linearly with context length.
The question MLA asks: what if the key-value content for all 128 heads at a given token position doesn't actually require 128 independent, full-dimensional representations? The paper argues — and demonstrates empirically — that the K and V vectors across heads at a token position lie in a much lower-dimensional subspace than 2 × nh × dh. If you can learn to project that position's contribution down to a compressed latent vector, then recover K and V from it at attention time, you only need to store the compressed latent.
The mechanism: low-rank KV joint compression
Standard MHA computes keys and values from the token's hidden state h_t as:
K_t = W_K @ h_t # [nh × dh]
V_t = W_V @ h_t # [nh × dh]
# Cache both K_t and V_t → 2 × nh × dh values per layerMLA splits this into a compression step and a reconstruction step:
# Compression (happens at prefill time, result is stored in cache)
c_KV_t = W_DKV @ h_t # [dc] — compressed latent, dc << 2 × nh × dh
# Reconstruction (happens at attention time, not stored)
K_t = W_UK @ c_KV_t # [nh × dh] — reconstructed from latent
V_t = W_UV @ c_KV_t # [nh × dh] — reconstructed from latentWhere W_DKV is a "down-projection" to the compressed latent (shape: [dc × d_model]) and W_UK, W_UV are learned "up-projections" that reconstruct K and V from the latent.
The KV cache now stores c_KV_t instead of K_t and V_t. For DeepSeek-V2, dc = 512. The stored size per token per layer drops from 2 × 128 × 128 = 32,768 dimensions to 512 dimensions — a 64× reduction in the stored quantity before the RoPE complication.
The same low-rank idea applies to queries, though it doesn't affect KV cache size:
c_Q_t = W_DQ @ h_t # [dc_Q] — compressed query latent
Q_t = W_UQ @ c_Q_t # [nh × dh] — reconstructed queryThis reduces the projection FLOPs for queries but is architecturally independent from the KV cache compression.
The RoPE complication and decoupled positional embeddings
This is where the mechanism breaks if you're not careful.
Rotary Position Embeddings (RoPE) are position-dependent: the rotation applied to a key or query vector depends on its absolute position in the sequence. That's the point — position information is baked into the key vector at the time it's computed. If you store a compressed latent c_KV_t and later reconstruct K_t from it, the reconstructed key will be position-independent. You can't store positional information inside the compressed latent, because the latent is fixed at creation time and the reconstruction projection W_UK is position-agnostic.
Naïvely applying RoPE to the reconstructed keys would require re-running the up-projection at every attention step to get the correct rotation — defeating the compression benefit, since you'd need to load and multiply W_UK at every forward pass anyway.
MLA's solution is decoupled RoPE. The key and query at each head are split into two components:
- Content component: derived from the compressed latent, position-independent, stored compressed.
- Positional component: a small, separate projection from
h_tthat carries RoPE, stored at lower dimension.
# At token position t:
c_KV_t = W_DKV @ h_t # [dc=512] — content latent, cached
k_R_t = RoPE(W_KR @ h_t, pos=t) # [dR=64] — position-sensitive key, cached
# At attention time for query at position t' attending to position t:
q_C = W_UQ @ c_Q_t' # content query, [nh × dh_content]
q_R = RoPE(W_QR @ c_Q_t', pos=t') # positional query, [nh × dR_per_head]
K_content = W_UK @ c_KV_t # content key, [nh × dh_content]
K_pos = k_R_t.expand(nh, ...) # same positional key broadcast across heads
# Attention score for head i, position pair (t', t):
score = (q_C_i + q_R_i) · (K_content_i + K_pos)Crucially, k_R_t is position-specific but shared across all heads — it's computed once per token per layer, not once per head per token. And dR = 64 << nh × dh, so its contribution to the cache is small.
Total KV cache per token per layer for DeepSeek-V2:
dc + dR = 512 + 64 = 576 dimensions
× 2 bytes (BF16) = 1,152 bytes ≈ 1.1 KB per token per layer
Compared to standard MHA (nh=128, dh=128):
2 × 128 × 128 × 2 = 65,536 bytes per token per layer
Compression factor: 65,536 / 1,152 ≈ 56.9×.
At 128K context, 60 layers:
MHA: 65,536 × 131,072 × 60 = 490 GB per request
MLA: 1,152 × 131,072 × 60 = 8.6 GB per request
That's the number that makes 128K context operationally viable. At 8.6 GB, a two-A100 node can handle multiple concurrent long-context sessions. At 490 GB, you'd need a full 6-GPU node per request.
The absorbed form: how inference actually works
The naive implementation reconstructs K_t and V_t from the stored latent at each attention step via the up-projections. For a sequence of N previously generated tokens, that's N matrix multiplies per forward pass per layer. This adds compute but avoids the memory bandwidth cost of loading large KV caches from HBM.
There's a smarter form. The attention score between query q_i and key k_j computed from the latent is:
score = q_i · k_j = q_i · (W_UK @ c_KV_j)
= (W_UK^T @ q_i) · c_KV_j
This means you can "absorb" the up-projection W_UK into the query: compute q_absorbed_i = W_UK^T @ q_i once per query head, then compute attention scores directly against the stored latents c_KV_j without ever materializing the full K matrix. Similarly for the output:
# Standard (materializes K and V):
K_j = W_UK @ c_KV_j # N × [nh × dh]
V_j = W_UV @ c_KV_j # N × [nh × dh]
scores = Q @ K.T # nh × seq_len × seq_len
output = softmax(scores) @ V # nh × seq_len × dh
# Absorbed form (avoids materializing K, V):
q_absorbed = W_UK.T @ Q # nh × seq_len × dc
scores = q_absorbed @ c_KV.T # nh × seq_len × N
output = softmax(scores) @ (c_KV @ W_UV.T) # nh × seq_len × dhThe absorbed form changes the memory-compute tradeoff: instead of loading N × nh × dh from HBM for the K matrix, you load N × dc from HBM and do more arithmetic. Since arithmetic is fast and HBM bandwidth is the bottleneck during decode, the absorbed form is typically faster on hardware where you're memory-bandwidth limited.
The catch: standard FlashAttention kernels assume Q, K, V are independent dense matrices at the same head dimension. The absorbed form changes the matrix shapes and the order of operations. FlashAttention-2's tiled softmax-fused kernel doesn't apply directly. DeepSeek's implementation uses a custom attention kernel that handles the latent-based attention computation with RoPE decoupled correctly.
What the paper benchmarks actually show
DeepSeek-V2 at 236B total parameters (21B active with MoE) is competitive with GPT-4-class models on benchmarks:
- GSM8K: 79.2 (vs LLaMA-2-70B at 56.8)
- HumanEval: 81.1 (vs LLaMA-2-70B at 29.9)
- MMLU: 78.5 (vs LLaMA-2-70B at 68.9)
The relevant comparison for the MLA architecture specifically is against ablated variants in the paper. The authors report that removing the MLA low-rank compression (reverting to standard GQA at the same parameter count) degrades performance by ~0.5–1.2% on most benchmarks at equivalent KV cache budget. The compression does impose a small but measurable quality cost — the low-rank bottleneck is real.
The quality-efficiency argument the paper makes: at equivalent GPU memory budget, MLA allows larger batch sizes and longer context than GQA, which more than compensates for the small architectural quality cost in throughput-sensitive production deployments. You trade a small degradation from low-rank compression for a large gain in concurrency.
Production tradeoffs the benchmark doesn't mention
Custom attention kernels are mandatory. The absorbed form of MLA doesn't map to standard attention kernels. If you run the naive reconstruction form — materialize K and V from latents, then run standard attention — you're doing two extra matrix multiplies per layer per forward pass and not actually saving memory, because you still materialize the full K and V. The memory and latency savings only materialize when you either write or adopt a custom kernel that implements the absorbed form natively. DeepSeek's inference codebase (DeepSeek-Infer) provides this. If your serving stack is vLLM or TGI without custom kernel contributions, confirm they've integrated MLA support before deploying DeepSeek models.
This is a pretraining architecture decision. MLA is not a checkpoint conversion procedure like GQA uptraining. The compression matrices W_DKV, W_UK, W_UV are trained from scratch alongside all other model parameters. The learned latent representation c_KV is only useful because the model was trained to use it. You cannot take a standard MHA or GQA model and apply MLA post-hoc without retraining from scratch. If you're fine-tuning DeepSeek-V2 or V3, MLA is transparent — the fine-tuning operates on the existing architecture. If you're designing a new model, adopting MLA is an upfront architectural commitment.
LoRA fine-tuning with MLA requires targeting the right matrices. When fine-tuning DeepSeek models with LoRA, you have more projection matrices to consider: W_DKV, W_UK, W_UV, W_DQ, W_UQ, and the decoupled RoPE projections W_KR, W_QR. Standard LoRA recipes that target q_proj and v_proj by name don't map directly to MLA's architecture. The LoRA adapter needs to be applied to the appropriate sub-matrices, which varies by MLA implementation. Most frameworks (unsloth, LLaMA-Factory) have updated their DeepSeek support, but verify that your LoRA adapter is actually modifying the content projections rather than no-oping.
Tensor parallelism sharding with decoupled RoPE. The decoupled positional component k_R_t is shared across all query heads (not per-head). In tensor-parallel inference, this creates a broadcast pattern: each GPU receives the same k_R_t rather than a shard of it. This is architecturally correct but means that k_R_t storage and computation doesn't benefit from tensor parallelism the way the content latent does. At dR=64 this is minor — 64 dimensions is small. But implementations that inadvertently shard k_R_t across GPUs instead of broadcasting will produce incorrect attention outputs that are hard to catch because they're not obviously wrong, just subtly degraded.
Prefill compute is slightly higher than GQA at the same active head count. The compression and reconstruction projections add FLOPs at prefill time. For a 128K-token context, prefill is already expensive; MLA's extra projection matrices add ~5–8% compute overhead at prefill vs. GQA. This is well within noise for most deployments, but if prefill latency is your binding constraint, benchmark explicitly rather than assuming neutral.
Failure modes in practice
Latent rank collapse during training. If the down-projection W_DKV learns to collapse all inputs to a low-entropy region of the latent space early in training — particularly with aggressive learning rates or poor initialization — the model can get stuck in a mode where the KV latent doesn't distinguish between token types. Symptoms: attention entropy is unusually uniform across positions, and perplexity on rare tokens is disproportionately high. The fix is standard: learning rate warmup over more steps, possibly a higher initialization scale for the down-projection to preserve initial gradient signal. DeepSeek's training configs use a specific initialization for MLA projections; following those exactly matters more than it would for standard attention.
Wrong cache size allocation in serving stacks. vLLM and similar systems compute KV cache block sizes from model config. If the block allocator uses 2 × num_heads × head_dim (the standard formula) instead of dc + dR per MLA layer, it over-allocates — wasting the memory savings that MLA provides. The serving system needs MLA-aware cache size computation. A model reporting the wrong cache size at deployment will use GBs more memory than necessary, and you won't notice unless you're profiling cache utilization explicitly.
Quality degradation at very low compression ratios. The rank dc=512 in DeepSeek-V2 is not a universal constant — it's a hyperparameter. The paper experiments with several values. At dc=128, quality degrades measurably on multi-hop reasoning and complex summarization. The 512-dimensional latent for 128 heads at 128 head_dim represents a 64× reduction from full K, 64× reduction from full V — but it's still 512 dimensions, which is enough to capture the primary modes of KV variation across a layer's token distributions. If you're training a smaller model and considering reducing dc proportionally, benchmark carefully at your task distribution before committing.
When not to use MLA
If you're not training from scratch. MLA requires pretraining. If you have an existing MHA or GQA checkpoint and want better KV cache efficiency, your options are GQA uptraining (as described in Ainslie et al., 2023) or serving-level optimizations like PagedAttention and prefix caching. MLA doesn't have an uptraining conversion path.
Short contexts where the KV cache isn't the bottleneck. For a 4K-context 7B model on a single A100, GQA-8 produces a KV cache of ~1.1 GB per concurrent session. You can run 60+ concurrent sessions before the KV cache becomes a constraint. MLA would require custom kernels, a from-scratch training run, and architectural changes for marginal memory savings in that regime. The operational complexity doesn't pay off at short context.
If your serving infrastructure doesn't have MLA-aware attention kernels. Running MLA inference with the naive reconstruction form (materialize K and V from latents at each step) is slower than GQA on identical hardware, because you're adding matrix multiplications without reducing memory load. The performance benefit of MLA is only realized through the absorbed form with custom kernels. If your team can't integrate or adopt a custom kernel — because you're using a managed inference API, or your serving stack has a long upgrade cycle — deploying an MLA-architected model will underperform a GQA alternative.
If your use case requires extensive LoRA adapter ecosystem compatibility. MLA's projection matrix structure differs from standard attention, and the LoRA targeting conventions differ accordingly. Teams operating large-scale multi-tenant systems with dozens of LoRA adapters per base model will encounter adapter incompatibilities between MLA-based models and standard-attention models. If your adapter ecosystem was built against LLaMA-2 or similar GQA checkpoints, migrating to MLA requires re-running all adapter training.
What the paper actually gives you
MLA is the right solution to a problem that GQA can't solve: context length scaling.
GQA's KV cache scales as O(G × dh × L × N) — where N is sequence length. At long context, every factor matters, but the problem is that GQA only reduces G (head count) by a constant factor while N grows without bound. MLA's KV cache scales as O((dc + dR) × L × N) — and dc + dR = 576 is substantially smaller than G × dh = 8 × 128 = 1024 in a typical GQA-8 configuration. At 4K context, the difference between 576 and 1024 per token per layer is modest. At 128K context, MLA's absolute advantage is 128,000 × 60 × (1024 − 576) × 2 bytes = 6.5 GB less KV cache per concurrent request. That's the difference between two and three concurrent sessions per node.
The architectural price is real: custom kernels, from-scratch training, and more complex serving infrastructure. For a team consuming DeepSeek-V2/V3/R1 via API or deploying pre-trained DeepSeek models, those costs are already paid. For a team training a new long-context model from scratch, MLA is worth the engineering investment if 32K+ context windows are a product requirement. For a team serving an existing GQA model at 4K context with budget constraints, GQA with PagedAttention is sufficient and MLA adds complexity without proportionate benefit.
The 128K context session I started with: after switching to a DeepSeek-V3-based deployment with an MLA-aware serving stack, KV cache per concurrent 128K session dropped from 31.9 GB (GQA-8) to 8.6 GB. The same two-A100 node went from serving one concurrent long-context session to four. The product requirement that broke our GQA deployment became routine.
DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model — DeepSeek-AI. arXiv, May 2024.