[Core] Pluggable sleep-mode backend abstraction (RFC #34303) (#44074)
Nils Matteson committed
Jun 30, 2026 at 23:00 UTC
93d8f834dd8acf33eb0e2a75b2711b628cb6e226
4 files changed
+309
-4
tests/v1/worker/test_sleep_mode_backend.py
new
+92
@@ -0,0 +1,92 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU-only unit tests for the sleep-mode backend abstraction (RFC #34303).
4
+
5
+These cover the registry/factory contract and capability flags. They do not
6
+touch CUDA - the ``cumem`` suspend/resume path is exercised end-to-end on GPU
7
+in ``tests/basic_correctness/test_cumem.py``.
8
+"""
9
+
10
+import pytest
11
+
12
+from vllm.device_allocator.sleep_mode_backend import (
13
+ CuMemBackend,
14
+ SleepModeBackend,
15
+ SleepModeBackendFactory,
16
+)
17
+
18
+
19
+def test_cumem_is_the_default_registered_backend():
20
+ backend_cls = SleepModeBackendFactory.get_backend_class("cumem")
21
+ assert backend_cls is CuMemBackend
22
+ assert issubclass(backend_cls, SleepModeBackend)
23
+
24
+
25
+def test_cumem_capability_flags():
26
+ # cumem leaves NCCL untouched but does not preserve compiled artifacts,
27
+ # graphs, or durable state - these flags are what the executor and /health
28
+ # introspect to decide reinit / persistence behavior.
29
+ assert CuMemBackend.is_supported() is True
30
+ assert CuMemBackend.preserves_nccl() is True
31
+ assert CuMemBackend.preserves_compiled_artifacts() is False
32
+ assert CuMemBackend.preserves_graphs_with_nccl() is False
33
+ assert CuMemBackend.supports_durable_storage() is False
34
+
35
+
36
+def test_new_backend_starts_in_running_state():
37
+ # Constructing a backend must not touch the GPU; only suspend/resume do.
38
+ assert CuMemBackend().state() == "RUNNING"
39
+
40
+
41
+def test_unknown_backend_raises():
42
+ with pytest.raises(ValueError, match="Unsupported sleep-mode backend"):
43
+ SleepModeBackendFactory.get_backend_class("does-not-exist")
44
+
45
+
46
+def test_duplicate_registration_raises():
47
+ with pytest.raises(ValueError, match="already registered"):
48
+ SleepModeBackendFactory.register_backend(
49
+ "cumem",
50
+ "vllm.device_allocator.sleep_mode_backend",
51
+ "CuMemBackend",
52
+ )
53
+
54
+
55
+def test_third_party_backend_registration_and_resolution():
56
+ """A plugin registers a backend by name; the factory resolves it lazily."""
57
+ name = "_pytest_dummy_backend"
58
+ try:
59
+ SleepModeBackendFactory.register_backend(
60
+ name,
61
+ "tests.v1.worker.test_sleep_mode_backend",
62
+ "DummyBackend",
63
+ )
64
+ resolved = SleepModeBackendFactory.get_backend_class(name)
65
+ assert resolved is DummyBackend
66
+ assert resolved.supports_durable_storage() is True
67
+ finally:
68
+ SleepModeBackendFactory._registry.pop(name, None)
69
+
70
+
71
+def test_suspend_resume_state_transitions():
72
+ """Lifecycle state advances RUNNING -> SUSPENDED -> RUNNING without GPU."""
73
+ backend = DummyBackend()
74
+ assert backend.state() == "RUNNING"
75
+ backend.suspend(level=1)
76
+ assert backend.state() == "SUSPENDED"
77
+ backend.resume()
78
+ assert backend.state() == "RUNNING"
79
+
80
+
81
+class DummyBackend(SleepModeBackend):
82
+ """A no-GPU backend used to exercise lifecycle + registration in CPU tests."""
83
+
84
+ def suspend(self, level: int = 1) -> None:
85
+ self._state = "SUSPENDED"
86
+
87
+ def resume(self, tags: list[str] | None = None) -> None:
88
+ self._state = "RUNNING"
89
+
90
+ @classmethod
91
+ def supports_durable_storage(cls) -> bool:
92
+ return True
vllm/config/model.py
+5
@@ -289,6 +289,11 @@ class ModelConfig:
289
enable_sleep_mode: bool = False
290
"""Enable sleep mode for the engine (only cuda and
291
hip platforms are supported)."""
292
+ sleep_mode_backend: str = "cumem"
293
+ """Mechanism used to free and restore GPU state for sleep mode. ``"cumem"``
294
+ (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible
295
+ with prior releases. Additional backends (CUDA checkpoint, CRIU, durable
296
+ snapshot) may be registered in-tree or by plugins (RFC #34303)."""
297
enable_cumem_allocator: bool = False
298
"""Enable the custom cumem allocator to leverage advanced GPU memory
299
allocation features such as multi-node NVLink support.
vllm/device_allocator/sleep_mode_backend.py
new
+195
@@ -0,0 +1,195 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Pluggable sleep-mode backends (RFC #34303).
4
+
5
+vLLM's sleep/wake-up today is hard-wired to ``CuMemAllocator``: the GPU worker
6
+calls ``allocator.sleep(...)`` / ``allocator.wake_up(...)`` directly. RFC #34303
7
+proposes additional mechanisms for freeing and restoring GPU state - CUDA
8
+process checkpoint, CRIU, durable snapshot/restore - that share the *dispatch*
9
+(``/sleep`` endpoint -> engine -> executor -> worker) but differ in *mechanism*
10
+and in which resources they preserve (NCCL communicators, compiled kernels,
11
+CUDA graphs, survival across process restart).
12
+
13
+This module introduces a thin backend abstraction so those mechanisms can be
14
+selected by name without changing the public API. The default ``cumem`` backend
15
+wraps today's ``CuMemAllocator`` path 1:1, so existing users see no behavior
16
+change. The factory mirrors ``KVConnectorFactory`` and lets third-party
17
+backends register through a ``vllm.general_plugins`` entry point at import time.
18
+"""
19
+
20
+from __future__ import annotations
21
+
22
+import importlib
23
+from abc import ABC, abstractmethod
24
+from collections.abc import Callable
25
+from typing import TYPE_CHECKING, Literal
26
+
27
+from vllm.logger import init_logger
28
+
29
+if TYPE_CHECKING:
30
+ from vllm.config.model import ModelConfig
31
+
32
+logger = init_logger(__name__)
33
+
34
+SleepModeState = Literal["RUNNING", "SUSPENDED", "RESUMING"]
35
+
36
+
37
+class SleepModeBackend(ABC):
38
+ """Interface for a mechanism that frees and restores GPU state.
39
+
40
+ A backend owns the *mechanism* of suspend/resume. The dispatch path
41
+ (``/sleep`` endpoint -> engine -> executor -> worker) is shared across all
42
+ backends and lives outside this class.
43
+
44
+ Capability flags are ``@classmethod`` so callers (executor, ``/health``,
45
+ AUTO selection) can introspect a backend without instantiating it, matching
46
+ the capability-flag convention used by attention backends.
47
+ """
48
+
49
+ def __init__(self) -> None:
50
+ self._state: SleepModeState = "RUNNING"
51
+
52
+ @abstractmethod
53
+ def suspend(self, level: int = 1) -> None:
54
+ """Free GPU state.
55
+
56
+ ``level`` follows existing sleep-mode semantics: level 1 offloads
57
+ weights to host RAM (restorable in-process); level 2 discards weights
58
+ (reloaded from the model source on resume).
59
+ """
60
+ raise NotImplementedError
61
+
62
+ @abstractmethod
63
+ def resume(self, tags: list[str] | None = None) -> None:
64
+ """Restore previously-suspended GPU state.
65
+
66
+ ``tags`` optionally limits which tagged allocations are restored
67
+ (e.g. ``["weights"]`` or ``["kv_cache"]``).
68
+ """
69
+ raise NotImplementedError
70
+
71
+ def state(self) -> SleepModeState:
72
+ """Current lifecycle state. Lets ``/health`` distinguish a healthy-idle
73
+ (suspended) engine from a healthy-serving one (see RFC #34303)."""
74
+ return self._state
75
+
76
+ # -- Capability introspection (no instance required) --
77
+
78
+ @classmethod
79
+ def is_supported(cls) -> bool:
80
+ """Whether this backend can run on the current platform/driver."""
81
+ return True
82
+
83
+ @classmethod
84
+ def preserves_nccl(cls) -> bool:
85
+ """If False, NCCL communicators are destroyed by ``suspend`` and the
86
+ executor must re-initialize them on ``resume``."""
87
+ return False
88
+
89
+ @classmethod
90
+ def preserves_compiled_artifacts(cls) -> bool:
91
+ """If True, torch.compile / JIT kernels survive suspend/resume and need
92
+ not be recompiled on resume."""
93
+ return False
94
+
95
+ @classmethod
96
+ def preserves_graphs_with_nccl(cls) -> bool:
97
+ """If True, CUDA graphs containing NCCL collectives stay valid after
98
+ resume. False when NCCL is rebuilt (embedded comm handles go stale)."""
99
+ return False
100
+
101
+ @classmethod
102
+ def supports_durable_storage(cls) -> bool:
103
+ """If True, suspended state can be persisted beyond the process
104
+ lifetime (disk or object storage) and restored in a new process."""
105
+ return False
106
+
107
+
108
+class CuMemBackend(SleepModeBackend):
109
+ """Default backend.
110
+
111
+ Wraps the platform sleep-mode allocator exactly as the GPU worker did
112
+ before this abstraction existed, so behavior is identical to vLLM's current
113
+ sleep/wake-up. ``get_mem_allocator_instance()`` resolves to
114
+ ``CuMemAllocator`` on CUDA and ``XpuMemAllocator`` on XPU; suspend offloads
115
+ per-allocation between GPU and host, with NCCL buffers left untouched (they
116
+ are allocated outside the allocator pool).
117
+ """
118
+
119
+ def suspend(self, level: int = 1) -> None:
120
+ from vllm.device_allocator import get_mem_allocator_instance
121
+
122
+ self._state = "SUSPENDED"
123
+ allocator = get_mem_allocator_instance()
124
+ allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
125
+
126
+ def resume(self, tags: list[str] | None = None) -> None:
127
+ from vllm.device_allocator import get_mem_allocator_instance
128
+
129
+ self._state = "RESUMING"
130
+ allocator = get_mem_allocator_instance()
131
+ allocator.wake_up(tags)
132
+ self._state = "RUNNING"
133
+
134
+ @classmethod
135
+ def preserves_nccl(cls) -> bool:
136
+ # NCCL buffers live outside CuMemAllocator's pool, so an allocator-level
137
+ # sleep leaves the communicators intact (no reinit needed on resume).
138
+ return True
139
+
140
+
141
+class SleepModeBackendFactory:
142
+ """Registry and resolver for sleep-mode backends.
143
+
144
+ Mirrors ``KVConnectorFactory``: lazy module/class registration and a
145
+ built-in registry populated at import time. Third-party backends register
146
+ the same way from a ``vllm.general_plugins`` entry point.
147
+ """
148
+
149
+ _registry: dict[str, Callable[[], type[SleepModeBackend]]] = {}
150
+
151
+ @classmethod
152
+ def register_backend(cls, name: str, module_path: str, class_name: str) -> None:
153
+ """Register a backend with a lazy-loading module and class name."""
154
+ if name in cls._registry:
155
+ raise ValueError(f"Sleep-mode backend '{name}' is already registered.")
156
+
157
+ def loader() -> type[SleepModeBackend]:
158
+ module = importlib.import_module(module_path)
159
+ return getattr(module, class_name)
160
+
161
+ cls._registry[name] = loader
162
+
163
+ @classmethod
164
+ def get_backend_class(cls, name: str) -> type[SleepModeBackend]:
165
+ """Resolve a registered backend class by name."""
166
+ if name not in cls._registry:
167
+ available = ", ".join(sorted(cls._registry)) or "<none>"
168
+ raise ValueError(
169
+ f"Unsupported sleep-mode backend '{name}'. "
170
+ f"Registered backends: {available}."
171
+ )
172
+ return cls._registry[name]()
173
+
174
+ @classmethod
175
+ def create_backend(cls, model_config: ModelConfig) -> SleepModeBackend:
176
+ """Instantiate the backend selected by ``model_config``."""
177
+ name = model_config.sleep_mode_backend
178
+ backend_cls = cls.get_backend_class(name)
179
+ if not backend_cls.is_supported():
180
+ raise ValueError(
181
+ f"Sleep-mode backend '{name}' is not supported on this platform."
182
+ )
183
+ logger.info("Using sleep-mode backend: %s", name)
184
+ return backend_cls()
185
+
186
+
187
+# Register built-in backends here. Registration is lazy: only the module for the
188
+# selected backend is imported. Third-party backends (CUDA checkpoint, CRIU,
189
+# durable snapshot) register the same way through a vllm.general_plugins entry
190
+# point, without changes to vLLM core.
191
+SleepModeBackendFactory.register_backend(
192
+ "cumem",
193
+ "vllm.device_allocator.sleep_mode_backend",
194
+ "CuMemBackend",
195
+)
vllm/v1/worker/gpu_worker.py
+17
-4
@@ -86,6 +86,7 @@ from .utils import request_memory
86
logger = init_logger(__name__)
87
88
if TYPE_CHECKING:
89
+ from vllm.device_allocator.sleep_mode_backend import SleepModeBackend
90
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
91
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
92
@@ -170,6 +171,20 @@ class Worker(WorkerBase):
171
# pending non-blocking PP send work from the previous iteration
172
self._pp_send_work: list[Handle] = []
173
174
+ # Resolved lazily on first sleep/wake; persists worker-process state.
175
+ self._sleep_mode_backend: SleepModeBackend | None = None
176
+
177
+ def _get_sleep_mode_backend(self) -> "SleepModeBackend":
178
+ if self._sleep_mode_backend is None:
179
+ from vllm.device_allocator.sleep_mode_backend import (
180
+ SleepModeBackendFactory,
181
+ )
182
+
183
+ self._sleep_mode_backend = SleepModeBackendFactory.create_backend(
184
+ self.vllm_config.model_config
185
+ )
186
+ return self._sleep_mode_backend
187
+
188
def sleep(self, level: int = 1) -> None:
189
torch.accelerator.synchronize()
190
free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]
@@ -181,8 +196,7 @@ class Worker(WorkerBase):
196
name: buffer.cpu().clone() for name, buffer in model.named_buffers()
197
}
198
184
- allocator = get_mem_allocator_instance()
185
- allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
199
+ self._get_sleep_mode_backend().suspend(level)
200
201
torch.accelerator.synchronize()
202
deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
@@ -202,8 +216,7 @@ class Worker(WorkerBase):
216
)
217
218
def wake_up(self, tags: list[str] | None = None) -> None:
205
- allocator = get_mem_allocator_instance()
206
- allocator.wake_up(tags)
219
+ self._get_sleep_mode_backend().resume(tags)
220
221
# Restore the buffers after level 2 sleep
222
if len(self._sleep_saved_buffers):