Fuse routed scaling factor in topk_reduce kernel (#6220)

This commit is contained in:
Xiaoyu Zhang
2025-06-08 02:06:50 +08:00
committed by GitHub
parent f5599ef124
commit 515ef4facb
10 changed files with 331 additions and 9 deletions

View File

@@ -346,7 +346,7 @@ class DeepseekV2MoE(nn.Module):
final_hidden_states = self.experts(
hidden_states=hidden_states, router_logits=router_logits
)
final_hidden_states *= self.routed_scaling_factor
if shared_output is not None:
final_hidden_states = final_hidden_states + shared_output
if self.tp_size > 1: