[Bugfix] Release CUDA graph profiling memory before KV cache allocation (#53955)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py

@@ -11,9 +11,14 @@ See https://github.com/vllm-project/vllm/issues/49224.
 """
 
 import contextlib
+import gc
 from types import SimpleNamespace
 from typing import Any
 
+import pytest
+import torch
+
+from tests.utils import create_new_process_for_each_test
 from vllm.compilation.counter import compilation_counter
 from vllm.config.compilation import CUDAGraphMode
 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
     assert runner.speculator.cudagraph_manager is None
     # The real global pool is restored afterwards.
     assert _FakePlatform._global_graph_pool == GLOBAL_POOL
+
+
+@create_new_process_for_each_test("spawn")
+@pytest.mark.skipif(not cgu.current_platform.is_cuda(), reason="requires CUDA")
+def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch):
+    """Profiling graph memory must be freed before teardown completes."""
+
+    @contextlib.contextmanager
+    def _fake_set_current_vllm_config(_cfg):
+        yield
+
+    runner = _make_profiling_runner(CUDAGraphMode.FULL_AND_PIECEWISE)
+    runner.compilation_config.static_forward_context = {}
+    runner.model_state = SimpleNamespace(supports_mm_inputs=False)
+    runner.cache_config = SimpleNamespace(num_gpu_blocks=1)
+    runner.kv_caches = []
+    runner.attn_groups = []
+    runner.kv_cache_config = SimpleNamespace()
+    runner.lora_config = None
+    runner.maybe_remove_all_loras = lambda _: None
+    allocation_bytes = 64 << 20
+    memory: dict[str, int] = {}
+
+    torch.accelerator.synchronize()
+    gc.collect()
+    torch.accelerator.empty_cache()
+    memory["before"] = torch.accelerator.memory_reserved()
+
+    def _init(r):
+        r.events.append("init")
+        kv_cache = torch.empty(allocation_bytes, dtype=torch.uint8, device="cuda")
+        r.kv_caches = [kv_cache]
+        r.compilation_config.static_forward_context = {
+            "layer": SimpleNamespace(kv_cache=kv_cache)
+        }
+        manager = cgu.CudaGraphManager.__new__(cgu.CudaGraphManager)
+        manager.pool = cgu.current_platform.get_global_graph_pool()
+        r.speculator = SimpleNamespace(cudagraph_manager=manager)
+
+    def _capture_model() -> int:
+        for owner in (
+            runner.cudagraph_manager,
+            runner.speculator.cudagraph_manager,
+        ):
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=owner.pool):
+                output = torch.empty(allocation_bytes, dtype=torch.uint8, device="cuda")
+                output.fill_(1)
+            owner.profiling_resources = (graph, output)
+        torch.accelerator.synchronize()
+        memory["captured"] = torch.accelerator.memory_reserved()
+        return memory["captured"] - memory["before"]
+
+    monkeypatch.setattr(cgu, "set_current_vllm_config", _fake_set_current_vllm_config)
+    monkeypatch.setattr(cgu, "_init_minimal_kv_cache_for_profiling", _init)
+    runner.capture_model = _capture_model
+
+    cgu.profile_cudagraph_memory(runner)
+    memory["after"] = torch.accelerator.memory_reserved()
+
+    assert memory["captured"] - memory["before"] >= 3 * allocation_bytes
+    assert memory["after"] == memory["before"]

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -753,7 +753,7 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
         all_wrappers: list[Any] = []
         original_pools: dict[int, Any] = {}
         speculator = getattr(runner, "speculator", None)
-        spec_managers: list[tuple[str, CudaGraphManager]] = []
+        spec_manager_names: list[str] = []
         try:
             if not manager.needs_capture():
                 return 0
@@ -769,8 +769,8 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
                 original_pools[id(wrapper)] = wrapper.graph_pool
                 wrapper.graph_pool = throwaway_pool
             if speculator is not None:
-                spec_managers = [
-                    (name, value)
+                spec_manager_names = [
+                    name
                     for name, value in vars(speculator).items()
                     if isinstance(value, CudaGraphManager)
                 ]
@@ -799,8 +799,11 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
             # Drop the speculator's cudagraph managers; the real
             # initialize_kv_cache re-creates them. Their profiling graphs
             # release the throwaway pool here rather than after the real init.
-            for name, _ in spec_managers:
+            for name in spec_manager_names:
                 setattr(speculator, name, None)
+            # Drop local references before teardown detaches the runner's
+            # manager and flushes the allocator.
+            del manager
             _teardown_profiling_state(runner)
     finally:
         platform_cls._global_graph_pool = saved_global_pool
@@ -869,6 +872,7 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
             layer.kv_cache = (
                 torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []
             )
+            del kv_cache
     runner.cache_config.num_gpu_blocks = None
     runner.maybe_remove_all_loras(runner.lora_config)
     gc.collect()