diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py index dbd72c4e5e..74636fb615 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py +++ b/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"] diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index bb08241dee..9311b966ee 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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()