Multi-Head Latent Attention: compressing KV into a shared low-rank latent

The problem, from first principles

[03-multi-query-attention] and [04-grouped-query-attention] reduce the KV cache by sharing K/V ACROSS heads. Multi-Head Latent Attention (MLA, used in DeepSeek-V2/V3) takes a different route: compress K and V for ALL heads into one small shared latent vector c of dimension d_latent (far smaller than n_heads * d_head), and reconstruct per-head keys/values from that latent on the fly. Implement a simplified version of this scheme.

From theory to code

Implement multi_head_latent_attention(X, W_Q, W_DKV, W_UK, W_UV, W_O, n_heads, mask=None):

c = X @ W_DKV                          # (seq_len, d_latent) -- this is what actually gets cached
K = c @ W_UK,   V = c @ W_UV            # up-project the latent into full per-head K, V
Q_h = X @ W_Q[h]                        # queries are computed directly, no compression
head_h = softmax(Q_h K_h^T / sqrt(d_head)) @ V_h;  output = concat(heads) @ W_O

Constraints

• Compute latent c = X @ W_DKV once (shape (seq_len, d_latent)).
• Reconstruct full K, V via K = c @ W_UK, V = c @ W_UV, then reshape into n_heads heads.
• Queries are computed directly from X per head (no compression on the query side).
• Return both the attention output and the latent cache c — what would actually be stored per token in a real KV cache.
• d_latent < n_heads * d_head (a real compression, not a no-op).

Hints