Multi-head latent attention

Grouped-query attention shrinks the cache by sharing key and value heads. DeepSeek-V2's multi-head latent attention (MLA) goes further. It caches one low-rank vector per token, shared by every head, and rebuilds the keys and values from it when needed.

Compress keys and values into one latent

Each token's hidden vector $h_t$ is projected down to a latent $c_t$ of rank $r$, much smaller than the model width. Keys and values for every head are projected back up from it:

Only $c_t$ is cached. For a model with many heads, $r$ numbers per token replace $2 \cdot H \cdot d_{\text{head}}$.

RoPE does not survive this. A rotation that depends on position cannot pass through a fixed up-projection, so MLA adds a separate small key, $k^{\text{rope}}_t$ of width $d_r$, shared across heads and rotated normally. The cache holds $c_t$ and $k^{\text{rope}}_t$: $r + d_r$ numbers per token per layer.

def _project_mla(
    self, x: Tensor, rope: tuple[Tensor, Tensor] | None
) -> tuple[Tensor, Tensor, Tensor]:
    """Content rides a shared low-rank latent, position rides a separate shared RoPE key."""
    rope_dim = self.config.mla_rope_dim
    content = self.kv_up(self.kv_down(x))
    key, value = content.chunk(2, dim=-1)
    key = self._heads(key, self.n_heads, self.head_size)
    value = self._heads(value, self.n_heads, self.head_size)
    query = self._heads(self.query(x), self.n_heads, self.head_size)
    cosine, sine = rope
    rope_slice = (cosine[..., : rope_dim // 2], sine[..., : rope_dim // 2])
    query_rope = apply_rope(self._heads(self.q_rope(x), self.n_heads, rope_dim), *rope_slice)
    key_rope = apply_rope(self.k_rope(x).unsqueeze(1), *rope_slice)
    key_rope = key_rope.expand(-1, self.n_heads, -1, -1)
    return torch.cat((query, query_rope), -1), torch.cat((key, key_rope), -1), value

A large win against MHA, a small one against GQA-2

PLAN.md and day-wise.md both carried MLA in as a large cache saving, and DeepSeek reports reductions past 90 percent. The note checked where that figure comes from. It is measured against multi-head attention with many heads. Day 2 had already taken most of that saving with two KV heads.

What happened?

• At width 256 with 8 heads, GQA-2 caches $2 \times 2 \times 32 = 128$ numbers per token per layer. MLA at rank 64 with a 16-wide RoPE key caches 80. That is 1.6 times smaller, not 10.
• MLA stops beating GQA-2 at rank 112 at this width.
• At width 4,096, MHA caches 8,192 numbers per token per layer, and MLA's savings look like the paper's. MLA pays off in wide models with many heads. The lab is neither.

This was written into the note as a conflict with both plans, before any code ran.

MLA loses to GQA-2 at this width

The useful number is the rank at which MLA's cache first falls below GQA-2's 128 numbers, and whether quality survives at that rank. Stop condition: ranks 32, 64 and 128 against the GQA-2 control, three seeds, reporting cache numbers per token, code and prose bits per byte, and seconds per step. The note marked the result experimental either way.

Built, tested, not scheduled

The plan changed before the run. The code stays behind a flag that defaults to off.

What we did not build, and why

• Absorbing the up-projection. At decode time MLA folds $W_{\text{up}}$ into the query and output matrices, so it never materializes full keys. That is an identity with no training effect, it belongs to serving work, and it is the part most likely to hide a silent bug.
• Low-rank queries. DeepSeek-V2 also compresses queries to save activation memory in training. It does not change the cache.