From 56834422a18af26479cd9b179b42d140053d0762 Mon Sep 17 00:00:00 2001 From: gaopengff Date: Mon, 24 Aug 2026 12:18:22 +0800 Subject: [PATCH] [Intel XPU] Add xpu pass for biased_topk and hash_topk (#33323) Co-authored-by: Ma Mingfei --- python/sglang/srt/layers/moe/hash_topk.py | 39 +++++- python/sglang/srt/layers/moe/topk.py | 44 +++++- test/registered/xpu/test_topk.py | 159 +++++++++++++++++++++- 3 files changed, 233 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index 972510be7..271344732 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/python/sglang/srt/layers/moe/hash_topk.py @@ -23,12 +23,13 @@ from sglang.srt.layers.moe.topk import ( ) from sglang.srt.layers.moe.utils import has_per_rank_fused_shared_slots from sglang.srt.runtime_context import get_exec -from sglang.srt.utils import is_hip, is_npu +from sglang.srt.utils import is_hip, is_npu, is_xpu logger = logging.getLogger(__name__) _is_hip = is_hip() _is_npu = is_npu() +_is_xpu = is_xpu() class HashTopK(nn.Module): @@ -177,6 +178,38 @@ class HashTopK(nn.Module): return topk_weights, topk_ids + def _forward_xpu( + self, router_logits: torch.Tensor, input_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + # The XPU 'hash_topk' kernel currently supports the 'sqrtsoftplus' score func only. + # Other score funcs fall back to the torch implementation; more will be supported in the future. + if self.score_func == "sqrtsoftplus": + from sgl_kernel import hash_topk + + num_tokens = router_logits.size(0) + topk_routed = self.tid2eid.size(1) + topk_fused = topk_routed + self.num_fused_shared_experts + topk_ids = torch.empty( + (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device + ) + topk_weights = torch.empty( + (num_tokens, topk_fused), + dtype=torch.float32, + device=router_logits.device, + ) + hash_topk( + router_logits, + input_ids, + self.tid2eid, + topk_weights, + topk_ids, + self.routed_scaling_factor, + self.score_func, + ) + return topk_weights, topk_ids + else: + return self._forward_torch(router_logits, input_ids) + def forward( self, hidden_states: torch.Tensor, @@ -189,7 +222,9 @@ class HashTopK(nn.Module): input_ids.shape[0] == hidden_states.shape[0] == router_logits.shape[0] ), f"{input_ids.shape=} {hidden_states.shape=} {router_logits.shape=}" - if envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.get(): + if _is_xpu: + topk_weights, topk_ids = self._forward_xpu(router_logits, input_ids) + elif envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.get(): from sglang.kernels.ops.attention.dsv4 import hash_topk topk_weights, topk_ids = hash_topk( diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 5fd00a5c0..bc8f14f52 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1291,6 +1291,47 @@ def biased_topk_jit_kernel_impl( return topk_weights, topk_ids +def biased_topk_xpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + renormalize: bool, + scoring_func: str = "sigmoid", + num_fused_shared_experts: int = 0, + routed_scaling_factor: Optional[float] = None, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + + num_rows, _ = gating_output.shape + device = gating_output.device + + output = torch.empty(num_rows, topk, dtype=torch.float32, device=device) + indices = torch.empty(num_rows, topk, dtype=torch.int32, device=device) + + from sgl_kernel import biased_topk + + biased_topk( + gating_output, + correction_bias, + output, + indices, + topk, + scoring_func, + num_fused_shared_experts, + renormalize, + routed_scaling_factor=(routed_scaling_factor if routed_scaling_factor else 1.0), + apply_routed_scaling_factor_on_output=bool( + apply_routed_scaling_factor_on_output + ), + ) + + return output, indices + + @torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) def biased_grouped_topk_impl( hidden_states: torch.Tensor, @@ -2212,7 +2253,8 @@ def select_experts( assert not apply_routed_scaling_factor_on_output, "Not implemented" if scoring_func == "sqrtsoftplus" or scoring_func == "sigmoid": - topk_weights, topk_ids = biased_topk_jit_kernel_impl( + _biased_topk = biased_topk_xpu if _is_xpu else biased_topk_jit_kernel_impl + topk_weights, topk_ids = _biased_topk( hidden_states=hidden_states, gating_output=router_logits, correction_bias=correction_bias, diff --git a/test/registered/xpu/test_topk.py b/test/registered/xpu/test_topk.py index ba4f6abc0..bf3365e44 100644 --- a/test/registered/xpu/test_topk.py +++ b/test/registered/xpu/test_topk.py @@ -3,6 +3,8 @@ from typing import Optional import torch +from sglang.srt.layers.moe.hash_topk import HashTopK + torch.use_deterministic_algorithms(True) from sglang.srt.layers.moe.topk import ( @@ -11,16 +13,26 @@ from sglang.srt.layers.moe.topk import ( from sglang.srt.layers.moe.topk import ( biased_grouped_topk_impl as native_biased_grouped_topk, ) +from sglang.srt.layers.moe.topk import biased_topk_impl as native_biased_topk +from sglang.srt.layers.moe.topk import ( + biased_topk_xpu, +) from sglang.srt.layers.moe.topk import grouped_topk_gpu as native_grouped_topk from sglang.srt.layers.moe.topk import ( grouped_topk_xpu, ) +from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_xpu_ci from sglang.test.test_utils import CustomTestCase register_xpu_ci(est_time=5, suite="stage-b-test-1-gpu-xpu") +def _set_seed_and_device(): + torch.manual_seed(1024) + return torch.device("xpu") + + def _scatter_by_expert( weights: torch.Tensor, indices: torch.Tensor, num_columns: int ) -> torch.Tensor: @@ -83,8 +95,7 @@ class TestBiasedGroupedTopK(CustomTestCase): bias_dtype, routed_scaling_factor, ): - torch.manual_seed(1024) - device = torch.device("xpu") + device = _set_seed_and_device() # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating hidden_states = torch.randn(M, 100, dtype=torch.bfloat16, device=device) @@ -162,8 +173,7 @@ class TestBiasedGroupedTopK(CustomTestCase): renormalize = True routed_scaling_factor = 2.5 - torch.manual_seed(1024) - device = torch.device("xpu") + device = _set_seed_and_device() bs = [1, 2, 4, 8] seq_len = 1024 @@ -224,8 +234,7 @@ class TestBiasedGroupedTopK(CustomTestCase): renormalize = True routed_scaling_factor = 2.5 - torch.manual_seed(1024) - device = torch.device("xpu") + device = _set_seed_and_device() bs = [1] seq_len = 1024 @@ -271,6 +280,144 @@ class TestBiasedGroupedTopK(CustomTestCase): seq_len=seq_len, ) + def test_biased_topk(self): + # DeepSeek-V4 style routing shape + E_num_list = [256, 384] + topk_value = 6 + gating_dtype = torch.float32 + bias_dtype = torch.float32 + renormalize = True + scoring_func_list = ["sqrtsoftplus", "sigmoid"] + routed_scaling_factor = 2.5 + + device = _set_seed_and_device() + + bs = [1] + seq_len = 1024 + num_tokens = [b * seq_len for b in bs] + num_fused_shared_experts_list = [0, 1] + + for E_num in E_num_list: + for M in num_tokens: + for scoring_func in scoring_func_list: + for num_fused_shared_experts in num_fused_shared_experts_list: + + topk_routed = topk_value - num_fused_shared_experts + hidden_states = torch.randn( + M, 100, dtype=gating_dtype, device=device + ) + gating_output = torch.randn( + M, E_num, dtype=gating_dtype, device=device + ) + correction_bias = torch.randn( + E_num, dtype=bias_dtype, device=device + ) + + ref_topk_weights, ref_topk_ids = native_biased_topk( + hidden_states, + gating_output, + correction_bias, + topk_value, + renormalize, + scoring_func, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output=True, + ) + + # fused version + topk_weights, topk_ids = biased_topk_xpu( + hidden_states, + gating_output, + correction_bias, + topk_value, + renormalize, + scoring_func, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output=True, + ) + + torch.testing.assert_close( + _scatter_by_expert( + topk_weights[:, :topk_routed], + topk_ids[:, :topk_routed], + E_num, + ), + _scatter_by_expert( + ref_topk_weights[:, :topk_routed], + ref_topk_ids[:, :topk_routed], + E_num, + ), + ) + + def test_hash_topk(self): + """Guard the XPU fused hash-topk path against math/ID drift from torch.""" + device = _set_seed_and_device() + + E_num_list = [256, 384] + topk = 6 + vocab_size = 128 + dtype = torch.float32 + + bs = [1] + seq_len = 1024 + num_tokens = [b * seq_len for b in bs] + num_fused_shared_experts_list = [0, 1] + + with get_context().override_server_args(enable_waterfill=False): + for E_num in E_num_list: + for M in num_tokens: + for num_fused_shared_experts in num_fused_shared_experts_list: + hidden_states = torch.randn( + M, 1, dtype=torch.float32, device=device + ) + router_logits = torch.randn( + M, E_num, dtype=dtype, device=device + ) + input_ids = torch.randint( + low=0, + high=vocab_size, + size=(M,), + dtype=torch.int64, + device=device, + ) + + hash_topk = HashTopK( + topk=topk, + num_experts=E_num, + num_fused_shared_experts=num_fused_shared_experts, + vocab_size=vocab_size, + scoring_func="sqrtsoftplus", + routed_scaling_factor=2.5, + ).to(device) + topk_routed = hash_topk.tid2eid.shape[1] + with torch.no_grad(): + hash_topk.tid2eid.copy_( + torch.randint( + low=0, + high=E_num, + size=(vocab_size, topk_routed), + dtype=torch.int32, + device=device, + ) + ) + + ref_topk_weights, ref_topk_ids = hash_topk._forward_torch( + router_logits, input_ids + ) + + output = hash_topk( + hidden_states=hidden_states, + router_logits=router_logits, + input_ids=input_ids, + ) + + torch.testing.assert_close(output.topk_ids, ref_topk_ids) + torch.testing.assert_close( + output.topk_weights, ref_topk_weights + ) + if __name__ == "__main__": unittest.main()