diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py index e35fc5cd40..c58ce31bd2 100644 --- a/vllm/compilation/passes/fusion/act_quant_fusion.py +++ b/vllm/compilation/passes/fusion/act_quant_fusion.py @@ -70,7 +70,11 @@ class ActivationQuantPattern(VllmPatternReplacement): self.silu_and_mul_matcher = MatcherSiluAndMul() def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor: - kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs} + kwargs = { + "dtype": self.quant_dtype, + "device": current_platform.device_type, + **kwargs, + } return torch.empty(*args, **kwargs) diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index cc986595d4..d47d825a16 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -54,19 +54,27 @@ def _rms_input_weight_dtype_match(match: pm.Match) -> bool: def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.bfloat16, device=current_platform.device_type + ) def empty_fp32(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.float32, device=current_platform.device_type + ) def empty_i32(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.int32, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.int32, device=current_platform.device_type + ) def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.int64, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.int64, device=current_platform.device_type + ) RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index fbf05f7753..fef494ca54 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -45,6 +45,10 @@ if current_platform.is_cuda(): from .fusion.allreduce_rms_fusion import AllReduceFusionPass from .fusion.collective_fusion import AsyncTPPass +if current_platform.is_xpu(): + from .fusion.act_quant_fusion import ActivationQuantFusionPass + from .fusion.rms_quant_fusion import RMSNormQuantFusionPass + from .inductor_pass import ( CustomGraphPass, InductorPass, diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index c2be7ff03a..f8b5c35a15 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -197,24 +197,25 @@ class XPUPlatform(Platform): ) # Disable fusion passes not yet supported on XPU. + from vllm.config.compilation import CompilationMode + pass_config = compilation_config.pass_config fusion_passes_to_disable = { "enable_sp": "Sequence parallelism", "fuse_gemm_comms": "Async TP", "fuse_allreduce_rms": "AllReduce + RMSNorm fusion", - "fuse_norm_quant": "RMSNorm + quant fusion", - "fuse_act_quant": "Activation + quant fusion", "fuse_attn_quant": "Attention + quant fusion", "fuse_act_padding": "Activation + padding fusion", "fuse_rope_kvcache": "RoPE + KV cache fusion", } - for flag, feature_name in fusion_passes_to_disable.items(): - if getattr(pass_config, flag): - logger.warning( - "Feature %r is not yet supported on XPU and will be disabled.", - feature_name, - ) - setattr(pass_config, flag, False) + if compilation_config.mode != CompilationMode.NONE: + for flag, feature_name in fusion_passes_to_disable.items(): + if getattr(pass_config, flag): + logger.warning( + "Feature %r is not yet supported on XPU and will be disabled.", + feature_name, + ) + setattr(pass_config, flag, False) # check and update parallel config parallel_config = vllm_config.parallel_config