First layer

Checkpoint configuration and implementation moonshotai/Kimi-K3 @ f831ab6, fla-org/flash-linear-attention @ 07ca1e4 · parameter count checked against the checkpoint’s tensor headers · text model only · How to read

93 layers › First layer

Stored sums

×

supplied · no parametersblock_residual

The sums kept from earlier groups of 12 layers, the embedding first.

shape
bb [B batch, T positions, m stored sums, 7168 features]
Sources

Contextual evidence from Token embedding

The token embedding is a learned table of vocabulary × width rows, looked up by token ID.

embedding · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 1096–1158

Read code excerpt
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx)

… lines 1098–1155 omitted …

        # Get inputs_embeds
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

Contextual evidence from Token embedding

The released text configuration: 93 layers of width 7,168, the lists of Kimi Delta Attention and full-attention layers, attention-residual blocks of 12 layers, one dense layer, 896 routed experts of which 16 are used, 2 shared experts, a routed-expert latent of 3,584 features, MLA ranks and head widths, the SiTU activation, vocabulary 163,840, untied output head.

config · moonshotai/Kimi-K3 / config.json
Pinned revision f831ab668142 ↗ · lines 18–254

Read code excerpt
  "text_config": {
    "_name_or_path": "",
    "activation_situ_beta": 4.0,
    "activation_situ_linear_beta": 25.0,
    "add_cross_attention": false,
    "architectures": [
      "KimiLinearForCausalLM"
    ],
    "attn_res_block_size": 12,

… lines 27–45 omitted …

    "first_k_dense_replace": 1,
    "forced_bos_token_id": null,
    "forced_eos_token_id": null,
    "hidden_act": "situ",

… lines 50–65 omitted …

    "linear_attn_config": {
      "full_attn_layers": [
        4,
        8,
        12,
        16,
        20,
        24,
        28,
        32,
        36,
        40,
        44,
        48,
        52,
        56,
        60,
        64,
        68,
        72,
        76,
        80,
        84,
        88,
        92,
        93
      ],
      "gate_lower_bound": -5.0,
      "head_dim": 128,
      "kda_layers": [
        1,
        2,
        3,
        5,
        6,
        7,
        9,
        10,
        11,
        13,
        14,
        15,
        17,
        18,
        19,
        21,
        22,
        23,
        25,
        26,
        27,
        29,
        30,
        31,
        33,
        34,
        35,
        37,
        38,
        39,
        41,
        42,
        43,
        45,
        46,
        47,
        49,
        50,
        51,
        53,
        54,
        55,
        57,
        58,
        59,
        61,
        62,
        63,
        65,
        66,
        67,
        69,
        70,
        71,
        73,
        74,
        75,
        77,
        78,
        79,
        81,
        82,
        83,
        85,
        86,
        87,
        89,
        90,
        91
      ],
      "num_heads": 96,
      "short_conv_kernel_size": 4,
      "use_full_rank_gate": true
    },
    "max_length": 20,
    "max_position_embeddings": 1048576,
    "min_length": 0,
    "mla_use_nope": true,

… lines 174–175 omitted …

    "moe_intermediate_size": 3072,
    "moe_layer_freq": 1,
    "moe_renormalize": true,
    "moe_router_activation_func": "sigmoid",
    "no_repeat_ngram_size": 0,
    "num_attention_heads": 96,
    "num_beam_groups": 1,
    "num_beams": 1,
    "num_expert_group": 1,
    "num_experts": 896,
    "num_experts_per_token": 16,
    "num_hidden_layers": 93,
    "num_key_value_heads": 96,
    "num_nextn_predict_layers": 0,
    "num_return_sequences": 1,
    "num_shared_experts": 2,
    "output_attentions": false,
    "output_hidden_states": false,
    "output_scores": false,
    "pad_token_id": 163839,
    "prefix": null,
    "problem_type": null,
    "pruned_heads": {},
    "q_lora_rank": 1536,
    "qk_nope_head_dim": 128,
    "qk_rope_head_dim": 64,

… lines 202–245 omitted …

    "routed_expert_hidden_size": 3584,
    "routed_scaling_factor": 1.0,
    "sep_token_id": null,
    "suppress_tokens": null,
    "task_specific_params": null,
    "temperature": 1.0,
    "tf_legacy_loss": false,
    "tie_encoder_decoder": false,
    "tie_word_embeddings": false,

Contextual evidence from Attention input mix

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)

