Multi-Head Latent Attention: DeepSeek's KV-Cache Design

Multi-Head Latent Attention: DeepSeek’s KV-Cache Design

DeepSeek’s Multi-Head Latent Attention (MLA) is not a smaller MHA: it is a different attention mechanism that trades projection work and a position-decoupling trick for a much smaller per-token KV cache. This lesson verifies every claim against the DeepSeek papers, computes an original toy cache example, and distinguishes the optimized inference path from the common “just reconstruct K/V” explanation.

Prerequisites: Read lesson 13 (kv-cache-explained) for cache mechanics and lesson 15 (mha-mqa-gqa) for head-grouping tradeoffs before working through the absorption and comparison sections below.

1. Why MLA exists: the MHA cache bottleneck

In standard multi-head attention (MHA), each layer stores a key and value for every attention head and every token. DeepSeek-V2 (arXiv:2405.04434, §2.1.1) writes the MHA projection as:

q_t = W^Q h_t,    k_t = W^K h_t,    v_t = W^V h_t

with W^Q, W^K, W^V ∈ ℝ^{d_h·n_h × d}. During generation, all n_h keys and values for every cached token must be read; the cache per token scales with 2 · n_h · d_h · l. DeepSeek-V2 reports a 93.3% KV-cache reduction relative to DeepSeek 67B (dense MHA), but that figure is a model-to-model comparison under the paper’s own hyper-parameters, not a universal ratio (arXiv:2405.04434, §2.1.4, Table 1; §3.2.3). We will derive a separate, explicitly toy example below and never treat the paper’s reported reduction as applying to arbitrary models.

2. Joint low-rank compression of K and V

The core of MLA (§2.1.2) is a single down-projection that produces a latent vector for both keys and values together:

c_t^{KV} = W^{DKV} h_t          # W^{DKV} ∈ ℝ^{d_c × d}, d_c ≪ d_h·n_h
k_t^C    = W^{UK} c_t^{KV}      # W^{UK} ∈ ℝ^{d_h·n_h × d_c}
v_t^C    = W^{UV} c_t^{KV}      # W^{UV} ∈ ℝ^{d_h·n_h × d_c}

Only c_t^{KV} is cached during generation. Its size is d_c · l elements per token, versus 2 · n_h · d_h · l for MHA. DeepSeek-V2 sets d_c = 4·d_h and the decoupled key dimension d_h^R = d_h / 2 (§2.1.3), so the total cached elements are (d_c + d_h^R)·l = 4.5·d_h·l, comparable to GQA with 2.25 groups (arXiv:2405.04434, Table 1). Note that the comparison in Table 1 uses model-specific settings (d_h, n_h, l); those dimensions are not universal specifications.

The projection matrices W^{DKV}, W^{UK}, and W^{UV} are learned parameters. The down-projection reduces activation memory during training; the up-projections reconstruct the full key/value shapes before attention. This is the first important qualification: compression is not free. There is extra compute at every layer, and the reconstructed tensors must still be formed (even if optimized away) before the attention scores are computed.

3. Decoupled RoPE and why ordinary RoPE breaks absorption

DeepSeek-V2 uses rotary position embedding (RoPE). If RoPE is applied directly to the reconstructed keys k_t^C, then W^{UK} is coupled with a position-sensitive RoPE matrix. Because matrix multiplication does not commute, W^{UK} can no longer be pre-combined with W^Q (§2.1.3): a position-dependent RoPE matrix would sit between them. The paper writes this clearly:

“RoPE is position-sensitive for both keys and queries. If we apply RoPE for the keys k_t^C, W^{UK} … will be coupled with a position-sensitive RoPE matrix. In this way, W^{UK} cannot be absorbed into W^Q any more during inference.” (arXiv:2405.04434, §2.1.3)

The solution is a decoupled RoPE strategy. A separate shared key k_t^R carries RoPE independently:

k_t^R = RoPE(W^{KR} h_t)       # W^{KR} ∈ ℝ^{d_h^R × d}

This k_t^R is also cached (together with c_t^{KV}), giving a total cached size of (d_c + d_h^R)·l elements (§2.1.3, §2.1.4). The concatenated attention key is:

k_t,i = [k_t,i^C; k_t^R]

The separate k_t^R is required, not optional. It carries the positional signal that the compressed latent key k_t^C cannot provide without destroying the absorption property. You cannot omit the RoPE key or treat it as absorbable; its separation is the structural reason the optimization works.

4. What the absorption trick actually does during decode

During inference (§2.1.3), the optimized path does not reconstruct full k_t and v_t explicitly for all cached tokens. Instead:

• W^{UK} is absorbed into the query projection: the query computation uses a modified W^Q that includes the up-projection for the latent key path.
• W^{UV} is absorbed into the output projection W^O: the value reconstruction is folded into the final linear map.

In practice, this means the model computes attention scores using q_t against the cached c_j^{KV} and k_j^R directly, without materializing full k_j and v_j tensors for the entire context. The paper states: “during inference, since W^{UK} can be absorbed into W^Q, and W^{UV} can be absorbed into W^O, we even do not need to compute keys and values out for attention.” (arXiv:2405.04434, §2.1.2; see also DeepSeek-V3, arXiv:2412.19437, §2.1.1, where the same equations are adopted with blue-boxed cached vectors c_t^{KV} and k_t^R).

This is the critical distinction from the simple explanation: the optimized path avoids full reconstruction of K and V, not by ignoring it, but by reordering operations so that the up-projections are applied at query-time and output-time rather than at key-time.

Qualification of the trick

The absorption is valid only when the matrix order and positional constraints are preserved:

• It requires W^{UK} to be applied without an intervening position-dependent matrix. If RoPE is applied to the reconstructed key, the absorption fails.
• It requires the separate k_t^R to be cached independently. If k_t^R is omitted, positional encoding is lost.
• It does not eliminate projection work; it relocates it. The down-projection W^{DKV} must still be computed for the new token, and the modified query/output projections are larger operations than standard MHA.

This is a constrained optimization, not a universal property of any compressed attention mechanism.

MLA architecture

Compression path (left), decoupled RoPE (center-top), concatenated attention (right), and the decode-time absorption annotations (bottom). Only c_t^KV and k_t^R are cached; W_UK and W_UV are absorbed into W_Q and W_O.

Read this diagram as text

A four-zone diagram. Zone 1: input h_t down-projected by W_DKV to latent c_t^KV (cached). Zone 2: c_t^KV up-projected by W_UK and W_UV to k_t^C and v_t^C; separate h_t path through W_KR + RoPE produces cached k_t^R. Zone 3: concatenated q_t = [q_t^C; q_t^R] and k_t = [k_t^C; k_t^R]. Zone 4: attention produces o_t; annotations note W_UK absorbed into W_Q and W_UV absorbed into W_O during decode.

5. Per-layer cache calculation: a toy, explicitly illustrative example

The following numbers are toy dimensions chosen to make ratios transparent. They are not DeepSeek-V2 or DeepSeek-V3 specifications; the paper’s actual hyper-parameters are larger (e.g., DeepSeek-V2: d_h = 128, n_h = 128, d_c = 512, d_h^R = 64, l = 60; DeepSeek-V3 uses the same MLA architecture with n_h = 128, d_h = 128, d_c = 512, d_h^R = 64, §2.1.1).

Stated toy dimensions: B = 2, S = 128, FP16 (b = 2 bytes), H_q = 8, d_head = 16, GQA H_kv = 2, MLA d_c = 24 and d_R = 8. These are arbitrary teaching values.

Formula from lesson 13 (per-layer, including both K and V):

bytes = 2 × H_kv × d_head × S × B × b     # MHA / GQA (full head storage)

