A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[XPU] Add xpu torch.compile support (#22609)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Aug 27, 2025 at 13:33 UTC
fce10dbed5441b4f918b23a2b63aae72bc00a2f6
8 files changed
+36
-11
.buildkite/scripts/hardware_ci/run-xpu-test.sh
+1
@@ -31,6 +31,7 @@ docker run \
31
set -e
32
echo $ZE_AFFINITY_MASK
33
VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
34
+ VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -O.cudagraph_mode=NONE
35
VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
36
VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
37
cd tests
vllm/attention/layer.py
+1
-2
@@ -190,8 +190,7 @@ class Attention(nn.Module, AttentionLayerBase):
190
# torch.compile works by registering the attention as one giant
191
# opaque custom op. For other platforms, we directly call them
192
# and let torch.compile handle them.
193
- self.use_direct_call = not current_platform.is_cuda_alike(
194
- ) and not current_platform.is_cpu()
193
+ self.use_direct_call = not current_platform.opaque_attention_op()
194
195
self.use_output = self.attn_backend.accept_output_buffer
196
compilation_config = get_current_vllm_config().compilation_config
vllm/compilation/fix_functionalization.py
+8
@@ -9,6 +9,7 @@ import torch
9
from torch._higher_order_ops.auto_functionalize import auto_functionalized
10
11
from vllm.logger import init_logger
12
+from vllm.platforms import current_platform
13
14
from .fx_utils import is_func
15
from .vllm_inductor_pass import VllmInductorPass
@@ -26,6 +27,13 @@ class FixFunctionalizationPass(VllmInductorPass):
27
"""
28
29
def __call__(self, graph: torch.fx.Graph):
30
+ # XPU does not support auto-functionalization yet.
31
+ # Will enable this when switch to vllm-xpu-kernels.
32
+ if current_platform.is_xpu():
33
+ logger.debug("XPU platform does not support fix functionalization"
34
+ "pass currently.")
35
+ return
36
+
37
self.begin()
38
self.dump_graph(graph, "before_fix_functionalization")
39
vllm/platforms/cpu.py
+4
@@ -335,3 +335,7 @@ class CpuPlatform(Platform):
335
return (cls.supports_v1(model_config)
336
and arch in (CpuArchEnum.X86, CpuArchEnum.POWERPC,
337
CpuArchEnum.ARM, CpuArchEnum.S390X))
338
+
339
+ @classmethod
340
+ def opaque_attention_op(cls) -> bool:
341
+ return True
vllm/platforms/cuda.py
+4
@@ -442,6 +442,10 @@ class CudaPlatformBase(Platform):
442
def use_custom_allreduce(cls) -> bool:
443
return True
444
445
+ @classmethod
446
+ def opaque_attention_op(cls) -> bool:
447
+ return True
448
+
449
@classmethod
450
def get_static_graph_wrapper_cls(cls) -> str:
451
return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
vllm/platforms/interface.py
+8
@@ -509,6 +509,14 @@ class Platform:
509
"""
510
return False
511
512
+ @classmethod
513
+ def opaque_attention_op(cls) -> bool:
514
+ """
515
+ Returns True if we register attention as one giant opaque custom op
516
+ on the current platform
517
+ """
518
+ return False
519
+
520
@classmethod
521
def validate_request(
522
cls,
vllm/platforms/rocm.py
+4
@@ -411,6 +411,10 @@ class RocmPlatform(Platform):
411
supported_archs = ['gfx94', 'gfx95']
412
return any(gfx in gcn_arch for gfx in supported_archs)
413
414
+ @classmethod
415
+ def opaque_attention_op(cls) -> bool:
416
+ return True
417
+
418
@classmethod
419
def get_cu_count(cls, device_id: int = 0) -> int:
420
return torch.cuda.get_device_properties(
vllm/platforms/xpu.py
+6
-9
@@ -90,21 +90,14 @@ class XPUPlatform(Platform):
90
if cache_config and cache_config.block_size is None:
91
cache_config.block_size = 64
92
93
- # FIXME: Temporarily forcing eager mode
94
- # remove after t.compile support stabilizes.
95
- if (envs.VLLM_USE_V1 and model_config is not None
96
- and not vllm_config.model_config.enforce_eager):
97
- from vllm.config import CompilationLevel
98
- vllm_config.compilation_config.level = CompilationLevel.NO_COMPILATION # noqa: E501
99
-
93
# lazy import to avoid circular import
94
from vllm.config import CUDAGraphMode
95
compilation_config = vllm_config.compilation_config
96
if compilation_config.cudagraph_mode is None or \
97
compilation_config.cudagraph_mode.max_cudagraph_mode() \
98
!= CUDAGraphMode.NONE:
106
- logger.info("[XPU] CUDA graph is not supported on XPU, "
107
- "disabling cudagraphs.")
99
+ logger.info("[XPU] CUDA graph is not supported on XPU, disabling "
100
+ "cudagraphs. Fallback to cudagraph_mode=NONE")
101
compilation_config.cudagraph_mode = CUDAGraphMode.NONE
102
103
# check and update parallel config
@@ -182,3 +175,7 @@ class XPUPlatform(Platform):
175
"Intel Arc A770 have bfloat16 accuracy known issue. "
176
"You can use float16 instead by explicitly setting the "
177
"`dtype` flag in CLI, for example: --dtype=half.")
178
+
179
+ @classmethod
180
+ def opaque_attention_op(cls) -> bool:
181
+ return True