[XPU] [MoE] add quant input when prepare for fusedmoe (#47122)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Signed-off-by: zofia <110436990+zufangzhu@users.noreply.github.com> Co-authored-by: mayuyuace <qiming1.zhang@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
zofia committed
Jul 20, 2026 at 15:47 UTC
df13b5aef55c5cf6a3111b28e8d48b4a71575d8d
3 files changed
+37
-4
vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
+28
-1
@@ -20,6 +20,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
kFp8StaticTensorSym,
21
kInt4Static,
22
kInt4Static32,
23
+ kMxfp4Dynamic,
24
kMxfp4Static,
25
kMxfp8Dynamic,
26
kMxfp8Static,
@@ -64,10 +65,16 @@ class XPUExperts(mk.FusedMoEExpertsModular):
65
)
66
self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
67
self.fused_moe_impl: XpuFusedMoe | None = None
68
+ is_xe2_or_xe3 = torch.ops._xpu_C.is_xe2_arch() or torch.ops._xpu_C.is_xe3_arch()
69
+ if not is_xe2_or_xe3:
70
+ raise NotImplementedError(
71
+ "XPUExperts is only supported on Intel Xe2/Xe3 GPUs"
72
+ )
73
+ self._expects_unquantized_inputs = is_xe2_or_xe3
74
75
@property
76
def expects_unquantized_inputs(self) -> bool:
70
- return True
77
+ return self._expects_unquantized_inputs
78
79
@staticmethod
80
def activation_format() -> mk.FusedMoEActivationFormat:
@@ -172,6 +179,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
179
hidden_states=hidden_states,
180
topk_weights=topk_weights,
181
topk_ids=topk_ids,
182
+ a1q_scale=a1q_scale,
183
)
184
185
@@ -309,6 +317,24 @@ class XPUExpertsMxFp4(XPUExperts):
317
num_dispatchers,
318
)
319
320
+ def workspace_shapes(
321
+ self,
322
+ M: int,
323
+ N: int,
324
+ K: int,
325
+ topk: int,
326
+ global_num_experts: int,
327
+ local_num_experts: int,
328
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
329
+ activation: MoEActivation,
330
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
331
+ # K = a1q.size(-1). When activations are pre-quantized packed mxfp4,
332
+ # K is the packed hidden_size (= logical / 2); the kernel output is at
333
+ # logical hidden_size (2 * K). When unquantized (bf16), K is already
334
+ # the logical size.
335
+ logical_K = K if self.expects_unquantized_inputs else 2 * K
336
+ return (0,), (0,), (M, logical_K)
337
+
338
@staticmethod
339
def _supports_quant_scheme(
340
weight_key: QuantKey | None,
@@ -316,5 +342,6 @@ class XPUExpertsMxFp4(XPUExperts):
342
) -> bool:
343
SUPPORTED_W_A = [
344
(kMxfp4Static, None),
345
+ (kMxfp4Static, kMxfp4Dynamic),
346
]
347
return (weight_key, activation_key) in SUPPORTED_W_A
vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py
+1
-1
@@ -71,7 +71,7 @@ class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):
71
assert output.size() == fused_expert_output.size(), (
72
"output shape is expected to match the fused_expert_output shape. "
73
f"But got output={output.size()}, "
74
- f"used_expert_output={fused_expert_output.size()}"
74
+ f"fused_expert_output={fused_expert_output.size()}"
75
)
76
output.copy_(fused_expert_output, non_blocking=True)
77
return output
vllm/model_executor/layers/fused_moe/utils.py
+8
-2
@@ -17,12 +17,14 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
17
)
18
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
19
quant_dequant_mxfp4,
20
+ xpu_mxfp4_quantize,
21
)
22
from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
23
quant_dequant_mxfp6,
24
)
25
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
26
mxfp8_e4m3_quantize,
27
+ xpu_mxfp8_quantize,
28
)
29
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
30
ref_nvfp4_quant_dequant,
@@ -195,6 +197,8 @@ def _mxfp4_quantize(
197
per_act_token_quant: bool,
198
block_shape: list[int] | None = None,
199
) -> tuple[torch.Tensor, None]:
200
+ if current_platform.is_xpu():
201
+ return xpu_mxfp4_quantize(A)
202
assert block_shape is None
203
# TODO: native mxfp4 is currently not integrated in vllm,
204
# so simulating even on devices supporting this data type natively.
@@ -223,6 +227,8 @@ def _mxfp8_e4m3_quantize(
227
is_sf_swizzled_layout: bool = False,
228
mx_alignment: int = 0,
229
) -> tuple[torch.Tensor, torch.Tensor]:
230
+ if current_platform.is_xpu():
231
+ return xpu_mxfp8_quantize(A)
232
assert A_scale is None
233
assert not per_act_token_quant
234
assert block_shape is None or block_shape == [1, 32]
@@ -309,7 +315,7 @@ def moe_kernel_quantize_input(
315
A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
316
return A, None
317
elif quant_dtype == "mxfp4":
312
- if not quantization_emulation:
318
+ if not current_platform.is_xpu() and not quantization_emulation:
319
raise NotImplementedError(
320
"moe_kernel_quantize_input should not be used for native"
321
" quant_dtype='mxfp4' MOE. Please open an issue."
@@ -318,7 +324,7 @@ def moe_kernel_quantize_input(
324
elif quant_dtype == "mxfp8":
325
# TODO: `quant_dtype == "mxfp8"` is ambiguous,
326
# should be fp8_e4m3. OCP MX also defines `fp8_e5m2`.
321
- if quantization_emulation:
327
+ if not current_platform.is_xpu() and quantization_emulation:
328
raise NotImplementedError(
329
"moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
330
"quantization emulation. Please open an issue."