Multi-Latent Attention (MLA)

Multi-latent attention (DeepSeek AI, 2024) replaces the per-token KV cache with a compressed latent vector. Instead of caching the full key and value projections for each token — 2 * n_kv_heads * d_head values per position — MLA caches a single low-dimensional vector c_t of dimension d_latent. At decode time, the full keys and values are reconstructed from c_t via learned up-projections. DeepSeek-V2 introduced this mechanism and DeepSeek-V3 refined it; both achieve KV cache compression ratios that exceed GQA and MQA.

The compression-reconstruction pipeline

For each input token x_t at position t, the standard KV cache stores:

$K_{t} = x_{t} @ W_{K}$ → shape (n_kv_heads * d_head,)

$V_{t} = x_{t} @ W_{V}$ → shape (n_kv_heads * d_head,)

Total cached per token: 2 * n_kv_heads * d_head values.

MLA replaces this with a two-stage process. In the encoding stage (during prefill), the input is projected into a compressed latent:

c_{t} = x_{t} @ W_{\text{DKV}}

where W_DKV has shape (d_model, d_latent) and d_latent is much smaller than 2 * n_kv_heads * d_head. Only c_t is cached — a vector of d_latent values per token.

In the reconstruction stage (during decode), keys and values are recovered from the latent:

$K_{t} = c_{t} @ W_{\text{UK}}$ → shape (n_kv_heads * d_head,)

$V_{t} = c_{t} @ W_{\text{UV}}$ → shape (n_kv_heads * d_head,)

W_UK and W_UV are learned up-projection matrices of shape (d_latent, n_kv_heads * d_head).

DeepSeek-V2 dimensions

DeepSeek-V2 (236B total parameters, 21B active via mixture of experts) uses:

$d_{\text{model}} = 5120$
$n_{\text{heads}} = 128$ (across all experts)
$d_{\text{head}} = 128$
$d_{\text{latent}} = 512$

Standard KV cache per token: $2 \cdot 128 \cdot 128 = 32, 768$ values.

MLA cache per token: 512 values (the latent c_t).

Compression ratio: $32, \frac{768}{512} = 64x$ .

For context, GQA-8 on the same model would cache $2 \cdot 8 \cdot 128 = 2048$ values per token — a 16x reduction from MHA. MLA achieves 4x more compression than GQA-8, and 64x more than MHA.

At sequence length 128K tokens in float16:

MHA: $128K \cdot 32, 768 \cdot 2 = 8 \text{GB}$
GQA-8: $128K \cdot 2, 048 \cdot 2 = 0.5 \text{GB}$
MLA: $128K \cdot 512 \cdot 2 = 0.125 \text{GB}$

Decoupled RoPE

Rotary position embedding (RoPE) applies a position-dependent rotation to query and key vectors so that the dot product q_t · k_s encodes the relative position t - s. RoPE is applied element-wise and depends on the vector dimension — it rotates pairs of adjacent dimensions by an angle proportional to the position.

The problem with MLA: RoPE must be applied to keys, but the keys are reconstructed from a compressed latent. If you apply RoPE to the latent c_t before compression, the rotation is lossy — the down-projection mixes dimensions that RoPE intended to keep separate. If you apply RoPE after reconstruction, the positional information isn't in the cache and must be recomputed, but this is feasible and is what MLA does — with a twist.

DeepSeek-V2 splits the key into two components:

Content key K_C_t: reconstructed from the latent. Contains semantic information, no positional encoding.
RoPE key K_R_t: a small additional vector (d_rope dimensions) projected directly from x_t, not compressed. RoPE is applied only to this component.

K_C_t = c_t @ W_UK                    # from latent, no RoPE
K_R_t = RoPE(x_t @ W_KR, position=t)  # separate small projection, with RoPE
K_t = concat(K_C_t, K_R_t)            # full key

The RoPE key K_R_t is also cached per token, but it's small: $d_{\text{rope}} = 64$ in DeepSeek-V2. Total cache per token: $d_{\text{latent}} + d_{\text{rope}} = 512 + 64 = 576$ values — still a 57x reduction from MHA.

The query is split symmetrically: a content query from the latent and a RoPE query with positional encoding, concatenated before the dot product.

The absorbed attention optimization

A key insight of MLA is that the up-projection matrices W_UK and W_UV can be algebraically absorbed into the query and output projections, eliminating the need to explicitly reconstruct K and V during attention.

Standard attention computes:

\text{score}_{t,s} = q_{t}^{T} @ K_{s} = q_{t}^{T} @ (c_{s} @ W_{\text{UK}})

This can be rewritten as:

\text{score}_{t,s} = (q_{t}^{T} @ W_{\text{UK}}^{T}) \cdot c_{s} = q_{\text{t,absorbed}}^{T} @ c_{s}

where $q_{\text{t,absorbed}} = W_{\text{UK}}^{T} @ q_{t}$ . The attention scores are computed directly between the transformed query and the raw latent c_s — the key is never materialized.

Similarly, the attention output for each head:

o_{t} = \text{sum}_{s}(\text{attn}_{t,s} \cdot V_{s}) = \text{sum}_{s}(\text{attn}_{t,s} \cdot c_{s} @ W_{\text{UV}}) = (\text{sum}_{s} \text{attn}_{t,s} \cdot c_{s}) @ W_{\text{UV}}

The weighted sum is computed over the raw latents, and W_UV is applied once to the result — not per token.

This means the full K and V tensors are never materialized in memory during decode. Attention operates directly in the d_latent-dimensional latent space.

PyTorch implementation

python

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiLatentAttention(nn.Module):
    def __init__(
        self,
        d_model: int,
        n_heads: int,
        d_head: int,
        d_latent: int,
        d_rope: int,
    ):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_head
        self.d_latent = d_latent
        self.d_rope = d_rope

        # Query projections
        self.W_Q = nn.Linear(d_model, n_heads * d_head, bias=False)
        self.W_QR = nn.Linear(d_model, n_heads * d_rope, bias=False)

        # KV down-projection (produces the cached latent)
        self.W_DKV = nn.Linear(d_model, d_latent, bias=False)

        # KV up-projections (reconstruct from latent)
        self.W_UK = nn.Linear(d_latent, n_heads * d_head, bias=False)
        self.W_UV = nn.Linear(d_latent, n_heads * d_head, bias=False)

        # RoPE key: separate small projection, not compressed
        self.W_KR = nn.Linear(d_model, d_rope, bias=False)

        self.W_O = nn.Linear(n_heads * d_head, d_model, bias=False)

    def forward(self, x: torch.Tensor, apply_rope_fn=None) -> torch.Tensor:
        B, S, _ = x.shape

        # Compress KV into latent (this is what gets cached)
        c = self.W_DKV(x)  # (B, S, d_latent)

        # Reconstruct full K and V from latent
        k_content = self.W_UK(c).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        v = self.W_UV(c).view(B, S, self.n_heads, self.d_head).transpose(1, 2)

        # RoPE component of keys (not compressed, cached separately)
        k_rope = self.W_KR(x)  # (B, S, d_rope)
        if apply_rope_fn is not None:
            k_rope = apply_rope_fn(k_rope)

        # Query: content + RoPE components
        q_content = self.W_Q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        q_rope = self.W_QR(x).view(B, S, self.n_heads, self.d_rope).transpose(1, 2)
        if apply_rope_fn is not None:
            q_rope = apply_rope_fn(q_rope)

        # Attention scores: content term + RoPE term
        content_scores = torch.matmul(q_content, k_content.transpose(-2, -1))
        rope_scores = torch.matmul(
            q_rope,
            k_rope.unsqueeze(1).expand(-1, self.n_heads, -1, -1).transpose(-2, -1),
        )

        scale = math.sqrt(self.d_head + self.d_rope)
        scores = (content_scores + rope_scores) / scale
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.W_O(out)

Cache footprint comparison

python

d_model, n_heads, d_head = 5120, 128, 128
d_latent, d_rope = 512, 64

mla = MultiLatentAttention(d_model, n_heads, d_head, d_latent, d_rope)

cache_per_token_mha = 2 * n_heads * d_head  # 32,768
cache_per_token_gqa8 = 2 * 8 * d_head       # 2,048
cache_per_token_mla = d_latent + d_rope      # 576

print(f"MHA cache/token:  {cache_per_token_mha:,} values")
print(f"GQA-8 cache/token: {cache_per_token_gqa8:,} values")
print(f"MLA cache/token:  {cache_per_token_mla:,} values")
print(f"MLA vs MHA: {cache_per_token_mha / cache_per_token_mla:.0f}x compression")
print(f"MLA vs GQA-8: {cache_per_token_gqa8 / cache_per_token_mla:.1f}x compression")

Output:

MHA cache/token:  32,768 values
GQA-8 cache/token: 2,048 values
MLA cache/token:  576 values
MLA vs MHA: 57x compression
MLA vs GQA-8: 3.6x compression

Quality and computational cost

MLA's compression is learned, not heuristic. The down-projection W_DKV and up-projections W_UK, W_UV are trained end-to-end with the rest of the model. The network learns to encode exactly the KV information that matters for attention into the d_latent-dimensional bottleneck.

DeepSeek-V2 reported performance matching or exceeding comparably-sized dense models (67B parameter dense baselines) on standard benchmarks — MMLU, HumanEval, GSM8K — despite the aggressive KV compression. The MoE architecture complicates direct comparison with dense GQA models, but the KV cache savings are unambiguous and orthogonal to the MoE decision.

The computational tradeoff: MLA adds two matrix multiplications per token during decode — the W_UK and W_UV up-projections from (d_latent) to (n_heads * d_head). For DeepSeek-V2, each is a (512, 16384) matrix multiply. In practice, the absorbed attention form eliminates these, computing attention directly in latent space. But even without absorption, the added compute is small relative to the FFN layers — and the memory savings during long-context decode far outweigh the compute cost.

Why MLA isn't universal yet

MLA requires training from scratch — there's no equivalent of GQA's mean-pooling uptraining trick to convert an existing MHA or GQA model. The compression bottleneck is learned jointly with all other model parameters, and retrofitting it would require substantial retraining. This makes MLA a commitment at architecture design time.

Additionally, the absorbed attention optimization requires custom CUDA kernels to be efficient. Standard FlashAttention implementations don't natively support the split content/RoPE score computation. DeepSeek wrote custom kernels for this; other frameworks are still catching up.

As of mid-2026, MLA is used by the DeepSeek model family (V2, V3, R1). Other major labs (Meta, Google, Mistral) continue to use GQA. Whether MLA becomes the next standard depends on whether the ecosystem builds the kernel support and training infrastructure to make it as turnkey as GQA.

← Previous