Kimi K2 · Key–value compression

61 transformer blocks › 1 dense block › Latent attention

Key–value compression

learnedself.kv_a_layernorm(compressed_kv)

Compress each position’s 7,168 features into a 512-feature key–value latent.

Why

The key content and the values of all 64 heads are rebuilt from one 512-feature latent per position. In the design of latent attention that latent is cached in their place: caching every head’s key and value for every token is what its authors name as the bottleneck of serving. Kimi K2’s pinned implementation caches the rebuilt keys and values (modeling_deepseek.py 807–814). DeepSeek-V2 §2.1.1–2.1.3

c^{KV}=xW_{DKV}

\hat c^{KV}=\operatorname{RMSNorm}(c^{KV})

in: $x$ [B batch, T positions, 7168 d]
out: $\hat c^{KV}$ [B batch, T positions, 512 kv latent]
parameters: 3,670,528

Parameters

3,670,528 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Key–value down-projection	$7{,}168 \times 512$	3,670,016
Key–value latent RMSNorm	$512$	512
Per-instance total		3,670,528

Across 1 independent dense blocks: $3{,}670{,}528 \times 1 = 3{,}670{,}528$ parameters.

Derived from the pinned configuration and implementation, and equal to the count of the checkpoint’s tensors read from its safetensors headers (FP8 block scales and stored rotary buffers excluded).

Explanation

Caching this latent and the rotary key would take 576 values per position; full per-head keys (192 features) and values (128) take 20,480. This implementation caches the full keys and values (modeling_deepseek.py 810–814).

One matrix, kv_a_proj_with_mqa, produces this latent and the 64-feature rotary key together, 576 output features in all; the implementation splits them after the product (modeling_deepseek.py 776–779).

Codemodeling_deepseek.py 662–688, 776–789

modeling_deepseek.py 662–688

            self.q_a_proj = nn.Linear(
                self.hidden_size, config.q_lora_rank, bias=config.attention_bias
            )
            self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)
            self.q_b_proj = nn.Linear(
                config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
            )

        self.kv_a_proj_with_mqa = nn.Linear(
            self.hidden_size,
            config.kv_lora_rank + config.qk_rope_head_dim,
            bias=config.attention_bias,
        )
        self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank)
        self.kv_b_proj = nn.Linear(
            config.kv_lora_rank,
            self.num_heads
            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
            bias=False,
        )

        self.o_proj = nn.Linear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=config.attention_bias,
        )
        self._init_rope()

modeling_deepseek.py 776–789

        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
        compressed_kv, k_pe = torch.split(
            compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
        )
        k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
        kv = (
            self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
            .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
            .transpose(1, 2)
        )

        k_nope, value_states = torch.split(
            kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
        )

Sources2

Direct evidence

MLA owns a query down-projection, query RMSNorm and up-projection, a joint key–value down-projection that also emits the shared rotary key, a latent RMSNorm, a key–value up-projection and an output projection.

mla-init · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 662–688

Read code excerpt

            self.q_a_proj = nn.Linear(
                self.hidden_size, config.q_lora_rank, bias=config.attention_bias
            )
            self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)
            self.q_b_proj = nn.Linear(
                config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
            )

        self.kv_a_proj_with_mqa = nn.Linear(
            self.hidden_size,
            config.kv_lora_rank + config.qk_rope_head_dim,
            bias=config.attention_bias,
        )
        self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank)
        self.kv_b_proj = nn.Linear(
            config.kv_lora_rank,
            self.num_heads
            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
            bias=False,
        )

        self.o_proj = nn.Linear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=config.attention_bias,
        )
        self._init_rope()

Direct evidence

The joint down-projection yields the key–value latent and one rotary key shared by all heads; the latent is normalized and up-projected into per-head content keys and values.

mla-kv · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 776–789

Read code excerpt

        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
        compressed_kv, k_pe = torch.split(
            compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
        )
        k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
        kv = (
            self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
            .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
            .transpose(1, 2)
        )

        k_nope, value_states = torch.split(
            kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
        )