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 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).

Source code in vllm/device_allocator/sleep_mode_backend.py 
116
117
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
153
154
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()
        # Runtime state is kept at every level; weights only at level 1.
        offload_tags = ("weights", "runtime") if level == 1 else ("runtime",)
        allocator.sleep(offload_tags=offload_tags)

    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:

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

    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 
76
77
78
79
80
81
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 
85
86
87
88
@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 
90
91
92
93
94
@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 
 96
 97
 98
 99
100
@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 
102
103
104
105
106
107
@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 
62
63
64
65
66
67
68
69
@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 
71
72
73
74
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 
109
110
111
112
113
@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(level=1) abstractmethod

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).

Source code in vllm/device_allocator/sleep_mode_backend.py 
52
53
54
55
56
57
58
59
60
@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

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 
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
199
200
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 
190
191
192
193
194
195
196
197
198
199
200
@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 
179
180
181
182
183
184
185
186
187
188
@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 
167
168
169
170
171
172
173
174
175
176
177
@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