[Feature] Release WorkspaceManager scratch on sleep (#59156)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/basic_correctness/memory/cumem/test_cumem.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import gc
+
 import pytest
 import torch
 
@@ -233,6 +235,98 @@ def test_discard_tags():
     assert torch.allclose(weights, torch.ones_like(weights))
 
 
+@pytest.mark.parametrize("tag", ["workspace", None], ids=["workspace", "default"])
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
+def test_selective_wake_restores_internal_tags(tag):
+    """A selective wake defers only weights/kv_cache; other tags always wake."""
+    allocator = get_mem_allocator_instance()
+    with allocator.use_memory_pool(tag):
+        internal = torch.empty(16 << 20, dtype=torch.uint8, device=DEVICE_TYPE)
+    with allocator.use_memory_pool("kv_cache"):
+        kv = torch.empty(32 << 20, dtype=torch.uint8, device=DEVICE_TYPE)
+
+    allocator.sleep(offload_tags=())
+    assert mapped_usage(allocator) == 0
+
+    allocator.wake_up(tags=["weights"])
+    assert mapped_usage(allocator) == internal.nbytes
+    allocator.wake_up(tags=["kv_cache"])
+    assert mapped_usage(allocator) == internal.nbytes + kv.nbytes
+
+
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
+@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
+def test_workspace_scratch_discarded_on_sleep():
+    """Workspace scratch lives in the sleep-mode pool: sleep discards it, any
+    (selective) wake remaps it at the same address so captured graphs replay,
+    and a buffer outgrown by a resize stays usable until it is released."""
+    from vllm.device_allocator.sleep_mode_backend import CuMemBackend
+    from vllm.v1.worker.workspace import WorkspaceManager
+
+    allocator = get_mem_allocator_instance()
+    manager = WorkspaceManager(
+        torch.device(DEVICE_TYPE),
+        alloc_context=lambda: allocator.use_memory_pool("workspace"),
+    )
+    (small,) = manager.get_simultaneous(((32 << 20,), torch.uint8))
+    # Resize while a view of the old buffer is alive: the tag is re-entered.
+    (scratch,) = manager.get_simultaneous(((64 << 20,), torch.uint8))
+    assert len(allocator.allocator_and_pools["workspace"]) == 2
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        scratch.add_(1)
+
+    backend = CuMemBackend()
+    backend.suspend(level=1)
+    assert mapped_usage(allocator) == 0
+    # Discarded, not copied.
+    assert all(
+        d.tag == "workspace" and d.cpu_backup_tensor is None
+        for d in allocator.pointer_to_data.values()
+    )
+
+    backend.resume(tags=["kv_cache"])
+    assert mapped_usage(allocator) == small.nbytes + scratch.nbytes
+    scratch.zero_()
+    graph.replay()
+    assert int(scratch.sum()) == scratch.numel()
+    small.fill_(1)
+
+    # The next resize releases the outgrown buffer; wake does not remap it.
+    del small
+    (big,) = manager.get_simultaneous(((128 << 20,), torch.uint8))
+    assert mapped_usage(allocator) == scratch.nbytes + big.nbytes
+    backend.suspend(level=1)
+    backend.resume(tags=["weights"])
+    assert mapped_usage(allocator) == scratch.nbytes + big.nbytes
+
+
+@pytest.mark.parametrize("free_x_early", [True, False], ids=["x-freed", "x-alive"])
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
+@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
+def test_reentered_tag(free_x_early):
+    """Enter a tag twice (two workspace ubatches): freeing X returns its memory
+    by the next entry, without touching Y or breaking sleep/wake."""
+    allocator = get_mem_allocator_instance()
+    nbytes = 64 << 20
+    with allocator.use_memory_pool("t"):
+        xs = [torch.empty(nbytes, dtype=torch.uint8, device=DEVICE_TYPE)]
+        if free_x_early:
+            xs.clear()
+    with allocator.use_memory_pool("t"):
+        y = torch.empty(nbytes, dtype=torch.uint8, device=DEVICE_TYPE)
+    xs.clear()
+    gc.collect()
+    with allocator.use_memory_pool("t"):
+        pass
+    assert mapped_usage(allocator) == nbytes
+    y.fill_(1)
+    allocator.sleep(offload_tags="t")
+    allocator.wake_up()
+    assert int(y.sum()) == y.numel()
+    allocator.release_pools()
+
+
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
 @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
 def test_level2_discards_ordinary_tensor_with_weights_tag():

vllm/device_allocator/__init__.py

@@ -13,6 +13,9 @@ from vllm.platforms import current_platform
 # py_handle has type list[int] on ROCm and int otherwise
 HandleType: TypeAlias = tuple[int, int, int, list[int] | int]
 
+# Selective wake may defer only these; any wake restores every other tag.
+DEFERRABLE_TAGS: tuple[str, ...] = ("weights", "kv_cache")
+
 
 @dataclasses.dataclass
 class AllocationData:

vllm/device_allocator/cumem.py

@@ -16,7 +16,7 @@ from typing import Any
 
 import torch
 
-from vllm.device_allocator import AllocationData, HandleType
+from vllm.device_allocator import DEFERRABLE_TAGS, AllocationData, HandleType
 from vllm.device_allocator.alloc_conf import (
     EXPANDABLE_SEGMENTS,
     conf_flag_enabled,
@@ -139,7 +139,7 @@ class CuMemAllocator:
     def __init__(self):
         self.pointer_to_data: dict[int, AllocationData] = {}
         self.current_tag: str = CuMemAllocator.default_tag
-        self.allocator_and_pools: dict[str, Any] = {}
+        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
@@ -167,7 +167,7 @@ class CuMemAllocator:
         if not self.allocator_and_pools:
             return
 
-        pool_entries = list(self.allocator_and_pools.values())
+        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]
@@ -334,9 +334,9 @@ class CuMemAllocator:
         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.
+            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()
@@ -345,7 +345,7 @@ class CuMemAllocator:
         for ptr, data in self.pointer_to_data.items():
             if not data.is_asleep:
                 continue
-            if tags is None or data.tag in tags:
+            if tags is None or data.tag in tags or data.tag not in DEFERRABLE_TAGS:
                 handle = data.handle
                 create_and_map(handle)
                 data.is_asleep = False
@@ -392,6 +392,10 @@ class CuMemAllocator:
         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 use_memory_pool_with_allocator(
                 self.python_malloc_callback, self.python_free_callback
             ) as data:
@@ -400,7 +404,7 @@ class CuMemAllocator:
                 # 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
+                self.allocator_and_pools.setdefault(tag, []).append(data)
                 yield
                 # PyTorch's bug, calling torch.cuda.empty_cache() will error
                 # when using pluggable allocator, see
@@ -413,16 +417,19 @@ class CuMemAllocator:
                 # 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)
+                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:
+                unmap_and_release(self._python_free_callback(allocation["address"]))
+
     def get_current_usage(self) -> int:
         """Get the total number of bytes allocated in the memory pool."""
         sum_bytes: int = 0

vllm/device_allocator/xpumem.py

@@ -9,7 +9,7 @@ from typing import Any
 
 import torch
 
-from vllm.device_allocator import AllocationData, HandleType
+from vllm.device_allocator import DEFERRABLE_TAGS, AllocationData, HandleType
 from vllm.logger import init_logger
 from vllm.utils.torch_utils import PIN_MEMORY
 
@@ -145,7 +145,7 @@ class XpuMemAllocator:
     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.allocator_and_pools: dict[str, list[Any]] = {}
         self.python_malloc_callback = self._python_malloc_callback
         self.python_free_callback = self._python_free_callback
 
@@ -267,7 +267,11 @@ class XpuMemAllocator:
         for ptr, data in self.pointer_to_data.items():
             if not data.is_asleep:
                 continue
-            if tags is not None and data.tag not in tags:
+            if (
+                tags is not None
+                and data.tag not in tags
+                and data.tag in DEFERRABLE_TAGS
+            ):
                 continue
             create_and_allocate(data.handle)
             data.is_asleep = False
@@ -300,7 +304,7 @@ class XpuMemAllocator:
         # 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())
+        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]
@@ -333,7 +337,8 @@ class XpuMemAllocator:
                 self.python_malloc_callback,
                 self.python_free_callback,
             ) as data:
