[Bugfix][NPU] Fix/Refactor routed scaling factor application in MoE routing

[xdtbynd]

Motivation

Currently, the routed_scaling_factor is applied inconsistently depending on whether renormalize is enabled. When renormalize=True, the scaling factor is not applied after weight normalization, leading to incorrect router weights and potentially degraded model performance. This PR unifies the application logic to ensure the scaling factor is always correctly applied, regardless of the renormalization setting.

Modifications

• topk_weights handling
  • When renormalize=True, call apply_routed_scaling_after_renorm(topk_weights, topk_config) after renormalization to apply the scaling factor.
  • When renormalize=False, multiply topk_weights by topk_config.routed_scaling_factor directly.
• MOE implementation class initialization
  • Pass routed_scaling_factor to the expert implementation constructor.
  • Set apply_routed_scaling_factor_on_output=True to ensure the scaling factor is applied on the final output, consistent with the intended behavior.

These changes eliminate the conditional omission of the scaling factor and align the code with the design expectation that routed_scaling_factor should always affect the router weights/output.

Accuracy Tests

Verified on Ascend910 NPU:

Metric	Before Fix	After Fix	Threshold
GSM8K Accuracy	0.62	0.93	0.88
Output Throughput	65.54 token/s	77.65 token/s	70
test_bs_1_speed	PASSED	PASSED	-
test_gsm8k	FAILED	PASSED	-
Total runtime	-	389s	400s est

Speed Tests and Profiling

No speed regression. Throughput restored from 65.54 to 77.65 token/s (matching the pre-regression baseline of ~78 token/s). The fix only restores a multiplication that was incorrectly removed, adding negligible overhead.

Checklist

• Format your code according to the contribution guide
• Add unit tests according to the contribution guide — existing test_npu_llada2_mini.py covers this regression
• Update documentation according to the contribution guide — no doc changes needed
• Provide accuracy and speed benchmark results — see Accuracy Tests section above
• Follow the SGLang code style — minimal 3-line diff, follows existing code patterns

---

CI States

Latest PR Test (Base): ❌ Run #29574948888
Latest PR Test (Extra): ❌ Run #29574948768

[gemini-code-assist]

Code Review

This pull request updates the NPU MoE top-k logic to scale topk_weights by routed_scaling_factor when renormalize is false, and adjusts the condition for applying this scaling factor on the output. The review feedback highlights potential runtime TypeError issues if topk_config.routed_scaling_factor is None, suggesting safety guards to handle None values during multiplication and when passing the factor to the custom NPU operator.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If topk_config.routed_scaling_factor is None (since it is an Optional[float]), multiplying topk_weights by it directly will raise a TypeError at runtime. We should guard this multiplication by checking if routed_scaling_factor is not None.

Suggested change

	else:
	topk_weights = topk_weights * topk_config.routed_scaling_factor
	elif topk_config.routed_scaling_factor is not None:
	topk_weights = topk_weights * topk_config.routed_scaling_factor

[gemini-code-assist]

If topk_config.routed_scaling_factor is None, passing it directly to the custom NPU op npu_moe_gating_top_k might cause a TypeError or unexpected behavior in the PyTorch C++ binding. It is safer to guard against None and default to 1.0.

                topk_config.routed_scaling_factor
                if (not renormalize or topk_config.apply_routed_scaling_factor_on_output)
                and topk_config.routed_scaling_factor is not None
                else 1.0

[sglang-npu-bot]

/tag-and-rerun-ci