For MHA (H_kv = H_q = 8):

2 × 8 × 16 × 128 × 2 × 2 = 131,072 bytes

For GQA (H_kv = 2):

2 × 2 × 16 × 128 × 2 × 2 = 32,768 bytes

For MLA, the cached elements are the latent vector c_t^{KV} (dimension d_c) and the shared RoPE key k_t^R (dimension d_R), per token, per layer, per batch item, in FP16:

MLA bytes = (d_c + d_R) × B × S × b
           = (24 + 8) × 2 × 128 × 2
           = 32 × 512
           = 16,384 bytes

Split by cached component:

• Latent c_t^{KV}: 24 × 2 × 128 × 2 = 12,288 bytes
• RoPE key k_t^R: 8 × 2 × 128 × 2 = 4,096 bytes

Ratio check (from these toy dimensions only):

• MHA / MLA = 131,072 / 16,384 = 8.0 → MLA is 8× smaller than this MHA cache.
• GQA / MLA = 32,768 / 16,384 = 2.0 → MLA is 2× smaller than this GQA cache.

These ratios depend entirely on the stated toy dimensions. Changing d_c, d_R, S, B, or b changes the ratios; they are not universal properties of MLA.

6. Cost, quality, and why cache ratios are not enough

MLA introduces three categories of overhead that must be evaluated before any architecture claim:

1. Projection work: The down-projection W^{DKV} is computed for every new token. The up-projections W^{UK} and W^{UV} must be applied (either explicitly or through the absorbed query/output paths) before attention. This is additional compute relative to MHA, where W^K and W^V are single linear maps.
2. Implementation and kernel costs: The absorption requires custom kernels or fused operations that combine query projection with latent up-projection, and output projection with value reconstruction. Not all serving frameworks implement this fusion; without it, the optimization is partially lost.
3. Quality impact requires measurement: DeepSeek-V2 reports that MLA achieves stronger performance than MHA in the paper’s evaluation (§2.1.1, §3.2.2), but that comparison is made under DeepSeek-V2’s specific training setup (8.1T tokens, MoE architecture, RL alignment). DeepSeek-V2’s ablation (Appendix D) compares MHA, GQA, MQA, and MLA under the same architecture; Appendix D.2 compares MLA and MHA specifically. Those results are model-specific. You should not assume MLA improves quality unconditionally; you should measure perplexity or task accuracy on your dataset and compute budget before adopting it.

The paper reports inference efficiency gains: DeepSeek-V2 achieves 5.76× maximum generation throughput and 93.3% KV-cache reduction relative to DeepSeek 67B (arXiv:2405.04434, Abstract; §3.2.3). These are comparisons between DeepSeek-V2 (MoE + MLA) and DeepSeek 67B (dense MHA), not between MHA and MLA in isolation. The throughput improvement includes the MoE architecture’s sparse activation (only 21B of 236B parameters activated), not only the attention change.

DeepSeek-V3 (arXiv:2412.19437, §2.1.1) adopts the same MLA equations (c_t^{KV}, k_t^C, k_t^R, v_t^C, absorption into W^Q and W^O), confirming the design is maintained at larger scale (n_h = 128, d_h = 128, d_c = 512, d_h^R = 64). DeepSeek-V3 does not claim universal superiority of MLA; it uses MLA because it was validated in DeepSeek-V2 (§2.1.1: “thoroughly validated by DeepSeek-V2”).

7. Sources, evidence limits, and what the papers actually claim

