Kimi K2 · Dense transformer block

61 transformer blocks

Dense transformer block

learnedDeepseekV3DecoderLayer(config, layer_idx)

MLA and a dense SwiGLU MLP, each added to the residual stream; block 1.

x'_n=x_n+\Delta_{\mathrm{attn}}

x_{n+1}=x'_n+\Delta_{\mathrm{ffn}}

All 28 steps

x_{\mathrm{attn}}=\operatorname{RMSNorm}(x_n)

\text{Multi-head latent attention:}

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

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

c^{KV}=xW_{DKV}

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

{\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}}}]

{\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}}}]

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

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

x'_n=x_n+\Delta_{\mathrm{attn}}

x_{\mathrm{ffn}}=\operatorname{RMSNorm}(x'_n)

\text{SwiGLU MLP:}

g=x_{\mathrm{ffn}}W_{1}

u=x_{\mathrm{ffn}}W_{3}

a=g\,\sigma(g)

h=a\odot u

\Delta_{\mathrm{ffn}}=hW_{2}

x_{n+1}=x'_n+\Delta_{\mathrm{ffn}}

in: $x_0$ [B batch, T positions, 7168 d] $R$ [T positions, 32 pairs] $M$ [T queries, T keys]
out: $x_{n+1}$ [B batch, T positions, 7168 d]
parameters: 497,500,160

Parameters

497,500,160 parameters per dense block, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Attention RMSNorm	$7{,}168$	7,168
Multi-head latent attention	101,124,096 across its components	101,124,096
MLP RMSNorm	$7{,}168$	7,168
SwiGLU MLP	396,361,728 across its components	396,361,728
Per-instance total		497,500,160

Across 1 independent dense blocks: $497{,}500{,}160 \times 1 = 497{,}500{,}160$ 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 1152–1166, 1198–1218

modeling_deepseek.py 1152–1166

        self.mlp = (
            DeepseekV3MoE(config)
            if (
                config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0
            )
            else DeepseekV3MLP(config)
        )
        self.input_layernorm = DeepseekV3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = DeepseekV3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

modeling_deepseek.py 1198–1218

        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

Sources1

Direct evidence

A block applies input RMSNorm, attention and a residual addition, then post-attention RMSNorm, the FFN and a second residual addition; the first first_k_dense_replace blocks use the dense MLP, the rest MoE.

block · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 1152–1218

Read code excerpt

        self.mlp = (
            DeepseekV3MoE(config)
            if (
                config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0
            )
            else DeepseekV3MLP(config)
        )
        self.input_layernorm = DeepseekV3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = DeepseekV3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

… lines 1167–1197 omitted …

        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states