- in
- [B batch, T positions, 7168 d]
- out
- [B batch, T positions, 7168 d]
- parameters
- 132,120,576 · × 68 KDA layers = 8,984,199,168
Parameters
132,120,576 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Gate projection | 44,040,192 | |
| Up projection | 44,040,192 | |
| Down projection | 44,040,192 | |
| Per-instance total | 132,120,576 | |
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
Its 6,144 hidden features are moe_intermediate_size × num_shared_experts (modeling_kimi_linear.py 798).
Codemodeling_kimi_linear.py 273–874
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_outclass KimiMLP(nn.Module):
def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
self.gate_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(
self.intermediate_size, self.hidden_size, bias=False)
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, x):
if self.config.hidden_act == "situ":
gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1)
down_proj = self.down_proj(self.act_fn(gate_up))
else:
down_proj = self.down_proj(self.act_fn(
self.gate_proj(x)) * self.up_proj(x))
return down_projSources3
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_outDirect evidence
The dense feed-forward network: down(SiTU(gate x, up x)).
mlp · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 273–301
Read code excerpt
class KimiMLP(nn.Module):
def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
self.gate_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(
self.intermediate_size, self.hidden_size, bias=False)
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, x):
if self.config.hidden_act == "situ":
gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1)
down_proj = self.down_proj(self.act_fn(gate_up))
else:
down_proj = self.down_proj(self.act_fn(
self.gate_proj(x)) * self.up_proj(x))
return down_proj