93 layers › 68 KDA layers › Latent mixture of experts › Router
Expert scores
×learnedF.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None)
One learned score vector per expert: 896 scores per position.
- in
- [B batch, 1 positions, 7168 d]
- out
- [B batch, 1 positions, 896 experts]
- parameters
- 6,422,528 · × 68 KDA layers = 436,731,904
In float32 (modeling_kimi_linear.py 707–710).
Parameters
6,422,528 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Weight matrix | 6,422,528 | |
| Per-instance total | 6,422,528 | |
This projection has no bias parameter.
The reference module holds the weight output by input, [out, in].
Across 68 independent KDA layers: parameters.
Derived from the configured matrix dimensions; the model total matches the checkpoint’s tensors.
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