Kimi K3 · Latent mixture of experts

93 layers › 24 MLA layers

Latent mixture of experts

learnedself.block_sparse_moe · KimiSparseMoeBlock

Each position runs 16 of 896 routed experts in a 3,584-feature latent, plus shared experts on the full width.

Why

Latent. In a conventional mixture each selected expert receives the full-width token, so communication and expert-weight traffic grow with the number of experts chosen. Here the routed experts work in a compact latent while the shared experts keep a full-width path for common transformations, which lets the layer use 896 routed experts with 16 active per token. Kimi K3 report §2.3

Stable. At that sparsity the authors name two failures: the routed path is a chain of nearly four matrix multiplications whose activations explode at this scale, and balancing the load of so many experts goes beyond what the earlier bias update handles. The norm before the up-projection, SiTU-GLU and the quantile rule answer them. Kimi K3 report §2.3

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

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

All 23 steps

\text{Router:}

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}

\text{Routed expert:}

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}

\text{Shared experts:}

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}}

in: $x_{\mathrm{ffn}}$ [B batch, 1 position, 7168 d]
out: $\Delta_{\mathrm{ffn}}$ [B batch, 1 position, 7168 d]
parameters: 29,784,936,832 · × 24 MLA layers = 714,838,483,968

Parameters

29,784,936,832 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Router	6,423,424 across its components	6,423,424
Latent down-projection	$7{,}168 \times 3{,}584$	25,690,112
Routed expert	$33{,}030{,}144 \times 896$	29,595,009,024
Latent RMSNorm	$3{,}584$	3,584
Latent up-projection	$3{,}584 \times 7{,}168$	25,690,112
Shared experts	132,120,576 across its components	132,120,576
Per-instance total		29,784,936,832

Across 24 independent MLA layers: $29{,}784{,}936{,}832 \times 24 = 714{,}838{,}483{,}968$ 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.

Explanation

The routed experts’ weights are stored as 4-bit MXFP4 with one scale per 32 weights; the rest of the layer’s weights are bfloat16 and the router’s selection bias is float32 (config.json, quantization_config; checkpoint-headers.json).

Codemodeling_kimi_linear.py 768–874

modeling_kimi_linear.py 768–813

    def __init__(self, config: KimiLinearConfig):
        super().__init__()
        self.config = config
        self.hidden_dim = config.hidden_size
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_token
        self.moe_renormalize = config.moe_renormalize

        self.use_latent_moe = getattr(config, "routed_expert_hidden_size", None) is not None
        self.moe_hidden_size = (
            config.routed_expert_hidden_size
            if self.use_latent_moe else config.hidden_size
        )
        self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False)

        self.ep_size = 1
        self.experts_per_rank = config.num_experts
        self.ep_rank = 0
        self.experts = nn.ModuleList(
            [
                KimiBlockSparseMLP(
                    config,
                    hidden_size=self.moe_hidden_size,
                    intermediate_size=config.moe_intermediate_size,
                )
                for _ in range(config.num_experts)
            ],
        )
        self.gate = KimiMoEGate(config)
        if config.num_shared_experts is not None:
            intermediate_size = config.moe_intermediate_size * config.num_shared_experts
            self.shared_experts = KimiMLP(
                config=config, intermediate_size=intermediate_size,
            )

        if self.use_latent_moe:
            self.routed_expert_down_proj = nn.Linear(
                config.hidden_size, self.moe_hidden_size, bias=False,
            )
            self.routed_expert_up_proj = nn.Linear(
                self.moe_hidden_size, config.hidden_size, bias=False,
            )
            if self.latent_moe_use_norm:
                self.routed_expert_norm = KimiRMSNorm(
                    self.moe_hidden_size, eps=config.rms_norm_eps,
                )

modeling_kimi_linear.py 815–838

    def forward(self, hidden_states):
        identity = hidden_states
        orig_shape = hidden_states.shape
        topk_idx, topk_weight = self.gate(hidden_states)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

        if self.use_latent_moe:
            hidden_states = self.routed_expert_down_proj(hidden_states)

        if not self.training:
            y = self.moe_infer(hidden_states, topk_idx, topk_weight)
        else:
            raise NotImplementedError("Training mode is not supported in KimiSparseMoeBlock")

        if self.use_latent_moe:
            if self.latent_moe_use_norm:
                y = self.routed_expert_norm(y)
            y = self.routed_expert_up_proj(y)

        y = y.view(*orig_shape)

        if self.config.num_shared_experts is not None:
            y = y + self.shared_experts(identity)
        return y

