| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # cumem-based pytorch pluggable allocator to implement sleep mode. |
| 5 | # other approaches tried but failed: |
| 6 | # - cuda-python package binding |
| 7 | # - custom libcuda driver ctypes wrapper |
| 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 atexit |
| 12 | import gc |
| 13 | from collections.abc import Callable, Iterator |
| 14 | from contextlib import contextmanager |
| 15 | from typing import Any |
| 16 | |
| 17 | import torch |
| 18 | |
| 19 | from vllm.device_allocator import DEFERRABLE_TAGS, AllocationData, HandleType |
| 20 | from vllm.device_allocator.alloc_conf import ( |
| 21 | EXPANDABLE_SEGMENTS, |
| 22 | conf_flag_enabled, |
| 23 | current_alloc_conf, |
| 24 | set_alloc_conf, |
| 25 | with_conf_flag, |
| 26 | ) |
| 27 | from vllm.logger import init_logger |
| 28 | from vllm.platforms import current_platform |
| 29 | from vllm.utils.system_utils import find_loaded_library |
| 30 | from vllm.utils.torch_utils import PIN_MEMORY |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | |
| 35 | cumem_available = False |
| 36 | libcudart: Any = None |
| 37 | try: |
| 38 | from vllm.cumem_allocator import ( |
| 39 | init_module, |
| 40 | python_create_and_map, |
| 41 | python_unmap_and_release, |
| 42 | ) |
| 43 | from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary |
| 44 | |
| 45 | lib_name = find_loaded_library("cumem_allocator") |
| 46 | libcudart = CudaRTLibrary() |
| 47 | cumem_available = True |
| 48 | except ModuleNotFoundError: |
| 49 | # only cuda and rocm platforms support cumem allocator |
| 50 | init_module = None |
| 51 | python_create_and_map = None |
| 52 | python_unmap_and_release = None |
| 53 | lib_name = None |
| 54 | |
| 55 | |
| 56 | def create_and_map(allocation_handle: HandleType) -> None: |
| 57 | python_create_and_map(*allocation_handle) |
| 58 | |
| 59 | |
| 60 | def unmap_and_release(allocation_handle: HandleType) -> None: |
| 61 | python_unmap_and_release(*allocation_handle) |
| 62 | |
| 63 | |
| 64 | def get_pluggable_allocator( |
| 65 | python_malloc_fn: Callable[[HandleType], None], |
| 66 | python_free_func: Callable[[int], HandleType], |
| 67 | ) -> torch.cuda.memory.CUDAPluggableAllocator: |
| 68 | init_module(python_malloc_fn, python_free_func) |
| 69 | new_alloc = torch.cuda.memory.CUDAPluggableAllocator( |
| 70 | lib_name, "my_malloc", "my_free" |
| 71 | ) |
| 72 | return new_alloc |
| 73 | |
| 74 | |
| 75 | @contextmanager |
| 76 | def use_memory_pool_with_allocator( |
| 77 | python_malloc_fn: Callable[[HandleType], None], |
| 78 | python_free_func: Callable[[int], HandleType], |
| 79 | ) -> Iterator[ |
| 80 | tuple[torch.cuda.memory.MemPool, torch.cuda.memory.CUDAPluggableAllocator] |
| 81 | ]: |
| 82 | new_alloc = get_pluggable_allocator(python_malloc_fn, python_free_func) |
| 83 | mem_pool = torch.cuda.memory.MemPool(new_alloc._allocator) |
| 84 | with torch.cuda.memory.use_mem_pool(mem_pool): |
| 85 | yield mem_pool, new_alloc |
| 86 | |
| 87 | |
| 88 | class CuMemAllocator: |
| 89 | """A singleton class that manages a memory pool for CUDA tensors. |
| 90 | The memory in this pool can be offloaded or discarded when the |
| 91 | allocator sleeps. |
| 92 | |
| 93 | Inside the `use_memory_pool(tag)` context, all tensors created will |
| 94 | be allocated in the memory pool, and has the same tag as the |
| 95 | tag passed to the context. |
| 96 | |
| 97 | When we call `sleep`, all tensors with the specified tag will be |
| 98 | offloaded to CPU memory, and the rest of the tensors will be discarded. |
| 99 | When we call `wake_up`, all tensors that are previously offloaded |
| 100 | will be loaded back to GPU memory, and the rest of the tensors will |
| 101 | have empty memory. |
| 102 | |
| 103 | Why it needs to be a singleton? |
| 104 | When allocated tensors are garbage collected, PyTorch will call |
| 105 | the free callback, which will call the `python_free_callback` method. |
| 106 | The C-extension uses a global variable to store the function of an |
| 107 | instance of this class. If we create multiple instances of this class, |
| 108 | the global variable will be overwritten and the free callback will |
| 109 | not work as expected. |
| 110 | """ |
| 111 | |
| 112 | instance: "CuMemAllocator | None" = None |
| 113 | default_tag: str = "default" |
| 114 | |
| 115 | @staticmethod |
| 116 | def get_instance() -> "CuMemAllocator": |
| 117 | """CuMemAllocator is a singleton class. |
| 118 | We cannot call the constructor directly. |
| 119 | Call this method to get the instance. |
| 120 | """ |
| 121 | assert cumem_available, "cumem allocator is not available" |
| 122 | if CuMemAllocator.instance is None: |
| 123 | CuMemAllocator.instance = CuMemAllocator() |
| 124 | # Ensure MemPool/allocator wrappers are released before interpreter |
| 125 | # finalization tears down PyTorch allocator internals. |
| 126 | atexit.register(CuMemAllocator._shutdown_singleton) |
| 127 | return CuMemAllocator.instance |
| 128 | |
| 129 | @staticmethod |
| 130 | def _shutdown_singleton() -> None: |
| 131 | instance = CuMemAllocator.instance |
| 132 | if instance is None: |
| 133 | return |
| 134 | try: |
| 135 | instance.release_pools() |
| 136 | except Exception: |
| 137 | logger.exception("CuMemAllocator singleton shutdown failed") |
| 138 | |
| 139 | def __init__(self): |
| 140 | self.pointer_to_data: dict[int, AllocationData] = {} |
| 141 | self.current_tag: str = CuMemAllocator.default_tag |
| 142 | self.allocator_and_pools: dict[str, list[Any]] = {} |
| 143 | # Creating strong references to the two callbacks here to prevent |
| 144 | # these ephemeral bound-method objects being garbage collected. |
| 145 | # See discussions in https://github.com/vllm-project/vllm/pull/22724 |
| 146 | self.python_malloc_callback = self._python_malloc_callback |
| 147 | self.python_free_callback = self._python_free_callback |
| 148 | |
| 149 | def release_pools(self) -> None: |
| 150 | """Drop Python references to MemPool/pluggable allocators eagerly. |
| 151 | |
| 152 | A cumem ``MemPool`` outlives the ``use_memory_pool`` context (a strong |
| 153 | reference is kept in ``allocator_and_pools`` to work around |
| 154 | pytorch/pytorch#146431), and a captured CUDA graph can keep it alive |
| 155 | longer still. ``MemPool`` only holds a non-owning pointer to the |
| 156 | allocator, whose owning reference lives in the Python |
| 157 | ``CUDAPluggableAllocator``. If both are instead dropped during |
| 158 | interpreter shutdown, GC may finalize the allocator first; the eventual |
| 159 | ``~MemPool`` -> ``emptyCache`` -> ``release_block`` then makes a virtual |
| 160 | call into the freed allocator -- aborting the process with "pure virtual |
| 161 | method called" (pytorch/pytorch#145168). |
| 162 | |
| 163 | Release the kept-alive pools before interpreter finalization, and keep |
| 164 | the pluggable allocator wrappers alive while MemPool destructors run. |
| 165 | This is safe to call more than once. |
| 166 | """ |
| 167 | if not self.allocator_and_pools: |
| 168 | return |
| 169 | |
| 170 | pool_entries = [e for es in self.allocator_and_pools.values() for e in es] |
| 171 | self.allocator_and_pools.clear() |
| 172 | |
| 173 | mem_pools = [entry[0] for entry in pool_entries] |
| 174 | allocators = [entry[1] for entry in pool_entries] |
| 175 | pool_entries.clear() |
| 176 | |
| 177 | # Phase 1: drop MemPool refs while allocators are still strongly held. |
| 178 | mem_pools.clear() |
| 179 | gc.collect() |
| 180 | |
| 181 | # Phase 2: now it is safe to release allocator wrappers. |
| 182 | allocators.clear() |
| 183 | |
| 184 | def close(self) -> None: |
| 185 | """Compatibility alias for deterministic pool release.""" |
| 186 | self.release_pools() |
| 187 | |
| 188 | def _python_malloc_callback(self, allocation_handle: HandleType) -> None: |
| 189 | """Internal method to store the allocation data |
| 190 | when memory is allocated in the memory pool. |
| 191 | """ |
| 192 | py_d_mem = allocation_handle[2] |
| 193 | self.pointer_to_data[py_d_mem] = AllocationData( |
| 194 | allocation_handle, self.current_tag |
| 195 | ) |
| 196 | logger.debug( |
| 197 | "Allocated %s bytes for %s with address %s from cumem allocator", |
| 198 | allocation_handle[1], |
| 199 | self.current_tag, |
| 200 | py_d_mem, |
| 201 | ) |
| 202 | return |
| 203 | |
| 204 | def _python_free_callback(self, ptr: int) -> HandleType: |
| 205 | """Internal method to look up the allocation data |
| 206 | when memory is freed in the memory pool. |
| 207 | """ |
| 208 | data = self.pointer_to_data.pop(ptr) |
| 209 | if data.cpu_backup_tensor is not None: |
| 210 | data.cpu_backup_tensor = None |
| 211 | if data.is_asleep and current_platform.is_rocm(): |
| 212 | # On ROCm, sleep() already unmapped and released this allocation's |
| 213 | # physical chunks and holds its virtual address as a placeholder |
| 214 | # reservation. Return a handle with an empty chunk list so the C |
| 215 | # extension skips unmap/release (avoiding a double-free) while |
| 216 | # still freeing the placeholder address. |
| 217 | device, size, d_mem, _ = data.handle |
| 218 | return (device, size, d_mem, []) |
| 219 | # Drain pending kernels before the C extension's cuMemUnmap. |
| 220 | # The pluggable allocator path doesn't defer reclaim like the |
| 221 | # regular caching allocator, so without this, in-flight work |
| 222 | # (e.g. quant helpers' transient tensors during weight loading) |
| 223 | # races the unmap and surfaces as CUDA_ERROR_ILLEGAL_ADDRESS. |
| 224 | torch.cuda.synchronize(data.handle[0]) |
| 225 | logger.debug( |
| 226 | "Freed %s bytes for %s with address %s from cumem allocator", |
| 227 | data.handle[1], |
| 228 | data.tag, |
| 229 | ptr, |
| 230 | ) |
| 231 | return data.handle |
| 232 | |
| 233 | def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: |
| 234 | """Put the allocator in sleep mode. |
| 235 | All data in the memory allocation with the specified tag will be |
| 236 | offloaded to CPU memory, and others will be discarded. |
| 237 | |
| 238 | Args: |
| 239 | offload_tags: The tags of the memory allocation that will be |
| 240 | offloaded. The rest of the memory allocation will be discarded. |
| 241 | |
| 242 | """ |
| 243 | if offload_tags is None: |
| 244 | # by default, allocated tensors are offloaded |
| 245 | # when the allocator sleeps |
| 246 | offload_tags = (CuMemAllocator.default_tag,) |
| 247 | elif isinstance(offload_tags, str): |
| 248 | offload_tags = (offload_tags,) |
| 249 | |
| 250 | assert isinstance(offload_tags, tuple) |
| 251 | |
| 252 | total_bytes = 0 |
| 253 | backup_bytes = 0 |
| 254 | has_policy_conflict = False |
| 255 | |
| 256 | for ptr, data in self.pointer_to_data.items(): |
| 257 | if data.is_asleep: |
| 258 | requests_offload = data.tag in offload_tags |
| 259 | was_offloaded = data.cpu_backup_tensor is not None |
| 260 | if requests_offload != was_offloaded: |
| 261 | has_policy_conflict = True |
| 262 | continue |
| 263 | handle = data.handle |
| 264 | total_bytes += handle[1] |
| 265 | if data.tag in offload_tags: |
| 266 | backup_bytes += handle[1] |
| 267 | size_in_bytes = handle[1] |
| 268 | cpu_backup_tensor = torch.empty( |
| 269 | size_in_bytes, |
| 270 | dtype=torch.uint8, |
| 271 | device="cpu", |
| 272 | pin_memory=PIN_MEMORY, |
| 273 | ) |
| 274 | cpu_ptr = cpu_backup_tensor.data_ptr() |
| 275 | libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes) |
| 276 | data.cpu_backup_tensor = cpu_backup_tensor |
| 277 | try: |
| 278 | unmap_and_release(handle) |
| 279 | finally: |
| 280 | data.is_asleep = True |
| 281 | |
| 282 | logger.info( |
| 283 | "CuMemAllocator: sleep freed %.2f GiB memory in total, of which " |
| 284 | "%.2f GiB is backed up in CPU and the rest %.2f GiB is discarded " |
| 285 | "directly.", |
| 286 | total_bytes / 1024**3, |
| 287 | backup_bytes / 1024**3, |
| 288 | (total_bytes - backup_bytes) / 1024**3, |
| 289 | ) |
| 290 | |
| 291 | if has_policy_conflict: |
| 292 | logger.warning( |
| 293 | "CuMemAllocator: sleep cannot change the policy of " |
| 294 | "already-asleep allocations; the existing policy was kept." |
| 295 | ) |
| 296 | |
| 297 | gc.collect() |
| 298 | torch.cuda.empty_cache() |
| 299 | |
| 300 | def discard(self, tags: tuple[str, ...] | str) -> None: |
| 301 | """Discard mapped allocations with the given tags without CPU backup.""" |
| 302 | if isinstance(tags, str): |
| 303 | tags = (tags,) |
| 304 | |
| 305 | discarded_bytes = 0 |
| 306 | has_policy_conflict = False |
| 307 | for data in self.pointer_to_data.values(): |
| 308 | if data.tag not in tags: |
| 309 | continue |
| 310 | if data.is_asleep: |
| 311 | if data.cpu_backup_tensor is not None: |
| 312 | has_policy_conflict = True |
| 313 | continue |
| 314 | torch.accelerator.synchronize(data.handle[0]) |
| 315 | unmap_and_release(data.handle) |
| 316 | data.is_asleep = True |
| 317 | discarded_bytes += data.handle[1] |
| 318 | |
| 319 | logger.info( |
| 320 | "CuMemAllocator: discarded %.2f GiB for tags %s.", |
| 321 | discarded_bytes / 1024**3, |
| 322 | tags, |
| 323 | ) |
| 324 | |
| 325 | if has_policy_conflict: |
| 326 | logger.warning( |
| 327 | "CuMemAllocator: discard cannot change the policy of " |
| 328 | "already-asleep allocations; the existing policy was kept." |
| 329 | ) |
| 330 | |
| 331 | def wake_up(self, tags: list[str] | None = None) -> None: |
| 332 | """Wake up the allocator from sleep mode. |
| 333 | All data that is previously offloaded will be loaded back to GPU |
| 334 | memory, and the rest of the data will have empty memory. |
| 335 | |
| 336 | Args: |
| 337 | tags: The deferrable tags (weights, kv_cache) to load back to GPU |
| 338 | memory; every other tag is always loaded back. If None, all |
| 339 | memory allocation will be loaded back to GPU memory. |
| 340 | |
| 341 | """ |
| 342 | gc.collect() |
| 343 | torch.accelerator.empty_cache() |
| 344 | |
| 345 | for ptr, data in self.pointer_to_data.items(): |
| 346 | if not data.is_asleep: |
| 347 | continue |
| 348 | if tags is None or data.tag in tags or data.tag not in DEFERRABLE_TAGS: |
| 349 | handle = data.handle |
| 350 | create_and_map(handle) |
| 351 | data.is_asleep = False |
| 352 | if data.cpu_backup_tensor is not None: |
| 353 | cpu_backup_tensor = data.cpu_backup_tensor |
| 354 | if cpu_backup_tensor is not None: |
| 355 | size_in_bytes = ( |
| 356 | cpu_backup_tensor.numel() * cpu_backup_tensor.element_size() |
| 357 | ) |
| 358 | cpu_ptr = cpu_backup_tensor.data_ptr() |
| 359 | libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) |
| 360 | data.cpu_backup_tensor = None |
| 361 | |
| 362 | @contextmanager |
| 363 | def use_memory_pool(self, tag: str | None = None): |
| 364 | """A context manager to use the memory pool. |
| 365 | All memory allocation created inside the context will be allocated |
| 366 | in the memory pool, and has the specified tag. |
| 367 | |
| 368 | Args: |
| 369 | tag: The tag of the memory allocation. If None, the default tag |
| 370 | will be used. |
| 371 | |
| 372 | """ |
| 373 | if tag is None: |
| 374 | tag = CuMemAllocator.default_tag |
| 375 | |
| 376 | assert isinstance(tag, str) |
| 377 | |
| 378 | # Expandable segments are incompatible with the memory pool used for |
| 379 | # sleep mode (see https://github.com/pytorch/pytorch/issues/147851). |
| 380 | # If the user has enabled expandable segments, temporarily disable |
| 381 | # them for the duration of the memory pool context and restore on |
| 382 | # exit. The whole config string is rewritten, not just this one |
| 383 | # field: torch resets every option that is absent from the string it |
| 384 | # is handed, so a bare "expandable_segments:False" would also drop |
| 385 | # max_split_size_mb, garbage_collection_threshold and |
| 386 | # roundup_power2_divisions for the rest of the process. |
| 387 | prev_conf = current_alloc_conf() |
| 388 | expandable_was_enabled = conf_flag_enabled(prev_conf, EXPANDABLE_SEGMENTS) |
| 389 | if expandable_was_enabled: |
| 390 | set_alloc_conf(with_conf_flag(prev_conf, EXPANDABLE_SEGMENTS, False)) |
| 391 | |
| 392 | old_tag = self.current_tag |
| 393 | self.current_tag = tag |
| 394 | try: |
| 395 | if tag != old_tag: |
| 396 | # Older pools of this tag are never reused, so trim them. |
| 397 | for pool, _ in self.allocator_and_pools.get(tag, []): |
| 398 | self._trim(pool) |
| 399 | with use_memory_pool_with_allocator( |
| 400 | self.python_malloc_callback, self.python_free_callback |
| 401 | ) as data: |
| 402 | # start to hit another PyTorch bug in PyTorch 2.6, |
| 403 | # possibly because of gc-related issue w.r.t. the allocator |
| 404 | # and the memory pool. |
| 405 | # to avoid the issue, we keep a reference of the data. |
| 406 | # see https://github.com/pytorch/pytorch/issues/146431 . |
| 407 | self.allocator_and_pools.setdefault(tag, []).append(data) |
| 408 | yield |
| 409 | # PyTorch's bug, calling torch.cuda.empty_cache() will error |
| 410 | # when using pluggable allocator, see |
| 411 | # https://github.com/pytorch/pytorch/issues/145168 . |
| 412 | # if we have some memory allocated and then freed, |
| 413 | # the memory will not be released, e.g. in online |
| 414 | # quantization, where the model is created in higher |
| 415 | # precision, and then quantized in lower precision. |
| 416 | # Find all unused allocations and manually release them. |
| 417 | # TODO: we should expose `empty_cache` method in the memory |
| 418 | # pool. |
| 419 | # TODO: ask for help from PyTorch team to expose this method. |
| 420 | self._trim(data[0]) |
| 421 | finally: |
| 422 | self.current_tag = old_tag |
| 423 | if expandable_was_enabled: |
| 424 | set_alloc_conf(prev_conf) |
| 425 | |
| 426 | def _trim(self, pool: Any) -> None: |
| 427 | """Release free segments of a pool that is never allocated from again.""" |
| 428 | for allocation in pool.snapshot(): |
| 429 | data = self.pointer_to_data.get(allocation["address"]) |
| 430 | if allocation["allocated_size"] == 0 and data and not data.is_asleep: |
| 431 | unmap_and_release(self._python_free_callback(allocation["address"])) |
| 432 | |
| 433 | def get_current_usage(self) -> int: |
| 434 | """Get the total number of bytes allocated in the memory pool.""" |
| 435 | sum_bytes: int = 0 |
| 436 | for ptr, data in self.pointer_to_data.items(): |
| 437 | handle = data.handle |
| 438 | sum_bytes += handle[1] |
| 439 | return sum_bytes |