Skip to content

vllm.device_allocator.sleep_mode_backend

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.

Classes:

CuMemBackend

Bases: SleepModeBackend

Default backend.

Wraps the platform sleep-mode allocator, CuMemAllocator on CUDA and XpuMemAllocator on XPU. Suspend offloads per allocation between GPU and host, with NCCL buffers left untouched (they are outside the allocator pool).

Source code in vllm/device_allocator/sleep_mode_backend.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
class CuMemBackend(SleepModeBackend):
    """Default backend.

    Wraps the platform sleep-mode allocator, ``CuMemAllocator`` on CUDA and
    ``XpuMemAllocator`` on XPU. Suspend offloads per allocation between GPU and
    host, with NCCL buffers left untouched (they are outside the allocator pool).
    """

    def suspend(self, discardable: Collection[int] = ()) -> None:
        from vllm.device_allocator import get_mem_allocator_instance

        self._state = "SUSPENDED"
        get_mem_allocator_instance().sleep(
            ("weights", "runtime", "cudagraph"), discardable
        )

    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"

    def discard(self, tags: tuple[str, ...]) -> None:
        from vllm.device_allocator import get_mem_allocator_instance

        allocator = get_mem_allocator_instance()
        allocator.discard(tags)
        self._state = "SUSPENDED"

    @classmethod
    def preserves_communicators(cls) -> bool:
        # Comm identity/topology survive memory suspension; no reinit needed.
        return True

SleepModeBackend

Bases: 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.

Methods:

  • discard –

    Discard selected tagged allocations without preserving contents.

  • is_supported –

    Whether this backend can run on the current platform/driver.

  • preserves_communicators –

    If False, collective communicators (e.g. NCCL) are destroyed by

  • preserves_compiled_artifacts –

    If True, torch.compile / JIT kernels survive suspend/resume and need

  • preserves_graphs_with_communicators –

    If True, CUDA graphs containing collective communicators (e.g. NCCL)

  • resume –

    Restore previously-suspended GPU state.

  • state –

    Current lifecycle state. Lets /health distinguish a healthy-idle

  • supports_durable_storage –

    If True, suspended state can be persisted beyond the process

  • suspend –

    Free GPU state, so that resume restores it at the same addresses.

Source code in vllm/device_allocator/sleep_mode_backend.py 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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, discardable: Collection[int] = ()) -> None:
        """Free GPU state, so that ``resume`` restores it at the same addresses.

        ``discardable`` holds the storage addresses of the tensors the caller
        rewrites after ``resume``, so they need no backup: the parameters not
        selected for preservation at level 2, nothing at level 1. A backend
        doesn't need to know the level. KV cache and workspace scratch are never
        restored.
        """
        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

    def discard(self, tags: tuple[str, ...]) -> None:
        """Discard selected tagged allocations without preserving contents."""
        raise RuntimeError(
            f"Sleep-mode backend {type(self).__name__!r} does not support "
            "selective discard."
        )

    # -- 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

discard(tags)

Discard selected tagged allocations without preserving contents.

Source code in vllm/device_allocator/sleep_mode_backend.py 
78
79
80
81
82
83
def discard(self, tags: tuple[str, ...]) -> None:
    """Discard selected tagged allocations without preserving contents."""
    raise RuntimeError(
        f"Sleep-mode backend {type(self).__name__!r} does not support "
        "selective discard."
    )

is_supported() classmethod

Whether this backend can run on the current platform/driver.

Source code in vllm/device_allocator/sleep_mode_backend.py 
87
88
89
90
@classmethod
def is_supported(cls) -> bool:
    """Whether this backend can run on the current platform/driver."""
    return True

preserves_communicators() classmethod

If False, collective communicators (e.g. NCCL) are destroyed by suspend and the executor must re-initialize them on resume.

Source code in vllm/device_allocator/sleep_mode_backend.py 
92
93
94
95
96
@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

preserves_compiled_artifacts() classmethod

If True, torch.compile / JIT kernels survive suspend/resume and need not be recompiled on resume.

Source code in vllm/device_allocator/sleep_mode_backend.py 
 98
 99
100
101
102
@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

preserves_graphs_with_communicators() classmethod

If True, CUDA graphs containing collective communicators (e.g. NCCL) stay valid after resume. False when communicators are rebuilt (embedded comm handles go stale).

Source code in vllm/device_allocator/sleep_mode_backend.py 
104
105
106
107
108
109
@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

resume(tags=None) abstractmethod

Restore previously-suspended GPU state.

tags optionally limits which tagged allocations are restored (e.g. ["weights"] or ["kv_cache"]).

Source code in vllm/device_allocator/sleep_mode_backend.py 
64
65
66
67
68
69
70
71
@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

state()

Current lifecycle state. Lets /health distinguish a healthy-idle (suspended) engine from a healthy-serving one (see RFC #34303).

Source code in vllm/device_allocator/sleep_mode_backend.py 
73
74
75
76
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

supports_durable_storage() classmethod

If True, suspended state can be persisted beyond the process lifetime (disk or object storage) and restored in a new process.

Source code in vllm/device_allocator/sleep_mode_backend.py 
111
112
113
114
115
@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

suspend(discardable=()) abstractmethod

Free GPU state, so that resume restores it at the same addresses.

discardable holds the storage addresses of the tensors the caller rewrites after resume, so they need no backup: the parameters not selected for preservation at level 2, nothing at level 1. A backend doesn't need to know the level. KV cache and workspace scratch are never restored.

Source code in vllm/device_allocator/sleep_mode_backend.py 
52
53
54
55
56
57
58
59
60
61
62
@abstractmethod
def suspend(self, discardable: Collection[int] = ()) -> None:
    """Free GPU state, so that ``resume`` restores it at the same addresses.

    ``discardable`` holds the storage addresses of the tensors the caller
    rewrites after ``resume``, so they need no backup: the parameters not
    selected for preservation at level 2, nothing at level 1. A backend
    doesn't need to know the level. KV cache and workspace scratch are never
    restored.
    """
    raise NotImplementedError

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.

Methods:

Source code in vllm/device_allocator/sleep_mode_backend.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
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()

create_backend(model_config) classmethod

Instantiate the backend selected by model_config.

Source code in vllm/device_allocator/sleep_mode_backend.py 
188
189
190
191
192
193
194
195
196
197
198
@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()

get_backend_class(name) classmethod

Resolve a registered backend class by name.

Source code in vllm/device_allocator/sleep_mode_backend.py 
177
178
179
180
181
182
183
184
185
186
@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]()

register_backend(name, module_path, class_name) classmethod

Register a backend with a lazy-loading module and class name.

Source code in vllm/device_allocator/sleep_mode_backend.py 
165
166
167
168
169
170
171
172
173
174
175
@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