AI Engineer Exhaustive
EasyPlusDeepSeek-V3

KV Compression

DeepSeek-V3

Easy · 🔒 Plus required

Low-rank down-projection for KV cache

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. With (n_{h}=128) and (d_{h}=128), that is (2n_{h}d_{h}=32768) numbers per token per layer. MLA’s first move is to replace those two full tensors with one joint latent. A single down-projection maps the residual into a vector whose width is a few hundred, and that vector is what generation stores (together with a 64-D rotary key). The paper calls this “low-rank joint compression for attention keys and values.” Joint means keys and values share the same latent; they are not compressed by two independent bottlenecks. V3 sets the latent width (d_{c}=512). As in V2, an RMSNorm sits on the latent before it is used or cached.

How it works

Let (\mathbf{h}_{t}\in\mathbb{R}^{d}) be the (already RMSNorm-ed) attention input. Equation (1) of the report is the entire compression:

\mathbf{c}_{t}^{KV}=W^{DKV}\mathbf{h}_{t},\qquad W^{DKV}\in\mathbb{R}^{d_{c}\times d},\quad \mathbf{c}_{t}^{KV}\in\mathbb{R}^{d_{c}}.

For 671B, (d=7168) and (d_{c}=512). In the official module the same linear also emits the decoupled rotary key, so the implemented map is slightly wider than the boxed symbol in the paper:

[\mathbf{c}_{t}^{KV};\tilde{\mathbf{k}}_{t}^{R}]=W_{kv,a}\mathbf{h}_{t}\in\mathbb{R}^{d_{c}+d_{h}^{R}}=\mathbb{R}^{576}.

The first 512 channels are the latent; the last 64 are (W^{KR}\mathbf{h}_{t}) before RoPE. After the split, the latent is normalized:

\hat{\mathbf{c}}_{t}^{KV}=\mathrm{RMSNorm}(\mathbf{c}_{t}^{KV}).

Sec. 4.2 states that V3, like V2, “employs additional RMSNorm layers after the compressed latent vectors.” The absorb inference path writes (\hat{\mathbf{c}}_{t}^{KV}) into kv_cache and the rotated 64-D vector into pe_cache. That is the KV cache: (512+64=576) elements per token per layer, about (57\times) smaller than MHA’s 32768.

Queries have a separate bottleneck (d_{c}'=1536) (q_lora_rank) that reduces training activation memory; it is not part of the KV cache. The demo ModelArgs sets q_lora_rank=0 (full-rank queries) while keeping kv_lora_rank=512, which shows the two ranks are independent knobs.

Compression is not optional at decode time. If you discard (\mathbf{c}_{t}^{KV}) you cannot rebuild keys or values for later tokens. Reconstruction (the up-projections (W^{UK},W^{UV})) is a later step and can be absorbed so the full tensors never exist.

Official code

MLA.__init__ builds

self.wkv_a   = Linear(dim, kv_lora_rank + qk_rope_head_dim)
self.kv_norm = RMSNorm(kv_lora_rank)

forward does kv, k_pe = split(wkv_a(x), [512, 64]), then kv_norm(kv). In attn_impl == "absorb" (the default) the normalized latent is stored in kv_cache with shape (max_batch, max_seq, kv_lora_rank). The naive path still computes the latent the same way; it just up-projects immediately and caches full heads instead.

Ranks are in config_671B.json: kv_lora_rank=512, qk_rope_head_dim=64.

Watch-outs

Sources