93 layers › 68 KDA layers › Latent mixture of experts
Input features
×tensorhidden_states · 815
The normalized features entering the MoE layer.
- shape
- [B batch, 1 position, 7168 d]
Sources
Contextual evidence from MoE RMSNorm
RMSNorm owns one gain per feature and computes in float32.
rmsnorm · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 226–236
Read code excerpt
class KimiRMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
dtype = hidden_states.dtype
x = hidden_states.float()
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
return self.weight * x.to(dtype)Contextual evidence from MoE RMSNorm
With attention residuals, a layer mixes its attention input from the stored sums and the running sum, stores the running sum and starts a new one when its index is divisible by the block size, adds the attention output to the running sum, mixes the feed-forward input the same way, and adds the feed-forward output.
block-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 973–1046
Read code excerpt
def _forward_attn_residual(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: tuple[torch.Tensor] | None = None,
output_attentions: bool | None = False,
use_cache: bool | None = False,
block_residual: torch.Tensor | None = None,
**kwargs: Unpack[FlashAttentionKwargs],
):
batch_size, seq_len, hidden_size = hidden_states.shape
prefix_sum = hidden_states
if block_residual is not None and block_residual.shape[1] > 0:
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.self_attention_res_proj,
self.self_attention_res_norm,
).view(batch_size, seq_len, hidden_size)
if self.layer_idx % self.attn_res_block_size == 0:
block_residual = torch.cat(
[block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1)
prefix_sum = None
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
if self.is_linear_attn is False:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
else:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
cache_params=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
if prefix_sum is not None:
prefix_sum = prefix_sum + hidden_states
else:
prefix_sum = hidden_states
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.mlp_res_proj,
self.mlp_res_norm,
).view(batch_size, seq_len, hidden_size)
hidden_states = self.post_attention_layernorm(hidden_states)
if hasattr(self, "block_sparse_moe"):
hidden_states = self.block_sparse_moe(hidden_states)
else:
hidden_states = self.mlp(hidden_states)
if prefix_sum is None:
prefix_sum = hidden_states
else:
prefix_sum = prefix_sum + hidden_states
return prefix_sum, block_residualContextual evidence from Router
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))Contextual evidence from Router
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_weightContextual evidence from Latent down-projection
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,
)Contextual evidence from Latent down-projection
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_outContextual evidence from Shared experts
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