diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index cde167e5d3..dfc578b441 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -20,6 +20,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kInt4Static, kInt4Static32, + kMxfp4Dynamic, kMxfp4Static, kMxfp8Dynamic, kMxfp8Static, @@ -64,10 +65,16 @@ class XPUExperts(mk.FusedMoEExpertsModular): ) self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit self.fused_moe_impl: XpuFusedMoe | None = None + is_xe2_or_xe3 = torch.ops._xpu_C.is_xe2_arch() or torch.ops._xpu_C.is_xe3_arch() + if not is_xe2_or_xe3: + raise NotImplementedError( + "XPUExperts is only supported on Intel Xe2/Xe3 GPUs" + ) + self._expects_unquantized_inputs = is_xe2_or_xe3 @property def expects_unquantized_inputs(self) -> bool: - return True + return self._expects_unquantized_inputs @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -172,6 +179,7 @@ class XPUExperts(mk.FusedMoEExpertsModular): hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, + a1q_scale=a1q_scale, ) @@ -309,6 +317,24 @@ class XPUExpertsMxFp4(XPUExperts): num_dispatchers, ) + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # K = a1q.size(-1). When activations are pre-quantized packed mxfp4, + # K is the packed hidden_size (= logical / 2); the kernel output is at + # logical hidden_size (2 * K). When unquantized (bf16), K is already + # the logical size. + logical_K = K if self.expects_unquantized_inputs else 2 * K + return (0,), (0,), (M, logical_K) + @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -316,5 +342,6 @@ class XPUExpertsMxFp4(XPUExperts): ) -> bool: SUPPORTED_W_A = [ (kMxfp4Static, None), + (kMxfp4Static, kMxfp4Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 837c149862..81c84d64be 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -71,7 +71,7 @@ class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce): assert output.size() == fused_expert_output.size(), ( "output shape is expected to match the fused_expert_output shape. " f"But got output={output.size()}, " - f"used_expert_output={fused_expert_output.size()}" + f"fused_expert_output={fused_expert_output.size()}" ) output.copy_(fused_expert_output, non_blocking=True) return output diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 6aa4eb0cd2..e3d6493dda 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -17,12 +17,14 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import ( ) from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( quant_dequant_mxfp4, + xpu_mxfp4_quantize, ) from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( quant_dequant_mxfp6, ) from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( mxfp8_e4m3_quantize, + xpu_mxfp8_quantize, ) from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( ref_nvfp4_quant_dequant, @@ -195,6 +197,8 @@ def _mxfp4_quantize( per_act_token_quant: bool, block_shape: list[int] | None = None, ) -> tuple[torch.Tensor, None]: + if current_platform.is_xpu(): + return xpu_mxfp4_quantize(A) assert block_shape is None # TODO: native mxfp4 is currently not integrated in vllm, # so simulating even on devices supporting this data type natively. @@ -223,6 +227,8 @@ def _mxfp8_e4m3_quantize( is_sf_swizzled_layout: bool = False, mx_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: + if current_platform.is_xpu(): + return xpu_mxfp8_quantize(A) assert A_scale is None assert not per_act_token_quant assert block_shape is None or block_shape == [1, 32] @@ -309,7 +315,7 @@ def moe_kernel_quantize_input( A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16) return A, None elif quant_dtype == "mxfp4": - if not quantization_emulation: + if not current_platform.is_xpu() and not quantization_emulation: raise NotImplementedError( "moe_kernel_quantize_input should not be used for native" " quant_dtype='mxfp4' MOE. Please open an issue." @@ -318,7 +324,7 @@ def moe_kernel_quantize_input( elif quant_dtype == "mxfp8": # TODO: `quant_dtype == "mxfp8"` is ambiguous, # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`. - if quantization_emulation: + if not current_platform.is_xpu() and quantization_emulation: raise NotImplementedError( "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE " "quantization emulation. Please open an issue."