Decoupled RoPE

DeepSeek-V3

Medium · 🔒 Plus required

Separate RoPE for compressed keys

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

MLA stores a low-rank latent (\mathbf{c}{t}^{KV}) instead of per-head keys. Rotary position embeddings (RoPE) are a position-dependent orthogonal transform. If you rotate the reconstructed full key (\mathbf{k}{t}^{C}=W^{UK}\mathbf{c}{t}^{KV}), that rotation cannot be folded into a position-independent matrix, so you lose the main inference trick of MLA: absorbing (W^{UK}) into the query. DeepSeek-V2/V3 therefore split each query and key into a content half that never sees RoPE and a short rotary half that does. Only the rotary key (\mathbf{k}{t}^{R}) is cached alongside (\mathbf{c}_{t}^{KV}). The paper applies YaRN exclusively to this shared rotary key when extending context from 4K to 128K.

How it works

After the content key is up-projected, a second linear map builds a rotary key from the uncompressed residual (not from (\mathbf{c}_{t}^{KV})):

\mathbf{k}_{t}^{R}=\operatorname{RoPE}(W^{KR}\mathbf{h}_{t})\in\mathbb{R}^{d_{h}^{R}},\qquad \mathbf{k}_{t,i}=[\mathbf{k}_{t,i}^{C};\mathbf{k}_{t}^{R}].

(W^{KR}\in\mathbb{R}^{d_{h}^{R}\times d}) with (d_{h}^{R}=64). The same (\mathbf{k}_{t}^{R}) is concatenated onto every head (i). Queries are compressed first, then split:

\mathbf{q}_{t}^{C}=W^{UQ}\mathbf{c}_{t}^{Q},\qquad \mathbf{q}_{t}^{R}=\operatorname{RoPE}(W^{QR}\mathbf{c}_{t}^{Q}),\qquad \mathbf{q}_{t,i}=[\mathbf{q}_{t,i}^{C};\mathbf{q}_{t,i}^{R}].

(W^{QR}\in\mathbb{R}^{n_{h}d_{h}^{R}\times d_{c}'}) produces a per-head rotary query. Attention (paper eq. 10) dots the concatenated vectors and scales by the sum of the two widths:

\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}.

Because the rotary inner product (\mathbf{q}{t,i}^{R\top}\mathbf{k}{j}^{R}) does not involve (W^{UK}), the content scores (\mathbf{q}{t,i}^{C\top}W^{UK}\mathbf{c}{j}^{KV}) can still absorb (W^{UK}) into the query. The official “absorb” path therefore caches only (\mathbf{c}^{KV}) and (\mathbf{k}^{R}).

RoPE itself is the usual complex multiply. With (\boldsymbol{\omega}{m}=\theta^{-2m/d{h}^{R}}) and YaRN-adjusted frequencies when (T>\texttt{original_seq_len}),

\operatorname{RoPE}(\mathbf{x})_{t}=\mathbf{x}_{t}\odot e^{i\,t\boldsymbol{\omega}}.

The 671B config uses (\theta=10^{4}), YaRN scale (s=40), and softmax mscale ((0.1\ln s+1)^{2}) once the sequence exceeds the pretrain length of 4096. Values are never rotated.

Official code

In inference/model.py:

• precompute_freqs_cis builds frequencies of length qk_rope_head_dim (64), with the YaRN ramp when max_seq_len > original_seq_len.
• apply_rotary_emb views the last axis as pairs, multiplies by freqs_cis, and writes real parts back.
• MLA packs (W^{DKV}) and (W^{KR}) into one linear wkv_a: output width kv_lora_rank + qk_rope_head_dim. The split is [512 | 64]. k_pe is unsqueezed to a singleton head axis, rotated, then broadcast across heads.
• Queries are produced at width qk_nope_head_dim + qk_rope_head_dim = 192 per head, then split and rotated on the last 64 channels only.

Watch-outs

• Rotating the compressed latent (\mathbf{c}_{t}^{KV}) (or the reconstructed (\mathbf{k}^{C})) is not what the paper does. RoPE is applied only to the decoupled 64-D streams.
• (\mathbf{k}^{R}) is shared across heads; (\mathbf{q}^{R}) is not. Expanding a single rotary query to every head, or allocating a per-head rotary key in the cache, mismatches both the math and the absorb kernel.
• Softmax scale must use (d_{h}+d_{h}^{R}=192). Scaling by (\sqrt{128}) or (\sqrt{64}) alone is wrong.
• YaRN in V3 is specified for the shared key (\mathbf{k}_{t}^{R}) only. Applying the frequency stretch to the content channels, or skipping the mscale on the softmax, will not match the long-context recipe in Sec. 4.3.

Sources

• Paper: DeepSeek-V3 Technical Report, arXiv:2412.19437 (Sec. 2.1.1 eqs. 3–4, 8–10; Sec. 4.3)
• Code: deepseek-ai/DeepSeek-V3 (inference/model.py: precompute_freqs_cis, apply_rotary_emb, MLA)