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 it needs them.

Compress keys and values into one latent

Each token's hidden vector $h_t$ gets 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 a lot of heads, $r$ numbers per token replace $2 \cdot H \cdot d_{\text{head}}$.

But RoPE doesn't survive this. A rotation that depends on position can't 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$, so $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 big win against MHA, a small one against GQA-2

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

So 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's 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. So MLA pays off in wide models with many heads, and this lab is neither.

This went 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 where 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 didn't build, and why

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