From 866f02db8a6944f87541ee9e77a6da6865c0066a Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Fri, 2 Oct 2026 10:14:37 +0800 Subject: [PATCH] [Feature] Release WorkspaceManager scratch on sleep (#59156) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .../memory/cumem/test_cumem.py | 94 +++++++++++++++++++ vllm/device_allocator/__init__.py | 3 + vllm/device_allocator/cumem.py | 33 ++++--- vllm/device_allocator/xpumem.py | 15 ++- vllm/v1/worker/gpu_worker.py | 2 + vllm/v1/worker/workspace.py | 16 +++- 6 files changed, 140 insertions(+), 23 deletions(-) diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index 135a3fc1ed..36a54ae219 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/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(): diff --git a/vllm/device_allocator/__init__.py b/vllm/device_allocator/__init__.py index 02d5160316..b70b7430b4 100644 --- a/vllm/device_allocator/__init__.py +++ b/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: diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 2e5612e5fa..6bf0ea71f8 100644 --- a/vllm/device_allocator/cumem.py +++ b/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 diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index 17fa69ee84..7d4b40a78e 100644 --- a/vllm/device_allocator/xpumem.py +++ b/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 diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index e9021bd142..15e2ee0611 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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: diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index a86110f2f8..4bae1f52be 100644 --- a/vllm/v1/worker/workspace.py +++ b/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: