93 layers › 68 KDA layers › Latent mixture of experts
Latent up-projection
×learnedself.routed_expert_up_proj(y) · 832
Project the normalized sum back to 7,168 features.
- in
- [B batch, T positions, 3584 latent]
- out
- [B batch, T positions, 7168 d]
- parameters
- 25,690,112 · × 68 KDA layers = 1,746,927,616
Parameters
25,690,112 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Weight matrix | 25,690,112 | |
| Per-instance total | 25,690,112 | |
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 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