93 layers › 68 KDA layers › Latent mixture of experts
learnedself.experts · KimiBlockSparseMLP
One of 896 independent gated MLPs on the latent; each position runs only the 16 its router selected.
All 6 steps
- in
- [B batch, 1 positions, 3584 latent]
- out
- [B batch, 1 positions, 3584]
- parameters
- 33,030,144 · × 68 KDA layers × 896 experts = 2,012,460,613,632
Parameters
33,030,144 parameters per expert, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Gate projection | 11,010,048 | |
| Up projection | 11,010,048 | |
| Down projection | 11,010,048 | |
| Per-instance total | 33,030,144 | |
Across 896 independent experts: 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 implementation sorts the positions by expert, runs each expert on its contiguous slice and scatters the outputs back (modeling_kimi_linear.py 841–874).
Codemodeling_kimi_linear.py 242–874
class KimiBlockSparseMLP(nn.Module):
def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None):
super().__init__()
self.config = config
self.ffn_dim = config.intermediate_size if intermediate_size is None else intermediate_size
self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size
self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate
self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down
self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up
if config.hidden_act == "situ":
beta, linear_beta = _get_situ_activation_params(config)
self.act_fn = SituAndMul(
beta=beta,
linear_beta=linear_beta,
)
else:
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, hidden_states):
if self.config.hidden_act == "situ":
gate_up = torch.cat([self.w1(hidden_states), self.w3(hidden_states)], dim=-1)
current_hidden_states = self.act_fn(gate_up)
else:
current_hidden_states = self.act_fn(
self.w1(hidden_states)) * self.w3(hidden_states)
current_hidden_states = self.w2(current_hidden_states)
return current_hidden_states 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,
) 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 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_outSources3
Direct evidence
Each routed expert is its own gated MLP of moe_intermediate_size hidden features on the routed-expert width, with the SiTU activation.
expert · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 242–270
Read code excerpt
class KimiBlockSparseMLP(nn.Module):
def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None):
super().__init__()
self.config = config
self.ffn_dim = config.intermediate_size if intermediate_size is None else intermediate_size
self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size
self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate
self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down
self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up
if config.hidden_act == "situ":
beta, linear_beta = _get_situ_activation_params(config)
self.act_fn = SituAndMul(
beta=beta,
linear_beta=linear_beta,
)
else:
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, hidden_states):
if self.config.hidden_act == "situ":
gate_up = torch.cat([self.w1(hidden_states), self.w3(hidden_states)], dim=-1)
current_hidden_states = self.act_fn(gate_up)
else:
current_hidden_states = self.act_fn(
self.w1(hidden_states)) * self.w3(hidden_states)
current_hidden_states = self.w2(current_hidden_states)
return current_hidden_statesDirect 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