| 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 DEFERRABLE_TAGS, AllocationData, HandleType |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.utils.torch_utils import PIN_MEMORY |
| 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 | |
| 122 | instance: "XpuMemAllocator | None" = None |
| 123 | default_tag: str = "default" |
| 124 | |
| 125 | @staticmethod |
| 126 | def get_instance() -> "XpuMemAllocator": |
| 127 | assert xpumem_available, "xpumem allocator is not available" |
| 128 | if XpuMemAllocator.instance is None: |
| 129 | XpuMemAllocator.instance = XpuMemAllocator() |
| 130 | # Ensure MemPool/allocator wrappers are released before interpreter |
| 131 | # finalization tears down XPU runtime internals. |
| 132 | atexit.register(XpuMemAllocator._shutdown_singleton) |
| 133 | return XpuMemAllocator.instance |
| 134 | |
| 135 | @staticmethod |
| 136 | def _shutdown_singleton() -> None: |
| 137 | instance = XpuMemAllocator.instance |
| 138 | if instance is None: |
| 139 | return |
| 140 | try: |
| 141 | instance.release_pools() |
| 142 | except Exception: |
| 143 | logger.exception("XpuMemAllocator singleton shutdown failed") |
| 144 | |
| 145 | def __init__(self): |
| 146 | self.pointer_to_data: dict[int, AllocationData] = {} |
| 147 | self.current_tag: str = XpuMemAllocator.default_tag |
| 148 | self.allocator_and_pools: dict[str, list[Any]] = {} |
| 149 | self.python_malloc_callback = self._python_malloc_callback |
| 150 | self.python_free_callback = self._python_free_callback |
| 151 | |
| 152 | def _python_malloc_callback(self, allocation_handle: HandleType) -> None: |
| 153 | ptr = allocation_handle[2] |
| 154 | self.pointer_to_data[ptr] = AllocationData(allocation_handle, self.current_tag) |
| 155 | logger.debug( |
| 156 | "Allocated %s bytes for %s at %s", |
| 157 | allocation_handle[1], |
| 158 | self.current_tag, |
| 159 | ptr, |
| 160 | ) |
| 161 | |
| 162 | def _python_free_callback(self, ptr: int) -> HandleType: |
| 163 | data = self.pointer_to_data.pop(ptr) |
| 164 | data.cpu_backup_tensor = None |
| 165 | logger.debug("Freed %s bytes for %s at %s", data.handle[1], data.tag, ptr) |
| 166 | return data.handle |
| 167 | |
| 168 | def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: |
| 169 | if offload_tags is None: |
| 170 | offload_tags = (XpuMemAllocator.default_tag,) |
| 171 | elif isinstance(offload_tags, str): |
| 172 | offload_tags = (offload_tags,) |
| 173 | |
| 174 | assert isinstance(offload_tags, tuple) |
| 175 | |
| 176 | total_bytes = 0 |
| 177 | backup_bytes = 0 |
| 178 | has_policy_conflict = False |
| 179 | |
| 180 | for ptr, data in self.pointer_to_data.items(): |
| 181 | if data.is_asleep: |
| 182 | requests_offload = data.tag in offload_tags |
| 183 | was_offloaded = data.cpu_backup_tensor is not None |
| 184 | if requests_offload != was_offloaded: |
| 185 | has_policy_conflict = True |
| 186 | continue |
| 187 | size_in_bytes = data.handle[1] |
| 188 | total_bytes += size_in_bytes |
| 189 | if data.tag not in offload_tags: |
| 190 | unmap_and_release(data.handle) |
| 191 | data.is_asleep = True |
| 192 | continue |
| 193 | |
| 194 | backup_bytes += size_in_bytes |
| 195 | device, _, _, _ = data.handle |
| 196 | cpu_backup_tensor = torch.empty( |
| 197 | size_in_bytes, |
| 198 | dtype=torch.uint8, |
| 199 | device="cpu", |
| 200 | pin_memory=PIN_MEMORY, |
| 201 | ) |
| 202 | cpu_ptr = cpu_backup_tensor.data_ptr() |
| 203 | _xpu_memcpy_sync( |
| 204 | cpu_ptr, |
| 205 | ptr, |
| 206 | size_in_bytes, |
| 207 | MEMCPY_DEVICE_TO_HOST, |
| 208 | device, |
| 209 | ) |
| 210 | data.cpu_backup_tensor = cpu_backup_tensor |
| 211 | |
| 212 | unmap_and_release(data.handle) |
| 213 | data.is_asleep = True |
| 214 | |
| 215 | logger.info( |
| 216 | "XpuMemAllocator: sleep freed %.2f GiB memory in total, of which " |
| 217 | "%.2f GiB is backed up in CPU and the rest %.2f GiB is discarded " |
| 218 | "directly.", |
| 219 | total_bytes / 1024**3, |
| 220 | backup_bytes / 1024**3, |
| 221 | (total_bytes - backup_bytes) / 1024**3, |
| 222 | ) |
| 223 | |
| 224 | if has_policy_conflict: |
| 225 | logger.warning( |
| 226 | "XpuMemAllocator: sleep cannot change the policy of " |
| 227 | "already-asleep allocations; the existing policy was kept." |
| 228 | ) |
| 229 | |
| 230 | gc.collect() |
| 231 | xpu_empty_cache = getattr(torch.xpu, "empty_cache", None) |
| 232 | if callable(xpu_empty_cache): |
| 233 | xpu_empty_cache() |
| 234 | |
| 235 | def discard(self, tags: tuple[str, ...] | str) -> None: |
| 236 | """Discard mapped allocations with the given tags without CPU backup.""" |
| 237 | if isinstance(tags, str): |
| 238 | tags = (tags,) |
| 239 | |
| 240 | discarded_bytes = 0 |
| 241 | has_policy_conflict = False |
| 242 | for data in self.pointer_to_data.values(): |
| 243 | if data.tag not in tags: |
| 244 | continue |
| 245 | if data.is_asleep: |
| 246 | if data.cpu_backup_tensor is not None: |
| 247 | has_policy_conflict = True |
| 248 | continue |
| 249 | torch.accelerator.synchronize(data.handle[0]) |
| 250 | unmap_and_release(data.handle) |
| 251 | data.is_asleep = True |
| 252 | discarded_bytes += data.handle[1] |
| 253 | |
| 254 | logger.info( |
| 255 | "XpuMemAllocator: discarded %.2f GiB for tags %s.", |
| 256 | discarded_bytes / 1024**3, |
| 257 | tags, |
| 258 | ) |
| 259 | |
| 260 | if has_policy_conflict: |
| 261 | logger.warning( |
| 262 | "XpuMemAllocator: discard cannot change the policy of " |
| 263 | "already-asleep allocations; the existing policy was kept." |
| 264 | ) |
| 265 | |
| 266 | def wake_up(self, tags: list[str] | None = None) -> None: |
| 267 | for ptr, data in self.pointer_to_data.items(): |
| 268 | if not data.is_asleep: |
| 269 | continue |
| 270 | if ( |
| 271 | tags is not None |
| 272 | and data.tag not in tags |
| 273 | and data.tag in DEFERRABLE_TAGS |
| 274 | ): |
| 275 | continue |
| 276 | create_and_allocate(data.handle) |
| 277 | data.is_asleep = False |
| 278 | |
| 279 | cpu_backup_tensor = data.cpu_backup_tensor |
| 280 | if cpu_backup_tensor is None: |
| 281 | continue |
| 282 | |
| 283 | device, size_in_bytes, _, _ = data.handle |
| 284 | _xpu_memcpy_sync( |
| 285 | ptr, |
| 286 | cpu_backup_tensor.data_ptr(), |
| 287 | size_in_bytes, |
| 288 | MEMCPY_HOST_TO_DEVICE, |
| 289 | device, |
| 290 | ) |
| 291 | data.cpu_backup_tensor = None |
| 292 | |
| 293 | def release_pools(self) -> None: |
| 294 | """Drop Python references to MemPool/pluggable allocators eagerly. |
| 295 | |
| 296 | This prevents pool destruction from being deferred to interpreter |
| 297 | finalization, which can happen after parts of XPU runtime are already |
| 298 | torn down. |
| 299 | """ |
| 300 | if not self.allocator_and_pools: |
| 301 | return |
| 302 | |
| 303 | # Note: keep allocators alive while MemPool objects are destroyed. |
| 304 | # MemPool teardown may invoke allocator virtual methods (e.g. raw_delete) |
| 305 | # when releasing cached blocks. If allocator wrappers are dropped first, |
| 306 | # C++ can hit "pure virtual method called" during shutdown. |
| 307 | pool_entries = [e for es in self.allocator_and_pools.values() for e in es] |
| 308 | self.allocator_and_pools.clear() |
| 309 | |
| 310 | mem_pools = [entry[0] for entry in pool_entries] |
| 311 | allocators = [entry[1] for entry in pool_entries] |
| 312 | pool_entries.clear() |
| 313 | |
| 314 | xpu_sync = getattr(torch.xpu, "synchronize", None) |
| 315 | if callable(xpu_sync): |
| 316 | try: |
| 317 | xpu_sync() |
| 318 | except Exception: |
| 319 | logger.debug("torch.xpu.synchronize() failed during release_pools") |
| 320 | |
| 321 | # Phase 1: drop MemPool refs while allocators are still strongly held. |
| 322 | mem_pools.clear() |
| 323 | gc.collect() |
| 324 | |
| 325 | # Phase 2: now it is safe to release allocator wrappers. |
| 326 | allocators.clear() |
| 327 | |
| 328 | @contextmanager |
| 329 | def use_memory_pool(self, tag: str | None = None): |
| 330 | if tag is None: |
| 331 | tag = XpuMemAllocator.default_tag |
| 332 | |
| 333 | old_tag = self.current_tag |
| 334 | self.current_tag = tag |
| 335 | try: |
| 336 | with use_memory_pool_with_allocator( |
| 337 | self.python_malloc_callback, |
| 338 | self.python_free_callback, |
| 339 | ) as data: |
| 340 | # Keep every entry's pool alive (pytorch/pytorch#146431). |
| 341 | self.allocator_and_pools.setdefault(tag, []).append(data) |
| 342 | yield |
| 343 | finally: |
| 344 | self.current_tag = old_tag |
| 345 | |
| 346 | def get_current_usage(self) -> int: |
| 347 | total = 0 |
| 348 | for data in self.pointer_to_data.values(): |
| 349 | total += data.handle[1] |
| 350 | return total |