DL0107 Multi-head Latent Attention MLA

Multi-head Latent Attention (MLA) attacks the KV-cache bottleneck of autoregressive decoding by low-rank joint compression of keys and values instead of by deleting KV heads. Each token is projected down to a single shared latent vector of dimension , and that latent is the only thing cached; per-head keys and values are reconstructed on the fly by up-projection matrices and , so every head still gets its own distinct K and V, unlike GQA and MQA which force heads to share. Because RoPE does not commute with the up-projection, MLA adds a small decoupled RoPE key of dimension that is shared across heads and cached directly, giving a total cache of values per token per layer against for the equivalent MHA configuration, a roughly 57x reduction. The reason this is not simply lossy compression is the absorption trick: at inference the up-projections fold into the query and output projections, so the model never materializes the full keys and values, and the attention score becomes a bilinear form on the two latents. The result reported in DeepSeek-V2 is unusual for a cache-reduction technique, since MLA matched or beat full MHA on benchmarks while caching about 1/57 as much, which is why DeepSeek-V3 kept it unchanged at 61 layers and 128 heads.

(1) Joint Low-Rank KV Compression: one down-projection produces a single latent per token that serves all heads and both K and V, so cache cost is rather than .
(2) Head Diversity Is Preserved: per-head up-projections mean MLA reduces cache without reducing the number of effective K/V heads, which is the quality cost GQA and MQA pay.
(3) Decoupled RoPE Path: position information travels in a separate small key dimension because applying RoPE inside the compressed path would make the up-projection position-dependent and destroy absorption.
(4) Absorption At Inference: folds into and folds into , so decoding reads only the cached latent and never reconstructs full K or V tensors.
(5) Query Compression Is A Separate Win: queries are also routed through a latent of dimension 1536, which cuts training activation memory but has no effect on the KV cache.
(6) Memory-Bandwidth Trade For FLOPs: the absorbed score is a dot product over 576 dimensions instead of 128, so MLA raises arithmetic intensity, exactly the right direction for memory-bound batch decoding.

Figure 1: Only the two boxes inside the dashed region enter the cache: the 512-dimensional KV latent and the 64-dimensional decoupled RoPE key. Per-head keys and values are reconstructed from the latent, and at inference those up-projections are absorbed so they are never computed explicitly.

The decoupled RoPE dimension is the part candidates most often miss, and it follows from a single algebraic fact. The score between query and key in the compressed path is , which lets you precompute the constant matrix . If RoPE were applied to , that matrix would become , which depends on the relative position and therefore differs for every cached token, so nothing can be precomputed and every prefix key would have to be rebuilt at each step. MLA sidesteps this by carrying position in a separate 64-dimensional key that is cached in rotated form and shared by all heads, so the query concatenates a 128-dimensional content part with a 64-dimensional positional part, giving an effective head dimension of 192 during training. A practical consequence is that the value head dimension (128) differs from the query and key dimension (192), which is why MLA needs purpose-built kernels such as FlashMLA rather than a stock FlashAttention call.

Mathematical Formulation:

Where:

• is the cached KV latent and the cached decoupled RoPE key; together they are the entire per-token cache.
• is the layer input ( in DeepSeek-V3), and is the query latent of dimension 1536.
• indexes the current token, ranges over cached prefix positions with , and indexes the heads.
• is the content head dimension, the positional dimension, so the concatenated query and key have dimension 192 while values keep 128.
• are down-projections, up-projections, the RoPE-path projections, and the output projection.
• is the absorbed score matrix, valid only because no rotation sits between the two up-projections; applies the rotary transform, and denotes concatenation.

Figure 2: Totals assume 61 layers, bf16, and a 128K-token context. MLA lands between MQA and 8-group GQA on cache size, but reaches that point by compressing rather than by removing KV heads, which is why it does not carry their quality penalty.