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, 05:15 UTC
[XPU] Enable ModelRunnerV2 on XPU (#36078)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Xinyu Chen committed
Mar 6, 2026 at 01:19 UTC
d8839ef7d964dd98b82e671e743b42754be3350c
2 files changed
+21
-2
vllm/v1/worker/xpu_model_runner.py
+18
@@ -8,6 +8,9 @@ import torch
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 (
12
+ GPUModelRunner as GPUModelRunnerV2,
13
+)
14
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
15
16
if TYPE_CHECKING:
@@ -30,6 +33,18 @@ class XPUModelRunner(GPUModelRunner):
33
self.cascade_attn_enabled = False
34
35
36
+class XPUModelRunnerV2(GPUModelRunnerV2):
37
+ """A model runner for XPU devices."""
38
+
39
+ def __init__(
40
+ self,
41
+ vllm_config: VllmConfig,
42
+ device: torch.device,
43
+ ):
44
+ with _torch_cuda_wrapper():
45
+ super().__init__(vllm_config, device)
46
+
47
+
48
@contextmanager
49
def _torch_cuda_wrapper():
50
try:
@@ -39,9 +54,12 @@ def _torch_cuda_wrapper():
54
torch.cuda.current_stream = torch.xpu.current_stream
55
torch.cuda.stream = torch.xpu.stream
56
torch.cuda.mem_get_info = torch.xpu.mem_get_info
57
+ torch.cuda.Event = torch.Event
58
+ torch.cuda.set_stream = torch.xpu.set_stream
59
if supports_xpu_graph():
60
torch.cuda.graph = torch.xpu.graph
61
torch.cuda.CUDAGraph = torch.xpu.XPUGraph
62
+ torch.cuda.graph_pool_handle = torch.xpu.graph_pool_handle
63
yield
64
finally:
65
pass
vllm/v1/worker/xpu_worker.py
+3
-2
@@ -15,7 +15,7 @@ from vllm.utils.torch_utils import set_random_seed
15
from vllm.v1.utils import report_usage_stats
16
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
17
from vllm.v1.worker.workspace import init_workspace_manager
18
-from vllm.v1.worker.xpu_model_runner import XPUModelRunner
18
+from vllm.v1.worker.xpu_model_runner import XPUModelRunner, XPUModelRunnerV2
19
20
from .utils import request_memory
21
@@ -105,7 +105,8 @@ class XPUWorker(Worker):
105
init_workspace_manager(self.device, num_ubatches)
106
107
# Construct the model runner
108
- self.model_runner = XPUModelRunner( # type: ignore
108
+ model_runner = XPUModelRunnerV2 if self.use_v2_model_runner else XPUModelRunner
109
+ self.model_runner = model_runner( # type: ignore
110
self.vllm_config, self.device
111
)
112