chore: import upstream snapshot with attribution

vllm/device_allocator/__init__.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import dataclasses
+from contextlib import AbstractContextManager
+from typing import Protocol, TypeAlias
+
+import torch
+
+from vllm.platforms import current_platform
+
+# py_device, py_size_or_aligned_size, py_ptr, py_handle
+# py_handle has type list[int] on ROCm and int otherwise
+HandleType: TypeAlias = tuple[int, int, int, list[int] | int]
+
+
+@dataclasses.dataclass
+class AllocationData:
+    handle: HandleType
+    tag: str
+    cpu_backup_tensor: torch.Tensor | None = None
+    is_asleep: bool = False
+
+
+class MemAllocator(Protocol):
+    def use_memory_pool(self, tag: str | None = None) -> AbstractContextManager: ...
+
+    def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ...
+
+    def wake_up(self, tags: list[str] | None = None) -> None: ...
+
+    def get_current_usage(self) -> int: ...
+
+
+def get_mem_allocator_instance() -> MemAllocator:
+    if current_platform.is_cuda_alike():
+        from vllm.device_allocator.cumem import CuMemAllocator
+
+        return CuMemAllocator.get_instance()
+
+    if current_platform.is_xpu():
+        from vllm.device_allocator.xpumem import XpuMemAllocator
+
+        return XpuMemAllocator.get_instance()
+
+    raise RuntimeError(
+        "Sleep mode allocator is not available on platform "
+        f"{type(current_platform).__name__} "
+        f"(device_type={current_platform.device_type})."
+    )

vllm/device_allocator/cumem.py

@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# cumem-based pytorch pluggable allocator to implement sleep mode.
+# other approaches tried but failed:
+# - cuda-python package binding
+# - custom libcuda driver ctypes wrapper
+# both of them failed because of cuda context mismatch.
+# not sure why, they are created from a different context.
+# the only successful approach is to call cuda driver API in C.
+import atexit
+import gc
+import os
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from typing import Any
+
+import torch
+
+from vllm.device_allocator import AllocationData, HandleType
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.system_utils import find_loaded_library
+from vllm.utils.torch_utils import PIN_MEMORY
+
+logger = init_logger(__name__)
+
+
+cumem_available = False
+libcudart: Any = None
+try:
+    from vllm.cumem_allocator import (
+        init_module,
+        python_create_and_map,
+        python_unmap_and_release,
+    )
+    from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
+
+    lib_name = find_loaded_library("cumem_allocator")
+    libcudart = CudaRTLibrary()
+    cumem_available = True
+except ModuleNotFoundError:
+    # only cuda and rocm platforms support cumem allocator
+    init_module = None
+    python_create_and_map = None
+    python_unmap_and_release = None
+    lib_name = None
+
+
+def create_and_map(allocation_handle: HandleType) -> None:
+    python_create_and_map(*allocation_handle)
+
+
+def unmap_and_release(allocation_handle: HandleType) -> None:
+    python_unmap_and_release(*allocation_handle)
+
+
+def get_pluggable_allocator(
+    python_malloc_fn: Callable[[HandleType], None],
+    python_free_func: Callable[[int], HandleType],
+) -> torch.cuda.memory.CUDAPluggableAllocator:
+    init_module(python_malloc_fn, python_free_func)
+    new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
+        lib_name, "my_malloc", "my_free"
+    )
+    return new_alloc
+
+
+@contextmanager
+def use_memory_pool_with_allocator(
+    python_malloc_fn: Callable[[HandleType], None],
+    python_free_func: Callable[[int], HandleType],
+) -> Iterator[
+    tuple[torch.cuda.memory.MemPool, torch.cuda.memory.CUDAPluggableAllocator]
+]:
+    new_alloc = get_pluggable_allocator(python_malloc_fn, python_free_func)
+    mem_pool = torch.cuda.memory.MemPool(new_alloc._allocator)
+    with torch.cuda.memory.use_mem_pool(mem_pool):
+        yield mem_pool, new_alloc
+
+
+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.
+
+    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, 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
+
+    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 = list(self.allocator_and_pools.values())
+        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:
+            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.cuda.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) -> 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.
+        """
+        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
+
+        for ptr, data in self.pointer_to_data.items():
+            handle = data.handle
+            total_bytes += handle[1]
+            if data.tag in offload_tags:
+                backup_bytes += handle[1]
+                size_in_bytes = handle[1]
+                cpu_backup_tensor = torch.empty(
+                    size_in_bytes,
+                    dtype=torch.uint8,
+                    device="cpu",
+                    pin_memory=PIN_MEMORY,
+                )
+                cpu_ptr = cpu_backup_tensor.data_ptr()
+                libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)
+                data.cpu_backup_tensor = cpu_backup_tensor
+            try:
+                unmap_and_release(handle)
+            finally:
+                data.is_asleep = True
+
+        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,
+        )
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    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 of the data will have empty memory.
+
+        Args:
+            tags: The tags of the memory allocation that will be loaded
+                back to GPU memory. If None, all memory allocation will be loaded
+                back to GPU memory.
+        """
+        for ptr, data in self.pointer_to_data.items():
+            if tags is None or data.tag in tags:
+                handle = data.handle
+                create_and_map(handle)
+                data.is_asleep = False
+                if data.cpu_backup_tensor is not None:
+                    cpu_backup_tensor = data.cpu_backup_tensor
+                    if cpu_backup_tensor is not None:
+                        size_in_bytes = (
+                            cpu_backup_tensor.numel() * cpu_backup_tensor.element_size()
+                        )
+                        cpu_ptr = cpu_backup_tensor.data_ptr()
+                        libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes)
+                        data.cpu_backup_tensor = None
+
+    @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 via
+        # PYTORCH_CUDA_ALLOC_CONF, temporarily disable them for the duration
+        # of the memory pool context and restore on exit.
+        conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
+        expandable_was_enabled = "expandable_segments:True" in conf
+        if expandable_was_enabled:
+            torch.cuda.memory._set_allocator_settings("expandable_segments:False")
+
+        old_tag = self.current_tag
+        self.current_tag = tag
+        try:
+            with use_memory_pool_with_allocator(
+                self.python_malloc_callback, self.python_free_callback
+            ) 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[tag] = data
+                yield
+                # PyTorch's bug, calling torch.cuda.empty_cache() will error
+                # when using 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.
+                allocations = data[0].snapshot()
+                for allocation in allocations:
+                    if allocation["allocated_size"] == 0:
+                        handle = self._python_free_callback(allocation["address"])
+                        unmap_and_release(handle)
+        finally:
+            self.current_tag = old_tag
+            if expandable_was_enabled:
+                torch.cuda.memory._set_allocator_settings("expandable_segments:True")
+
+    def get_current_usage(self) -> int:
+        """
+        Get the total number of bytes allocated in the memory pool.
+        """
+        sum_bytes: int = 0
+        for ptr, data in self.pointer_to_data.items():
+            handle = data.handle
+            sum_bytes += handle[1]
+        return sum_bytes

vllm/device_allocator/sleep_mode_backend.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pluggable sleep-mode backends (RFC #34303).
+
+vLLM's sleep/wake-up today is hard-wired to ``CuMemAllocator``: the GPU worker
+calls ``allocator.sleep(...)`` / ``allocator.wake_up(...)`` directly. RFC #34303
+proposes additional mechanisms for freeing and restoring GPU state - CUDA
+process checkpoint, CRIU, durable snapshot/restore - that share the *dispatch*
+(``/sleep`` endpoint -> engine -> executor -> worker) but differ in *mechanism*
+and in which resources they preserve (NCCL communicators, compiled kernels,
+CUDA graphs, survival across process restart).
+
+This module introduces a thin backend abstraction so those mechanisms can be
+selected by name without changing the public API. The default ``cumem`` backend
+wraps today's ``CuMemAllocator`` path 1:1, so existing users see no behavior
+change. The factory mirrors ``KVConnectorFactory`` and lets third-party
+backends register through a ``vllm.general_plugins`` entry point at import time.
+"""
+
+from __future__ import annotations
+
+import importlib
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Literal
+
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config.model import ModelConfig
+
+logger = init_logger(__name__)
+
+SleepModeState = Literal["RUNNING", "SUSPENDED", "RESUMING"]
+
+
+class SleepModeBackend(ABC):
+    """Interface for a mechanism that frees and restores GPU state.
+
+    A backend owns the *mechanism* of suspend/resume. The dispatch path
+    (``/sleep`` endpoint -> engine -> executor -> worker) is shared across all
+    backends and lives outside this class.
+
+    Capability flags are ``@classmethod`` so callers (executor, ``/health``,
+    AUTO selection) can introspect a backend without instantiating it, matching
+    the capability-flag convention used by attention backends.
+    """
+
+    def __init__(self) -> None:
+        self._state: SleepModeState = "RUNNING"
+
+    @abstractmethod
+    def suspend(self, level: int = 1) -> None:
+        """Free GPU state.
+
+        ``level`` follows existing sleep-mode semantics: level 1 offloads
+        weights to host RAM (restorable in-process); level 2 discards weights
+        (reloaded from the model source on resume).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def resume(self, tags: list[str] | None = None) -> None:
+        """Restore previously-suspended GPU state.
+
+        ``tags`` optionally limits which tagged allocations are restored
+        (e.g. ``["weights"]`` or ``["kv_cache"]``).
+        """
+        raise NotImplementedError
+
+    def state(self) -> SleepModeState:
+        """Current lifecycle state. Lets ``/health`` distinguish a healthy-idle
+        (suspended) engine from a healthy-serving one (see RFC #34303)."""
+        return self._state
+
+    # -- Capability introspection (no instance required) --
+
+    @classmethod
+    def is_supported(cls) -> bool:
+        """Whether this backend can run on the current platform/driver."""
+        return True
+
+    @classmethod
+    def preserves_communicators(cls) -> bool:
+        """If False, collective communicators (e.g. NCCL) are destroyed by
+        ``suspend`` and the executor must re-initialize them on ``resume``."""
+        return False
+
+    @classmethod
+    def preserves_compiled_artifacts(cls) -> bool:
+        """If True, torch.compile / JIT kernels survive suspend/resume and need
+        not be recompiled on resume."""
+        return False
+
+    @classmethod
+    def preserves_graphs_with_communicators(cls) -> bool:
+        """If True, CUDA graphs containing collective communicators (e.g. NCCL)
+        stay valid after resume. False when communicators are rebuilt (embedded
+        comm handles go stale)."""
+        return False
+
+    @classmethod
+    def supports_durable_storage(cls) -> bool:
+        """If True, suspended state can be persisted beyond the process
+        lifetime (disk or object storage) and restored in a new process."""
+        return False
+
+
+class CuMemBackend(SleepModeBackend):
+    """Default backend.
+
+    Wraps the platform sleep-mode allocator exactly as the GPU worker did
+    before this abstraction existed, so behavior is identical to vLLM's current
+    sleep/wake-up. ``get_mem_allocator_instance()`` resolves to
+    ``CuMemAllocator`` on CUDA and ``XpuMemAllocator`` on XPU; suspend offloads
+    per-allocation between GPU and host, with NCCL buffers left untouched (they
+    are allocated outside the allocator pool).
+    """
+
+    def suspend(self, level: int = 1) -> None:
+        from vllm.device_allocator import get_mem_allocator_instance
+
+        self._state = "SUSPENDED"
+        allocator = get_mem_allocator_instance()
+        allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
+
+    def resume(self, tags: list[str] | None = None) -> None:
+        from vllm.device_allocator import get_mem_allocator_instance
+
+        self._state = "RESUMING"
+        allocator = get_mem_allocator_instance()
+        allocator.wake_up(tags)
+        self._state = "RUNNING"
+
+    @classmethod
+    def preserves_communicators(cls) -> bool:
+        # Communicator buffers (e.g. NCCL) live outside CuMemAllocator's pool, so
+        # an allocator-level sleep leaves them intact (no reinit needed on resume).
+        return True
+
+
+class SleepModeBackendFactory:
+    """Registry and resolver for sleep-mode backends.
+
+    Mirrors ``KVConnectorFactory``: lazy module/class registration and a
+    built-in registry populated at import time. Third-party backends register
+    the same way from a ``vllm.general_plugins`` entry point.
+    """
+
+    _registry: dict[str, Callable[[], type[SleepModeBackend]]] = {}
+
+    @classmethod
+    def register_backend(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a backend with a lazy-loading module and class name."""
+        if name in cls._registry:
+            raise ValueError(f"Sleep-mode backend '{name}' is already registered.")
+
+        def loader() -> type[SleepModeBackend]:
+            module = importlib.import_module(module_path)
+            return getattr(module, class_name)
+
+        cls._registry[name] = loader
+
+    @classmethod
+    def get_backend_class(cls, name: str) -> type[SleepModeBackend]:
+        """Resolve a registered backend class by name."""
+        if name not in cls._registry:
+            available = ", ".join(sorted(cls._registry)) or "<none>"
+            raise ValueError(
+                f"Unsupported sleep-mode backend '{name}'. "
+                f"Registered backends: {available}."
+            )
+        return cls._registry[name]()
+
+    @classmethod
+    def create_backend(cls, model_config: ModelConfig) -> SleepModeBackend:
+        """Instantiate the backend selected by ``model_config``."""
+        name = model_config.sleep_mode_backend
+        backend_cls = cls.get_backend_class(name)
+        if not backend_cls.is_supported():
+            raise ValueError(
+                f"Sleep-mode backend '{name}' is not supported on this platform."
+            )
+        logger.info("Using sleep-mode backend: %s", name)
+        return backend_cls()
+
+
+# Register built-in backends here. Registration is lazy: only the module for the
+# selected backend is imported. Third-party backends (CUDA checkpoint, CRIU,
+# durable snapshot) register the same way through a vllm.general_plugins entry
+# point, without changes to vLLM core.
+SleepModeBackendFactory.register_backend(
+    "cumem",
+    "vllm.device_allocator.sleep_mode_backend",
+    "CuMemBackend",
+)

vllm/device_allocator/xpumem.py

@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import atexit
+import gc
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from typing import Any
+
+import torch
+
+from vllm.device_allocator import AllocationData, HandleType
+from vllm.logger import init_logger
+from vllm.utils.torch_utils import PIN_MEMORY
+
+logger = init_logger(__name__)
+
+MEMCPY_HOST_TO_DEVICE = 0
+MEMCPY_DEVICE_TO_HOST = 1
+MEMCPY_DEVICE_TO_DEVICE = 2
+
+xpumem_available = False
+xpumem_allocator: Any = None
+
+try:
+    from vllm_xpu_kernels import xpumem_allocator as _xpumem_allocator
+
+    xpumem_allocator = _xpumem_allocator
+    xpumem_available = True
+except ImportError:
+    xpumem_allocator = None
+
+
+def _xpu_memory_module() -> Any:
+    mem_mod = getattr(torch.xpu, "memory", None)
+    if mem_mod is None:
+        raise RuntimeError("torch.xpu.memory is not available")
+    return mem_mod
+
+
+def _supports_xpu_mem_pool(mem_mod: Any) -> bool:
+    return hasattr(mem_mod, "MemPool") and hasattr(mem_mod, "use_mem_pool")
+
+
+def _xpu_memcpy_sync(
+    dst_ptr: int,
+    src_ptr: int,
+    n_bytes: int,
+    kind: int,
+    device: int,
+) -> None:
+    def _to_i64_ptr(ptr: int) -> int:
+        # torch custom-op `int` arguments are signed int64.
+        # data_ptr() may return a uint64 value above 2^63-1, so normalize it.
+        return ptr if ptr < (1 << 63) else ptr - (1 << 64)
+
+    torch.ops._C.xpu_memcpy_sync(
+        _to_i64_ptr(dst_ptr),
+        _to_i64_ptr(src_ptr),
+        n_bytes,
+        kind,
+        device,
+    )
+
+
+def get_pluggable_allocator(
+    python_malloc_fn: Callable[[HandleType], None],
+    python_free_func: Callable[[int], HandleType],
+) -> Any:
+    if not xpumem_available or xpumem_allocator is None:
+        raise RuntimeError("xpumem allocator extension is not available")
+
+    xpumem_allocator.init_module(python_malloc_fn, python_free_func)
+    mem_mod = _xpu_memory_module()
+    alloc_cls = getattr(mem_mod, "XPUPluggableAllocator", None)
+    if alloc_cls is None:
+        raise RuntimeError("torch.xpu.memory.XPUPluggableAllocator is not available")
+
+    lib_name = xpumem_allocator.__file__
+    return alloc_cls(lib_name, "my_malloc", "my_free")
+
+
+def create_and_allocate(allocation_handle: HandleType) -> None:
+    if not xpumem_available or xpumem_allocator is None:
+        raise RuntimeError("xpumem allocator extension is not available")
+    xpumem_allocator.python_create_and_allocate(*allocation_handle)
+
+
+def unmap_and_release(allocation_handle: HandleType) -> None:
+    if not xpumem_available or xpumem_allocator is None:
+        raise RuntimeError("xpumem allocator extension is not available")
+    xpumem_allocator.python_unmap_and_release(*allocation_handle)
+
+
+@contextmanager
+def use_memory_pool_with_allocator(
+    python_malloc_fn: Callable[[HandleType], None],
+    python_free_func: Callable[[int], HandleType],
+) -> Iterator[tuple[Any, Any]]:
+    mem_mod = _xpu_memory_module()
+    if not _supports_xpu_mem_pool(mem_mod):
+        raise RuntimeError(
+            "torch.xpu.memory MemPool APIs are not available "
+            "(need MemPool and use_mem_pool)."
+        )
+    new_alloc = get_pluggable_allocator(python_malloc_fn, python_free_func)
+    mem_pool = mem_mod.MemPool(new_alloc._allocator)
+    with mem_mod.use_mem_pool(mem_pool):
+        yield mem_pool, new_alloc
+
+
+class XpuMemAllocator:
+    """A singleton pluggable allocator helper for XPU.
+
+    Note:
+    Sleep will offload selected payloads to CPU or discard and unmap XPU
+    physical memory. Wake-up remaps physical memory back to the same
+    reserved virtual address and restores payload.
+    """
+
+    instance: "XpuMemAllocator | None" = None
+    default_tag: str = "default"
+
+    @staticmethod
+    def get_instance() -> "XpuMemAllocator":
+        assert xpumem_available, "xpumem allocator is not available"
+        if XpuMemAllocator.instance is None:
+            XpuMemAllocator.instance = XpuMemAllocator()
+            # Ensure MemPool/allocator wrappers are released before interpreter
+            # finalization tears down XPU runtime internals.
+            atexit.register(XpuMemAllocator._shutdown_singleton)
+        return XpuMemAllocator.instance
+
+    @staticmethod
+    def _shutdown_singleton() -> None:
+        instance = XpuMemAllocator.instance
+        if instance is None:
+            return
+        try:
+            instance.release_pools()
+        except Exception:
+            logger.exception("XpuMemAllocator singleton shutdown failed")
+
+    def __init__(self):
+        self.pointer_to_data: dict[int, AllocationData] = {}
+        self.current_tag: str = XpuMemAllocator.default_tag
+        self.allocator_and_pools: dict[str, Any] = {}
+        self.python_malloc_callback = self._python_malloc_callback
+        self.python_free_callback = self._python_free_callback
+
+    def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
+        ptr = allocation_handle[2]
+        self.pointer_to_data[ptr] = AllocationData(allocation_handle, self.current_tag)
+        logger.debug(
+            "Allocated %s bytes for %s at %s",
+            allocation_handle[1],
+            self.current_tag,
+            ptr,
+        )
+
+    def _python_free_callback(self, ptr: int) -> HandleType:
+        data = self.pointer_to_data.pop(ptr)
+        data.cpu_backup_tensor = None
+        logger.debug("Freed %s bytes for %s at %s", data.handle[1], data.tag, ptr)
+        return data.handle
+
+    def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
+        if offload_tags is None:
+            offload_tags = (XpuMemAllocator.default_tag,)
+        elif isinstance(offload_tags, str):
+            offload_tags = (offload_tags,)
+
+        assert isinstance(offload_tags, tuple)
+
+        total_bytes = 0
+        backup_bytes = 0
+
+        for ptr, data in self.pointer_to_data.items():
+            size_in_bytes = data.handle[1]
+            total_bytes += size_in_bytes
+            if data.tag not in offload_tags:
+                unmap_and_release(data.handle)
+                continue
+
+            backup_bytes += size_in_bytes
+            device, _, _, _ = data.handle
+            cpu_backup_tensor = torch.empty(
+                size_in_bytes,
+                dtype=torch.uint8,
+                device="cpu",
+                pin_memory=PIN_MEMORY,
+            )
+            cpu_ptr = cpu_backup_tensor.data_ptr()
+            _xpu_memcpy_sync(
+                cpu_ptr,
+                ptr,
+                size_in_bytes,
+                MEMCPY_DEVICE_TO_HOST,
+                device,
+            )
+            data.cpu_backup_tensor = cpu_backup_tensor
+
+            unmap_and_release(data.handle)
+
+        logger.info(
+            "XpuMemAllocator: 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,
+        )
+
+        gc.collect()
+        xpu_empty_cache = getattr(torch.xpu, "empty_cache", None)
+        if callable(xpu_empty_cache):
+            xpu_empty_cache()
+
+    def wake_up(self, tags: list[str] | None = None) -> None:
+        for ptr, data in self.pointer_to_data.items():
+            if tags is not None and data.tag not in tags:
+                continue
+            create_and_allocate(data.handle)
+
+            cpu_backup_tensor = data.cpu_backup_tensor
+            if cpu_backup_tensor is None:
+                continue
+
+            device, size_in_bytes, _, _ = data.handle
+            _xpu_memcpy_sync(
+                ptr,
+                cpu_backup_tensor.data_ptr(),
+                size_in_bytes,
+                MEMCPY_HOST_TO_DEVICE,
+                device,
+            )
+            data.cpu_backup_tensor = None
+
+    def release_pools(self) -> None:
+        """Drop Python references to MemPool/pluggable allocators eagerly.
+
+        This prevents pool destruction from being deferred to interpreter
+        finalization, which can happen after parts of XPU runtime are already
+        torn down.
+        """
+        if not self.allocator_and_pools:
+            return
+
+        # Note: keep allocators alive while MemPool objects are destroyed.
+        # MemPool teardown may invoke allocator virtual methods (e.g. raw_delete)
+        # when releasing cached blocks. If allocator wrappers are dropped first,
+        # C++ can hit "pure virtual method called" during shutdown.
+        pool_entries = list(self.allocator_and_pools.values())
+        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()
+
+        xpu_sync = getattr(torch.xpu, "synchronize", None)
+        if callable(xpu_sync):
+            try:
+                xpu_sync()
+            except Exception:
+                logger.debug("torch.xpu.synchronize() failed during release_pools")
+
+        # 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()
+
+    @contextmanager
+    def use_memory_pool(self, tag: str | None = None):
+        if tag is None:
+            tag = XpuMemAllocator.default_tag
+
+        old_tag = self.current_tag
+        self.current_tag = tag
+        try:
+            with use_memory_pool_with_allocator(
+                self.python_malloc_callback,
+                self.python_free_callback,
+            ) as data:
+                self.allocator_and_pools[tag] = data
+                yield
+        finally:
+            self.current_tag = old_tag
+
+    def get_current_usage(self) -> int:
+        total = 0
+        for data in self.pointer_to_data.values():
+            total += data.handle[1]
+        return total