A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 01:45 UTC
[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>
Ma Jian committed
May 22, 2026 at 18:22 UTC
d3d1cf6972607c53327b5ce1748e56a95fc41c37
2 files changed
+74
-1
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+13
-1
@@ -297,6 +297,8 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
297
"""
298
if current_platform.is_rocm():
299
return [Mxfp4MoeBackend.AITER_MXFP4_BF16]
300
+ if current_platform.is_xpu():
301
+ return [Mxfp4MoeBackend.XPU]
302
_AVAILABLE_BACKENDS = [
303
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
304
Mxfp4MoeBackend.DEEPGEMM_MXFP4,
@@ -1528,10 +1530,20 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1530
w13_bias,
1531
w2_bias,
1532
)
1533
+ elif mxfp4_backend == Mxfp4MoeBackend.XPU:
1534
+ # No additional transformation needed for XPU backend
1535
+ return (
1536
+ w13_weight,
1537
+ w2_weight,
1538
+ w13_weight_scale,
1539
+ w2_weight_scale,
1540
+ w13_bias,
1541
+ w2_bias,
1542
+ )
1543
else:
1544
raise ValueError(
1545
f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
1534
- f"Expected TRTLLM, Triton, or AITER backend."
1546
+ f"Expected TRTLLM, Triton, AITER, or XPU backend."
1547
)
1548
1549
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+61
@@ -57,6 +57,52 @@ def vllm_topk_sigmoid(
57
return topk_weights, topk_indices
58
59
60
+def _topk_softplus_sqrt_torch(
61
+ topk_weights: torch.Tensor,
62
+ topk_indices: torch.Tensor,
63
+ token_expert_indices: torch.Tensor,
64
+ gating_output: torch.Tensor,
65
+ renormalize: bool = False,
66
+ e_score_correction_bias: torch.Tensor | None = None,
67
+ input_tokens: torch.Tensor | None = None,
68
+ hash_indices_table: torch.Tensor | None = None,
69
+ routed_scaling_factor: float = 1.0,
70
+) -> tuple[torch.Tensor, ...]:
71
+ """Pure PyTorch fallback for topk_softplus_sqrt (XPU/CPU)."""
72
+ # scores = sqrt(softplus(gating_output))
73
+ scores = torch.sqrt(F.softplus(gating_output.float()))
74
+
75
+ # Bias is used for expert SELECTION only, not for weight computation.
76
+ # Using biased scores as weights flattens the distribution when the bias
77
+ # is near-uniform (e.g., DSv4-Flash where all biases ≈ 8.08).
78
+ if e_score_correction_bias is not None:
79
+ scores_for_choice = scores + e_score_correction_bias.float()
80
+ else:
81
+ scores_for_choice = scores
82
+
83
+ topk = topk_weights.shape[-1]
84
+
85
+ if hash_indices_table is not None and input_tokens is not None:
86
+ # Hash MoE: expert indices predetermined by lookup table
87
+ # hash_indices_table: [vocab_size, topk] mapping token_id -> expert_ids
88
+ expert_ids = hash_indices_table[input_tokens.long()] # [M, topk]
89
+ topk_indices.copy_(expert_ids)
90
+ # Gather weights from unbiased scores
91
+ weights = scores.gather(1, expert_ids.long())
92
+ else:
93
+ # Standard topk selection using biased scores
94
+ _, indices = torch.topk(scores_for_choice, k=topk, dim=-1)
95
+ topk_indices.copy_(indices)
96
+ # Gather weights from unbiased scores
97
+ weights = scores.gather(1, indices)
98
+
99
+ if renormalize:
100
+ weights = weights / (weights.sum(dim=-1, keepdim=True).clamp(min=1e-20))
101
+
102
+ topk_weights.copy_(weights * routed_scaling_factor)
103
+ return topk_weights, topk_indices
104
+
105
+
106
def vllm_topk_softplus_sqrt(
107
topk_weights: torch.Tensor,
108
topk_indices: torch.Tensor,
@@ -68,6 +114,21 @@ def vllm_topk_softplus_sqrt(
114
hash_indices_table: torch.Tensor | None = None,
115
routed_scaling_factor: float = 1.0,
116
) -> tuple[torch.Tensor, ...]:
117
+ from vllm.platforms import current_platform
118
+
119
+ if current_platform.is_xpu():
120
+ return _topk_softplus_sqrt_torch(
121
+ topk_weights,
122
+ topk_indices,
123
+ token_expert_indices,
124
+ gating_output,
125
+ renormalize,
126
+ e_score_correction_bias,
127
+ input_tokens,
128
+ hash_indices_table,
129
+ routed_scaling_factor,
130
+ )
131
+
132
ops.topk_hash_softplus_sqrt(
133
topk_weights,
134
topk_indices,