Kimi K2 · Multi-head latent attention

61 transformer blocks › 1 dense block

Multi-head latent attention

learnedself.self_attn

Queries, keys and values come from compressed latents; every query attends to the current and all earlier positions.

Why

Latent. Standard attention caches a full key and value for every head at every position, and the authors of latent attention call that cache a large bottleneck that limits batch size and sequence length when a model is served. Their design caches one latent per position and rebuilds keys and values from it. Kimi K2’s latent has 512 features; its pinned implementation caches the rebuilt keys and values of all 64 heads (modeling_deepseek.py 807–814). DeepSeek-V2 §2.1.1–2.1.2

Kimi K2 adopted latent attention from DeepSeek-V2, as did Kimi K2.5; Kimi K3 keeps it in its global-attention layers. Kimi K3 report §2.1.2

O=AV

\Delta_{\mathrm{attn}}=oW_{O}

All 30 steps

\text{Query compression:}

c^{{\htmlClass{vl-stream-q}{Q}}}=xW_{DQ}

\hat c^{{\htmlClass{vl-stream-q}{Q}}}=\operatorname{RMSNorm}(c^{{\htmlClass{vl-stream-q}{Q}}})

\text{Key–value compression:}

c^{KV}=xW_{DKV}

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

\text{Query heads:}

{\htmlClass{vl-stream-q}{q^{C}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{UQ}

{\htmlClass{vl-stream-q}{Q^{C}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{C}}})

{\htmlClass{vl-stream-q}{q^{\mathrm{pe}}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{QR}

{\htmlClass{vl-stream-q}{Q^{\mathrm{pe}}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{\mathrm{pe}}}})

{\htmlClass{vl-stream-q}{Q^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-q}{Q^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})

{\htmlClass{vl-stream-q}{Q}}=[{\htmlClass{vl-stream-q}{Q^{C}}};\,{\htmlClass{vl-stream-q}{Q^{R}}}]

\text{Key heads:}

{\htmlClass{vl-stream-k}{k^{C}}}=\hat c^{KV}W_{UK}

{\htmlClass{vl-stream-k}{K^{C}}}=\operatorname{reshape}({\htmlClass{vl-stream-k}{k^{C}}})

{\htmlClass{vl-stream-k}{k^{\mathrm{pe}}}}=xW_{KR}

{\htmlClass{vl-stream-k}{k^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-k}{k^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})

{\htmlClass{vl-stream-k}{K_i}}=[{\htmlClass{vl-stream-k}{K^{C}_i}};\,{\htmlClass{vl-stream-k}{k^{R}}}]

\text{Value heads:}

{\htmlClass{vl-stream-v}{v}}=\hat c^{KV}W_{UV}

{\htmlClass{vl-stream-v}{V}}=\operatorname{reshape}({\htmlClass{vl-stream-v}{v}})

\text{Causal attention:}

S=QK^{\mathsf T}

S'=S\cdot\frac{m^2}{\sqrt{192}}

m=1+0.1\ln 32

\tilde S=S'+{\htmlClass{vl-stream-context}{M}}

A[b,i,t,s]=\frac{\exp(\tilde S[b,i,t,s])}{\sum_{u=0}^{T-1}\exp(\tilde S[b,i,t,u])}

O=AV

o=\operatorname{reshape}(O)

\Delta_{\mathrm{attn}}=oW_{O}

in: $x_{\mathrm{attn}}$ [B batch, T positions, 7168 d] $R$ [T positions, 32 pairs] $M$ [T queries, T keys]
out: $\Delta_{\mathrm{attn}}$ [B batch, T positions, 7168 d]
parameters: 101,124,096

Parameters

101,124,096 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Query compression	11,011,584 across its components	11,011,584
Key–value compression	3,670,528 across its components	3,670,528
Query heads	18,874,368 across its components	18,874,368
Key heads	4,653,056 across its components	4,653,056
Value heads	4,194,304 across its components	4,194,304
Causal attention	58,720,256 across its components	58,720,256
Per-instance total		101,124,096

Across 1 independent dense blocks: $101{,}124{,}096 \times 1 = 101{,}124{,}096$ 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).

Codemodeling_deepseek.py 662–852

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 799–840

        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)

        q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)

        query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
        query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
        query_states[:, :, :, self.qk_nope_head_dim :] = q_pe

        key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
        key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
        key_states[:, :, :, self.qk_nope_head_dim :] = k_pe
        if past_key_value is not None:
            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
            key_states, value_states = past_key_value.update(
                key_states, value_states, self.layer_idx, cache_kwargs
            )

        attn_weights = (
            torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale
        )

        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
                f" {attn_weights.size()}"
            )
        assert attention_mask is not None
        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights + attention_mask

        # upcast attention to fp32
        attn_weights = nn.functional.softmax(
            attn_weights, dim=-1, dtype=torch.float32
        ).to(query_states.dtype)
        attn_weights = nn.functional.dropout(
            attn_weights, p=self.attention_dropout, training=self.training
        )
        attn_output = torch.matmul(attn_weights, value_states)

modeling_deepseek.py 848–852

        attn_output = attn_output.transpose(1, 2).contiguous()

        attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim)

        attn_output = self.o_proj(attn_output)

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 rotary parts are rotated; each head’s query and key join their non-rotated and rotated parts, the key’s rotated part being the shared rotary key; scores are scaled, masked by the causal mask, normalized and used to mix values.

mla-prefill · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 799–840

Read code excerpt

        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)

        q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)

        query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
        query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
        query_states[:, :, :, self.qk_nope_head_dim :] = q_pe

        key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
        key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
        key_states[:, :, :, self.qk_nope_head_dim :] = k_pe
        if past_key_value is not None:
            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
            key_states, value_states = past_key_value.update(
                key_states, value_states, self.layer_idx, cache_kwargs
            )

        attn_weights = (
            torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale
        )

        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
                f" {attn_weights.size()}"
            )
        assert attention_mask is not None
        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights + attention_mask

        # upcast attention to fp32
        attn_weights = nn.functional.softmax(
            attn_weights, dim=-1, dtype=torch.float32
        ).to(query_states.dtype)
        attn_weights = nn.functional.dropout(
            attn_weights, p=self.attention_dropout, training=self.training
        )
        attn_output = torch.matmul(attn_weights, value_states)

Direct evidence

Head outputs are flattened and projected back to the model width.

mla-output · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 848–852

Read code excerpt

        attn_output = attn_output.transpose(1, 2).contiguous()

        attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim)

        attn_output = self.o_proj(attn_output)