Key–value compression

Checkpoint configuration and implementation moonshotai/Kimi-K2-Instruct @ fd1984e · parameter count checked against the checkpoint’s tensor headers · How to read

61 transformer blocks › 60 MoE blocks › Latent attention › Key–value compression

Key–value down-projection

×

learnedself.kv_a_proj_with_mqa(hidden_states)

Map 7,168 features to the 512-feature key–value latent.

cKV=xWDKV{\htmlData{light=raw down}{c^{KV}}}={\htmlData{light=input}{x}}W_{DKV}
in
xx [B batch, T positions, 7168 d]
out
cKVc^{KV} [B batch, T positions, 512 kv latent]
parameters
3,670,016 · × 60 MoE blocks = 220,200,960
Parameters

3,670,016 parameters per instance.

Exact per-instance accounting
ComponentCalculationParameters
Weight matrix7,168×5127{,}168 \times 5123,670,016
Per-instance total3,670,016

This projection has no bias parameter.

The reference module holds the weight output by input, [out, in].

Across 60 independent MoE blocks: 3,670,016×60=220,200,9603{,}670{,}016 \times 60 = 220{,}200{,}960 parameters.

Derived from the configured matrix dimensions; the model total matches the checkpoint’s tensors.

Explanation

WDKVW_{DKV} is the first 512 output features of kv_a_proj_with_mqa.

Codemodeling_deepseek.py 662–688, 776–789 · config.json 38–46
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
        )
config.json 38–46
  "quantization_config": {
    "activation_scheme": "dynamic",
    "fmt": "e4m3",
    "quant_method": "fp8",
    "weight_block_size": [
      128,
      128
    ]
  },
Sources3

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
        )

Direct evidence

Weights are stored in FP8 (e4m3) with one scale per 128 × 128 block; the scales are quantization metadata and are excluded from parameter counts.

linear · moonshotai/Kimi-K2-Instruct / config.json
Pinned revision fd1984e2b7a3 ↗ · lines 38–46

Read code excerpt
  "quantization_config": {
    "activation_scheme": "dynamic",
    "fmt": "e4m3",
    "quant_method": "fp8",
    "weight_block_size": [
      128,
      128
    ]
  },