KV Reconstruction

DeepSeek-V3

Easy · 🔒 Plus required

Up-project compressed KV back to full size

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

Compression writes a 512-D latent. Attention still needs per-head content keys and values. Reconstruction is the pair of up-projections that expand that latent back to (n_{h}d_{h}) channels. The paper’s eqs. (2) and (5) are this step. Values have no rotary half; keys get the shared (\mathbf{k}^{R}) concatenated after reconstruction. You do not have to materialize the expanded tensors at decode time: because reconstruction is a position-independent linear map, (W^{UK}) can be absorbed into the query and (W^{UV}) into the attention output. The official repo implements both the literal “naive” expansion and the absorbed path.

How it works

Starting from the (RMSNorm-ed) latent (\mathbf{c}{t}^{KV}\in\mathbb{R}^{d{c}}),

\mathbf{k}_{t}^{C}=W^{UK}\mathbf{c}_{t}^{KV},\qquad \mathbf{v}_{t}^{C}=W^{UV}\mathbf{c}_{t}^{KV},

with (W^{UK},W^{UV}\in\mathbb{R}^{n_{h}d_{h}\times d_{c}}). Split (\mathbf{k}{t}^{C}) and (\mathbf{v}{t}^{C}) into (n_{h}) heads of width (d_{h}=128). Then form the key that attention actually uses:

\mathbf{k}_{t,i}=[\mathbf{k}_{t,i}^{C};\mathbf{k}_{t}^{R}]\in\mathbb{R}^{d_{h}+d_{h}^{R}}=\mathbb{R}^{192}.

(\mathbf{v}{t,i}^{C}) stays 128-D. For 671B, (n{h}=128), so each up-projection is (16384\times 512). In code the two maps are one matrix (W_{kv,b}) of shape ((n_{h}(d_{h}+d_{h}),,d_{c})=(n_{h}\cdot 256,,512)), then a split into qk_nope_head_dim and v_head_dim.

Naive path. Apply (W_{kv,b}), split, concat (\mathbf{k}^{R}), cache full (\mathbf{k},\mathbf{v}), and run ordinary attention. This matches the paper’s equations literally and wastes the cache reduction.

Absorb path (default). Rewrite the content score and the value mix so (W^{UK}) and (W^{UV}) never hit the sequence axis of the cache:

\mathbf{q}_{t,i}^{C\top}\mathbf{k}_{j,i}^{C} =\bigl((W^{UK}_{i})^{\top}\mathbf{q}_{t,i}^{C}\bigr)^{\top}\mathbf{c}_{j}^{KV},

\sum_{j}\alpha_{tj}\mathbf{v}_{j,i}^{C} =W^{UV}_{i}\sum_{j}\alpha_{tj}\mathbf{c}_{j}^{KV}.

Rotary scores (\mathbf{q}{t,i}^{R\top}\mathbf{k}{j}^{R}) are added separately. Cached tensors stay (\mathbf{c}^{KV}) and (\mathbf{k}^{R}). Absorption is valid only for the content half; that is why RoPE was decoupled.

Paper eq. (10) still uses the reconstructed view in the math. The two implementations are algebraically the same when (W^{UK},W^{UV}) are applied as above and the softmax scale is (1/\sqrt{d_{h}+d_{h}^{R}}).

What is specified versus inferred: eqs. (2) and (5) specify the up-projections. Fusing them into one wkv_b and the absorb einsums are implementation choices in inference/model.py. The report discusses recomputing “MLA up-projections” in the backward pass (Sec. 3) to save activation memory; that is a training tactic, not a change to the forward algebra.

Official code

MLA:

self.wkv_b = ColumnParallelLinear(kv_lora_rank, n_heads * (qk_nope_head_dim + v_head_dim))

attn_impl == "naive": wkv_b(kv_norm(kv)), view as heads, split into k_nope and v, cat with expanded k_pe, cache k and v.
attn_impl == "absorb": reshape wkv_b.weight to (n_local_heads, 256, kv_lora_rank); einsum q_nope with the first 128 rows (this is (W^{UK})); after softmax, einsum the context vector with the last 128 rows (this is (W^{UV})).

Output projection wo always maps n_heads * v_head_dim back to dim.

Watch-outs

Reconstructing keys and then applying RoPE to the full 192-D vector double-applies position (the rotary half is already rotated) and prevents absorption.
(W^{UK}) and (W^{UV}) have the same shape but are not interchangeable. Swapping the two slices of wkv_b mixes keys into values.
Absorbing (W^{UV}) before the weighted sum (W^{UV} c then attend) is mathematically fine but returns you to the naive cache. Absorb after the sum (attend on c, then (W^{UV})) is the cheap path.
Under column parallelism, wkv_b is sharded on heads. Viewing it as (n_heads, 256, rank) without using n_local_heads breaks multi-GPU inference.

Sources

Paper: DeepSeek-V3 Technical Report, arXiv:2412.19437 (Sec. 2.1.1 eqs. 2, 5, 10–11; Sec. 3 up-projection recomputation)
Code: deepseek-ai/DeepSeek-V3 (inference/model.py MLA.forward, naive vs absorb)