[Bugfix][NPU] Fix Hunyuan3 model where MoE's routing_scaling_ratio is missing on NPU (#29909)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Jinyan Yi
2026-07-13 09:48:41 +03:00
committed by GitHub
co-authored by github-actions[bot] gemini-code-assist[bot] Even Zhou ronnie_zheng
parent 82e7cdcff9
commit 08d6d297e5
@@ -16,6 +16,20 @@ if TYPE_CHECKING:
from sglang.srt.layers.moe.topk import TopKConfig, TopKOutput
def _apply_routed_scaling_after_renorm(
topk_weights: torch.Tensor,
topk_config: "TopKConfig",
) -> torch.Tensor:
"""Mirror GPU post-renorm scaling when apply_routed_scaling_factor_on_output is set."""
if (
topk_config.renormalize
and topk_config.apply_routed_scaling_factor_on_output
and topk_config.routed_scaling_factor is not None
):
return topk_weights * topk_config.routed_scaling_factor
return topk_weights
def fused_topk_npu(
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
@@ -60,8 +74,7 @@ def fused_topk_npu(
topk_weights = scores.gather(1, topk_ids)
if renormalize:
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
else:
topk_weights = topk_weights * topk_config.routed_scaling_factor
topk_weights = _apply_routed_scaling_after_renorm(topk_weights, topk_config)
topk_weights = topk_weights.to(torch.float32)
# Support grouped top-k or correction bias or sigmoid or routed_scaling_factor
@@ -86,7 +99,9 @@ def fused_topk_npu(
# 1 for sigmoid, 0 for softmax
norm_type=1,
routed_scaling_factor=(
1 if renormalize else topk_config.routed_scaling_factor
topk_config.routed_scaling_factor
if topk_config.apply_routed_scaling_factor_on_output
else 1
),
eps=float(1e-20),
)