Contextual evidence from Attention input mix

With attention residuals, a layer mixes its attention input from the stored sums and the running sum, stores the running sum and starts a new one when its index is divisible by the block size, adds the attention output to the running sum, mixes the feed-forward input the same way, and adds the feed-forward output.

block-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 973–1046

Read code excerpt
    def _forward_attn_residual(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: tuple[torch.Tensor] | None = None,
        output_attentions: bool | None = False,
        use_cache: bool | None = False,
        block_residual: torch.Tensor | None = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ):
        batch_size, seq_len, hidden_size = hidden_states.shape
        prefix_sum = hidden_states

        if block_residual is not None and block_residual.shape[1] > 0:
            hidden_states = _apply_attn_res(
                prefix_sum.view(-1, hidden_size),
                block_residual,
                self.self_attention_res_proj,
                self.self_attention_res_norm,
            ).view(batch_size, seq_len, hidden_size)

        if self.layer_idx % self.attn_res_block_size == 0:
            block_residual = torch.cat(
                [block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1)
            prefix_sum = None

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        if self.is_linear_attn is False:
            hidden_states = self.self_attn(
                hidden_states=hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                output_attentions=output_attentions,
                use_cache=use_cache,
                **kwargs,
            )
        else:
            hidden_states = self.self_attn(
                hidden_states=hidden_states,
                attention_mask=attention_mask,
                cache_params=past_key_values,
                output_attentions=output_attentions,
                use_cache=use_cache,
                **kwargs,
            )

        if prefix_sum is not None:
            prefix_sum = prefix_sum + hidden_states
        else:
            prefix_sum = hidden_states

        hidden_states = _apply_attn_res(
            prefix_sum.view(-1, hidden_size),
            block_residual,
            self.mlp_res_proj,
            self.mlp_res_norm,
        ).view(batch_size, seq_len, hidden_size)

        hidden_states = self.post_attention_layernorm(hidden_states)
        if hasattr(self, "block_sparse_moe"):
            hidden_states = self.block_sparse_moe(hidden_states)
        else:
            hidden_states = self.mlp(hidden_states)

        if prefix_sum is None:
            prefix_sum = hidden_states
        else:
            prefix_sum = prefix_sum + hidden_states

        return prefix_sum, block_residual

Contextual evidence from Attention input mix

A layer has Kimi Delta Attention or latent attention by its index, a dense MLP in the first first_k_dense_replace layers and a MoE in the rest, two pre-norms, and for each of its two attention-residual mixes an RMSNorm and a one-row projection.

block-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 878–917

Read code excerpt
    def __init__(self, config: KimiLinearConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.config = config
        self.layer_idx = layer_idx
        if config.is_kda_layer(layer_idx):
            self.is_linear_attn = True
            self.self_attn = KimiDeltaAttention(
                config=config, layer_idx=layer_idx)
        elif config.is_mla:
            self.is_linear_attn = False
            self.self_attn = KimiMLAAttention(
                config=config, layer_idx=layer_idx)
        else:
            raise NotImplementedError
        if (
            config.num_experts is not None
            and layer_idx >= config.first_k_dense_replace
            and layer_idx % getattr(config, "moe_layer_freq", 1) == 0
        ):
            self.block_sparse_moe = KimiSparseMoeBlock(config)
        else:
            self.mlp = KimiMLP(config)
        self.input_layernorm = KimiRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = KimiRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)

        # Attention residual
        self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None
        if self.use_attn_residuals:
            self.attn_res_block_size = config.attn_res_block_size
            self.self_attention_res_norm = KimiRMSNorm(
                config.hidden_size, eps=config.rms_norm_eps)
            self.mlp_res_norm = KimiRMSNorm(
                config.hidden_size, eps=config.rms_norm_eps)
            self.self_attention_res_proj = nn.Linear(
                config.hidden_size, 1, bias=False)
            self.mlp_res_proj = nn.Linear(
                config.hidden_size, 1, bias=False)