vllm.device_allocator.cumem

Classes:

• CuMemAllocator –

  A singleton class that manages a memory pool for CUDA tensors.

CuMemAllocator

A singleton class that manages a memory pool for CUDA tensors. The memory in this pool can be offloaded or discarded when the allocator sleeps.

Inside the use_memory_pool(tag) context, all tensors created will be allocated in the memory pool, and has the same tag as the tag passed to the context.

When we call sleep, all tensors with the specified tag will be offloaded to CPU memory, and the rest of the tensors will be discarded. When we call wake_up, all tensors that are previously offloaded will be loaded back to GPU memory, and the rest of the tensors will have empty memory (zeroed on ROCm).

The _create_and_map, _unmap_and_release, _create_cpu_backup, _release_cpu_backup, _get_pluggable_allocator, and _use_memory_pool_with_allocator methods are subclass hooks for out-of-tree device allocators. Their names and signatures form the subclassing contract; renaming them or changing their signatures is a breaking change.

Why it needs to be a singleton? When allocated tensors are garbage collected, PyTorch will call the free callback, which will call the python_free_callback method. The C-extension uses a global variable to store the function of an instance of this class. If we create multiple instances of this class, the global variable will be overwritten and the free callback will not work as expected.

Methods:

• close –

  Compatibility alias for deterministic pool release.

• cudagraph_pool –

  Tag CUDA graph capture allocations and yield the graph pool id.

• discard –

  Discard mapped allocations with the given tags without CPU backup.

• get_current_usage –

  Get the total number of bytes allocated in the memory pool.

• get_instance –

  CuMemAllocator is a singleton class.

• release_pools –

  Drop Python references to MemPool/pluggable allocators eagerly.

• sleep –

  Put the allocator in sleep mode.

• use_memory_pool –

  A context manager to use the memory pool.

• wake_up –

  Wake up the allocator from sleep mode.

Source code in vllm/device_allocator/cumem.py

class CuMemAllocator:
    """A singleton class that manages a memory pool for CUDA tensors.
    The memory in this pool can be offloaded or discarded when the
    allocator sleeps.

    Inside the `use_memory_pool(tag)` context, all tensors created will
    be allocated in the memory pool, and has the same tag as the
    tag passed to the context.

    When we call `sleep`, all tensors with the specified tag will be
    offloaded to CPU memory, and the rest of the tensors will be discarded.
    When we call `wake_up`, all tensors that are previously offloaded
    will be loaded back to GPU memory, and the rest of the tensors will
    have empty memory (zeroed on ROCm).

    The `_create_and_map`, `_unmap_and_release`, `_create_cpu_backup`,
    `_release_cpu_backup`, `_get_pluggable_allocator`, and
    `_use_memory_pool_with_allocator` methods are subclass hooks for
    out-of-tree device allocators. Their names and signatures form the
    subclassing contract; renaming them or changing their signatures is
    a breaking change.

    Why it needs to be a singleton?
    When allocated tensors are garbage collected, PyTorch will call
    the free callback, which will call the `python_free_callback` method.
    The C-extension uses a global variable to store the function of an
    instance of this class. If we create multiple instances of this class,
    the global variable will be overwritten and the free callback will
    not work as expected.
    """

    instance: "CuMemAllocator | None" = None
    default_tag: str = "default"

    @staticmethod
    def get_instance() -> "CuMemAllocator":
        """CuMemAllocator is a singleton class.
        We cannot call the constructor directly.
        Call this method to get the instance.
        """
        assert cumem_available, "cumem allocator is not available"
        if CuMemAllocator.instance is None:
            CuMemAllocator.instance = CuMemAllocator()
            # Ensure MemPool/allocator wrappers are released before interpreter
            # finalization tears down PyTorch allocator internals.
            atexit.register(CuMemAllocator._shutdown_singleton)
        return CuMemAllocator.instance

    @staticmethod
    def _shutdown_singleton() -> None:
        instance = CuMemAllocator.instance
        if instance is None:
            return
        try:
            instance.release_pools()
        except Exception:
            logger.exception("CuMemAllocator singleton shutdown failed")

    def __init__(self):
        self.pointer_to_data: dict[int, AllocationData] = {}
        self.current_tag: str = CuMemAllocator.default_tag
        self.allocator_and_pools: dict[str, list[Any]] = {}
        # Creating strong references to the two callbacks here to prevent
        # these ephemeral bound-method objects being garbage collected.
        # See discussions in https://github.com/vllm-project/vllm/pull/22724
        self.python_malloc_callback = self._python_malloc_callback
        self.python_free_callback = self._python_free_callback

    @staticmethod
    def _create_and_map(allocation_handle: HandleType) -> None:
        python_create_and_map(*allocation_handle)

    @staticmethod
    def _unmap_and_release(allocation_handle: HandleType) -> None:
        python_unmap_and_release(*allocation_handle)

    @staticmethod
    def _create_cpu_backup(size_in_bytes: int) -> torch.Tensor:
        backup = torch.frombuffer(
            mmap.mmap(-1, size_in_bytes, flags=mmap.MAP_PRIVATE), dtype=torch.uint8
        )
        if PIN_MEMORY:
            error = libcudart.cudaHostRegister(backup.data_ptr(), size_in_bytes)
            if error:
                libcudart.cudaGetLastError()
                logger.warning_once(
                    "cudaHostRegister failed (error %d); sleep mode CPU backups "
                    "fall back to pageable memory.",
                    error,
                )
        return backup

    @staticmethod
    def _release_cpu_backup(backup: torch.Tensor) -> None:
        if backup.is_pinned() and libcudart.cudaHostUnregister(backup.data_ptr()):
            libcudart.cudaGetLastError()

    def _get_pluggable_allocator(self) -> torch.cuda.memory.CUDAPluggableAllocator:
        init_module(self.python_malloc_callback, self.python_free_callback)
        return torch.cuda.memory.CUDAPluggableAllocator(
            lib_name, "my_malloc", "my_free"
        )

    @contextmanager
    def _use_memory_pool_with_allocator(
        self,
    ) -> Iterator[
        tuple[torch.cuda.memory.MemPool, torch.cuda.memory.CUDAPluggableAllocator]
    ]:
        new_alloc = self._get_pluggable_allocator()
        mem_pool = torch.cuda.memory.MemPool(new_alloc._allocator)
        with torch.cuda.memory.use_mem_pool(mem_pool):
            yield mem_pool, new_alloc

    def release_pools(self) -> None:
        """Drop Python references to MemPool/pluggable allocators eagerly.

        A cumem `MemPool` outlives the `use_memory_pool` context (a strong
        reference is kept in `allocator_and_pools` to work around
        pytorch/pytorch#146431), and a captured CUDA graph can keep it alive
        longer still. `MemPool` only holds a non-owning pointer to the
        allocator, whose owning reference lives in the Python
        `CUDAPluggableAllocator`. If both are instead dropped during
        interpreter shutdown, GC may finalize the allocator first; the eventual
        `~MemPool` -> `emptyCache` -> `release_block` then makes a virtual
        call into the freed allocator -- aborting the process with "pure virtual
        method called" (pytorch/pytorch#145168).

        Release the kept-alive pools before interpreter finalization, and keep
        the pluggable allocator wrappers alive while MemPool destructors run.
        This is safe to call more than once.
        """
        if not self.allocator_and_pools:
            return

        pool_entries = [e for es in self.allocator_and_pools.values() for e in es]
        self.allocator_and_pools.clear()

        mem_pools = [entry[0] for entry in pool_entries]
        allocators = [entry[1] for entry in pool_entries]
        pool_entries.clear()

        # Phase 1: drop MemPool refs while allocators are still strongly held.
        mem_pools.clear()
        gc.collect()

        # Phase 2: now it is safe to release allocator wrappers.
        allocators.clear()

    def close(self) -> None:
        """Compatibility alias for deterministic pool release."""
        self.release_pools()

    def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
        """Internal method to store the allocation data
        when memory is allocated in the memory pool.
        """
        py_d_mem = allocation_handle[2]
        self.pointer_to_data[py_d_mem] = AllocationData(
            allocation_handle, self.current_tag
        )
        logger.debug(
            "Allocated %s bytes for %s with address %s from cumem allocator",
            allocation_handle[1],
            self.current_tag,
            py_d_mem,
        )
        return

    def _python_free_callback(self, ptr: int) -> HandleType:
        """Internal method to look up the allocation data
        when memory is freed in the memory pool.
        """
        data = self.pointer_to_data.pop(ptr)
        if data.cpu_backup_tensor is not None:
            self._release_cpu_backup(data.cpu_backup_tensor)
            data.cpu_backup_tensor = None
        if data.is_asleep and current_platform.is_rocm():
            # On ROCm, sleep() already unmapped and released this allocation's
            # physical chunks and holds its virtual address as a placeholder
            # reservation. Return a handle with an empty chunk list so the C
            # extension skips unmap/release (avoiding a double-free) while
            # still freeing the placeholder address.
            device, size, d_mem, _ = data.handle
            return (device, size, d_mem, [])
        # Drain pending kernels before the C extension's cuMemUnmap.
        # The pluggable allocator path doesn't defer reclaim like the
        # regular caching allocator, so without this, in-flight work
        # (e.g. quant helpers' transient tensors during weight loading)
        # races the unmap and surfaces as CUDA_ERROR_ILLEGAL_ADDRESS.
        torch.accelerator.synchronize(data.handle[0])
        logger.debug(
            "Freed %s bytes for %s with address %s from cumem allocator",
            data.handle[1],
            data.tag,
            ptr,
        )
        return data.handle

    def sleep(
        self,
        offload_tags: tuple[str, ...] | str | None = None,
        discardable: Collection[int] = (),
    ) -> None:
        """Put the allocator in sleep mode.
        All data in the memory allocation with the specified tag will be
        offloaded to CPU memory, and others will be discarded.

        Args:
            offload_tags: The tags of the memory allocation that will be
                offloaded. The rest of the memory allocation will be discarded.
            discardable: Addresses of tensors the caller rewrites after waking.
                Weights allocations holding nothing else are discarded too.

        """
        if offload_tags is None:
            # by default, allocated tensors are offloaded
            # when the allocator sleeps
            offload_tags = (CuMemAllocator.default_tag,)
        elif isinstance(offload_tags, str):
            offload_tags = (offload_tags,)

        assert isinstance(offload_tags, tuple)

        total_bytes = 0
        backup_bytes = 0
        has_policy_conflict = False
        discarded = discardable_allocations(self.allocator_and_pools, discardable)

        for ptr, data in self.pointer_to_data.items():
            offload = data.tag in offload_tags and ptr not in discarded
            if data.is_asleep:
                if offload != (data.cpu_backup_tensor is not None):
                    has_policy_conflict = True
                continue
            total_bytes += data.handle[1]
            if offload:
                backup_bytes += data.handle[1]
            self._put_to_sleep(ptr, data, offload)

        logger.info(
            "CuMemAllocator: sleep freed %.2f GiB memory in total, of which "
            "%.2f GiB is backed up in CPU and the rest %.2f GiB is discarded "
            "directly.",
            total_bytes / 1024**3,
            backup_bytes / 1024**3,
            (total_bytes - backup_bytes) / 1024**3,
        )

        if has_policy_conflict:
            logger.warning(
                "CuMemAllocator: sleep cannot change the policy of "
                "already-asleep allocations; the existing policy was kept."
            )

        gc.collect()
        torch.accelerator.empty_cache()

    def _put_to_sleep(self, ptr: int, data: AllocationData, offload: bool) -> None:
        if offload:
            if data.cpu_backup_tensor is None:
                data.cpu_backup_tensor = self._create_cpu_backup(data.handle[1])
            cpu_ptr = data.cpu_backup_tensor.data_ptr()
            libcudart.cudaMemcpy(cpu_ptr, ptr, data.handle[1])
        elif data.cpu_backup_tensor is not None:
            self._release_cpu_backup(data.cpu_backup_tensor)
            data.cpu_backup_tensor = None
        try:
            self._unmap_and_release(data.handle)
        finally:
            data.is_asleep = True

    def discard(self, tags: tuple[str, ...] | str) -> None:
        """Discard mapped allocations with the given tags without CPU backup."""
        if isinstance(tags, str):
            tags = (tags,)

        discarded_bytes = 0
        has_policy_conflict = False
        for ptr, data in self.pointer_to_data.items():
            if data.tag not in tags:
                continue
            if data.is_asleep:
                if data.cpu_backup_tensor is not None:
                    has_policy_conflict = True
                continue
            torch.accelerator.synchronize(data.handle[0])
            self._put_to_sleep(ptr, data, offload=False)
            discarded_bytes += data.handle[1]

        logger.info(
            "CuMemAllocator: discarded %.2f GiB for tags %s.",
            discarded_bytes / 1024**3,
            tags,
        )

        if has_policy_conflict:
            logger.warning(
                "CuMemAllocator: discard cannot change the policy of "
                "already-asleep allocations; the existing policy was kept."
            )

    def wake_up(self, tags: list[str] | None = None) -> None:
        """Wake up the allocator from sleep mode.
        All data that is previously offloaded will be loaded back to GPU
        memory, and the rest will have empty memory (zeroed on ROCm).

        Args:
            tags: The deferrable tags (weights, kv_cache) to load back to GPU
                memory; every other tag is always loaded back. If None, all
                memory allocation will be loaded back to GPU memory.

        """
        gc.collect()
        torch.accelerator.empty_cache()

        for ptr, data in self.pointer_to_data.items():
            if not data.is_asleep:
                continue
            if tags is None or data.tag in tags or data.tag not in DEFERRABLE_TAGS:
                handle = data.handle
                self._create_and_map(handle)
                data.is_asleep = False
                if data.cpu_backup_tensor is not None:
                    cpu_ptr = data.cpu_backup_tensor.data_ptr()
                    libcudart.cudaMemcpy(ptr, cpu_ptr, handle[1])
                elif current_platform.is_rocm():
                    # amdgpu <6.14 may return stale VRAM (#44972); drop with its support
                    libcudart.cudaMemset(ptr, 0, handle[1])
                    torch.accelerator.synchronize()

    @contextmanager
    def use_memory_pool(self, tag: str | None = None):
        """A context manager to use the memory pool.
        All memory allocation created inside the context will be allocated
        in the memory pool, and has the specified tag.

        Args:
            tag: The tag of the memory allocation. If None, the default tag
                will be used.

        """
        if tag is None:
            tag = CuMemAllocator.default_tag

        assert isinstance(tag, str)

        # Expandable segments are incompatible with the memory pool used for
        # sleep mode (see https://github.com/pytorch/pytorch/issues/147851).
        # If the user has enabled expandable segments, temporarily disable
        # them for the duration of the memory pool context and restore on
        # exit. The whole config string is rewritten, not just this one
        # field: torch resets every option that is absent from the string it
        # is handed, so a bare "expandable_segments:False" would also drop
        # max_split_size_mb, garbage_collection_threshold and
        # roundup_power2_divisions for the rest of the process.
        prev_conf = current_alloc_conf()
        expandable_was_enabled = conf_flag_enabled(prev_conf, EXPANDABLE_SEGMENTS)
        if expandable_was_enabled:
            set_alloc_conf(with_conf_flag(prev_conf, EXPANDABLE_SEGMENTS, False))

        old_tag = self.current_tag
        self.current_tag = tag
        try:
            if tag != old_tag:
                # Older pools of this tag are never reused, so trim them.
                for pool, _ in self.allocator_and_pools.get(tag, []):
                    self._trim(pool)
            with self._use_memory_pool_with_allocator() as data:
                # start to hit another PyTorch bug in PyTorch 2.6,
                # possibly because of gc-related issue w.r.t. the allocator
                # and the memory pool.
                # to avoid the issue, we keep a reference of the data.
                # see https://github.com/pytorch/pytorch/issues/146431 .
                self.allocator_and_pools.setdefault(tag, []).append(data)
                yield
                # PyTorch's bug, torch.accelerator.empty_cache() does not
                # release memory held by a pluggable allocator, see
                # https://github.com/pytorch/pytorch/issues/145168 .
                # if we have some memory allocated and then freed,
                # the memory will not be released, e.g. in online
                # quantization, where the model is created in higher
                # precision, and then quantized in lower precision.
                # Find all unused allocations and manually release them.
                # TODO: we should expose `empty_cache` method in the memory
                # pool.
                # TODO: ask for help from PyTorch team to expose this method.
                self._trim(data[0])
        finally:
            self.current_tag = old_tag
            if expandable_was_enabled:
                set_alloc_conf(prev_conf)

    def _trim(self, pool: Any) -> None:
        """Release free segments of a pool that is never allocated from again."""
        for allocation in pool.snapshot():
            data = self.pointer_to_data.get(allocation["address"])
            if allocation["allocated_size"] == 0 and data and not data.is_asleep:
                self._unmap_and_release(
                    self._python_free_callback(allocation["address"])
                )

    @contextmanager
    def cudagraph_pool(self) -> Iterator[tuple[int, int]]:
        """Tag CUDA graph capture allocations and yield the graph pool id."""
        if "cudagraph" not in self.allocator_and_pools:
            allocator = self._get_pluggable_allocator()
            mem_pool = torch.cuda.memory.MemPool(allocator._allocator)
            self.allocator_and_pools["cudagraph"] = [(mem_pool, allocator)]
        old_tag = self.current_tag
        self.current_tag = "cudagraph"
        try:
            # capture_begin routes allocations to this pool; no use_mem_pool.
            yield self.allocator_and_pools["cudagraph"][0][0].id
        finally:
            self.current_tag = old_tag

    def get_current_usage(self) -> int:
        """Get the total number of bytes allocated in the memory pool."""
        sum_bytes: int = 0
        for data in self.pointer_to_data.values():
            handle = data.handle
            sum_bytes += handle[1]
        return sum_bytes

_python_free_callback(ptr)

Internal method to look up the allocation data when memory is freed in the memory pool.

Source code in vllm/device_allocator/cumem.py

def _python_free_callback(self, ptr: int) -> HandleType:
    """Internal method to look up the allocation data
    when memory is freed in the memory pool.
    """
    data = self.pointer_to_data.pop(ptr)
    if data.cpu_backup_tensor is not None:
        self._release_cpu_backup(data.cpu_backup_tensor)
        data.cpu_backup_tensor = None
    if data.is_asleep and current_platform.is_rocm():
        # On ROCm, sleep() already unmapped and released this allocation's
        # physical chunks and holds its virtual address as a placeholder
        # reservation. Return a handle with an empty chunk list so the C
        # extension skips unmap/release (avoiding a double-free) while
        # still freeing the placeholder address.
        device, size, d_mem, _ = data.handle
        return (device, size, d_mem, [])
    # Drain pending kernels before the C extension's cuMemUnmap.
    # The pluggable allocator path doesn't defer reclaim like the
    # regular caching allocator, so without this, in-flight work
    # (e.g. quant helpers' transient tensors during weight loading)
    # races the unmap and surfaces as CUDA_ERROR_ILLEGAL_ADDRESS.
    torch.accelerator.synchronize(data.handle[0])
    logger.debug(
        "Freed %s bytes for %s with address %s from cumem allocator",
        data.handle[1],
        data.tag,
        ptr,
    )
    return data.handle

_python_malloc_callback(allocation_handle)

Internal method to store the allocation data when memory is allocated in the memory pool.

Source code in vllm/device_allocator/cumem.py

def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
    """Internal method to store the allocation data
    when memory is allocated in the memory pool.
    """
    py_d_mem = allocation_handle[2]
    self.pointer_to_data[py_d_mem] = AllocationData(
        allocation_handle, self.current_tag
    )
    logger.debug(
        "Allocated %s bytes for %s with address %s from cumem allocator",
        allocation_handle[1],
        self.current_tag,
        py_d_mem,
    )
    return

_trim(pool)

Release free segments of a pool that is never allocated from again.

Source code in vllm/device_allocator/cumem.py

def _trim(self, pool: Any) -> None:
    """Release free segments of a pool that is never allocated from again."""
    for allocation in pool.snapshot():
        data = self.pointer_to_data.get(allocation["address"])
        if allocation["allocated_size"] == 0 and data and not data.is_asleep:
            self._unmap_and_release(
                self._python_free_callback(allocation["address"])
            )

close()

Compatibility alias for deterministic pool release.

Source code in vllm/device_allocator/cumem.py

def close(self) -> None:
    """Compatibility alias for deterministic pool release."""
    self.release_pools()

cudagraph_pool()

Tag CUDA graph capture allocations and yield the graph pool id.

