From a15dbbd63daaba913b3302a858ad7e20e6dcc219 Mon Sep 17 00:00:00 2001
From: liuzhenwei <zhenweiliu@habana.ai>
Date: Mon, 21 Sep 2026 13:15:29 +0800
Subject: [PATCH] [XPU][CI] fallback to 7-args of moe_align_block_size (#57855)

Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
---
 vllm/_custom_ops.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
index 1d38bcade0..dcecce71b3 100644
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -2274,6 +2274,19 @@ def moe_align_block_size(
     expert_map: torch.Tensor | None = None,
     scatter_idx: torch.Tensor | None = None,
 ) -> None:
+    if current_platform.is_xpu():
+        if scatter_idx is not None:
+            raise NotImplementedError("scatter_idx is not supported on XPU")
+        torch.ops._moe_C.moe_align_block_size(
+            topk_ids,
+            num_experts,
+            block_size,
+            sorted_token_ids,
+            experts_ids,
+            num_tokens_post_pad,
+            expert_map,
+        )
+        return
     torch.ops._moe_C.moe_align_block_size(
         topk_ids,
         num_experts,