- in
- [B batch, T positions, 7168 d]
- out
- [B batch, T positions, 896 experts]
- parameters
- 6,423,424 · × 68 KDA layers = 436,792,832
Parameters
6,423,424 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Expert scores | 6,422,528 | |
| Select experts | 896 | |
| Per-instance total | 6,423,424 | |
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 router returns each position’s 16 weights with their expert indices (modeling_kimi_linear.py 759); spreads them over all 896 experts, zero for the other 880.
The pinned router runs only in evaluation mode (modeling_kimi_linear.py 721); how the bias is balanced in training is in the technical report, §2.3.3.
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