• DeepSeek-V2, arXiv:2405.04434 (§2.1.1-2.1.4, Table 1, Appendix C, Appendix D): defines MLA equations, explains the RoPE incompatibility, states the absorption trick, and provides ablation comparisons against MHA, GQA, and MQA under the same model configuration. The cache comparison in Table 1 uses d_c = 4·d_h and d_h^R = d_h / 2.
• DeepSeek-V2, Appendix C: full formulas of MLA (referenced for the complete derivation, not reproduced here).
• DeepSeek-V2, Appendix D.2: direct comparison between MLA and MHA. The paper reports MLA performs stronger than MHA in its setup; this is a model-specific result.
• DeepSeek-V3, arXiv:2412.19437 (§2.1.1): adopts the same MLA equations with the same cached vectors (c_t^{KV} and k_t^R marked in blue), confirming continuity but making no new universal claims.
• The reported 5.76× throughput and 93.3% KV reduction (arXiv:2405.04434, Abstract) compare DeepSeek-V2 (236B total, 21B active, MLA + DeepSeekMoE) against DeepSeek 67B (dense MHA, no MoE). These are not MHA-vs-MLA ratios in isolation.

If you implement MLA, verify: (a) the cached tensors are exactly c_t^{KV} and k_t^R, (b) W^{UK} and W^{UV} are either reconstructed or properly absorbed, (c) RoPE is applied only through W^{QR} on c_t^Q (for the query RoPE component) and through W^{KR} on h_t (for k_t^R), and (d) quality is measured on your dataset, not inferred from the paper’s benchmark scores.

8. Exit exercise

Given the toy dimensions from section 5 (B = 2, S = 128, FP16, H_q = 8, d_head = 16, GQA H_kv = 2, MLA d_c = 24, d_R = 8), answer the following. The answers can be checked from the arithmetic in section 5; they do not require external sources.

1. Compute the per-layer MHA, GQA, and MLA cache bytes. Show the formula substitution for each.
2. State the MLA-to-MHA and MLA-to-GQA ratios. Confirm that MLA is 8× smaller than MHA and 2× smaller than GQA for these dimensions.
3. Split the MLA cache into the latent c_t^{KV} component and the RoPE-key k_t^R component. Verify the split sums to the total MLA bytes.
4. If you increase S from 128 to 256 (same dimensions otherwise), by what factor does each cache change? Explain why the ratio between variants stays the same.
5. State one reason why a model with 8× smaller KV cache does not necessarily have 8× faster decode latency. Reference either the projection work, kernel fusion, or the fact that the paper’s 5.76× throughput compares DeepSeek-V2 (MoE + MLA) against DeepSeek 67B (dense MHA), not MHA-only to MLA-only.

Exercise: MLA cache arithmetic and qualified claims

1. 01Using B=2, S=128, FP16 (b=2), H_q=8, d_head=16, compute per-layer MHA (H_kv=8), GQA (H_kv=2), and MLA (d_c=24, d_R=8) cache bytes. Show each formula and result.

2. 02State the MLA-to-MHA and MLA-to-GQA ratios from these results, and split MLA into its latent (c_t^KV) and RoPE-key (k_t^R) components.

3. 03If S doubles to 256 with all other dimensions unchanged, how does each cache change? Why does the variant ratio stay constant?

4. 04Give one specific reason (from projection work, kernel fusion, or the paper comparison scope) why 8x cache reduction does not guarantee 8x faster decode.

9. Related lessons

• Lesson 6 (multi-head-attention): projection, splitting, merging, and output projection mechanics that MLA modifies.
• Lesson 13 (kv-cache-explained): per-layer cache growth and the raw byte formula used in this lesson’s arithmetic.
• Lesson 15 (mha-mqa-gqa): grouping rules and the quality-evaluation framework needed before adopting any attention variant.
• Lesson 20 (architecture-comparison): a later decision framework; this lesson does not link it since it is not yet published.

Related lessons

• LLM Internals & Architecture · Attention mechanisms from scratchMulti-Head Attention: Splitting, Merging, and Projecting14 min
• LLM Internals & Architecture · Generation, KV cache, and contextKV Cache Implementation and Memory Accounting14 min
• LLM Internals & Architecture · Modern LLM architecturesMHA vs. MQA vs. GQA: Attention Architecture Tradeoffs14 min