-                self.allocator_and_pools[tag] = data
+                # Keep every entry's pool alive (pytorch/pytorch#146431).
+                self.allocator_and_pools.setdefault(tag, []).append(data)
                 yield
         finally:
             self.current_tag = old_tag

vllm/v1/worker/gpu_worker.py

@@ -532,6 +532,8 @@ class Worker(WorkerBase):
             self.device,
             num_ubatches,
             _num_workspace_lanes(self.vllm_config, self.use_v2_model_runner),
+            # Scratch holds no state across steps: discard it on sleep.
+            alloc_context=lambda: self._maybe_get_memory_pool_context("workspace"),
         )
         self.model_runner: GPUModelRunner = self._make_model_runner()
         if self.rank == 0:

vllm/v1/worker/workspace.py

@@ -4,7 +4,7 @@
 import inspect
 import os
 from collections.abc import Callable, Hashable, Iterator
-from contextlib import contextmanager
+from contextlib import AbstractContextManager, contextmanager, nullcontext
 from contextvars import ContextVar
 from itertools import accumulate
 from math import prod
@@ -58,8 +58,10 @@ class WorkspaceManager:
         device: torch.device,
         num_ubatches: int | None = None,
         num_lanes: int = 1,
+        alloc_context: Callable[[], AbstractContextManager] = nullcontext,
     ):
         self._device = device
+        self._alloc_context = alloc_context
         # Cache num ubatches at init based on configuration (default to 1)
         self._num_ubatches = num_ubatches if num_ubatches is not None else 1
         if num_lanes < 1:
@@ -261,9 +263,11 @@ class WorkspaceManager:
             # dead segment in reserved memory which can cause higher peak
             # memory usage.
             torch.accelerator.empty_cache()
-            self._current_workspaces[workspace_id] = torch.empty(
-                (required_bytes,), dtype=torch.uint8, device=self._device
-            )
+            # In sleep mode, entering the pool also trims outgrown buffers.
+            with self._alloc_context():
+                self._current_workspaces[workspace_id] = torch.empty(
+                    (required_bytes,), dtype=torch.uint8, device=self._device
+                )
             current_workspace = self._current_workspaces[workspace_id]
 
             if envs.VLLM_DEBUG_WORKSPACE:
@@ -308,6 +312,7 @@ def init_workspace_manager(
     device: torch.device,
     num_ubatches: int | None = None,
     num_lanes: int = 1,
+    alloc_context: Callable[[], AbstractContextManager] = nullcontext,
 ) -> None:
     """Initialize the workspace manager with a device.
 
@@ -318,6 +323,7 @@ def init_workspace_manager(
         device: The device to allocate workspace on.
         num_ubatches: Number of workspace ubatch slots. Defaults to 1.
         num_lanes: Number of independent execution lanes per ubatch. Defaults to 1.
+        alloc_context: Factory for the context scratch is allocated in.
 
     """
     global _manager
@@ -328,7 +334,7 @@ def init_workspace_manager(
             _manager._device,
             device,
         )
-    _manager = WorkspaceManager(device, num_ubatches, num_lanes)
+    _manager = WorkspaceManager(device, num_ubatches, num_lanes, alloc_context)
 
 
 def lock_workspace() -> None: