add xpu torch.compile support

Signed-off-by: Kunshang Ji <[email protected]>

Author: jikunshang

vllm/attention/layer.py

@@ -190,7 +190,7 @@ def __init__(
         # opaque custom op. For other platforms, we directly call them
         # and let torch.compile handle them.
         self.use_direct_call = not current_platform.is_cuda_alike(
-        ) and not current_platform.is_cpu()
+        ) and not current_platform.is_cpu() and not current_platform.is_xpu()
 
         self.use_output = self.attn_backend.accept_output_buffer
         compilation_config = get_current_vllm_config().compilation_config

vllm/compilation/fix_functionalization.py

@@ -9,6 +9,7 @@
 from torch._higher_order_ops.auto_functionalize import auto_functionalized
 
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 
 from .fx_utils import is_func
 from .vllm_inductor_pass import VllmInductorPass
@@ -32,6 +33,10 @@ def __call__(self, graph: torch.fx.Graph):
         self.nodes_to_remove: list[torch.fx.Node] = []
         count = 0
         for node in graph.nodes:
+            # XPU does not support auto-functionalization yet.
+            # Will enable this when switch to vllm-xpu-kernels.
+            if current_platform.is_xpu():
+                continue
             if not is_func(node, auto_functionalized):
                 continue  # Avoid deep if-elif nesting
 

vllm/platforms/xpu.py

@@ -78,6 +78,10 @@ def get_device_total_memory(cls, device_id: int = 0) -> int:
     def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
         return True
 
+    @classmethod
+    def get_piecewise_backend_cls(cls) -> str:
+        return "vllm.compilation.cuda_piecewise_backend.CUDAPiecewiseBackend"  # noqa
+
     @classmethod
     def inference_mode(cls):
         return torch.no_grad()