@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 866f02db8a

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

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

aoshen02 committed Oct 2, 2026 at 10:14 UTC 866f02db8a6944f87541ee9e77a6da6865c0066a

6 files changed +140 -23

tests/basic_correctness/memory/cumem/test_cumem.py

+94

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import gc

5 +

6 import pytest

7 import torch

@@ -233,6 +235,98 @@ def test_discard_tags():

235 assert torch.allclose(weights, torch.ones_like(weights))

236

237

238 +@pytest.mark.parametrize("tag", ["workspace", None], ids=["workspace", "default"])

239 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")

240 +def test_selective_wake_restores_internal_tags(tag):

241 + """A selective wake defers only weights/kv_cache; other tags always wake."""

242 + allocator = get_mem_allocator_instance()

243 + with allocator.use_memory_pool(tag):

244 + internal = torch.empty(16 << 20, dtype=torch.uint8, device=DEVICE_TYPE)

245 + with allocator.use_memory_pool("kv_cache"):

246 + kv = torch.empty(32 << 20, dtype=torch.uint8, device=DEVICE_TYPE)

247 +

248 + allocator.sleep(offload_tags=())

249 + assert mapped_usage(allocator) == 0

250 +

251 + allocator.wake_up(tags=["weights"])

252 + assert mapped_usage(allocator) == internal.nbytes

253 + allocator.wake_up(tags=["kv_cache"])

254 + assert mapped_usage(allocator) == internal.nbytes + kv.nbytes

255 +

256 +

257 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")

258 +@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")

259 +def test_workspace_scratch_discarded_on_sleep():

260 + """Workspace scratch lives in the sleep-mode pool: sleep discards it, any

261 + (selective) wake remaps it at the same address so captured graphs replay,

262 + and a buffer outgrown by a resize stays usable until it is released."""

263 + from vllm.device_allocator.sleep_mode_backend import CuMemBackend

264 + from vllm.v1.worker.workspace import WorkspaceManager

265 +

266 + allocator = get_mem_allocator_instance()

267 + manager = WorkspaceManager(

268 + torch.device(DEVICE_TYPE),

269 + alloc_context=lambda: allocator.use_memory_pool("workspace"),

270 + )

271 + (small,) = manager.get_simultaneous(((32 << 20,), torch.uint8))

272 + # Resize while a view of the old buffer is alive: the tag is re-entered.

273 + (scratch,) = manager.get_simultaneous(((64 << 20,), torch.uint8))

274 + assert len(allocator.allocator_and_pools["workspace"]) == 2

275 + graph = torch.cuda.CUDAGraph()

276 + with torch.cuda.graph(graph):

277 + scratch.add_(1)

278 +

279 + backend = CuMemBackend()

280 + backend.suspend(level=1)

281 + assert mapped_usage(allocator) == 0

282 + # Discarded, not copied.

283 + assert all(

284 + d.tag == "workspace" and d.cpu_backup_tensor is None

285 + for d in allocator.pointer_to_data.values()

286 + )

287 +

288 + backend.resume(tags=["kv_cache"])

289 + assert mapped_usage(allocator) == small.nbytes + scratch.nbytes

290 + scratch.zero_()

291 + graph.replay()

292 + assert int(scratch.sum()) == scratch.numel()

293 + small.fill_(1)

294 +

295 + # The next resize releases the outgrown buffer; wake does not remap it.

296 + del small

297 + (big,) = manager.get_simultaneous(((128 << 20,), torch.uint8))

298 + assert mapped_usage(allocator) == scratch.nbytes + big.nbytes

299 + backend.suspend(level=1)

300 + backend.resume(tags=["weights"])

301 + assert mapped_usage(allocator) == scratch.nbytes + big.nbytes

302 +

303 +

304 +@pytest.mark.parametrize("free_x_early", [True, False], ids=["x-freed", "x-alive"])

305 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")

306 +@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")

307 +def test_reentered_tag(free_x_early):

308 + """Enter a tag twice (two workspace ubatches): freeing X returns its memory

309 + by the next entry, without touching Y or breaking sleep/wake."""

310 + allocator = get_mem_allocator_instance()

311 + nbytes = 64 << 20

312 + with allocator.use_memory_pool("t"):

313 + xs = [torch.empty(nbytes, dtype=torch.uint8, device=DEVICE_TYPE)]

314 + if free_x_early:

315 + xs.clear()

316 + with allocator.use_memory_pool("t"):

317 + y = torch.empty(nbytes, dtype=torch.uint8, device=DEVICE_TYPE)

318 + xs.clear()

319 + gc.collect()

320 + with allocator.use_memory_pool("t"):

321 + pass

322 + assert mapped_usage(allocator) == nbytes

323 + y.fill_(1)

324 + allocator.sleep(offload_tags="t")

325 + allocator.wake_up()

326 + assert int(y.sum()) == y.numel()

327 + allocator.release_pools()

328 +

329 +

330 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")

331 @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")

332 def test_level2_discards_ordinary_tensor_with_weights_tag():

vllm/device_allocator/__init__.py

