[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
|
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(
|
def fused_topk_npu(
|
||||||
hidden_states: torch.Tensor,
|
hidden_states: torch.Tensor,
|
||||||
router_logits: torch.Tensor,
|
router_logits: torch.Tensor,
|
||||||
@@ -60,8 +74,7 @@ def fused_topk_npu(
|
|||||||
topk_weights = scores.gather(1, topk_ids)
|
topk_weights = scores.gather(1, topk_ids)
|
||||||
if renormalize:
|
if renormalize:
|
||||||
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||||
else:
|
topk_weights = _apply_routed_scaling_after_renorm(topk_weights, topk_config)
|
||||||
topk_weights = topk_weights * topk_config.routed_scaling_factor
|
|
||||||
topk_weights = topk_weights.to(torch.float32)
|
topk_weights = topk_weights.to(torch.float32)
|
||||||
|
|
||||||
# Support grouped top-k or correction bias or sigmoid or routed_scaling_factor
|
# 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
|
# 1 for sigmoid, 0 for softmax
|
||||||
norm_type=1,
|
norm_type=1,
|
||||||
routed_scaling_factor=(
|
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),
|
eps=float(1e-20),
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user