[XPU] Route grouped_topk to the fused _moe_C kernel on XPU (#53580)
Signed-off-by: Marceli Fylcek <marceli.fylcek@intel.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Marceli Fylcek committed
Sep 7, 2026 at 11:43 UTC
9b85112e130d8beffe9928ec0c86a4029d6ba7a8
2 files changed
+23
-3
vllm/_custom_ops.py
+22
-2
@@ -2500,9 +2500,9 @@ def grouped_topk(
2500
bias: Bias tensor (e_score_correction_bias). Always fused in kernel.
2501
scoring_func: 0=none (no activation), 1=sigmoid
2502
"""
2503
- if not current_platform.is_cuda():
2503
+ if not (current_platform.is_cuda() or current_platform.is_xpu()):
2504
raise NotImplementedError(
2505
- "The fused grouped_topk kernel is only available on CUDA platforms"
2505
+ "The fused grouped_topk kernel is only available on CUDA and XPU platforms"
2506
)
2507
return torch.ops._moe_C.grouped_topk(
2508
scores,
@@ -2516,6 +2516,26 @@ def grouped_topk(
2516
)
2517
2518
2519
+if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "grouped_topk"):
2520
+
2521
+ @register_fake("_moe_C::grouped_topk")
2522
+ def _grouped_topk_fake(
2523
+ scores: torch.Tensor,
2524
+ num_expert_group: int,
2525
+ topk_group: int,
2526
+ topk: int,
2527
+ renormalize: bool,
2528
+ routed_scaling_factor: float,
2529
+ bias: torch.Tensor,
2530
+ scoring_func: int = 0,
2531
+ ) -> tuple[torch.Tensor, torch.Tensor]:
2532
+ num_tokens = scores.size(0)
2533
+ return (
2534
+ scores.new_empty((num_tokens, topk), dtype=torch.float32),
2535
+ scores.new_empty((num_tokens, topk), dtype=torch.int32),
2536
+ )
2537
+
2538
+
2539
def moe_wna16_marlin_gemm(
2540
input: torch.Tensor,
2541
output: torch.Tensor | None,
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
+1
-1
@@ -90,7 +90,7 @@ def grouped_topk(
90
) -> tuple[torch.Tensor, torch.Tensor]:
91
if (
92
envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK
93
- and current_platform.is_cuda()
93
+ and (current_platform.is_cuda() or current_platform.is_xpu())
94
and num_expert_group <= 32
95
and topk <= 32
96
and e_score_correction_bias is not None