@@ -13,6 +13,9 @@ from vllm.platforms import current_platform

13 # py_handle has type list[int] on ROCm and int otherwise

14 HandleType: TypeAlias = tuple[int, int, int, list[int] | int]

16 +# Selective wake may defer only these; any wake restores every other tag.

17 +DEFERRABLE_TAGS: tuple[str, ...] = ("weights", "kv_cache")

18 +

20 @dataclasses.dataclass

21 class AllocationData:

vllm/device_allocator/cumem.py

+20 -13

@@ -16,7 +16,7 @@ from typing import Any

17 import torch

19 -from vllm.device_allocator import AllocationData, HandleType

19 +from vllm.device_allocator import DEFERRABLE_TAGS, AllocationData, HandleType

20 from vllm.device_allocator.alloc_conf import (

21 EXPANDABLE_SEGMENTS,

22 conf_flag_enabled,

@@ -139,7 +139,7 @@ class CuMemAllocator:

139 def __init__(self):

140 self.pointer_to_data: dict[int, AllocationData] = {}

141 self.current_tag: str = CuMemAllocator.default_tag

142 - self.allocator_and_pools: dict[str, Any] = {}

142 + self.allocator_and_pools: dict[str, list[Any]] = {}

143 # Creating strong references to the two callbacks here to prevent

144 # these ephemeral bound-method objects being garbage collected.

145 # See discussions in https://github.com/vllm-project/vllm/pull/22724

@@ -167,7 +167,7 @@ class CuMemAllocator:

167 if not self.allocator_and_pools:

168 return

169

170 - pool_entries = list(self.allocator_and_pools.values())

170 + pool_entries = [e for es in self.allocator_and_pools.values() for e in es]

171 self.allocator_and_pools.clear()

172

173 mem_pools = [entry[0] for entry in pool_entries]

@@ -334,9 +334,9 @@ class CuMemAllocator:

334 memory, and the rest of the data will have empty memory.

335

336 Args:

337 - tags: The tags of the memory allocation that will be loaded

338 - back to GPU memory. If None, all memory allocation will be loaded

339 - back to GPU memory.

337 + tags: The deferrable tags (weights, kv_cache) to load back to GPU

338 + memory; every other tag is always loaded back. If None, all

339 + memory allocation will be loaded back to GPU memory.

340

341 """

342 gc.collect()

@@ -345,7 +345,7 @@ class CuMemAllocator:

345 for ptr, data in self.pointer_to_data.items():

346 if not data.is_asleep:

347 continue

348 - if tags is None or data.tag in tags:

348 + if tags is None or data.tag in tags or data.tag not in DEFERRABLE_TAGS:

349 handle = data.handle

350 create_and_map(handle)

351 data.is_asleep = False

@@ -392,6 +392,10 @@ class CuMemAllocator:

392 old_tag = self.current_tag

393 self.current_tag = tag

394 try:

395 + if tag != old_tag:

396 + # Older pools of this tag are never reused, so trim them.

397 + for pool, _ in self.allocator_and_pools.get(tag, []):

398 + self._trim(pool)

399 with use_memory_pool_with_allocator(

400 self.python_malloc_callback, self.python_free_callback

401 ) as data:

@@ -400,7 +404,7 @@ class CuMemAllocator:

404 # and the memory pool.

405 # to avoid the issue, we keep a reference of the data.

406 # see https://github.com/pytorch/pytorch/issues/146431 .

403 - self.allocator_and_pools[tag] = data

407 + self.allocator_and_pools.setdefault(tag, []).append(data)

408 yield

409 # PyTorch's bug, calling torch.cuda.empty_cache() will error

410 # when using pluggable allocator, see

@@ -413,16 +417,19 @@ class CuMemAllocator:

417 # TODO: we should expose `empty_cache` method in the memory

418 # pool.

419 # TODO: ask for help from PyTorch team to expose this method.

416 - allocations = data[0].snapshot()

417 - for allocation in allocations:

418 - if allocation["allocated_size"] == 0:

419 - handle = self._python_free_callback(allocation["address"])

420 - unmap_and_release(handle)

420 + self._trim(data[0])

421 finally:

422 self.current_tag = old_tag

423 if expandable_was_enabled:

424 set_alloc_conf(prev_conf)

425

426 + def _trim(self, pool: Any) -> None:

427 + """Release free segments of a pool that is never allocated from again."""

428 + for allocation in pool.snapshot():

429 + data = self.pointer_to_data.get(allocation["address"])

430 + if allocation["allocated_size"] == 0 and data and not data.is_asleep:

431 + unmap_and_release(self._python_free_callback(allocation["address"]))

432 +

433 def get_current_usage(self) -> int:

434 """Get the total number of bytes allocated in the memory pool."""

435 sum_bytes: int = 0

vllm/device_allocator/xpumem.py

+10 -5

@@ -9,7 +9,7 @@ from typing import Any

10 import torch

12 -from vllm.device_allocator import AllocationData, HandleType

12 +from vllm.device_allocator import DEFERRABLE_TAGS, AllocationData, HandleType

13 from vllm.logger import init_logger

14 from vllm.utils.torch_utils import PIN_MEMORY

@@ -145,7 +145,7 @@ class XpuMemAllocator:

145 def __init__(self):

146 self.pointer_to_data: dict[int, AllocationData] = {}

147 self.current_tag: str = XpuMemAllocator.default_tag

148 - self.allocator_and_pools: dict[str, Any] = {}

148 + self.allocator_and_pools: dict[str, list[Any]] = {}

149 self.python_malloc_callback = self._python_malloc_callback

150 self.python_free_callback = self._python_free_callback

151

@@ -267,7 +267,11 @@ class XpuMemAllocator:

267 for ptr, data in self.pointer_to_data.items():

268 if not data.is_asleep:

269 continue

270 - if tags is not None and data.tag not in tags:

270 + if (

271 + tags is not None

272 + and data.tag not in tags

273 + and data.tag in DEFERRABLE_TAGS

274 + ):

275 continue

276 create_and_allocate(data.handle)

277 data.is_asleep = False

@@ -300,7 +304,7 @@ class XpuMemAllocator:

304 # MemPool teardown may invoke allocator virtual methods (e.g. raw_delete)

305 # when releasing cached blocks. If allocator wrappers are dropped first,

306 # C++ can hit "pure virtual method called" during shutdown.

303 - pool_entries = list(self.allocator_and_pools.values())

307 + pool_entries = [e for es in self.allocator_and_pools.values() for e in es]

308 self.allocator_and_pools.clear()

309

310 mem_pools = [entry[0] for entry in pool_entries]

@@ -333,7 +337,8 @@ class XpuMemAllocator:

337 self.python_malloc_callback,

338 self.python_free_callback,

339 ) as data:

336 - self.allocator_and_pools[tag] = data

340 + # Keep every entry's pool alive (pytorch/pytorch#146431).

341 + self.allocator_and_pools.setdefault(tag, []).append(data)

342 yield

343 finally:

344 self.current_tag = old_tag

vllm/v1/worker/gpu_worker.py

@@ -532,6 +532,8 @@ class Worker(WorkerBase):

532 self.device,

533 num_ubatches,

534 _num_workspace_lanes(self.vllm_config, self.use_v2_model_runner),

535 + # Scratch holds no state across steps: discard it on sleep.

536 + alloc_context=lambda: self._maybe_get_memory_pool_context("workspace"),

537 )

538 self.model_runner: GPUModelRunner = self._make_model_runner()

539 if self.rank == 0:

vllm/v1/worker/workspace.py

+11 -5

@@ -4,7 +4,7 @@

4 import inspect

5 import os

6 from collections.abc import Callable, Hashable, Iterator

7 -from contextlib import contextmanager

7 +from contextlib import AbstractContextManager, contextmanager, nullcontext

8 from contextvars import ContextVar

9 from itertools import accumulate

10 from math import prod

@@ -58,8 +58,10 @@ class WorkspaceManager:

58 device: torch.device,

59 num_ubatches: int | None = None,

60 num_lanes: int = 1,

61 + alloc_context: Callable[[], AbstractContextManager] = nullcontext,

62 ):

63 self._device = device

64 + self._alloc_context = alloc_context

65 # Cache num ubatches at init based on configuration (default to 1)

66 self._num_ubatches = num_ubatches if num_ubatches is not None else 1

67 if num_lanes < 1:

@@ -261,9 +263,11 @@ class WorkspaceManager:

263 # dead segment in reserved memory which can cause higher peak

264 # memory usage.

265 torch.accelerator.empty_cache()

264 - self._current_workspaces[workspace_id] = torch.empty(

265 - (required_bytes,), dtype=torch.uint8, device=self._device

266 - )

266 + # In sleep mode, entering the pool also trims outgrown buffers.

267 + with self._alloc_context():

268 + self._current_workspaces[workspace_id] = torch.empty(

269 + (required_bytes,), dtype=torch.uint8, device=self._device

270 + )

271 current_workspace = self._current_workspaces[workspace_id]

272

273 if envs.VLLM_DEBUG_WORKSPACE:

@@ -308,6 +312,7 @@ def init_workspace_manager(

312 device: torch.device,

313 num_ubatches: int | None = None,

314 num_lanes: int = 1,

315 + alloc_context: Callable[[], AbstractContextManager] = nullcontext,

316 ) -> None:

317 """Initialize the workspace manager with a device.

318

@@ -318,6 +323,7 @@ def init_workspace_manager(

323 device: The device to allocate workspace on.

324 num_ubatches: Number of workspace ubatch slots. Defaults to 1.

325 num_lanes: Number of independent execution lanes per ubatch. Defaults to 1.

326 + alloc_context: Factory for the context scratch is allocated in.

327

328 """

329 global _manager

@@ -328,7 +334,7 @@ def init_workspace_manager(

334 _manager._device,

335 device,

336 )

331 - _manager = WorkspaceManager(device, num_ubatches, num_lanes)

337 + _manager = WorkspaceManager(device, num_ubatches, num_lanes, alloc_context)

338

339

340 def lock_workspace() -> None: