Latent Attention (MLA)

During autoregressive decoding, standard Multi-Head Attention (MHA) caches the Key and Value vectors of every previous token. For long contexts, this KV cache dominates GPU memory and limits batch size and sequence length. How can we shrink the cache without hurting attention quality?

Latent attention (introduced in DeepSeek-V2) does not cache \(K\) and \(V\) directly. Instead, each token's hidden state \(x\) is first compressed by a down-projection \(W_{DKV}\) into a small shared latent vector \(c_{KV}\). Only this latent vector is stored in the cache. When Keys and Values are needed, they are reconstructed from the latent with two up-projection matrices, \(W_{UK}\) and \(W_{UV}\).

Because \(K\) and \(V\) are both derived from the same low-dimensional latent, the cache stores one compact vector per token instead of two full-size vectors per head.

import torch
import torch.nn as nn

class LatentKVCompression(nn.Module):
    """Compress hidden states into a latent and rebuild K, V from it."""
    def __init__(self, d_model: int = 8, d_latent: int = 4, d_head: int = 4):
        super().__init__()
        self.W_dkv = nn.Linear(d_model, d_latent, bias=False)   # down-projection
        self.W_uk = nn.Linear(d_latent, d_head, bias=False)     # up-projection for K
        self.W_uv = nn.Linear(d_latent, d_head, bias=False)     # up-projection for V

    def forward(self, x: torch.Tensor):
        c_kv = self.W_dkv(x)      # (T, d_latent)  <- only this is cached
        k = self.W_uk(c_kv)       # (T, d_head)
        v = self.W_uv(c_kv)       # (T, d_head)
        return c_kv, k, v
    

Trace the full forward pass of latent attention for the input sentence "The next day is", using sample dimensions: \(T = 4\) tokens, model dimension \(d_{model} = 8\), latent dimension \(d_c = 4\) and head dimension \(d_h = 4\). Compute every tensor from the input \(X\) to the context matrix.

The sample dimensions used in the diagram are:

Input \(X\): \((4, 8)\), one 8-dimensional feature row per token (The, next, day, is).
\(W_q\): \((8, 4)\), \(W_{DKV}\): \((8, 4)\), \(W_{UK}\): \((4, 4)\), \(W_{UV}\): \((4, 4)\).

Step 1: Query. Project the input directly to queries: \(Q = X W_q\), giving shape \((4, 8) \times (8, 4) = (4, 4)\).

Step 2: Latent (KV compression). Down-project the input: \(C_{KV} = X W_{DKV}\), giving shape \((4, 8) \times (8, 4) = (4, 4)\). This is the only tensor stored in the KV cache.

Step 3: Key and Value reconstruction. Up-project the latent: \(K = C_{KV} W_{UK}\) and \(V = C_{KV} W_{UV}\), each of shape \((4, 4) \times (4, 4) = (4, 4)\).

Step 4: Attention scores. \(S = Q K^{T}\), shape \((4, 4) \times (4, 4) = (4, 4)\).

Step 5: Scaling, causal mask and softmax. Divide by \(\sqrt{d_h}\), set the upper triangle (future tokens) to \(-\infty\), then apply softmax row-wise:

\[ A = \text{softmax}\left(\frac{Q K^{T}}{\sqrt{d_h}} + M_{causal}\right) \]

The attention weights \(A\) have shape \((4, 4)\), and each row sums to 1. For example, the token day (row 3) can only attend to The, next and itself.

Step 6: Context matrix. Weight the Values by the attention weights: \(\text{Context} = A V\), shape \((4, 4) \times (4, 4) = (4, 4)\). Each row is the context vector of one token.

Cache comparison for this example: standard attention caches \(K\) and \(V\), which is \(2 \times 4 \times 4 = 32\) values. Latent attention caches only \(C_{KV}\), which is \(4 \times 4 = 16\) values. In real models \(d_c\) is much smaller than \(n_h \cdot d_h\), so the saving is far larger.

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

T, d_model, d_c, d_h = 4, 8, 4, 4
X = torch.randn(T, d_model)            # (4, 8)  "The next day is"

