[XPU] Enable rms_norm/act quant fusions (#43963)

Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

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)
 
 

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

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,

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