Multi-Head Latent Attention

§ 1 · Premise

KV cache is the autoregressive tax

In a decoder Transformer of $L$ layers with $H$ heads of dimension $d_h$, each generated token appends one key and one value vector per head per layer to a cache that the next step has to read. For DeepSeek-67B’s geometry — $L = 95$, $H = 64$, $d_h = 128$, fp16 — the cache cost is $2 \cdot L \cdot H \cdot d_h \cdot 2\text{ B} \approx 3.1\text{ MB per token}$. A 32K context holds about 100 GB of cache before any model parameters are loaded, and the cache, not the FFN, dominates inference-time memory bandwidth at long context. DeepSeek-V2 quantifies the same problem on its 236B-parameter geometry: 1.6 MB of cache per token at $L = 60$, $H = 128$, $d_h = 128$ (DeepSeek-V2, Table 1).

The three predecessors all chase this number, with diminishing returns.

• MHA (Vaswani et al., 2017) gives every query head its own K and V; the cache scales with $H \cdot d_h$. Quality is the reference, memory is the worst.
• MQA (Shazeer, 2019) keeps a single shared K, V across all heads. Cache shrinks by $H\times$, but quality drops measurably on multi-task evals: GQA’s authors report MQA losing 0.7 points of average score vs. MHA at the 11B scale (Ainslie et al., 2023, Table 2).
• GQA (Ainslie et al., 2023) interpolates: $G$ groups of query heads share a K, V pair. At $G = 8$ on Llama-2-70B the cache is $8\times$ smaller than MHA while quality is within noise.

The shared move in MQA and GQA is to throw away K, V diversity. MLA keeps the diversity but stores a compressed representation: at each token, cache one low-rank latent and reconstruct every head’s K and V on the fly. The reconstruction matrices are folded into adjacent projections at inference, so the per-token memory cost is the latent — not the per-head vectors.

Two complications keep this from being a one-line description. First, the natural place to apply RoPE — to each head’s K — is incompatible with the latent representation, because rotating after up-projection forces the up-projection to be re-applied per cached token. Second, the same trick that lets MLA evade that cost (absorbing the K up-projection into the query at inference) only works when the query side is also restructured to keep position out of the content channel. Both points are derived below.

§ 2 · Derivation

From MHA to a cached latent

MHA baseline. Let $\mathbf{h}_t \in \mathbb{R}^{d}$ be the residual-stream vector at position $t$. A standard MHA layer with $H$ heads of dimension $d_h$ computes, for head $i \in \{1, \dots, H\}$:

with $W^{(i)}_Q, W^{(i)}_K, W^{(i)}_V \in \mathbb{R}^{d_h \times d}$. The per-token cache footprint is $2 \cdot H \cdot d_h$ scalars per layer — the two vectors $\mathbf{k}^{(i)}_t$ and $\mathbf{v}^{(i)}_t$ stacked across heads. For $H = 128$, $d_h = 128$, $L = 60$ (DeepSeek-V2) this is $1{,}966{,}080$ scalars or about 3.84 MB at fp16 per token.

Low-rank latent. MLA replaces the per-head $W^{(i)}_K, W^{(i)}_V$ with a shared down-projection to a single latent $\mathbf{c}^{KV}_t$ of dimension $d_c \ll H \cdot d_h$, and per-head up-projections that reconstruct K and V at attention time (DeepSeek-V2, eq. 9–11):

with $W_{DKV} \in \mathbb{R}^{d_c \times d}$ and $W^{(i)}_{UK}, W^{(i)}_{UV} \in \mathbb{R}^{d_h \times d_c}$. The superscript $C$ marks the “content” part of the key, distinguished below from a position-aware part. Only $\mathbf{c}^{KV}_t$ enters the cache: $d_c$ scalars per token per layer instead of $2 H d_h$. DeepSeek-V2 picks $d_c = 512$, so the latent is $\tfrac{1}{64}$ the size of the stacked per-head K and V.