Source code in vllm/device_allocator/cumem.py

@contextmanager
def cudagraph_pool(self) -> Iterator[tuple[int, int]]:
    """Tag CUDA graph capture allocations and yield the graph pool id."""
    if "cudagraph" not in self.allocator_and_pools:
        allocator = self._get_pluggable_allocator()
        mem_pool = torch.cuda.memory.MemPool(allocator._allocator)
        self.allocator_and_pools["cudagraph"] = [(mem_pool, allocator)]
    old_tag = self.current_tag
    self.current_tag = "cudagraph"
    try:
        # capture_begin routes allocations to this pool; no use_mem_pool.
        yield self.allocator_and_pools["cudagraph"][0][0].id
    finally:
        self.current_tag = old_tag

discard(tags)

Discard mapped allocations with the given tags without CPU backup.

Source code in vllm/device_allocator/cumem.py

def discard(self, tags: tuple[str, ...] | str) -> None:
    """Discard mapped allocations with the given tags without CPU backup."""
    if isinstance(tags, str):
        tags = (tags,)

    discarded_bytes = 0
    has_policy_conflict = False
    for ptr, data in self.pointer_to_data.items():
        if data.tag not in tags:
            continue
        if data.is_asleep:
            if data.cpu_backup_tensor is not None:
                has_policy_conflict = True
            continue
        torch.accelerator.synchronize(data.handle[0])
        self._put_to_sleep(ptr, data, offload=False)
        discarded_bytes += data.handle[1]

    logger.info(
        "CuMemAllocator: discarded %.2f GiB for tags %s.",
        discarded_bytes / 1024**3,
        tags,
    )

    if has_policy_conflict:
        logger.warning(
            "CuMemAllocator: discard cannot change the policy of "
            "already-asleep allocations; the existing policy was kept."
        )