W_q   = torch.randn(d_model, d_h)      # (8, 4)
W_dkv = torch.randn(d_model, d_c)      # (8, 4)
W_uk  = torch.randn(d_c, d_h)          # (4, 4)
W_uv  = torch.randn(d_c, d_h)          # (4, 4)

Q    = X @ W_q                         # (4, 4)
C_kv = X @ W_dkv                       # (4, 4)  latent, cached
K    = C_kv @ W_uk                     # (4, 4)
V    = C_kv @ W_uv                     # (4, 4)

scores = Q @ K.T / math.sqrt(d_h)      # (4, 4)
mask = torch.triu(torch.ones(T, T), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))
weights = F.softmax(scores, dim=-1)    # (4, 4)

context = weights @ V                  # (4, 4)
print(context.shape)                   # torch.Size([4, 4])
    

At inference time, rebuilding \(K\) and \(V\) from the latent for every cached token on every decoding step is wasteful. Can we avoid materialising \(K\) and \(V\) at all, and compute attention directly on the cached latent \(C_{KV}\)?

Matrix multiplication is associative, so the up-projection matrices can be absorbed into neighbouring weights. This is done once, offline, and creates the absorption matrices.

Absorbing \(W_{UK}\) into the query side. Substitute \(K = C_{KV} W_{UK}\) into the score:

\[ Q K^{T} = (X W_q)(C_{KV} W_{UK})^{T} = X \underbrace{(W_q W_{UK}^{T})}_{W_{absorb}} C_{KV}^{T} \]

The merged matrix \(W_{absorb} = W_q W_{UK}^{T}\) has shape \((8, 4) \times (4, 4) = (8, 4)\). The scores are now computed as \((X W_{absorb}) C_{KV}^{T}\), directly against the cached latent, so \(K\) is never built.

Absorbing \(W_{UV}\) into the output projection. Substitute \(V = C_{KV} W_{UV}\) into the context and apply the output projection \(W_O\):

\[ \text{Out} = (A\, C_{KV} W_{UV})\, W_O = (A\, C_{KV}) \underbrace{(W_{UV} W_O)}_{W_{absorb}'} \]

Attention weights are applied to the latent first, and the merged matrix \(W_{UV} W_O\) maps the result to the model dimension, so \(V\) is never built either.

Important notes:

1. Absorption is an inference-time optimisation. Training keeps the separate matrices.
2. It pays off when \(d_c < n_h \cdot d_h\) (many heads, small latent). In this toy example \(d_c = d_h\) so there is no real saving, but the algebra is identical.
3. RoPE position encoding breaks absorption, because the rotation sits between \(W_q\) and \(W_{UK}^{T}\) and does not commute with them. DeepSeek solves this with a small separate decoupled RoPE key/query component that is not absorbed.

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

T, d_model, d_c, d_h = 4, 8, 4, 4
X     = torch.randn(T, d_model)
W_q   = torch.randn(d_model, d_h)
W_dkv = torch.randn(d_model, d_c)
W_uk  = torch.randn(d_c, d_h)
W_uv  = torch.randn(d_c, d_h)
W_o   = torch.randn(d_h, d_model)

C_kv = X @ W_dkv                        # only this is cached

# Offline: merge weights (absorption matrices)
W_absorb_q = W_q @ W_uk.T               # (8, 4)
W_absorb_o = W_uv @ W_o                 # (4, 8)

# Attention directly on the latent (K and V never materialised)
scores = (X @ W_absorb_q) @ C_kv.T / math.sqrt(d_h)   # (4, 4)
mask = torch.triu(torch.ones(T, T), diagonal=1).bool()
weights = F.softmax(scores.masked_fill(mask, float('-inf')), dim=-1)
out_absorbed = (weights @ C_kv) @ W_absorb_o          # (4, 8)

# Reference: standard path with explicit K and V
Q, K, V = X @ W_q, C_kv @ W_uk, C_kv @ W_uv
ref_scores = Q @ K.T / math.sqrt(d_h)
ref_w = F.softmax(ref_scores.masked_fill(mask, float('-inf')), dim=-1)
out_ref = (ref_w @ V) @ W_o

print(torch.allclose(out_absorbed, out_ref, atol=1e-4))  # True