From 9b85112e130d8beffe9928ec0c86a4029d6ba7a8 Mon Sep 17 00:00:00 2001 From: Marceli Fylcek Date: Mon, 7 Sep 2026 11:43:38 +0200 Subject: [PATCH] [XPU] Route grouped_topk to the fused _moe_C kernel on XPU (#53580) Signed-off-by: Marceli Fylcek Co-authored-by: Claude Opus 5 --- vllm/_custom_ops.py | 24 +++++++++++++++++-- .../fused_moe/router/grouped_topk_router.py | 2 +- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 450a7a57da..93b5eec0f2 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2500,9 +2500,9 @@ def grouped_topk( bias: Bias tensor (e_score_correction_bias). Always fused in kernel. scoring_func: 0=none (no activation), 1=sigmoid """ - if not current_platform.is_cuda(): + if not (current_platform.is_cuda() or current_platform.is_xpu()): raise NotImplementedError( - "The fused grouped_topk kernel is only available on CUDA platforms" + "The fused grouped_topk kernel is only available on CUDA and XPU platforms" ) return torch.ops._moe_C.grouped_topk( scores, @@ -2516,6 +2516,26 @@ def grouped_topk( ) +if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "grouped_topk"): + + @register_fake("_moe_C::grouped_topk") + def _grouped_topk_fake( + scores: torch.Tensor, + num_expert_group: int, + topk_group: int, + topk: int, + renormalize: bool, + routed_scaling_factor: float, + bias: torch.Tensor, + scoring_func: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = scores.size(0) + return ( + scores.new_empty((num_tokens, topk), dtype=torch.float32), + scores.new_empty((num_tokens, topk), dtype=torch.int32), + ) + + def moe_wna16_marlin_gemm( input: torch.Tensor, output: torch.Tensor | None, diff --git a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py index 1d4e4a8b5e..2df3489358 100644 --- a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py @@ -90,7 +90,7 @@ def grouped_topk( ) -> tuple[torch.Tensor, torch.Tensor]: if ( envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK - and current_platform.is_cuda() + and (current_platform.is_cuda() or current_platform.is_xpu()) and num_expert_group <= 32 and topk <= 32 and e_score_correction_bias is not None