Kimi K3 · KDA layer · OpenModelAtlas

93 layers

KDA layer

learnedKimiDecoderLayer(config, layer_idx)

Kimi Delta Attention and a latent mixture of experts, each reading a mix over depth and adding to the running sum; the first three of every four layers, layer 1 apart.

s'_n=\bar{s}_n+\Delta_{\mathrm{attn}}

s_{n+1}=s'_n+\Delta_{\mathrm{ffn}}

All 45 steps

h_{\mathrm{attn}}=\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},\;s_n\,]

b'=[\,b_0,\ldots,b_{m-1},\;s_n\,]

s_n\leftarrow\varnothing

\text{Kimi Delta Attention (KDA):}

{\htmlClass{vl-stream-q}{q^{\mathrm p}}}=x_{\mathrm{attn}}W_{{\htmlClass{vl-stream-q}{Q}}}

{\htmlClass{vl-stream-q}{q^{\mathrm c}_t}}=\operatorname{SiLU}\!\Big(\sum_{j=0}^{3}w_j\odot {\htmlClass{vl-stream-q}{q^{\mathrm p}_{t-j}}}\Big)

\hat {\htmlClass{vl-stream-q}{q}}={\htmlClass{vl-stream-q}{q^{\mathrm h}}}/\lVert {\htmlClass{vl-stream-q}{q^{\mathrm h}}}\rVert_2

{\htmlClass{vl-stream-k}{k^{\mathrm p}}}=x_{\mathrm{attn}}W_{{\htmlClass{vl-stream-k}{K}}}

{\htmlClass{vl-stream-k}{k^{\mathrm c}_t}}=\operatorname{SiLU}\!\Big(\sum_{j=0}^{3}w_j\odot {\htmlClass{vl-stream-k}{k^{\mathrm p}_{t-j}}}\Big)

\hat {\htmlClass{vl-stream-k}{k}}={\htmlClass{vl-stream-k}{k^{\mathrm h}}}/\lVert {\htmlClass{vl-stream-k}{k^{\mathrm h}}}\rVert_2

{\htmlClass{vl-stream-v}{v^{\mathrm p}}}=x_{\mathrm{attn}}W_{{\htmlClass{vl-stream-v}{V}}}

{\htmlClass{vl-stream-v}{v^{\mathrm c}_t}}=\operatorname{SiLU}\!\Big(\sum_{j=0}^{3}w_j\odot {\htmlClass{vl-stream-v}{v^{\mathrm p}_{t-j}}}\Big)

z^{\downarrow}=x_{\mathrm{attn}}W^{\downarrow}_{\alpha}

z=z^{\downarrow}W^{\uparrow}_{\alpha}

\alpha=\exp\!\big(-5\,\sigma\big(e^{A}(z+b_{\alpha})\big)\big)

w=x_{\mathrm{attn}}W_{\beta}

\beta=\sigma(w)

S\leftarrow\operatorname{Diag}(\alpha)\,S

S\leftarrow S+\beta\,\hat {\htmlClass{vl-stream-k}{k}}\,({\htmlClass{vl-stream-v}{v}}-S^{\top}\hat {\htmlClass{vl-stream-k}{k}})^{\top}

\tilde o=S^{\top}\hat {\htmlClass{vl-stream-q}{q}}/\sqrt{128}

g=x_{\mathrm{attn}}W_{g}

o^{g}=\sigma(g)\odot\operatorname{RMSNorm}(\tilde o)

\Delta_{\mathrm{attn}}=oW_{O}

s'_n=\bar{s}_n+\Delta_{\mathrm{attn}}

h_{\mathrm{ffn}}=\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},\;s'_n\,]

\text{Latent mixture of experts:}

z=x_{\mathrm{ffn}}W_{r}

s=\sigma(z)

E^{\mathrm{sel}}_i=\big[\,i\in\operatorname{top}_{16}(s+b)\,\big]

G_{i}=\frac{s_i\,E^{\mathrm{sel}}_i}{\sum_{j}s_j\,E^{\mathrm{sel}}_j}

x^{c}=x_{\mathrm{ffn}}W_{\downarrow}

g=x^{c}W_{\mathrm{gate}}

u=x^{c}W_{\mathrm{up}}

a=4\tanh\!\big(g/4\big)\,\sigma(g)

\bar u=25\tanh\!\big(u/25\big)

h=a\odot\bar u

E_i=hW_{\mathrm{down}}

u=\sum_{i=1}^{896}G_i\,E_i

y^{\mathrm{routed}}=\hat uW_{\uparrow}

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

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

a=4\tanh\!\big(g/4\big)\,\sigma(g)

\bar u=25\tanh\!\big(u/25\big)

h=a\odot\bar u

y^{\mathrm{shared}}=hW_{\mathrm{down}}

\Delta_{\mathrm{ffn}}=y^{\mathrm{routed}}+y^{\mathrm{shared}}

s_{n+1}=s'_n+\Delta_{\mathrm{ffn}}

in: $s_n$ [B batch, T positions, 7168 d] $b$ [B batch, T positions, m stored sums, 7168 features]
out: $s_{n+1}$ [B batch, T positions, 7168 d]
parameters: 30,228,720,256 · × 68 KDA layers = 2,055,552,977,408

Each layer of this kind receives the running sum $s_n$ from the layer before it: $s_1$ in layer 2. The kinds alternate as the stack’s period shows.

Parameters

30,228,720,256 parameters per KDA layer, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Attention input mix	$7{,}168 + 7{,}168 \times 1$	14,336
Attention RMSNorm	$7{,}168$	7,168
Kimi Delta Attention (KDA)	443,740,416 across its components	443,740,416
Feed-forward input mix	$7{,}168 + 7{,}168 \times 1$	14,336
MoE RMSNorm	$7{,}168$	7,168
Latent mixture of experts	29,784,936,832 across its components	29,784,936,832
Per-instance total		30,228,720,256

Across 68 independent KDA layers: $30{,}228{,}720{,}256 \times 68 = 2{,}055{,}552{,}977{,}408$ parameters.

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.

Codemodeling_kimi_linear.py 878–917, 973–1046 · configuration_kimi_k3.py 152–156

modeling_kimi_linear.py 878–917

    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)

modeling_kimi_linear.py 973–1046

    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

configuration_kimi_k3.py 152–156

    def is_kda_layer(self, layer_idx: int):
        return (
            self.linear_attn_config is not None
            and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
        )

Sources3

Direct evidence

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)

Direct evidence

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

Direct evidence

A layer uses Kimi Delta Attention when its 1-based index is in kda_layers; the lists are read 1-based.

layer-kinds · moonshotai/Kimi-K3 / configuration_kimi_k3.py
Pinned revision f831ab668142 ↗ · lines 152–156

Read code excerpt

    def is_kda_layer(self, layer_idx: int):
        return (
            self.linear_attn_config is not None
            and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
        )