[Feature] Offload KV-init runtime state on sleep (#59158)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
aoshen02 committed
Oct 2, 2026 at 17:49 UTC
8e58ac22aff66f06f49b5ed0ce95ec9c92f8d022
4 files changed
+45
-4
tests/basic_correctness/memory/cumem/test_cumem.py
+29
@@ -327,6 +327,35 @@ def test_reentered_tag(free_x_early):
327
allocator.release_pools()
328
329
330
+@pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"])
331
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
332
+@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
333
+def test_runtime_state_survives_sleep(level):
334
+ """Runtime state is offloaded at both levels and restored by the first
335
+ selective wake; a nested kv_cache pool keeps its own tag."""
336
+ from vllm.device_allocator.sleep_mode_backend import CuMemBackend
337
+
338
+ allocator = get_mem_allocator_instance()
339
+ with allocator.use_memory_pool("runtime"):
340
+ state = torch.arange(1 << 20, device=DEVICE_TYPE)
341
+ with allocator.use_memory_pool("kv_cache"):
342
+ kv = torch.zeros(1 << 20, device=DEVICE_TYPE)
343
+ assert {d.tag for d in allocator.pointer_to_data.values()} == {
344
+ "runtime",
345
+ "kv_cache",
346
+ }
347
+
348
+ backend = CuMemBackend()
349
+ backend.suspend(level=level)
350
+ assert mapped_usage(allocator) == 0
351
+
352
+ backend.resume(tags=["weights"])
353
+ assert torch.equal(state, torch.arange(1 << 20, device=DEVICE_TYPE))
354
+ backend.resume(tags=["kv_cache"])
355
+ kv.fill_(1)
356
+ assert int(kv.sum()) == kv.numel()
357
+
358
+
359
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
360
@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
361
def test_level2_discards_ordinary_tensor_with_weights_tag():
vllm/device_allocator/sleep_mode_backend.py
+3
-1
@@ -129,7 +129,9 @@ class CuMemBackend(SleepModeBackend):
129
130
self._state = "SUSPENDED"
131
allocator = get_mem_allocator_instance()
132
- allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
132
+ # Runtime state is kept at every level; weights only at level 1.
133
+ offload_tags = ("weights", "runtime") if level == 1 else ("runtime",)
134
+ allocator.sleep(offload_tags=offload_tags)
135
136
def resume(self, tags: list[str] | None = None) -> None:
137
from vllm.device_allocator import get_mem_allocator_instance
vllm/v1/worker/gpu_model_runner.py
+3
@@ -6531,6 +6531,9 @@ class GPUModelRunner(
6531
if hasattr(self, "kv_cache_config"):
6532
delattr(self, "kv_cache_config")
6533
self.cache_config.num_gpu_blocks = None
6534
+ # Profiling may have rebuilt the InputBatch outside the worker's
6535
+ # "runtime" pool; the real initialize_kv_cache rebuilds it.
6536
+ self._init_block_sizes = []
6537
6538
for layer in self.compilation_config.static_forward_context.values():
6539
if hasattr(layer, "kv_cache"):
vllm/v1/worker/gpu_worker.py
+10
-3
@@ -841,10 +841,17 @@ class Worker(WorkerBase):
841
if mem_pool_context is None:
842
mem_pool_context = self._maybe_get_memory_pool_context(tag="kv_cache")
843
844
- self.model_runner.initialize_kv_cache(
845
- kv_cache_config,
846
- kv_cache_allocation_context=mem_pool_context,
844
+ # Offload KV-init state, except on XPU, whose outermost pool would win.
845
+ runtime_pool = (
846
+ nullcontext()
847
+ if current_platform.is_xpu()
848
+ else self._maybe_get_memory_pool_context(tag="runtime")
849
)
850
+ with runtime_pool:
851
+ self.model_runner.initialize_kv_cache(
852
+ kv_cache_config,
853
+ kv_cache_allocation_context=mem_pool_context,
854
+ )
855
856
# Build KV-zero metadata outside the CuMem pool so the bookkeeping
857
# GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch