[Intel XPU] Enable (biased) grouped topk for xpu (#31126)

test/registered/xpu/test_topk.py

@@ -8,12 +8,31 @@ 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 grouped_topk_gpu as native_grouped_topk
+from sglang.srt.layers.moe.topk import (
+    grouped_topk_xpu,
+)
 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 _scatter_by_expert(
+    weights: torch.Tensor, indices: torch.Tensor, num_columns: int
+) -> torch.Tensor:
+    """Scatter (weight, id) pairs into a dense ``[M, num_columns]`` tensor.
+
+    Makes the comparison independent of the per-row slot order, so the test does
+    not depend on how ties between equal scores are broken.
+    """
+    dense = torch.zeros(
+        (weights.shape[0], num_columns), dtype=torch.float32, device=weights.device
+    )
+    dense.scatter_(1, indices.long(), weights.float())
+    return dense
+
+
 # Nemotron-3 uses biased_grouped_topk
 class TestBiasedGroupedTopK(CustomTestCase):
     def _run_single_test(
@@ -96,6 +115,128 @@ class TestBiasedGroupedTopK(CustomTestCase):
                 routed_scaling_factor,
             )
 
+    def test_biased_grouped_topk(self):
+        # DeepSeek-V3 style grouped routing shape
+        E_num = 256
+        num_expert_group = 8
+        topk_value = 8
+        topk_group = 4
+        gating_dtype = torch.bfloat16
+        bias_dtype = torch.float32
+        renormalize = True
+        routed_scaling_factor = 2.5
+
+        torch.manual_seed(1024)
+        device = torch.device("xpu")
+
+        bs = [1, 2, 4, 8]
+        seq_len = 1024
+        num_tokens = [b * seq_len for b in bs]
+        num_fused_shared_experts_list = [0, 1]
+
+        for M in num_tokens:
+            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=torch.bfloat16, 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_grouped_topk(
+                    hidden_states.float(),
+                    gating_output.float(),
+                    correction_bias,
+                    topk_value,
+                    renormalize,
+                    num_expert_group,
+                    topk_group,
+                    num_fused_shared_experts,
+                    routed_scaling_factor=routed_scaling_factor,
+                )
+
+                # fused version
+                topk_weights, topk_ids = biased_grouped_topk_gpu(
+                    hidden_states,
+                    gating_output,
+                    correction_bias,
+                    topk_value,
+                    renormalize,
+                    num_expert_group,
+                    topk_group,
+                    num_fused_shared_experts,
+                    routed_scaling_factor,
+                )
+
+                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_grouped_topk(self):
+        # DeepSeek-V3 style grouped routing shape
+        E_num = 256
+        num_expert_group = 8
+        topk_value = 8
+        topk_group = 4
+        gating_dtype = torch.bfloat16
+        renormalize = True
+        routed_scaling_factor = 2.5
+
+        torch.manual_seed(1024)
+        device = torch.device("xpu")
+
+        bs = [1]
+        seq_len = 1024
+        num_tokens = [b * seq_len for b in bs]
+        num_fused_shared_experts_list = [0, 1]
+
+        for M in num_tokens:
+            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=torch.bfloat16, device=device)
+                gating_output = torch.randn(M, E_num, dtype=gating_dtype, device=device)
+
+                ref_topk_weights, ref_topk_ids = native_grouped_topk(
+                    hidden_states.float(),
+                    gating_output.float(),
+                    topk_value,
+                    renormalize,
+                    num_expert_group,
+                    topk_group,
+                    num_fused_shared_experts,
+                    routed_scaling_factor=routed_scaling_factor,
+                )
+
+                # fused version
+                topk_weights, topk_ids = grouped_topk_xpu(
+                    hidden_states,
+                    gating_output,
+                    topk_value,
+                    renormalize,
+                    num_expert_group,
+                    topk_group,
+                    num_fused_shared_experts,
+                    routed_scaling_factor,
+                )
+
+                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,
+                    ),
+                )
+
 
 if __name__ == "__main__":
     unittest.main()