Kimi K3 · Final RMSNorm · OpenModelAtlas

Final RMSNorm

learnedself.norm(hidden_states) · 1219

Rescale each position’s features using their root mean square, then apply learned gain.

\begin{aligned}m_2(x)&=\frac{1}{d}\sum_{j=1}^d x_j^2\\y&=\gamma_{\mathrm{final}}\odot\frac{x}{\sqrt{m_2(x)+\varepsilon}}\\\varepsilon&=10^{-5}\end{aligned}

in: $h_{\mathrm{out}}$ [B batch, T positions, 7168 d]
out: $z$ [B batch, T positions, 7168 d]
parameters: 7,168

Parameters

7,168 parameters per instance.

Exact per-instance accounting
Component	Calculation	Parameters
Learned gain	$7{,}168$	7,168
Per-instance total		7,168

Each normalization instance owns its learned gain vector.

Derived from the configured feature width; the model total matches the checkpoint’s tensors.

Codemodeling_kimi_linear.py 226–1233

modeling_kimi_linear.py 226–236

class KimiRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        dtype = hidden_states.dtype
        x = hidden_states.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
        return self.weight * x.to(dtype)

modeling_kimi_linear.py 1098–1108

        self.layers = nn.ModuleList([KimiDecoderLayer(
            config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
        self.norm = KimiRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)

        self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None
        if self.use_attn_residuals:
            self.output_attn_res_norm = KimiRMSNorm(
                config.hidden_size, eps=config.rms_norm_eps)
            self.output_attn_res_proj = nn.Linear(
                config.hidden_size, 1, bias=False)

modeling_kimi_linear.py 1173–1219

        causal_mask = create_causal_mask(
            config=self.config,
            input_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_position=cache_position,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )
        linear_attn_mask = self._update_linear_attn_mask(
            attention_mask, cache_position)

        hidden_states = inputs_embeds
        if past_key_values is not None:
            assert isinstance(past_key_values, KimiDynamicCache)

        block_residual = None
        if self.use_attn_residuals:
            block_residual = hidden_states.new_zeros(
                hidden_states.shape[0] * hidden_states.shape[1], 0,
                hidden_states.shape[2])

        for decoder_layer in self.layers:
            layer_mask = linear_attn_mask if decoder_layer.is_linear_attn else causal_mask

            if self.use_attn_residuals:
                hidden_states, block_residual = decoder_layer(
                    hidden_states,
                    attention_mask=layer_mask,
                    past_key_values=past_key_values,
                    cache_position=cache_position,
                    block_residual=block_residual,
                    **kwargs,
                )
            else:
                hidden_states = decoder_layer(
                    hidden_states,
                    attention_mask=layer_mask,
                    past_key_values=past_key_values,
                    cache_position=cache_position,
                    **kwargs,
                )

        if self.use_attn_residuals:
            hidden_states = self._apply_output_attn_res(
                hidden_states, block_residual)

        hidden_states = self.norm(hidden_states)

modeling_kimi_linear.py 1226–1233

    def _apply_output_attn_res(self, hidden_states, block_residual):
        batch_size, seq_len, hidden_size = hidden_states.shape
        return _apply_attn_res(
            hidden_states.view(-1, hidden_size),
            block_residual,
            self.output_attn_res_proj,
            self.output_attn_res_norm,
        ).view(batch_size, seq_len, hidden_size)

Sources2

Direct evidence

RMSNorm owns one gain per feature and computes in float32.

rmsnorm · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 226–236

Read code excerpt

class KimiRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        dtype = hidden_states.dtype
        x = hidden_states.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
        return self.weight * x.to(dtype)

Direct evidence

The model embeds tokens, calls create_causal_mask for the latent-attention layers (with flash attention it yields no matrix) and passes the padding mask to the Kimi Delta Attention layers, starts with no stored sums, applies every layer, mixes the output from the stored sums and the last running sum, and applies the final RMSNorm.

transformer · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 1098–1233

Read code excerpt

        self.layers = nn.ModuleList([KimiDecoderLayer(
            config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
        self.norm = KimiRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)

        self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None
        if self.use_attn_residuals:
            self.output_attn_res_norm = KimiRMSNorm(
                config.hidden_size, eps=config.rms_norm_eps)
            self.output_attn_res_proj = nn.Linear(
                config.hidden_size, 1, bias=False)

… lines 1109–1172 omitted …

        causal_mask = create_causal_mask(
            config=self.config,
            input_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_position=cache_position,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )
        linear_attn_mask = self._update_linear_attn_mask(
            attention_mask, cache_position)

        hidden_states = inputs_embeds
        if past_key_values is not None:
            assert isinstance(past_key_values, KimiDynamicCache)

        block_residual = None
        if self.use_attn_residuals:
            block_residual = hidden_states.new_zeros(
                hidden_states.shape[0] * hidden_states.shape[1], 0,
                hidden_states.shape[2])

        for decoder_layer in self.layers:
            layer_mask = linear_attn_mask if decoder_layer.is_linear_attn else causal_mask

            if self.use_attn_residuals:
                hidden_states, block_residual = decoder_layer(
                    hidden_states,
                    attention_mask=layer_mask,
                    past_key_values=past_key_values,
                    cache_position=cache_position,
                    block_residual=block_residual,
                    **kwargs,
                )
            else:
                hidden_states = decoder_layer(
                    hidden_states,
                    attention_mask=layer_mask,
                    past_key_values=past_key_values,
                    cache_position=cache_position,
                    **kwargs,
                )

        if self.use_attn_residuals:
            hidden_states = self._apply_output_attn_res(
                hidden_states, block_residual)

        hidden_states = self.norm(hidden_states)

… lines 1220–1225 omitted …

    def _apply_output_attn_res(self, hidden_states, block_residual):
        batch_size, seq_len, hidden_size = hidden_states.shape
        return _apply_attn_res(
            hidden_states.view(-1, hidden_size),
            block_residual,
            self.output_attn_res_proj,
            self.output_attn_res_norm,
        ).view(batch_size, seq_len, hidden_size)