Kimi K2 · Transformer blocks

in: $x_0$ [B batch, T positions, 7168 d] $R$ [T positions, 32 pairs] $M$ [T queries, T keys]
out: $x_{61}$ [B batch, T positions, 7168 d]
parameters: 1,024,059,415,040

Parameters

1,024,059,415,040 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Dense transformer block	$497{,}500{,}160 \times 1$	497,500,160
MoE transformer block	$17{,}059{,}365{,}248 \times 60$	1,023,561,914,880
Per-instance total		1,024,059,415,040

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

The implementation uses DeepseekV3MoE when config.n_routed_experts is not None, layer_idx >= config.first_k_dense_replace and layer_idx % config.moe_layer_freq == 0, and DeepseekV3MLP otherwise (modeling_deepseek.py 1152–1160). With moe_layer_freq = 1, every layer after the dense ones qualifies.

Codemodeling_deepseek.py 1152–1616

modeling_deepseek.py 1360–1370

        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                DeepseekV3DecoderLayer(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
        self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

modeling_deepseek.py 1423–1438

        past_key_values_length = 0
        if use_cache:
            use_legacy_cache = not isinstance(past_key_values, Cache)
            if use_legacy_cache:
                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
            past_key_values_length = past_key_values.get_usable_length(seq_length)

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length,
                seq_length + past_key_values_length,
                dtype=torch.long,
                device=device,
            )
            position_ids = position_ids.unsqueeze(0)

modeling_deepseek.py 1443–1457

        if self._use_flash_attention_2:
            # 2d mask is passed through the layers
            attention_mask = (
                attention_mask
                if (attention_mask is not None and 0 in attention_mask)
                else None
            )
        else:
            # 4d mask is passed through the layers
            attention_mask = _prepare_4d_causal_attention_mask(
                attention_mask,
                (batch_size, seq_length),
                inputs_embeds,
                past_key_values_length,
            )

modeling_deepseek.py 1522

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

modeling_deepseek.py 1614–1616

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        logits = logits.float()

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

Sources2

Direct evidence

The model embeds tokens, numbers positions from the cached length, builds a 4D causal mask, applies all blocks, normalizes and projects to vocabulary logits with an untied head, computing every position’s logits in float32.

transformer · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 1360–1616

Read code excerpt

        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                DeepseekV3DecoderLayer(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
        self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

… lines 1371–1422 omitted …

        past_key_values_length = 0
        if use_cache:
            use_legacy_cache = not isinstance(past_key_values, Cache)
            if use_legacy_cache:
                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
            past_key_values_length = past_key_values.get_usable_length(seq_length)

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length,
                seq_length + past_key_values_length,
                dtype=torch.long,
                device=device,
            )
            position_ids = position_ids.unsqueeze(0)

… lines 1439–1442 omitted …

        if self._use_flash_attention_2:
            # 2d mask is passed through the layers
            attention_mask = (
                attention_mask
                if (attention_mask is not None and 0 in attention_mask)
                else None
            )
        else:
            # 4d mask is passed through the layers
            attention_mask = _prepare_4d_causal_attention_mask(
                attention_mask,
                (batch_size, seq_length),
                inputs_embeds,
                past_key_values_length,
            )

… lines 1458–1521 omitted …

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

… lines 1523–1613 omitted …

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        logits = logits.float()

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