From d3d1cf6972607c53327b5ce1748e56a95fc41c37 Mon Sep 17 00:00:00 2001
From: Ma Jian <jian1.ma@intel.com>
Date: Fri, 22 May 2026 18:22:45 +0800
Subject: [PATCH] [XPU]feat: add XPU fallback for MoE topk routing and MXFP4
 backend (#42951)

Signed-off-by: Ma Jian <jian1.ma@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
---
 .../layers/fused_moe/oracle/mxfp4.py          | 14 ++++-
 .../router/fused_topk_bias_router.py          | 61 +++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
index a69bae1b05..3ecc4f018f 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
@@ -297,6 +297,8 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
     """
     if current_platform.is_rocm():
         return [Mxfp4MoeBackend.AITER_MXFP4_BF16]
+    if current_platform.is_xpu():
+        return [Mxfp4MoeBackend.XPU]
     _AVAILABLE_BACKENDS = [
         Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
         Mxfp4MoeBackend.DEEPGEMM_MXFP4,
@@ -1528,10 +1530,20 @@ def convert_weight_to_mxfp4_moe_kernel_format(
             w13_bias,
             w2_bias,
         )
+    elif mxfp4_backend == Mxfp4MoeBackend.XPU:
+        # No additional transformation needed for XPU backend
+        return (
+            w13_weight,
+            w2_weight,
+            w13_weight_scale,
+            w2_weight_scale,
+            w13_bias,
+            w2_bias,
+        )
     else:
         raise ValueError(
             f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
-            f"Expected TRTLLM, Triton, or AITER backend."
+            f"Expected TRTLLM, Triton, AITER, or XPU backend."
         )
 
 
diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
index 0ca5c3f979..31a75c860d 100644
--- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
@@ -57,6 +57,52 @@ def vllm_topk_sigmoid(
     return topk_weights, topk_indices
 
 
+def _topk_softplus_sqrt_torch(
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    token_expert_indices: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    e_score_correction_bias: torch.Tensor | None = None,
+    input_tokens: torch.Tensor | None = None,
+    hash_indices_table: torch.Tensor | None = None,
+    routed_scaling_factor: float = 1.0,
+) -> tuple[torch.Tensor, ...]:
+    """Pure PyTorch fallback for topk_softplus_sqrt (XPU/CPU)."""
+    # scores = sqrt(softplus(gating_output))
+    scores = torch.sqrt(F.softplus(gating_output.float()))
+
+    # Bias is used for expert SELECTION only, not for weight computation.
+    # Using biased scores as weights flattens the distribution when the bias
+    # is near-uniform (e.g., DSv4-Flash where all biases ≈ 8.08).
+    if e_score_correction_bias is not None:
+        scores_for_choice = scores + e_score_correction_bias.float()
+    else:
+        scores_for_choice = scores
+
+    topk = topk_weights.shape[-1]
+
+    if hash_indices_table is not None and input_tokens is not None:
+        # Hash MoE: expert indices predetermined by lookup table
+        # hash_indices_table: [vocab_size, topk] mapping token_id -> expert_ids
+        expert_ids = hash_indices_table[input_tokens.long()]  # [M, topk]
+        topk_indices.copy_(expert_ids)
+        # Gather weights from unbiased scores
+        weights = scores.gather(1, expert_ids.long())
+    else:
+        # Standard topk selection using biased scores
+        _, indices = torch.topk(scores_for_choice, k=topk, dim=-1)
+        topk_indices.copy_(indices)
+        # Gather weights from unbiased scores
+        weights = scores.gather(1, indices)
+
+    if renormalize:
+        weights = weights / (weights.sum(dim=-1, keepdim=True).clamp(min=1e-20))
+
+    topk_weights.copy_(weights * routed_scaling_factor)
+    return topk_weights, topk_indices
+
+
 def vllm_topk_softplus_sqrt(
     topk_weights: torch.Tensor,
     topk_indices: torch.Tensor,
@@ -68,6 +114,21 @@ def vllm_topk_softplus_sqrt(
     hash_indices_table: torch.Tensor | None = None,
     routed_scaling_factor: float = 1.0,
 ) -> tuple[torch.Tensor, ...]:
+    from vllm.platforms import current_platform
+
+    if current_platform.is_xpu():
+        return _topk_softplus_sqrt_torch(
+            topk_weights,
+            topk_indices,
+            token_expert_indices,
+            gating_output,
+            renormalize,
+            e_score_correction_bias,
+            input_tokens,
+            hash_indices_table,
+            routed_scaling_factor,
+        )
+
     ops.topk_hash_softplus_sqrt(
         topk_weights,
         topk_indices,