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 7, 19:00 UTC
[XPU]Support CUDAGraph on XPU Platform (#34482)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com> Co-authored-by: chzhang <chaojun.zhang@intel.com> Co-authored-by: zhenwei-intel <zhenwei.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Xinyu Chen committed
Feb 25, 2026 at 14:22 UTC
35d44b45570396b28ce94186b2438dbc608fd6c0
3 files changed
+45
-4
vllm/platforms/xpu.py
+33
-4
@@ -13,6 +13,7 @@ import vllm_xpu_kernels._moe_C # noqa
13
import vllm_xpu_kernels._xpu_C # noqa
14
15
from vllm.logger import init_logger
16
+from vllm.utils.torch_utils import supports_xpu_graph
17
from vllm.v1.attention.backends.registry import AttentionBackendEnum
18
19
from .interface import DeviceCapability, Platform, PlatformEnum
@@ -151,10 +152,15 @@ class XPUPlatform(Platform):
152
def inference_mode(cls):
153
return torch.no_grad()
154
155
+ @classmethod
156
+ def get_static_graph_wrapper_cls(cls) -> str:
157
+ return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
158
+
159
@classmethod
160
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
161
cache_config = vllm_config.cache_config
162
model_config = vllm_config.model_config
163
+ parallel_config = vllm_config.parallel_config
164
# in V1(or with chunked prefill) block_size is 64
165
if cache_config and cache_config.block_size is None:
166
cache_config.block_size = 64
@@ -166,9 +172,32 @@ class XPUPlatform(Platform):
172
if compilation_config.compile_sizes is None:
173
compilation_config.compile_sizes = []
174
169
- assert compilation_config.cudagraph_mode == CUDAGraphMode.NONE, (
170
- "CUDA graph mode should be NONE on XPU"
171
- )
175
+ attention_config = vllm_config.attention_config
176
+ if attention_config.backend is None:
177
+ attention_config.backend = AttentionBackendEnum.FLASH_ATTN
178
+ if not supports_xpu_graph():
179
+ compilation_config.cudagraph_mode = CUDAGraphMode.NONE
180
+ logger.warning(
181
+ "XPU Graph is not supported in the current PyTorch version, "
182
+ "disabling cudagraph_mode."
183
+ )
184
+ elif parallel_config.world_size_across_dp > 1:
185
+ compilation_config.cudagraph_mode = CUDAGraphMode.NONE
186
+ logger.warning(
187
+ "XPU Graph doesn't support capture communication ops, "
188
+ "disabling cudagraph_mode."
189
+ )
190
+ else:
191
+ if (
192
+ attention_config.backend == AttentionBackendEnum.FLASH_ATTN
193
+ and compilation_config.cudagraph_mode
194
+ not in {CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE}
195
+ ):
196
+ compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
197
+ logger.warning(
198
+ "FMHA sycl-tla kernels cannot be captured with XPU graphs, "
199
+ "falling back to PIECEWISE graph mode on XPU platform."
200
+ )
201
202
if vllm_config.lora_config is not None:
203
compilation_config.mode = CompilationMode.NONE
@@ -201,7 +230,7 @@ class XPUPlatform(Platform):
230
231
@classmethod
232
def support_static_graph_mode(cls) -> bool:
204
- return False
233
+ return True
234
235
@classmethod
236
def is_pin_memory_available(cls):
vllm/utils/torch_utils.py
+5
@@ -745,6 +745,11 @@ def supports_xccl() -> bool:
745
return torch.distributed.is_xccl_available()
746
747
748
+# Supports XPU Graph with PyTorch versions >= 2.11.0.dev for XPU platform
749
+def supports_xpu_graph() -> bool:
750
+ return is_torch_equal_or_newer("2.11.0.dev")
751
+
752
+
753
# create a library to hold the custom op
754
vllm_lib = Library("vllm", "FRAGMENT") # noqa
755
vllm/v1/worker/xpu_model_runner.py
+7
@@ -7,6 +7,7 @@ import torch
7
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
+from vllm.utils.torch_utils import supports_xpu_graph
11
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
12
13
if TYPE_CHECKING:
@@ -40,6 +41,12 @@ def _torch_cuda_wrapper():
41
torch.cuda.default_stream = torch.xpu.current_stream
42
torch.cuda.current_stream = torch.xpu.current_stream
43
torch.cuda.stream = torch.xpu.stream
44
+ torch.cuda.mem_get_info = torch.xpu.mem_get_info
45
+ torch.cuda.synchronize = torch.xpu.synchronize
46
+ if supports_xpu_graph():
47
+ torch.cuda.graph = torch.xpu.graph
48
+ torch.cuda.CUDAGraph = torch.xpu.XPUGraph
49
+ torch.cuda.empty_cache = torch.xpu.empty_cache
50
yield
51
finally:
52
pass