[Bugfix] Release CUDA graph profiling memory before KV cache allocation (#53955)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Aug 26, 2026 at 21:35 UTC
94d96e2446d68a949541f2a458a748e08ccbe717
2 files changed
+75
-4
tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py
+67
@@ -11,9 +11,14 @@ See https://github.com/vllm-project/vllm/issues/49224.
11
"""
12
13
import contextlib
14
+import gc
15
from types import SimpleNamespace
16
from typing import Any
17
18
+import pytest
19
+import torch
20
+
21
+from tests.utils import create_new_process_for_each_test
22
from vllm.compilation.counter import compilation_counter
23
from vllm.config.compilation import CUDAGraphMode
24
from vllm.v1.worker.gpu import cudagraph_utils as cgu
@@ -335,3 +340,65 @@ def test_profile_cudagraph_memory_swaps_and_drops_speculator_managers(monkeypatc
340
assert runner.speculator.cudagraph_manager is None
341
# The real global pool is restored afterwards.
342
assert _FakePlatform._global_graph_pool == GLOBAL_POOL
343
+
344
+
345
+@create_new_process_for_each_test("spawn")
346
+@pytest.mark.skipif(not cgu.current_platform.is_cuda(), reason="requires CUDA")
347
+def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch):
348
+ """Profiling graph memory must be freed before teardown completes."""
349
+
350
+ @contextlib.contextmanager
351
+ def _fake_set_current_vllm_config(_cfg):
352
+ yield
353
+
354
+ runner = _make_profiling_runner(CUDAGraphMode.FULL_AND_PIECEWISE)
355
+ runner.compilation_config.static_forward_context = {}
356
+ runner.model_state = SimpleNamespace(supports_mm_inputs=False)
357
+ runner.cache_config = SimpleNamespace(num_gpu_blocks=1)
358
+ runner.kv_caches = []
359
+ runner.attn_groups = []
360
+ runner.kv_cache_config = SimpleNamespace()
361
+ runner.lora_config = None
362
+ runner.maybe_remove_all_loras = lambda _: None
363
+ allocation_bytes = 64 << 20
364
+ memory: dict[str, int] = {}
365
+
366
+ torch.accelerator.synchronize()
367
+ gc.collect()
368
+ torch.accelerator.empty_cache()
369
+ memory["before"] = torch.accelerator.memory_reserved()
370
+
371
+ def _init(r):
372
+ r.events.append("init")
373
+ kv_cache = torch.empty(allocation_bytes, dtype=torch.uint8, device="cuda")
374
+ r.kv_caches = [kv_cache]
375
+ r.compilation_config.static_forward_context = {
376
+ "layer": SimpleNamespace(kv_cache=kv_cache)
377
+ }
378
+ manager = cgu.CudaGraphManager.__new__(cgu.CudaGraphManager)
379
+ manager.pool = cgu.current_platform.get_global_graph_pool()
380
+ r.speculator = SimpleNamespace(cudagraph_manager=manager)
381
+
382
+ def _capture_model() -> int:
383
+ for owner in (
384
+ runner.cudagraph_manager,
385
+ runner.speculator.cudagraph_manager,
386
+ ):
387
+ graph = torch.cuda.CUDAGraph()
388
+ with torch.cuda.graph(graph, pool=owner.pool):
389
+ output = torch.empty(allocation_bytes, dtype=torch.uint8, device="cuda")
390
+ output.fill_(1)
391
+ owner.profiling_resources = (graph, output)
392
+ torch.accelerator.synchronize()
393
+ memory["captured"] = torch.accelerator.memory_reserved()
394
+ return memory["captured"] - memory["before"]
395
+
396
+ monkeypatch.setattr(cgu, "set_current_vllm_config", _fake_set_current_vllm_config)
397
+ monkeypatch.setattr(cgu, "_init_minimal_kv_cache_for_profiling", _init)
398
+ runner.capture_model = _capture_model
399
+
400
+ cgu.profile_cudagraph_memory(runner)
401
+ memory["after"] = torch.accelerator.memory_reserved()
402
+
403
+ assert memory["captured"] - memory["before"] >= 3 * allocation_bytes
404
+ assert memory["after"] == memory["before"]
vllm/v1/worker/gpu/cudagraph_utils.py
+8
-4
@@ -753,7 +753,7 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
753
all_wrappers: list[Any] = []
754
original_pools: dict[int, Any] = {}
755
speculator = getattr(runner, "speculator", None)
756
- spec_managers: list[tuple[str, CudaGraphManager]] = []
756
+ spec_manager_names: list[str] = []
757
try:
758
if not manager.needs_capture():
759
return 0
@@ -769,8 +769,8 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
769
original_pools[id(wrapper)] = wrapper.graph_pool
770
wrapper.graph_pool = throwaway_pool
771
if speculator is not None:
772
- spec_managers = [
773
- (name, value)
772
+ spec_manager_names = [
773
+ name
774
for name, value in vars(speculator).items()
775
if isinstance(value, CudaGraphManager)
776
]
@@ -799,8 +799,11 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
799
# Drop the speculator's cudagraph managers; the real
800
# initialize_kv_cache re-creates them. Their profiling graphs
801
# release the throwaway pool here rather than after the real init.
802
- for name, _ in spec_managers:
802
+ for name in spec_manager_names:
803
setattr(speculator, name, None)
804
+ # Drop local references before teardown detaches the runner's
805
+ # manager and flushes the allocator.
806
+ del manager
807
_teardown_profiling_state(runner)
808
finally:
809
platform_cls._global_graph_pool = saved_global_pool
@@ -869,6 +872,7 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
872
layer.kv_cache = (
873
torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []
874
)
875
+ del kv_cache
876
runner.cache_config.num_gpu_blocks = None
877
runner.maybe_remove_all_loras(runner.lora_config)
878
gc.collect()