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 9, 10:00 UTC
[Intel GPU] refine xpu worker (#32894)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Jan 29, 2026 at 20:26 UTC
8bb6271c77520c8df3bd7d17899c50225bc42e0a
2 files changed
+27
-90
vllm/platforms/__init__.py
+8
-6
@@ -138,16 +138,18 @@ def xpu_platform_plugin() -> str | None:
138
139
if supports_xccl():
140
dist_backend = "xccl"
141
- else:
142
- dist_backend = "ccl"
143
- import oneccl_bindings_for_pytorch # noqa: F401
144
-
145
- if hasattr(torch, "xpu") and torch.xpu.is_available():
146
- is_xpu = True
141
from vllm.platforms.xpu import XPUPlatform
142
143
XPUPlatform.dist_backend = dist_backend
144
logger.debug("Confirmed %s backend is available.", XPUPlatform.dist_backend)
145
+ else:
146
+ logger.warning(
147
+ "xccl is not enabled in this torch build, "
148
+ "communication is not available."
149
+ )
150
+
151
+ if hasattr(torch, "xpu") and torch.xpu.is_available():
152
+ is_xpu = True
153
logger.debug("Confirmed XPU platform is available.")
154
except Exception as e:
155
logger.debug("XPU platform is not available because: %s", str(e))
vllm/v1/worker/xpu_worker.py
+19
-84
@@ -7,14 +7,18 @@ import torch
7
import torch.distributed
8
9
from vllm.config import VllmConfig
10
-from vllm.distributed import get_world_group
10
from vllm.logger import init_logger
11
from vllm.platforms import current_platform
12
from vllm.profiler.wrapper import TorchProfilerWrapper
13
+from vllm.utils.mem_utils import MemorySnapshot, format_gib
14
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
19
20
+from .utils import request_memory
21
+
22
logger = init_logger(__name__)
23
24
@@ -48,86 +52,6 @@ class XPUWorker(Worker):
52
activities=["CPU", "XPU"],
53
)
54
51
- # we provide this function due to `torch.xpu.mem_get_info()` doesn't
52
- # return correct free_gpu_memory on intel client GPU. We need to
53
- # calculate/estiamte it.
54
- def xpu_get_mem_info(self):
55
- if current_platform.is_data_center_gpu():
56
- return torch.xpu.mem_get_info()
57
- else:
58
- _, total_gpu_memory = torch.xpu.mem_get_info()
59
- # FIXME: memory_allocated() doesn't count non-torch allocations,
60
- # and we don't have any API to get it. so we mark it as 128MB.
61
- used_memory = torch.xpu.memory_allocated()
62
- non_torch_allocations = 128 * 1024 * 1024
63
- free_gpu_memory = total_gpu_memory - (used_memory + non_torch_allocations)
64
- return free_gpu_memory, total_gpu_memory
65
-
66
- @torch.inference_mode()
67
- def determine_available_memory(self) -> int:
68
- """Profiles the peak memory usage of the model to determine how many
69
- KV blocks may be allocated without OOMs.
70
- The engine will first conduct a profiling of the existing memory usage.
71
- Then, it calculates the maximum possible number of GPU and CPU blocks
72
- that can be allocated with the remaining free memory.
73
- .. tip::
74
- You may limit the usage of GPU memory
75
- by adjusting the `gpu_memory_utilization` parameter.
76
- """
77
- # Profile the memory usage of the model and get the maximum number of
78
- # cache blocks that can be allocated with the remaining free memory.
79
- torch.xpu.empty_cache()
80
- torch.xpu.reset_peak_memory_stats()
81
-
82
- free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info()
83
- current_allocated_bytes = torch.xpu.memory_allocated()
84
- msg = (
85
- "Before memory profiling run, "
86
- f"total GPU memory: {total_gpu_memory / 1024**2:.2f} MB, "
87
- f"model load takes {current_allocated_bytes / 1024**2:.2f} MB, "
88
- f"free gpu memory is {free_gpu_memory / 1024**2:.2f} MB."
89
- )
90
- logger.info(msg)
91
- # Execute a forward pass with dummy inputs to profile the memory usage
92
- # of the model.
93
- self.model_runner.profile_run()
94
-
95
- free_gpu_memory, _ = self.xpu_get_mem_info()
96
- # NOTE(woosuk): Here we assume that the other processes using the same
97
- # GPU did not change their memory usage during the profiling.
98
- assert self.init_gpu_memory > free_gpu_memory, (
99
- "Error in memory profiling. "
100
- f"Initial free memory {self.init_gpu_memory}, current free memory"
101
- f" {free_gpu_memory}. This happens when the GPU memory was "
102
- "not properly cleaned up before initializing the vLLM instance."
103
- )
104
-
105
- # Get the peak memory allocation recorded by torch
106
- peak_memory = torch.xpu.memory_stats()["allocated_bytes.all.peak"]
107
-
108
- torch.xpu.empty_cache()
109
- torch_allocated_bytes = torch.xpu.memory_stats()["allocated_bytes.all.current"]
110
- free_mem, total_mem = self.xpu_get_mem_info()
111
- total_allocated_bytes = total_mem - free_mem
112
-
113
- non_torch_allocations = total_allocated_bytes - torch_allocated_bytes
114
- if non_torch_allocations > 0:
115
- peak_memory += non_torch_allocations
116
- available_kv_cache_memory = (
117
- total_gpu_memory * self.cache_config.gpu_memory_utilization - peak_memory
118
- )
119
-
120
- msg = (
121
- "After memory profiling run, "
122
- f"peak memory usage is {peak_memory / 1024**2:.2f} MB,"
123
- f"torch mem is {torch_allocated_bytes / 1024**2:.2f} MB, "
124
- f"non-torch mem is {non_torch_allocations / 1024**2:.2f} MB, "
125
- f"free gpu memory is {free_gpu_memory / 1024**2:.2f} MB."
126
- )
127
- logger.info(msg)
128
-
129
- return int(available_kv_cache_memory)
130
-
55
def init_device(self):
56
device = self.device_config.device
57
if (
@@ -161,15 +85,26 @@ class XPUWorker(Worker):
85
current_platform.dist_backend,
86
)
87
164
- # global all_reduce needed for overall oneccl warm up
165
- torch.distributed.all_reduce(
166
- torch.zeros(1).xpu(), group=get_world_group().device_group
88
+ torch.xpu.empty_cache()
89
+ self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
90
+ self.requested_memory = request_memory(init_snapshot, self.cache_config)
91
+ logger.debug("worker init memory snapshot: %r", self.init_snapshot)
92
+ logger.debug(
93
+ "worker requested memory: %sGiB", format_gib(self.requested_memory)
94
)
95
96
# Set random seed.
97
set_random_seed(self.model_config.seed)
98
99
+ # Initialize workspace manager
100
+ num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
101
+ init_workspace_manager(self.device, num_ubatches)
102
+
103
# Construct the model runner
104
self.model_runner = XPUModelRunner( # type: ignore
105
self.vllm_config, self.device
106
)
107
+
108
+ if self.rank == 0:
109
+ # If usage stat is enabled, collect relevant info.
110
+ report_usage_stats(self.vllm_config)