93 layers › 68 KDA layers › Latent mixture of experts › Router
Select experts
×learnedtorch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False)
Keep the top 16 of 896 experts, ranked with a selection bias.
- in
- [B batch, 1 positions, 896 experts]
- out
- [B batch, 1 positions, 896 experts]
- parameters
- 896 · × 68 KDA layers = 60,928
Parameters
896 parameters owned by this component
| Component | Calculation | Parameters |
|---|---|---|
| learned vector | 896 | |
| Per-instance total | 896 | |
Across 68 independent KDA layers: 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 configuration sets one expert group, so the group stage keeps all 896 experts.
The bias (896 values) is added only when choosing experts; the weights use the unbiased affinities (modeling_kimi_linear.py 723, 750).
Codemodeling_kimi_linear.py 676–701, 703–759
def __init__(self, config: KimiLinearConfig):
super().__init__()
self.config = config
self.top_k = config.num_experts_per_token
self.num_experts = config.num_experts
self.routed_scaling_factor = config.routed_scaling_factor
self.moe_router_activation_func = config.moe_router_activation_func
self.num_expert_group = getattr(config, "num_expert_group", 1)
self.topk_group = getattr(config, "topk_group", 1)
# topk selection algorithm
self.moe_renormalize = config.moe_renormalize
self.gating_dim = config.hidden_size
self.weight = nn.Parameter(
torch.empty((self.num_experts, self.gating_dim)),
)
self.e_score_correction_bias = nn.Parameter(
torch.empty(self.num_experts),
)
self.reset_parameters()
def reset_parameters(self) -> None:
import torch.nn.init as init
init.kaiming_uniform_(self.weight, a=math.sqrt(5)) def forward(self, hidden_states):
bsz, seq_len, h = hidden_states.shape
# compute gating score
hidden_states = hidden_states.view(-1, h)
logits = F.linear(
hidden_states.type(torch.float32), self.weight.type(
torch.float32), None,
)
if self.moe_router_activation_func == "sigmoid":
scores = logits.sigmoid()
elif self.moe_router_activation_func == "softmax":
scores = logits.softmax(dim=1)
else:
raise NotImplementedError(
f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}",
)
# select top-k experts
assert not self.training
scores = scores.view(bsz * seq_len, -1)
scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
if self.num_expert_group > 1 and self.num_expert_group > self.topk_group:
group_scores = (
scores_for_choice.view(
bsz * seq_len, self.num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
) # [n, num_expert_group]
group_idx = torch.topk(
group_scores, k=self.topk_group, dim=-1, sorted=False,
)[
1
] # [n, top_k_group]
group_mask = torch.zeros_like(group_scores) # [n, num_expert_group]
group_mask.scatter_(1, group_idx, 1) # [n, num_expert_group]
score_mask = (
group_mask.unsqueeze(-1)
.expand(
bsz * seq_len, self.num_expert_group, self.num_experts // self.num_expert_group,
)
.reshape(bsz * seq_len, -1)
) # [n, e]
tmp_scores = scores_for_choice.masked_fill(
~score_mask.bool(), float("-inf")) # [n, e]
else:
tmp_scores = scores_for_choice
_, topk_idx = torch.topk(
tmp_scores, k=self.top_k, dim=-1, sorted=False,
)
topk_weight = scores.gather(1, topk_idx)
# norm gate to sum 1
if self.top_k > 1 and self.moe_renormalize:
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
topk_weight = topk_weight / denominator
# must multiply the scaling factor
topk_weight = topk_weight * self.routed_scaling_factor
return topk_idx, topk_weightSources2
Direct evidence
The router owns one score vector per routed expert on the full width and a per-expert correction bias.
gate-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 676–701
Read code excerpt
def __init__(self, config: KimiLinearConfig):
super().__init__()
self.config = config
self.top_k = config.num_experts_per_token
self.num_experts = config.num_experts
self.routed_scaling_factor = config.routed_scaling_factor
self.moe_router_activation_func = config.moe_router_activation_func
self.num_expert_group = getattr(config, "num_expert_group", 1)
self.topk_group = getattr(config, "topk_group", 1)
# topk selection algorithm
self.moe_renormalize = config.moe_renormalize
self.gating_dim = config.hidden_size
self.weight = nn.Parameter(
torch.empty((self.num_experts, self.gating_dim)),
)
self.e_score_correction_bias = nn.Parameter(
torch.empty(self.num_experts),
)
self.reset_parameters()
def reset_parameters(self) -> None:
import torch.nn.init as init
init.kaiming_uniform_(self.weight, a=math.sqrt(5))Direct evidence
Routing: sigmoid scores in float32, the bias added for selection, the top 16 experts chosen (one expert group), their unbiased scores renormalized and scaled by the route scale.
gate-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 703–759
Read code excerpt
def forward(self, hidden_states):
bsz, seq_len, h = hidden_states.shape
# compute gating score
hidden_states = hidden_states.view(-1, h)
logits = F.linear(
hidden_states.type(torch.float32), self.weight.type(
torch.float32), None,
)
if self.moe_router_activation_func == "sigmoid":
scores = logits.sigmoid()
elif self.moe_router_activation_func == "softmax":
scores = logits.softmax(dim=1)
else:
raise NotImplementedError(
f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}",
)
# select top-k experts
assert not self.training
scores = scores.view(bsz * seq_len, -1)
scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
if self.num_expert_group > 1 and self.num_expert_group > self.topk_group:
group_scores = (
scores_for_choice.view(
bsz * seq_len, self.num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
) # [n, num_expert_group]
group_idx = torch.topk(
group_scores, k=self.topk_group, dim=-1, sorted=False,
)[
1
] # [n, top_k_group]
group_mask = torch.zeros_like(group_scores) # [n, num_expert_group]
group_mask.scatter_(1, group_idx, 1) # [n, num_expert_group]
score_mask = (
group_mask.unsqueeze(-1)
.expand(
bsz * seq_len, self.num_expert_group, self.num_experts // self.num_expert_group,
)
.reshape(bsz * seq_len, -1)
) # [n, e]
tmp_scores = scores_for_choice.masked_fill(
~score_mask.bool(), float("-inf")) # [n, e]
else:
tmp_scores = scores_for_choice
_, topk_idx = torch.topk(
tmp_scores, k=self.top_k, dim=-1, sorted=False,
)
topk_weight = scores.gather(1, topk_idx)
# norm gate to sum 1
if self.top_k > 1 and self.moe_renormalize:
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
topk_weight = topk_weight / denominator
# must multiply the scaling factor
topk_weight = topk_weight * self.routed_scaling_factor
return topk_idx, topk_weight