Kimi K3 · Output mix · OpenModelAtlas

Output mix

learnedself._apply_output_attn_res(hidden_states, block_residual) · 1215–1217

Mix the model’s output from the stored sums and the last running sum.

Why

Over depth. A standard residual connection compresses everything earlier into one state over depth, which the authors compare to a recurrent network’s bottleneck over time. Here each sublayer selects what it reads from the earlier representations with attention weights computed from a learned query, as attention does over positions. Kimi K3 report §2.2

The RMSNorm inside the weights keeps a layer with large outputs from dominating them. Kimi K3 report §2.2

{\htmlData{light=output-mix}{h_{\mathrm{out}}}}=\sum_{i}\operatorname{softmax}_i\!\big(w^{\top}\operatorname{RMSNorm}({\htmlClass{vl-stream-v}{v_i}})\big)\,{\htmlClass{vl-stream-v}{v_i}}

{\htmlClass{vl-stream-v}{v}}=[\,b_0,\ldots,b_{m-1},\;{\htmlData{light=final-block-state blocks}{s_{93}}}\,]

in: $b$ [B batch, T positions, m stored sums, 7168 features] $s_{93}$ [B batch, T positions, 7168 d]
out: $h_{\mathrm{out}}$ [B batch, T positions, 7168 d]
parameters: 14,336

Parameters

14,336 parameters owned by this component

Exact per-instance accounting
Component	Calculation	Parameters
RMSNorm gain	$7{,}168$	7,168
weights, no bias	$7{,}168 \times 1$	7,168
Per-instance total		14,336

Derived from the pinned configuration and implementation, with the per-head log-scale at the checkpoint’s 128 elements where the code creates 96; equal to the checkpoint’s own tensors. The quantisation scales are not counted.

An RMSNorm gain and one scoring vector, 7,168 values each.

Explanation

The stored sums are the embedding and the 7 finished groups of 12 layers; the last running sum is the unfinished group’s.

Codemodeling_kimi_linear.py 1075–1233

modeling_kimi_linear.py 1075–1088

def _apply_attn_res(prefix_sum, block_residual, proj, norm):
    """
    prefix_sum:     (num_tokens, hidden_size)
    block_residual: (num_tokens, num_blocks, hidden_size)
    """
    v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1)
    v_float = v.float()
    variance = v_float.pow(2).mean(-1, keepdim=True)
    k = v_float * torch.rsqrt(variance + norm.variance_epsilon)
    score_weight = norm.weight.float() * proj.weight.squeeze(0).float()
    scores = (k * score_weight).sum(-1)
    probs = scores.softmax(-1).unsqueeze(1)
    hidden_states = torch.matmul(probs, v_float).squeeze(1)
    return hidden_states.to(v.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

The mix: the stored sums and the running sum are RMS-normalized, scored by one learned vector, and averaged with softmax weights, in float32.

attn-res · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 1075–1088

Read code excerpt

def _apply_attn_res(prefix_sum, block_residual, proj, norm):
    """
    prefix_sum:     (num_tokens, hidden_size)
    block_residual: (num_tokens, num_blocks, hidden_size)
    """
    v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1)
    v_float = v.float()
    variance = v_float.pow(2).mean(-1, keepdim=True)
    k = v_float * torch.rsqrt(variance + norm.variance_epsilon)
    score_weight = norm.weight.float() * proj.weight.squeeze(0).float()
    scores = (k * score_weight).sum(-1)
    probs = scores.softmax(-1).unsqueeze(1)
    hidden_states = torch.matmul(probs, v_float).squeeze(1)
    return hidden_states.to(v.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)