modeling_kimi_linear.py 841–874

    def moe_infer(self, x, topk_ids, topk_weight):
        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
        cnts.scatter_(1, topk_ids, 1)
        tokens_per_expert = cnts.sum(dim=0)
        idxs = topk_ids.view(-1).argsort()
        sorted_tokens = x[idxs // topk_ids.shape[1]]

        tokens_per_expert = tokens_per_expert.cpu().numpy()

        outputs = []
        start_idx = 0
        for i, num_tokens in enumerate(tokens_per_expert):
            end_idx = start_idx + num_tokens
            if num_tokens == 0:
                continue
            expert = self.experts[i + self.ep_rank * self.experts_per_rank]
            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
            expert_out = expert(tokens_for_this_expert)
            outputs.append(expert_out)
            start_idx = end_idx

        outs = torch.cat(outputs, dim=0) if len(
            outputs) else sorted_tokens.new_empty(0)

        new_x = torch.empty_like(outs)
        new_x[idxs] = outs
        final_out = (
            new_x.view(*topk_ids.shape, -1)
            .type(topk_weight.dtype)
            .mul_(topk_weight.unsqueeze(dim=-1))
            .sum(dim=1)
            .type(new_x.dtype)
        )
        return final_out

Sources2

Direct evidence

The MoE layer holds the routed experts on the latent width, the router, one shared-expert MLP of num_shared_experts × moe_intermediate_size hidden features on the full width, a down-projection into the latent, an up-projection out of it and an RMSNorm before the up-projection.

moe-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 768–813

Read code excerpt

    def __init__(self, config: KimiLinearConfig):
        super().__init__()
        self.config = config
        self.hidden_dim = config.hidden_size
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_token
        self.moe_renormalize = config.moe_renormalize

        self.use_latent_moe = getattr(config, "routed_expert_hidden_size", None) is not None
        self.moe_hidden_size = (
            config.routed_expert_hidden_size
            if self.use_latent_moe else config.hidden_size
        )
        self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False)

        self.ep_size = 1
        self.experts_per_rank = config.num_experts
        self.ep_rank = 0
        self.experts = nn.ModuleList(
            [
                KimiBlockSparseMLP(
                    config,
                    hidden_size=self.moe_hidden_size,
                    intermediate_size=config.moe_intermediate_size,
                )
                for _ in range(config.num_experts)
            ],
        )
        self.gate = KimiMoEGate(config)
        if config.num_shared_experts is not None:
            intermediate_size = config.moe_intermediate_size * config.num_shared_experts
            self.shared_experts = KimiMLP(
                config=config, intermediate_size=intermediate_size,
            )

        if self.use_latent_moe:
            self.routed_expert_down_proj = nn.Linear(
                config.hidden_size, self.moe_hidden_size, bias=False,
            )
            self.routed_expert_up_proj = nn.Linear(
                self.moe_hidden_size, config.hidden_size, bias=False,
            )
            if self.latent_moe_use_norm:
                self.routed_expert_norm = KimiRMSNorm(
                    self.moe_hidden_size, eps=config.rms_norm_eps,
                )

Direct evidence

The router scores the full-width input; the input is projected down, each selected expert runs on its tokens and their outputs are weighted and summed, the sum is normalized and projected up, and the shared experts’ output on the full-width input is added.

moe-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 815–874

Read code excerpt

    def forward(self, hidden_states):
        identity = hidden_states
        orig_shape = hidden_states.shape
        topk_idx, topk_weight = self.gate(hidden_states)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

        if self.use_latent_moe:
            hidden_states = self.routed_expert_down_proj(hidden_states)

        if not self.training:
            y = self.moe_infer(hidden_states, topk_idx, topk_weight)
        else:
            raise NotImplementedError("Training mode is not supported in KimiSparseMoeBlock")

        if self.use_latent_moe:
            if self.latent_moe_use_norm:
                y = self.routed_expert_norm(y)
            y = self.routed_expert_up_proj(y)

        y = y.view(*orig_shape)

        if self.config.num_shared_experts is not None:
            y = y + self.shared_experts(identity)
        return y

… lines 839–840 omitted …

    def moe_infer(self, x, topk_ids, topk_weight):
        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
        cnts.scatter_(1, topk_ids, 1)
        tokens_per_expert = cnts.sum(dim=0)
        idxs = topk_ids.view(-1).argsort()
        sorted_tokens = x[idxs // topk_ids.shape[1]]

        tokens_per_expert = tokens_per_expert.cpu().numpy()

        outputs = []
        start_idx = 0
        for i, num_tokens in enumerate(tokens_per_expert):
            end_idx = start_idx + num_tokens
            if num_tokens == 0:
                continue
            expert = self.experts[i + self.ep_rank * self.experts_per_rank]
            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
            expert_out = expert(tokens_for_this_expert)
            outputs.append(expert_out)
            start_idx = end_idx

        outs = torch.cat(outputs, dim=0) if len(
            outputs) else sorted_tokens.new_empty(0)

        new_x = torch.empty_like(outs)
        new_x[idxs] = outs
        final_out = (
            new_x.view(*topk_ids.shape, -1)
            .type(topk_weight.dtype)
            .mul_(topk_weight.unsqueeze(dim=-1))
            .sum(dim=1)
            .type(new_x.dtype)
        )
        return final_out