[Intel XPU] Add xpu pass for biased_topk and hash_topk (#33323)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

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(

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,

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 = _set_seed_and_device()
-        device = torch.device("xpu")
 
         # 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 = _set_seed_and_device()
-        device = torch.device("xpu")
 
         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 = _set_seed_and_device()
-        device = torch.device("xpu")
 
         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()