- in
- [B batch, 1 positions, 7168 d]
- out
- [B batch, 1 positions, 7168 d]
- parameters
- 29,784,936,832 · × 68 KDA layers = 2,025,375,704,576
- mode
- decode · one new token after P = 1,000 positions
Parameters
29,784,936,832 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Router | 6,423,424 across its components | 6,423,424 |
| Latent down-projection | 25,690,112 | |
| Routed expert | 29,595,009,024 | |
| Latent RMSNorm | 3,584 | |
| Latent up-projection | 25,690,112 | |
| Shared experts | 132,120,576 across its components | 132,120,576 |
| Per-instance total | 29,784,936,832 | |
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 routed experts’ weights are stored as 4-bit MXFP4 with one scale per 32 weights; the rest of the layer’s weights are bfloat16 and the router’s selection bias is float32 (config.json, quantization_config; checkpoint-headers.json).
Codemodeling_kimi_linear.py 768–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_outSources2
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_out