Why this is not just GQA with $G = 1$. MQA also caches one K and one V per layer, of dimension $d_h$ each — $2 d_h = 256$ scalars at DeepSeek-V2’s $d_h = 128$. The MLA latent is 512, twice as large in raw scalars, but expresses a $d_c$-dimensional subspace from which every head’s K and V are independently reconstructed by its own $W^{(i)}_{UK}, W^{(i)}_{UV}$. MQA collapses head diversity at the cache level; MLA collapses only the rank, leaving each head free to read the latent through its own linear map. The empirical consequence is in § 4.

Query-side compression. DeepSeek-V2 also down-projects queries to $\mathbf{c}^{Q}_t = W_{DQ} \mathbf{h}_t \in \mathbb{R}^{d_c'}$ with $d_c' = 1536$, then up-projects per head: $\mathbf{q}^{C,(i)}_t = W^{(i)}_{UQ} \mathbf{c}^{Q}_t$ (DeepSeek-V2, eq. 6–8). The query compression saves training-time activation memory; it is not cached, since queries are recomputed at every step. The choice $d_c' = 3 d_c$ is reported without a sensitivity study.

The RoPE obstruction. Rotary position embeddings (Su et al., 2021) multiply $\mathbf{q}^{(i)}_t$ and $\mathbf{k}^{(i)}_t$ by a position-dependent rotation $R_t \in \mathbb{R}^{d_h \times d_h}$ before the inner product. With latent compression, the natural attention score for head $i$ between positions $t$ (query) and $s$ (key) is

The trick that makes MLA cheap at inference is to absorb $W^{(i)}_{UK}{}^{\top}$ into the query-side matrix $W^{(i)}_{UQ}$ once, offline: define $\tilde{W}^{(i)}_Q = W^{(i)}_{UK} W^{(i)}_{UQ}$ and read the cached $\mathbf{c}^{KV}_s$ directly. But if RoPE rotates the per-head K after the up-projection, the score becomes $(R_t \mathbf{q}^{C,(i)}_t)^{\top} R_s W^{(i)}_{UK} \mathbf{c}^{KV}_s$ — the $R_s$ depends on $s$, so the up-projection cannot be precomposed: every cached token would need its $W^{(i)}_{UK}$ reapplied on the fly, defeating the purpose. Rotating before the down-projection mixes positions into the latent in a way that cannot be unrotated per head.

Decoupled RoPE head. MLA’s resolution is to split each key into a content part (no RoPE, reconstructed from the latent) and a small position-aware part (RoPE applied, cached separately) (DeepSeek-V2, eq. 12–13):

where $W_{KR} \in \mathbb{R}^{d_R \times d}$ produces a single $d_R$-dimensional rotary key shared across all heads, with $d_R = 64$ in DeepSeek-V2. The query is split symmetrically: $\mathbf{q}^{(i)}_t = [W^{(i)}_{UQ} \mathbf{c}^{Q}_t ;\, R_t W^{(i)}_{QR} \mathbf{c}^{Q}_t]$. The content channel carries the absorbable up-projection; the rotary channel carries position with a small extra cost.

Per-token cache after decoupling. The cache stores $\mathbf{c}^{KV}_t$ and the single shared $\mathbf{k}^{R}_t \in \mathbb{R}^{d_R}$, i.e., $d_c + d_R = 512 + 64 = 576$ scalars per layer per token. Across 60 layers that is $34{,}560$ scalars or ~67.5 KB per token at fp16, versus 3.84 MB for MHA on the same geometry — a $\sim 57\times$ reduction at the layer-cache level (DeepSeek-V2, Table 1 reports the equivalent figure as 4.5% of MHA cache after grouping by representation precision).

Parameter count. Setting aside biases, an MLA layer’s KV-side weights total

For DeepSeek-V2 ($d = 5120$, $H = 128$, $d_h = 128$, $d_c = 512$, $d_R = 64$) this is $5120 \cdot 512 + 2 \cdot 128 \cdot 128 \cdot 512 + 5120 \cdot 64 \approx 19.4\text{M}$ parameters per layer for the K/V path. Standard MHA at the same geometry uses $2 \cdot d \cdot H \cdot d_h = 2 \cdot 5120 \cdot 128 \cdot 128 \approx 167.8\text{M}$, so MLA’s KV-side weight count is ~12% of MHA’s. The query-side compression adds another $d \cdot d_c' + H \cdot d_h \cdot d_c' \approx 33.6\text{M}$ parameters per layer.

Computational complexity. Per layer, MLA’s attention compute for sequence length $T$ is $\mathcal{O}(T^2 H d_h)$ for the score and value-application — identical to MHA in the asymptotic sense — plus $\mathcal{O}(T \cdot d \cdot d_c)$ for the latent down-projection and $\mathcal{O}(T \cdot H \cdot d_h \cdot d_c)$ for the per-head up-projections. The absorption trick hides the up-projection behind the query: at inference, the cached $\mathbf{c}^{KV}_s$ is read directly as a $d_c$-vector and the merged $\tilde{W}^{(i)}_Q$ produces head-$i$ queries that score against it without ever materializing $\mathbf{k}^{C,(i)}_s$. The flop count is unchanged; the memory traffic is what shrinks.

§ 3 · Reference implementation

Sketch

# Shapes: B batch, T tokens, d model dim, H heads, d_h head dim,
#         d_c KV latent dim, d_c_q Q latent dim, d_R RoPE head dim.

# --- projections (per layer) ---
c_kv = x @ W_DKV                          # [B, T, d_c]            <- cached
k_rope = rope(x @ W_KR, pos)              # [B, T, d_R]            <- cached
c_q = x @ W_DQ                            # [B, T, d_c_q]          (not cached)

# Up-projections used at attention time. At inference these get folded:
# tilde_W_Q[i] = W_UK[i] @ W_UQ[i], so queries can score against c_kv directly.
k_content = c_kv @ W_UK                   # [B, T, H, d_h]         (logical only)
v         = c_kv @ W_UV                   # [B, T, H, d_h]

q_content = c_q  @ W_UQ                   # [B, T, H, d_h]
q_rope    = rope(c_q @ W_QR, pos)         # [B, T, H, d_R]

# Stack the content + rope channels into a (d_h + d_R)-wide key/query.
k = concat([k_content, broadcast(k_rope, H)], dim=-1)   # [B, T, H, d_h + d_R]
q = concat([q_content, q_rope], dim=-1)                 # [B, T, H, d_h + d_R]

# Standard scaled dot-product over the concatenated dims.
attn = softmax(q @ k.transpose(-1, -2) / sqrt(d_h + d_R))   # [B, H, T, T]
out  = (attn @ v) @ W_O                                     # [B, T, d]

§ 4 · Empirical evidence

Ablations and scaling

Quality vs. MHA, MQA, GQA at fixed scale. DeepSeek-V2’s Table 9 holds the model at 7B activated parameters and varies only the attention variant. MLA reports 50.7 average across the listed benchmarks vs. 50.0 for MHA, with MLA winning on BBH (+1.5), MMLU (+0.4), and C-Eval (+1.5) while losing 0.3 on the GSM8K subset; MQA at 47.5 average and GQA at 48.6 both trail MHA on this setup (DeepSeek-V2, Table 9). The cache reduction at the same table is from 110.6 KB/token (MHA) to 33.75 KB/token (MLA), or 4.5% of the MHA budget after DeepSeek’s representation-precision accounting — a 3.3× reduction at the comparable-quality point.

Throughput. DeepSeek-V2 reports a 5.76× maximum generation throughput improvement over DeepSeek-67B (which used GQA at $G = 8$) on the same H800 cluster (DeepSeek-V2, §1 and Figure 1). The throughput gain conflates MLA with the DeepSeekMoE FFN; the paper does not isolate the MLA-only fraction.

Carry-over to V3 and V3.1. DeepSeek-V3 keeps $d_c = 512$, $d_R = 64$, $d_c' = 1536$, $H = 128$, $d_h = 128$ across 61 layers with $d = 7168$ (DeepSeek-V3 Technical Report, §2.1.1). No re-ablation against MHA at V3 scale is published; the V3 report cites the V2 ablations as the basis for keeping MLA. DeepSeek-V3.1’s HF model card reports the same attention configuration carried over.

Scaling to V3.2 + sparse attention. DeepSeek-V3.2 (arXiv 2509.04559) keeps MLA intact and adds a DeepSeek Sparse Attention (DSA) module that uses a Lightning Indexer to select a top-K subset of the cached latents (k=2048) for each query. The V3.2 report’s Table 2 shows DSA on top of MLA matching V3.1’s quality on most benchmarks while cutting prefill cost materially at long context — the MLA cache is what makes the sparse indexer cheap to score, since each candidate is a single $d_c + d_R$ vector rather than $H$ per-head K vectors. This is the clearest scaling evidence: the same MLA shape extends from 236B to 671B parameters and composes with a sparse selector layered above it.

Reproductions and independent reports. Kimi-K2 (Moonshot AI repo) adopts the MLA pattern with the decoupled RoPE head; the public model card describes the same $d_c = 512, d_R = 64$ shape. Kimi-Linear-48B (arXiv 2510.26692, Table 1) uses MLA on 7 of 27 layers in a hybrid stack with KDA linear attention on the rest, reporting that those 7 MLA layers contribute the bulk of long-context recall accuracy on RULER. No independent academic reproduction at the DeepSeek-V2 scale is public as of 2026-05-12; the available evidence is the introducing labs and their successors.

Training-time vs. inference-time cost. The empirical reports emphasize per-token cache and generation throughput, both of which are inference-time quantities. The training-time picture is different: MLA’s per-step compute matches MHA asymptotically (still $\mathcal{O}(T^2 H d_h)$) and adds a low-rank projection that costs a fraction of a percent of total FLOPs at $H d_h \gg d_c$. The DeepSeek-V2 report does not provide a training-time ablation isolating MLA’s wall-clock cost from the DeepSeekMoE FFN that runs alongside it; the V3 report similarly treats the two as a single joint design choice. Practitioners reproducing MLA from the equations should expect training-step time comparable to MHA on the same geometry; the wins are exclusively on the autoregressive side.

What is not in the public record. The sensitivity of quality to $d_c$ (would $d_c = 256$ work? $d_c = 1024$?) is not published. The sensitivity to $d_R$ — and whether the decoupled head can be removed entirely with a NoPE-style positional scheme — is also not studied in the DeepSeek reports. At kernel level, the inference fusion of $W^{(i)}_{UK}$ into $\tilde{W}^{(i)}_Q$ that makes MLA cache-efficient is described in DeepSeek-V2 §2.1.2 as a recipe but without throughput numbers isolated from the rest of the inference stack. The choice of $d_c' = 3 d_c$ for the query latent is likewise unmotivated in print. “I don’t know” is the honest answer to all three.

@article{arxiv2405_04434,
  title  = {DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model},
  author = {DeepSeek-AI},
  year   = {2024},
  eprint = {2405.04434},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2405.04434}
}

@article{arxiv_2405_04434,
  title         = {DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model},
  author        = {DeepSeek-AI},
  year          = {2024},
  eprint        = {2405.04434},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2405.04434}
}

{
  "id": "arxiv_2405_04434",
  "type": "article-journal",
  "title": "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model",
  "author": [
    {
      "literal": "DeepSeek-AI"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2024
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/2405.04434",
  "number": "2405.04434",
  "source": "arXiv"
}

TY  - JOUR
TI  - DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
AU  - DeepSeek-AI
PY  - 2024
JO  - arXiv
AN  - arXiv:2405.04434
UR  - https://arxiv.org/abs/2405.04434
ER  -