[Bugfix][NPU] Fix/Refactor routed scaling factor application in MoE routing (#31449)

2026-07-18 10:59:13 +08:00
parent 67e7f8d13a
commit 359009fa00
3 changed files with 4 additions and 1 deletions
@@ -74,7 +74,8 @@ def fused_topk_npu(
        topk_weights = scores.gather(1, topk_ids)
        if renormalize:
            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
-            topk_weights = _apply_routed_scaling_after_renorm(topk_weights, topk_config)
+        else:
+            topk_weights = topk_weights * topk_config.routed_scaling_factor
        topk_weights = topk_weights.to(torch.float32)

    # Support grouped top-k or correction bias or sigmoid or routed_scaling_factor
@@ -265,6 +265,7 @@ class LLaDA2MoeSparseMoeBlock(nn.Module):
            correction_bias=self.correction_bias,
            scoring_func=self.score_function,
            routed_scaling_factor=self.routed_scaling_factor,
+            apply_routed_scaling_factor_on_output=_is_npu,
        )

        self.experts = get_moe_impl_class(quant_config)(