K3: remove the add operation for megamoe path (#51146)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Aug 6, 2026 at 09:48 UTC
f85c1d2f84f55b061e9b27f1a64e7c2ab5dea82c
1 file changed
+27
-8
vllm/models/kimi_k3/nvidia/model.py
+27
-8
@@ -242,9 +242,23 @@ class KimiRoutedOutputTransform(nn.Module):
242
self.norm = norm
243
self.up_proj = up_proj
244
245
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
245
+ def forward(
246
+ self,
247
+ hidden_states: torch.Tensor,
248
+ residual: torch.Tensor | None = None,
249
+ ) -> torch.Tensor:
250
+ """Project the routed latent back to the hidden dim.
251
+
252
+ Args:
253
+ hidden_states: Routed expert output in latent space.
254
+ residual: Optional tensor of the up-projection's output shape to
255
+ accumulate into. It is consumed in the GEMM's beta-add
256
+ epilogue, so adding it costs no extra kernel.
257
+ """
258
if self.norm is not None:
259
hidden_states = self.norm(hidden_states)
260
+ if residual is not None:
261
+ return residual.addmm_(hidden_states, self.up_proj.weight.t())
262
hidden_states, _ = self.up_proj(hidden_states)
263
return hidden_states
264
@@ -523,8 +537,8 @@ class KimiMoE(nn.Module):
537
if self.num_shared_experts is not None:
538
shared_intermediate_size = moe_intermediate_size * self.num_shared_experts
539
self.shared_experts = KimiMLP(
526
- hidden_size=config.hidden_size, # 7618
527
- intermediate_size=shared_intermediate_size, # 3072*2
540
+ hidden_size=config.hidden_size,
541
+ intermediate_size=shared_intermediate_size,
542
hidden_act=config.hidden_act,
543
quant_config=quant_config,
544
reduce_results=False,
@@ -721,11 +735,16 @@ class KimiMoE(nn.Module):
735
topk_ids,
736
activation_clamp=None,
737
)
724
- final_hidden_states = self.routed_output_transform(final_hidden_states)
725
- if self.shared_experts is not None:
726
- final_hidden_states = final_hidden_states + self.shared_experts(
727
- hidden_states
728
- )
738
+ # The shared output is folded into the up-projection GEMM's beta-add
739
+ # epilogue, so combining the two branches costs no extra kernel.
740
+ shared_output = (
741
+ self.shared_experts(hidden_states)
742
+ if self.shared_experts is not None
743
+ else None
744
+ )
745
+ final_hidden_states = self.routed_output_transform(
746
+ final_hidden_states, residual=shared_output
747
+ )
748
else:
749
# Routed experts consume the down-projected latent; shared experts
750
# (inside MoERunner) get the original hidden states via