get_current_usage()

Get the total number of bytes allocated in the memory pool.

Source code in vllm/device_allocator/cumem.py

def get_current_usage(self) -> int:
    """Get the total number of bytes allocated in the memory pool."""
    sum_bytes: int = 0
    for data in self.pointer_to_data.values():
        handle = data.handle
        sum_bytes += handle[1]
    return sum_bytes

get_instance() staticmethod

CuMemAllocator is a singleton class. We cannot call the constructor directly. Call this method to get the instance.

Source code in vllm/device_allocator/cumem.py

@staticmethod
def get_instance() -> "CuMemAllocator":
    """CuMemAllocator is a singleton class.
    We cannot call the constructor directly.
    Call this method to get the instance.
    """
    assert cumem_available, "cumem allocator is not available"
    if CuMemAllocator.instance is None:
        CuMemAllocator.instance = CuMemAllocator()
        # Ensure MemPool/allocator wrappers are released before interpreter
        # finalization tears down PyTorch allocator internals.
        atexit.register(CuMemAllocator._shutdown_singleton)
    return CuMemAllocator.instance

release_pools()

Drop Python references to MemPool/pluggable allocators eagerly.

A cumem MemPool outlives the use_memory_pool context (a strong reference is kept in allocator_and_pools to work around pytorch/pytorch#146431), and a captured CUDA graph can keep it alive longer still. MemPool only holds a non-owning pointer to the allocator, whose owning reference lives in the Python CUDAPluggableAllocator. If both are instead dropped during interpreter shutdown, GC may finalize the allocator first; the eventual ~MemPool -> emptyCache -> release_block then makes a virtual call into the freed allocator -- aborting the process with "pure virtual method called" (pytorch/pytorch#145168).

Release the kept-alive pools before interpreter finalization, and keep the pluggable allocator wrappers alive while MemPool destructors run. This is safe to call more than once.

Source code in vllm/device_allocator/cumem.py

def release_pools(self) -> None:
    """Drop Python references to MemPool/pluggable allocators eagerly.

    A cumem `MemPool` outlives the `use_memory_pool` context (a strong
    reference is kept in `allocator_and_pools` to work around
    pytorch/pytorch#146431), and a captured CUDA graph can keep it alive
    longer still. `MemPool` only holds a non-owning pointer to the
    allocator, whose owning reference lives in the Python
    `CUDAPluggableAllocator`. If both are instead dropped during
    interpreter shutdown, GC may finalize the allocator first; the eventual
    `~MemPool` -> `emptyCache` -> `release_block` then makes a virtual
    call into the freed allocator -- aborting the process with "pure virtual
    method called" (pytorch/pytorch#145168).

    Release the kept-alive pools before interpreter finalization, and keep
    the pluggable allocator wrappers alive while MemPool destructors run.
    This is safe to call more than once.
    """
    if not self.allocator_and_pools:
        return

    pool_entries = [e for es in self.allocator_and_pools.values() for e in es]
    self.allocator_and_pools.clear()

    mem_pools = [entry[0] for entry in pool_entries]
    allocators = [entry[1] for entry in pool_entries]
    pool_entries.clear()

    # Phase 1: drop MemPool refs while allocators are still strongly held.
    mem_pools.clear()
    gc.collect()

    # Phase 2: now it is safe to release allocator wrappers.
    allocators.clear()

sleep(offload_tags=None, discardable=())

Put the allocator in sleep mode. All data in the memory allocation with the specified tag will be offloaded to CPU memory, and others will be discarded.

Parameters:

• offload_tags

  (tuple[str, ...] | str | None, default: None ) –

  The tags of the memory allocation that will be offloaded. The rest of the memory allocation will be discarded.

• discardable

  (Collection[int], default: () ) –

  Addresses of tensors the caller rewrites after waking. Weights allocations holding nothing else are discarded too.

Source code in vllm/device_allocator/cumem.py

def sleep(
    self,
    offload_tags: tuple[str, ...] | str | None = None,
    discardable: Collection[int] = (),
) -> None:
    """Put the allocator in sleep mode.
    All data in the memory allocation with the specified tag will be
    offloaded to CPU memory, and others will be discarded.

    Args:
        offload_tags: The tags of the memory allocation that will be
            offloaded. The rest of the memory allocation will be discarded.
        discardable: Addresses of tensors the caller rewrites after waking.
            Weights allocations holding nothing else are discarded too.

    """
    if offload_tags is None:
        # by default, allocated tensors are offloaded
        # when the allocator sleeps
        offload_tags = (CuMemAllocator.default_tag,)
    elif isinstance(offload_tags, str):
        offload_tags = (offload_tags,)

    assert isinstance(offload_tags, tuple)

    total_bytes = 0
    backup_bytes = 0
    has_policy_conflict = False
    discarded = discardable_allocations(self.allocator_and_pools, discardable)

    for ptr, data in self.pointer_to_data.items():
        offload = data.tag in offload_tags and ptr not in discarded
        if data.is_asleep:
            if offload != (data.cpu_backup_tensor is not None):
                has_policy_conflict = True
            continue
        total_bytes += data.handle[1]
        if offload:
            backup_bytes += data.handle[1]
        self._put_to_sleep(ptr, data, offload)

    logger.info(
        "CuMemAllocator: sleep freed %.2f GiB memory in total, of which "
        "%.2f GiB is backed up in CPU and the rest %.2f GiB is discarded "
        "directly.",
        total_bytes / 1024**3,
        backup_bytes / 1024**3,
        (total_bytes - backup_bytes) / 1024**3,
    )

    if has_policy_conflict:
        logger.warning(
            "CuMemAllocator: sleep cannot change the policy of "
            "already-asleep allocations; the existing policy was kept."
        )

    gc.collect()
    torch.accelerator.empty_cache()

use_memory_pool(tag=None)

A context manager to use the memory pool. All memory allocation created inside the context will be allocated in the memory pool, and has the specified tag.

Parameters:

• tag

  (str | None, default: None ) –

  The tag of the memory allocation. If None, the default tag will be used.

Source code in vllm/device_allocator/cumem.py

@contextmanager
def use_memory_pool(self, tag: str | None = None):
    """A context manager to use the memory pool.
    All memory allocation created inside the context will be allocated
    in the memory pool, and has the specified tag.

    Args:
        tag: The tag of the memory allocation. If None, the default tag
            will be used.

    """
    if tag is None:
        tag = CuMemAllocator.default_tag

    assert isinstance(tag, str)

    # Expandable segments are incompatible with the memory pool used for
    # sleep mode (see https://github.com/pytorch/pytorch/issues/147851).
    # If the user has enabled expandable segments, temporarily disable
    # them for the duration of the memory pool context and restore on
    # exit. The whole config string is rewritten, not just this one
    # field: torch resets every option that is absent from the string it
    # is handed, so a bare "expandable_segments:False" would also drop
    # max_split_size_mb, garbage_collection_threshold and
    # roundup_power2_divisions for the rest of the process.
    prev_conf = current_alloc_conf()
    expandable_was_enabled = conf_flag_enabled(prev_conf, EXPANDABLE_SEGMENTS)
    if expandable_was_enabled:
        set_alloc_conf(with_conf_flag(prev_conf, EXPANDABLE_SEGMENTS, False))

    old_tag = self.current_tag
    self.current_tag = tag
    try:
        if tag != old_tag:
            # Older pools of this tag are never reused, so trim them.
            for pool, _ in self.allocator_and_pools.get(tag, []):
                self._trim(pool)
        with self._use_memory_pool_with_allocator() as data:
            # start to hit another PyTorch bug in PyTorch 2.6,
            # possibly because of gc-related issue w.r.t. the allocator
            # and the memory pool.
            # to avoid the issue, we keep a reference of the data.
            # see https://github.com/pytorch/pytorch/issues/146431 .
            self.allocator_and_pools.setdefault(tag, []).append(data)
            yield
            # PyTorch's bug, torch.accelerator.empty_cache() does not
            # release memory held by a pluggable allocator, see
            # https://github.com/pytorch/pytorch/issues/145168 .
            # if we have some memory allocated and then freed,
            # the memory will not be released, e.g. in online
            # quantization, where the model is created in higher
            # precision, and then quantized in lower precision.
            # Find all unused allocations and manually release them.
            # TODO: we should expose `empty_cache` method in the memory
            # pool.
            # TODO: ask for help from PyTorch team to expose this method.
            self._trim(data[0])
    finally:
        self.current_tag = old_tag
        if expandable_was_enabled:
            set_alloc_conf(prev_conf)

wake_up(tags=None)

Wake up the allocator from sleep mode. All data that is previously offloaded will be loaded back to GPU memory, and the rest will have empty memory (zeroed on ROCm).

Parameters:

• tags

  (list[str] | None, default: None ) –

  The deferrable tags (weights, kv_cache) to load back to GPU memory; every other tag is always loaded back. If None, all memory allocation will be loaded back to GPU memory.

Source code in vllm/device_allocator/cumem.py

def wake_up(self, tags: list[str] | None = None) -> None:
    """Wake up the allocator from sleep mode.
    All data that is previously offloaded will be loaded back to GPU
    memory, and the rest will have empty memory (zeroed on ROCm).

    Args:
        tags: The deferrable tags (weights, kv_cache) to load back to GPU
            memory; every other tag is always loaded back. If None, all
            memory allocation will be loaded back to GPU memory.

    """
    gc.collect()
    torch.accelerator.empty_cache()

    for ptr, data in self.pointer_to_data.items():
        if not data.is_asleep:
            continue
        if tags is None or data.tag in tags or data.tag not in DEFERRABLE_TAGS:
            handle = data.handle
            self._create_and_map(handle)
            data.is_asleep = False
            if data.cpu_backup_tensor is not None:
                cpu_ptr = data.cpu_backup_tensor.data_ptr()
                libcudart.cudaMemcpy(ptr, cpu_ptr, handle[1])
            elif current_platform.is_rocm():
                # amdgpu <6.14 may return stale VRAM (#44972); drop with its support
                libcudart.cudaMemset(ptr, 0, handle[1])
                torch.accelerator.synchronize()