AI Engineer Exhaustive
HardPlusDeepSeek-V3

Multi-head Latent Attention

DeepSeek-V3

Hard

Full MLA combining compression and decoupled RoPE

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Standard multi-head attention caches a key and a value per head per token. At 128 heads and d_h=128 that is a large decode-time cache. Multi-head Latent Attention (MLA), carried from DeepSeek-V2 into V3, compresses keys and values through a joint low-rank latent and caches that latent plus a short RoPE-carrying vector instead of the full per-head KV. Queries are also low-rank compressed to cut training activations. Position information is decoupled: RoPE is applied only to a dedicated slice, so the compressed latent itself is position-free and can be reused across heads. V3 keeps the V2 MLA equations and adds YaRN on the shared RoPE key for the 128K context extension.

How it works

Let \mathbf{h}_t\in\mathbb{R}^{d} be the attention input (d=7168 on the 671B model). KV compression produces a latent of width $$ and a single shared RoPE key of width d_h^R=64:

\mathbf{c}_t^{KV}=W^{DKV}\mathbf{h}_t,\qquad \mathbf{k}_t^{C}=W^{UK}\mathbf{c}_t^{KV},\qquad \mathbf{k}_t^{R}=\operatorname{RoPE}(W^{KR}\mathbf{h}_t).

Head i then concatenates its content key with the same shared \mathbf{k}_t^{R}:

\mathbf{k}_{t,i}=[\mathbf{k}_{t,i}^{C};\mathbf{k}_t^{R}],\qquad \mathbf{v}_t^{C}=W^{UV}\mathbf{c}_t^{KV}.

Only the blue-boxed caches in the paper — \mathbf{c}_t^{KV} and \mathbf{k}_t^{R} — are stored at generation time. Queries use a wider bottleneck d_c'=1536:

\mathbf{c}_t^{Q}=W^{DQ}\mathbf{h}_t,\qquad \mathbf{q}_t^{C}=W^{UQ}\mathbf{c}_t^{Q},\qquad \mathbf{q}_t^{R}=\operatorname{RoPE}(W^{QR}\mathbf{c}_t^{Q}),

\mathbf{q}_{t,i}=[\mathbf{q}_{t,i}^{C};\mathbf{q}_{t,i}^{R}].

Attention is ordinary scaled-dot-product on the concatenated widths d_h+d_h^R=192, values taken from the uncompressed \mathbf{v}_{j,i}^{C} only:

\mathbf{o}_{t,i}=\sum_{j=1}^{t}\operatorname{Softmax}_j\!\left(\frac{\mathbf{q}_{t,i}^{\top}\mathbf{k}_{j,i}}{\sqrt{d_h+d_h^{R}}}\right)\mathbf{v}_{j,i}^{C}, \qquad \mathbf{u}_t=W^{O}[\mathbf{o}_{t,1};\ldots;\mathbf{o}_{t,n_h}].

The paper states that V3, like V2, applies RMSNorm after each compressed latent and extra scales on the width bottlenecks. On 671B: n_h=128, d_h=128, d_c=512, d_c'=1536, d_h^R=64. YaRN (s=40) is applied only to the decoupled shared key \mathbf{k}_t^{R} during the 32K/128K extension stages.

The official inference kernel has two numerically equivalent layouts. attn_impl="naive" materializes per-head \mathbf{k},\mathbf{v} and caches them. attn_impl="absorb" (default) absorbs $$ into the query so the content score is a dot of a projected q^{C} with the cached \mathbf{c}^{KV}, plus a separate RoPE score q^{R}\cdot k^{R}. That is the cache-friendly form the paper's blue boxes describe.

Shapes for one token on 671B (before tensor-parallel head split): q^{C}\in\mathbb{R}^{n_h\times 128}, q^{R}\in\mathbb{R}^{n_h\times 64}, c^{KV}\in\mathbb{R}^{512}, k^{R}\in\mathbb{R}^{64}, v^{C}\in\mathbb{R}^{n_h\times 128}.

Official code

All of this is class MLA in inference/model.py. Named weights: wq_a / q_norm / wq_b when q_lora_rank>0 (671B uses 1536); otherwise a single wq. Joint KV down-proj is wkv_a of width kv_lora_rank + qk_rope_head_dim, then kv_norm on the d_c slice only, then wkv_b up to $$ for (k^{C},v^{C}). Output is wo. Rotary frequencies are built in precompute_freqs_cis (YaRN ramp when max_seq_len > original_seq_len) and applied by apply_rotary_emb to $$ and k^{R} only. Widths match inference/configs/config_671B.json.

Watch-outs

Sources