[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:
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),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user