diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py index 9aef68476..d59498038 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/topk.py +++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py @@ -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 diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 94aec91b8..5fa3f64d8 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -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)( diff --git a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py index 467b73dfc..b8216893a 100644 --- a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py +++ b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py @@ -29,6 +29,7 @@ class TestLLaDA2Mini(GSM8KAscendMixin, CustomTestCase): "ascend", "--dllm-algorithm", "LowConfidence", # TODO: Add dLLM configurations + "--no-dllm-fdfo", # FDFO (PR #27551) halves single-batch speed on NPU; use sync mode ] env = { **os.environ,