@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3bb46975bd
Fork 0

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 1, 18:45 UTC

[XPU][Feature] transparent sleep mode support for XPU platform (#37149)

Signed-off-by: Yan Ma <yan.ma@intel.com>

Yan Ma committed Jun 7, 2026 at 10:45 UTC 3bb46975bd316bdd87e083758b86e2588cec8934
10 files changed +410 -47
.buildkite/intel_jobs/basic_correctness.yaml new
+22
@@ -0,0 +1,22 @@
1 +group: Basic Correctness
2 +depends_on:
3 + - image-build-xpu
4 +steps:
5 +- label: XPU Sleep Mode
6 + timeout_in_minutes: 30
7 + device: intel_gpu
8 + no_plugin: true
9 + working_dir: "."
10 + env:
11 + REGISTRY: "public.ecr.aws/q9t5s3a7"
12 + REPO: "vllm-ci-test-repo"
13 + VLLM_TEST_DEVICE: "xpu"
14 + source_file_dependencies:
15 + - vllm/
16 + - tests/basic_correctness/test_cumem.py
17 + commands:
18 + - >-
19 + bash .buildkite/scripts/hardware_ci/run-intel-test.sh
20 + 'cd tests &&
21 + export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
22 + pytest -v -s basic_correctness/test_mem.py::test_end_to_end'
.buildkite/test-amd.yaml
+2 -2
@@ -908,10 +908,10 @@ steps:
908 - vllm/
909 - tests/basic_correctness/test_basic_correctness
910 - tests/basic_correctness/test_cpu_offload
911 - - tests/basic_correctness/test_cumem.py
911 + - tests/basic_correctness/test_mem.py
912 commands:
913 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
914 - - pytest -v -s basic_correctness/test_cumem.py
914 + - pytest -v -s basic_correctness/test_mem.py
915 - pytest -v -s basic_correctness/test_basic_correctness.py
916 - pytest -v -s basic_correctness/test_cpu_offload.py
917
.buildkite/test_areas/basic_correctness.yaml
+2 -2
@@ -10,9 +10,9 @@ steps:
10 - vllm/
11 - tests/basic_correctness/test_basic_correctness
12 - tests/basic_correctness/test_cpu_offload
13 - - tests/basic_correctness/test_cumem.py
13 + - tests/basic_correctness/test_mem.py
14 commands:
15 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
16 - - pytest -v -s basic_correctness/test_cumem.py
16 + - pytest -v -s basic_correctness/test_mem.py
17 - pytest -v -s basic_correctness/test_basic_correctness.py
18 - pytest -v -s basic_correctness/test_cpu_offload.py
tests/basic_correctness/test_mem.py renamed
+22 -21
@@ -7,7 +7,7 @@ import pytest
7 import torch
8
9 from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
10 -from vllm.device_allocator.cumem import CuMemAllocator
10 +from vllm.device_allocator import get_mem_allocator_instance
11 from vllm.platforms import current_platform
12 from vllm.utils.mem_constants import GiB_bytes
13
@@ -16,14 +16,14 @@ from ..utils import create_new_process_for_each_test, requires_fp8
16 DEVICE_TYPE = current_platform.device_type
17
18
19 -@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
19 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
20 def test_python_error():
21 """
22 Test if Python error occurs when there's low-level
23 error happening from the C++ side.
24 """
25 - allocator = CuMemAllocator.get_instance()
26 - total_bytes = torch.cuda.mem_get_info()[1]
25 + allocator = get_mem_allocator_instance()
26 + total_bytes = current_platform.mem_get_info()[1]
27 alloc_bytes = int(total_bytes * 0.7)
28 tensors = []
29 with allocator.use_memory_pool():
@@ -42,7 +42,7 @@ def test_python_error():
42 allocator.wake_up()
43
44
45 -@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
45 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
46 def test_basic_cumem():
47 # some tensors from default memory pool
48 shape = (1024, 1024)
@@ -50,7 +50,7 @@ def test_basic_cumem():
50 x.zero_()
51
52 # some tensors from custom memory pool
53 - allocator = CuMemAllocator.get_instance()
53 + allocator = get_mem_allocator_instance()
54 with allocator.use_memory_pool():
55 # custom memory pool
56 y = torch.empty(shape, device=DEVICE_TYPE)
@@ -64,9 +64,9 @@ def test_basic_cumem():
64 output = x + y + z
65 assert torch.allclose(output, torch.ones_like(output) * 3)
66
67 - free_bytes = torch.cuda.mem_get_info()[0]
67 + free_bytes = current_platform.mem_get_info()[0]
68 allocator.sleep()
69 - free_bytes_after_sleep = torch.cuda.mem_get_info()[0]
69 + free_bytes_after_sleep = current_platform.mem_get_info()[0]
70 assert free_bytes_after_sleep > free_bytes
71 allocator.wake_up()
72
@@ -75,9 +75,10 @@ def test_basic_cumem():
75 assert torch.allclose(output, torch.ones_like(output) * 3)
76
77
78 -@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
78 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
79 +@pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU")
80 def test_cumem_with_cudagraph():
80 - allocator = CuMemAllocator.get_instance()
81 + allocator = get_mem_allocator_instance()
82 with allocator.use_memory_pool():
83 weight = torch.eye(1024, device=DEVICE_TYPE)
84 with allocator.use_memory_pool(tag="discard"):
@@ -98,9 +99,9 @@ def test_cumem_with_cudagraph():
99 with torch.cuda.graph(model_graph):
100 y = model(x)
101
101 - free_bytes = torch.cuda.mem_get_info()[0]
102 + free_bytes = current_platform.mem_get_info()[0]
103 allocator.sleep()
103 - free_bytes_after_sleep = torch.cuda.mem_get_info()[0]
104 + free_bytes_after_sleep = current_platform.mem_get_info()[0]
105 assert free_bytes_after_sleep > free_bytes
106 allocator.wake_up()
107
@@ -120,7 +121,7 @@ def test_cumem_with_cudagraph():
121 assert torch.allclose(y, x + 1)
122
123
123 -@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
124 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
125 @pytest.mark.parametrize(
126 "model",
127 [
@@ -131,7 +132,7 @@ def test_cumem_with_cudagraph():
132 ],
133 )
134 def test_end_to_end(model: str):
134 - free, total = torch.cuda.mem_get_info()
135 + free, total = current_platform.mem_get_info()
136 used_bytes_baseline = total - free # in case other process is running
137 llm = LLM(model, enable_sleep_mode=True)
138 prompt = "How are you?"
@@ -143,7 +144,7 @@ def test_end_to_end(model: str):
144 # test sleep level 1 here.
145 llm.sleep(level=1)
146
146 - free_gpu_bytes_after_sleep, total = torch.cuda.mem_get_info()
147 + free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()
148 used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
149 # now the memory usage is mostly cudagraph memory pool,
150 # and it should be less than the model weights (1B model, 2GiB weights)
@@ -163,7 +164,7 @@ def test_end_to_end(model: str):
164 llm.sleep(level=1)
165 llm.wake_up(tags=["weights"])
166
166 - free_gpu_bytes_wake_up_w, total = torch.cuda.mem_get_info()
167 + free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
168 used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
169
170 # should just reallocate memory for weights (1B model, ~2GiB weights)
@@ -180,7 +181,7 @@ def test_end_to_end(model: str):
181 @create_new_process_for_each_test()
182 def test_deep_sleep():
183 model = "hmellor/tiny-random-LlamaForCausalLM"
183 - free, total = torch.cuda.mem_get_info()
184 + free, total = current_platform.mem_get_info()
185 used_bytes_baseline = total - free # in case other process is running
186 llm = LLM(model, enable_sleep_mode=True)
187 prompt = "How are you?"
@@ -190,13 +191,13 @@ def test_deep_sleep():
191 # Put the engine to deep sleep
192 llm.sleep(level=2)
193
193 - free_gpu_bytes_after_sleep, total = torch.cuda.mem_get_info()
194 + free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()
195 used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
196 assert used_bytes < 3 * GiB_bytes
197
198 llm.wake_up(tags=["weights"])
199 llm.collective_rpc("reload_weights")
199 - free_gpu_bytes_wake_up_w, total = torch.cuda.mem_get_info()
200 + free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
201 used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
202 assert used_bytes < 4 * GiB_bytes
203
@@ -212,7 +213,7 @@ def test_deep_sleep():
213 def test_deep_sleep_async():
214 async def test():
215 model = "hmellor/tiny-random-LlamaForCausalLM"
215 - free, total = torch.cuda.mem_get_info()
216 + free, total = current_platform.mem_get_info()
217 used_bytes_baseline = total - free # in case other process is running
218 engine_args = AsyncEngineArgs(
219 model=model,
@@ -231,7 +232,7 @@ def test_deep_sleep_async():
232
233 await llm.wake_up(tags=["weights"])
234 await llm.collective_rpc("reload_weights")
234 - free_gpu_bytes_wake_up_w, total = torch.cuda.mem_get_info()
235 + free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
236 used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
237 assert used_bytes < 4 * GiB_bytes
238
vllm/config/model.py
+1 -1
@@ -528,7 +528,7 @@ class ModelConfig:
528 if self.enable_sleep_mode:
529 if not current_platform.is_sleep_mode_available():
530 raise ValueError("Sleep mode is not supported on current platform.")
531 - if not self.enable_cumem_allocator:
531 + if current_platform.is_cuda_alike() and not self.enable_cumem_allocator:
532 logger.info_once(
533 "Enabling cumem allocator because sleep mode requires it."
534 )
vllm/device_allocator/__init__.py
+48
@@ -0,0 +1,48 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import dataclasses
5 +from contextlib import AbstractContextManager
6 +from typing import Protocol
7 +
8 +import torch
9 +
10 +from vllm.platforms import current_platform
11 +
12 +# py_device, py_size_or_aligned_size, py_ptr, py_handle
13 +HandleType = tuple[int, int, int, int]
14 +
15 +
16 +@dataclasses.dataclass
17 +class AllocationData:
18 + handle: HandleType
19 + tag: str
20 + cpu_backup_tensor: torch.Tensor | None = None
21 +
22 +
23 +class MemAllocator(Protocol):
24 + def use_memory_pool(self, tag: str | None = None) -> AbstractContextManager: ...
25 +
26 + def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ...
27 +
28 + def wake_up(self, tags: list[str] | None = None) -> None: ...
29 +
30 + def get_current_usage(self) -> int: ...
31 +
32 +
33 +def get_mem_allocator_instance() -> MemAllocator:
34 + if current_platform.is_cuda_alike():
35 + from vllm.device_allocator.cumem import CuMemAllocator
36 +
37 + return CuMemAllocator.get_instance()
38 +
39 + if current_platform.is_xpu():
40 + from vllm.device_allocator.xpumem import XpuMemAllocator
41 +
42 + return XpuMemAllocator.get_instance()
43 +
44 + raise RuntimeError(
45 + "Sleep mode allocator is not available on platform "
46 + f"{type(current_platform).__name__} "
47 + f"(device_type={current_platform.device_type})."
48 + )
vllm/device_allocator/cumem.py
+1 -11
@@ -8,7 +8,6 @@
8 # both of them failed because of cuda context mismatch.
9 # not sure why, they are created from a different context.
10 # the only successful approach is to call cuda driver API in C.
11 -import dataclasses
11 import gc
12 import os
13 from collections.abc import Callable, Iterator
@@ -17,6 +16,7 @@ from typing import Any
16
17 import torch
18
19 +from vllm.device_allocator import AllocationData, HandleType
20 from vllm.logger import init_logger
21 from vllm.utils.platform_utils import is_pin_memory_available
22 from vllm.utils.system_utils import find_loaded_library
@@ -44,16 +44,6 @@ except ModuleNotFoundError:
44 python_unmap_and_release = None
45 lib_name = None
46
47 -# py_device, py_alignedSize, py_d_mem, py_p_memHandle
48 -HandleType = tuple[int, int, int, int]
49 -
50 -
51 -@dataclasses.dataclass
52 -class AllocationData:
53 - handle: HandleType
54 - tag: str
55 - cpu_backup_tensor: torch.Tensor | None = None
56 -
47
48 def create_and_map(allocation_handle: HandleType) -> None:
49 python_create_and_map(*allocation_handle)
vllm/device_allocator/xpumem.py new
+295
@@ -0,0 +1,295 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import atexit
5 +import gc
6 +from collections.abc import Callable, Iterator
7 +from contextlib import contextmanager
8 +from typing import Any
9 +
10 +import torch
11 +
12 +from vllm.device_allocator import AllocationData, HandleType
13 +from vllm.logger import init_logger
14 +from vllm.utils.platform_utils import is_pin_memory_available
15 +
16 +logger = init_logger(__name__)
17 +
18 +MEMCPY_HOST_TO_DEVICE = 0
19 +MEMCPY_DEVICE_TO_HOST = 1
20 +MEMCPY_DEVICE_TO_DEVICE = 2
21 +
22 +xpumem_available = False
23 +xpumem_allocator: Any = None
24 +
25 +try:
26 + from vllm_xpu_kernels import xpumem_allocator as _xpumem_allocator
27 +
28 + xpumem_allocator = _xpumem_allocator
29 + xpumem_available = True
30 +except ImportError:
31 + xpumem_allocator = None
32 +
33 +
34 +def _xpu_memory_module() -> Any:
35 + mem_mod = getattr(torch.xpu, "memory", None)
36 + if mem_mod is None:
37 + raise RuntimeError("torch.xpu.memory is not available")
38 + return mem_mod
39 +
40 +
41 +def _supports_xpu_mem_pool(mem_mod: Any) -> bool:
42 + return hasattr(mem_mod, "MemPool") and hasattr(mem_mod, "use_mem_pool")
43 +
44 +
45 +def _xpu_memcpy_sync(
46 + dst_ptr: int,
47 + src_ptr: int,
48 + n_bytes: int,
49 + kind: int,
50 + device: int,
51 +) -> None:
52 + def _to_i64_ptr(ptr: int) -> int:
53 + # torch custom-op `int` arguments are signed int64.
54 + # data_ptr() may return a uint64 value above 2^63-1, so normalize it.
55 + return ptr if ptr < (1 << 63) else ptr - (1 << 64)
56 +
57 + torch.ops._C.xpu_memcpy_sync(
58 + _to_i64_ptr(dst_ptr),
59 + _to_i64_ptr(src_ptr),
60 + n_bytes,
61 + kind,
62 + device,
63 + )
64 +
65 +
66 +def get_pluggable_allocator(
67 + python_malloc_fn: Callable[[HandleType], None],
68 + python_free_func: Callable[[int], HandleType],
69 +) -> Any:
70 + if not xpumem_available or xpumem_allocator is None:
71 + raise RuntimeError("xpumem allocator extension is not available")
72 +
73 + xpumem_allocator.init_module(python_malloc_fn, python_free_func)
74 + mem_mod = _xpu_memory_module()
75 + alloc_cls = getattr(mem_mod, "XPUPluggableAllocator", None)
76 + if alloc_cls is None:
77 + raise RuntimeError("torch.xpu.memory.XPUPluggableAllocator is not available")
78 +
79 + lib_name = xpumem_allocator.__file__
80 + return alloc_cls(lib_name, "my_malloc", "my_free")
81 +
82 +
83 +def create_and_allocate(allocation_handle: HandleType) -> None:
84 + if not xpumem_available or xpumem_allocator is None:
85 + raise RuntimeError("xpumem allocator extension is not available")
86 + xpumem_allocator.python_create_and_allocate(*allocation_handle)
87 +
88 +
89 +def unmap_and_release(allocation_handle: HandleType) -> None:
90 + if not xpumem_available or xpumem_allocator is None:
91 + raise RuntimeError("xpumem allocator extension is not available")
92 + xpumem_allocator.python_unmap_and_release(*allocation_handle)
93 +
94 +
95 +@contextmanager
96 +def use_memory_pool_with_allocator(
97 + python_malloc_fn: Callable[[HandleType], None],
98 + python_free_func: Callable[[int], HandleType],
99 +) -> Iterator[tuple[Any, Any]]:
100 + mem_mod = _xpu_memory_module()
101 + if not _supports_xpu_mem_pool(mem_mod):
102 + raise RuntimeError(
103 + "torch.xpu.memory MemPool APIs are not available "
104 + "(need MemPool and use_mem_pool)."
105 + )
106 + new_alloc = get_pluggable_allocator(python_malloc_fn, python_free_func)
107 + mem_pool = mem_mod.MemPool(new_alloc._allocator)
108 + with mem_mod.use_mem_pool(mem_pool):
109 + yield mem_pool, new_alloc
110 +
111 +
112 +class XpuMemAllocator:
113 + """A singleton pluggable allocator helper for XPU.
114 +
115 + Note:
116 + Sleep will offload selected payloads to CPU or discard and unmap XPU
117 + physical memory. Wake-up remaps physical memory back to the same
118 + reserved virtual address and restores payload.
119 + """
120 +
121 + instance: "XpuMemAllocator | None" = None
122 + default_tag: str = "default"
123 +
124 + @staticmethod
125 + def get_instance() -> "XpuMemAllocator":
126 + assert xpumem_available, "xpumem allocator is not available"
127 + if XpuMemAllocator.instance is None:
128 + XpuMemAllocator.instance = XpuMemAllocator()
129 + # Ensure MemPool/allocator wrappers are released before interpreter
130 + # finalization tears down XPU runtime internals.
131 + atexit.register(XpuMemAllocator._shutdown_singleton)
132 + return XpuMemAllocator.instance
133 +
134 + @staticmethod
135 + def _shutdown_singleton() -> None:
136 + instance = XpuMemAllocator.instance
137 + if instance is None:
138 + return
139 + try:
140 + instance.release_pools()
141 + except Exception:
142 + logger.exception("XpuMemAllocator singleton shutdown failed")
143 +
144 + def __init__(self):
145 + self.pointer_to_data: dict[int, AllocationData] = {}
146 + self.current_tag: str = XpuMemAllocator.default_tag
147 + self.allocator_and_pools: dict[str, Any] = {}
148 + self.python_malloc_callback = self._python_malloc_callback
149 + self.python_free_callback = self._python_free_callback
150 +
151 + def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
152 + ptr = allocation_handle[2]
153 + self.pointer_to_data[ptr] = AllocationData(allocation_handle, self.current_tag)
154 + logger.debug(
155 + "Allocated %s bytes for %s at %s",
156 + allocation_handle[1],
157 + self.current_tag,
158 + ptr,
159 + )
160 +
161 + def _python_free_callback(self, ptr: int) -> HandleType:
162 + data = self.pointer_to_data.pop(ptr)
163 + data.cpu_backup_tensor = None
164 + logger.debug("Freed %s bytes for %s at %s", data.handle[1], data.tag, ptr)
165 + return data.handle
166 +
167 + def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
168 + if offload_tags is None:
169 + offload_tags = (XpuMemAllocator.default_tag,)
170 + elif isinstance(offload_tags, str):
171 + offload_tags = (offload_tags,)
172 +
173 + assert isinstance(offload_tags, tuple)
174 +
175 + total_bytes = 0
176 + backup_bytes = 0
177 +
178 + for ptr, data in self.pointer_to_data.items():
179 + size_in_bytes = data.handle[1]
180 + total_bytes += size_in_bytes
181 + if data.tag not in offload_tags:
182 + unmap_and_release(data.handle)
183 + continue
184 +
185 + backup_bytes += size_in_bytes
186 + device, _, _, _ = data.handle
187 + cpu_backup_tensor = torch.empty(
188 + size_in_bytes,
189 + dtype=torch.uint8,
190 + device="cpu",
191 + pin_memory=is_pin_memory_available(),
192 + )
193 + cpu_ptr = cpu_backup_tensor.data_ptr()
194 + _xpu_memcpy_sync(
195 + cpu_ptr,
196 + ptr,
197 + size_in_bytes,
198 + MEMCPY_DEVICE_TO_HOST,
199 + device,
200 + )
201 + data.cpu_backup_tensor = cpu_backup_tensor
202 +
203 + unmap_and_release(data.handle)
204 +
205 + logger.info(
206 + "XpuMemAllocator: sleep freed %.2f GiB memory in total, of which "
207 + "%.2f GiB is backed up in CPU and the rest %.2f GiB is discarded "
208 + "directly.",
209 + total_bytes / 1024**3,
210 + backup_bytes / 1024**3,
211 + (total_bytes - backup_bytes) / 1024**3,
212 + )
213 +
214 + gc.collect()
215 + xpu_empty_cache = getattr(torch.xpu, "empty_cache", None)
216 + if callable(xpu_empty_cache):
217 + xpu_empty_cache()
218 +
219 + def wake_up(self, tags: list[str] | None = None) -> None:
220 + for ptr, data in self.pointer_to_data.items():
221 + if tags is not None and data.tag not in tags:
222 + continue
223 + create_and_allocate(data.handle)
224 +
225 + cpu_backup_tensor = data.cpu_backup_tensor
226 + if cpu_backup_tensor is None:
227 + continue
228 +
229 + device, size_in_bytes, _, _ = data.handle
230 + _xpu_memcpy_sync(
231 + ptr,
232 + cpu_backup_tensor.data_ptr(),
233 + size_in_bytes,
234 + MEMCPY_HOST_TO_DEVICE,
235 + device,
236 + )
237 + data.cpu_backup_tensor = None
238 +
239 + def release_pools(self) -> None:
240 + """Drop Python references to MemPool/pluggable allocators eagerly.
241 +
242 + This prevents pool destruction from being deferred to interpreter
243 + finalization, which can happen after parts of XPU runtime are already
244 + torn down.
245 + """
246 + if not self.allocator_and_pools:
247 + return
248 +
249 + # Note: keep allocators alive while MemPool objects are destroyed.
250 + # MemPool teardown may invoke allocator virtual methods (e.g. raw_delete)
251 + # when releasing cached blocks. If allocator wrappers are dropped first,
252 + # C++ can hit "pure virtual method called" during shutdown.
253 + pool_entries = list(self.allocator_and_pools.values())
254 + self.allocator_and_pools.clear()
255 +
256 + mem_pools = [entry[0] for entry in pool_entries]
257 + allocators = [entry[1] for entry in pool_entries]
258 + pool_entries.clear()
259 +
260 + xpu_sync = getattr(torch.xpu, "synchronize", None)
261 + if callable(xpu_sync):
262 + try:
263 + xpu_sync()
264 + except Exception:
265 + logger.debug("torch.xpu.synchronize() failed during release_pools")
266 +
267 + # Phase 1: drop MemPool refs while allocators are still strongly held.
268 + mem_pools.clear()
269 + gc.collect()
270 +
271 + # Phase 2: now it is safe to release allocator wrappers.
272 + allocators.clear()
273 +
274 + @contextmanager
275 + def use_memory_pool(self, tag: str | None = None):
276 + if tag is None:
277 + tag = XpuMemAllocator.default_tag
278 +
279 + old_tag = self.current_tag
280 + self.current_tag = tag
281 + try:
282 + with use_memory_pool_with_allocator(
283 + self.python_malloc_callback,
284 + self.python_free_callback,
285 + ) as data:
286 + self.allocator_and_pools[tag] = data
287 + yield
288 + finally:
289 + self.current_tag = old_tag
290 +
291 + def get_current_usage(self) -> int:
292 + total = 0
293 + for data in self.pointer_to_data.values():
294 + total += data.handle[1]
295 + return total
vllm/platforms/interface.py
+1 -1
@@ -197,7 +197,7 @@ class Platform:
197 # for ROCm, but currently we don't have a way to detect the
198 # exact GPU model statelessly here. So we return True for
199 # all ROCm platforms for now.
200 - return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
200 + return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM, PlatformEnum.XPU)
201
202 def is_cumem_allocator_available(self) -> bool:
203 try:
vllm/v1/worker/gpu_worker.py
+16 -9
@@ -18,6 +18,7 @@ import torch.nn as nn
18 import vllm.envs as envs
19 from vllm.config import CUDAGraphMode, VllmConfig, set_current_vllm_config
20 from vllm.config.compilation import CompilationMode
21 +from vllm.device_allocator import get_mem_allocator_instance
22 from vllm.distributed import (
23 ensure_model_parallel_initialized,
24 init_distributed_environment,
@@ -162,8 +163,6 @@ class Worker(WorkerBase):
163 self._pp_send_work: list[Handle] = []
164
165 def sleep(self, level: int = 1) -> None:
165 - from vllm.device_allocator.cumem import CuMemAllocator
166 -
166 free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
167
168 # Save the buffers before level 2 sleep
@@ -173,7 +172,7 @@ class Worker(WorkerBase):
172 name: buffer.cpu().clone() for name, buffer in model.named_buffers()
173 }
174
176 - allocator = CuMemAllocator.get_instance()
175 + allocator = get_mem_allocator_instance()
176 allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
177 free_bytes_after_sleep, total = torch.cuda.mem_get_info()
178 freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
@@ -186,9 +185,7 @@ class Worker(WorkerBase):
185 )
186
187 def wake_up(self, tags: list[str] | None = None) -> None:
189 - from vllm.device_allocator.cumem import CuMemAllocator
190 -
191 - allocator = CuMemAllocator.get_instance()
188 + allocator = get_mem_allocator_instance()
189 allocator.wake_up(tags)
190
191 # Restore the buffers after level 2 sleep
@@ -203,12 +200,22 @@ class Worker(WorkerBase):
200 self.model_runner.post_kv_cache_wake_up()
201
202 def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
206 - if not self.vllm_config.model_config.enable_cumem_allocator:
203 + if (
204 + current_platform.is_cuda_alike()
205 + and not self.vllm_config.model_config.enable_cumem_allocator
206 + ):
207 return nullcontext()
208
209 - from vllm.device_allocator.cumem import CuMemAllocator
209 + if (
210 + current_platform.is_xpu()
211 + and not self.vllm_config.model_config.enable_sleep_mode
212 + ):
213 + return nullcontext()
214 +
215 + if current_platform.is_cpu():
216 + return nullcontext()
217
211 - allocator = CuMemAllocator.get_instance()
218 + allocator = get_mem_allocator_instance()
219 if tag == "weights":
220 assert allocator.get_current_usage() == 0, (
221 "CuMem allocator can only be used for one instance per process."