@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e6cc0899b6

[Bugfix][MRV2] Run cudagraph memory profiling in a throwaway graph pool (#53682)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Aug 25, 2026 at 11:52 UTC e6cc0899b6e29ad613efae572eb98e0d776afe4d

2 files changed +131 -58

tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py

+56 -5

@@ -74,17 +74,28 @@ def _make_profiling_runner(

74 return runner

77 +class _FakePlatform:

78 + """Stands in for current_platform; the global graph pool is a class attr,

79 + matching vllm.platforms.Platform's lazy singleton."""

80 +

81 + _global_graph_pool: Any = GLOBAL_POOL

82 +

83 + @staticmethod

84 + def graph_pool_handle() -> Any:

85 + return THROWAWAY_POOL

86 +

87 + def get_global_graph_pool(self) -> Any:

88 + return type(self)._global_graph_pool

89 +

90 +

91 def _patch_module(monkeypatch) -> None:

92 @contextlib.contextmanager

93 def _fake_set_current_vllm_config(_cfg):

94 yield

96 + _FakePlatform._global_graph_pool = GLOBAL_POOL

97 monkeypatch.setattr(cgu, "set_current_vllm_config", _fake_set_current_vllm_config)

83 - monkeypatch.setattr(

84 - cgu,

85 - "current_platform",

86 - SimpleNamespace(graph_pool_handle=lambda: THROWAWAY_POOL),

87 - )

98 + monkeypatch.setattr(cgu, "current_platform", _FakePlatform())

99 monkeypatch.setattr(

100 cgu, "_init_minimal_kv_cache_for_profiling", lambda r: r.events.append("init")

101 )

@@ -284,3 +295,43 @@ def test_profile_cudagraph_memory_redirects_wrapper_pools(monkeypatch):

295 assert wrapper.graph_pool == GLOBAL_POOL

296 finally:

297 cgu.CUDAGraphWrapper._all_instances.discard(wrapper)

298 +

299 +

300 +def test_profile_cudagraph_memory_swaps_and_drops_speculator_managers(monkeypatch):

301 + """Speculator cudagraph managers must also capture into the throwaway pool.

302 +

303 + They are created during the profiling KV-cache bootstrap, binding the

304 + (swapped) global graph pool, and are re-created by the real

305 + initialize_kv_cache; profiling-captured graphs must be dropped at

306 + teardown rather than released against the persistent global pool, which

307 + would trip the c10 create_or_incref_pool assert at the real capture.

308 + """

309 + _patch_module(monkeypatch)

310 + runner = _make_profiling_runner(CUDAGraphMode.FULL_AND_PIECEWISE)

311 +

312 + def _init(r):

313 + r.events.append("init")

314 + # Mirror production: the speculator's cudagraph managers are created

315 + # during the profiling KV-cache bootstrap and bind the global pool

316 + # (which profiling has already pointed at the throwaway pool).

317 + manager = cgu.CudaGraphManager.__new__(cgu.CudaGraphManager)

318 + manager.pool = cgu.current_platform.get_global_graph_pool()

319 + r.speculator = SimpleNamespace(cudagraph_manager=manager)

320 +

321 + monkeypatch.setattr(cgu, "_init_minimal_kv_cache_for_profiling", _init)

322 +

323 + pools_seen: list[Any] = []

324 + capture_model = runner.capture_model

325 +

326 + def _capture_model() -> int:

327 + pools_seen.append(runner.speculator.cudagraph_manager.pool)

328 + return capture_model()

329 +

330 + runner.capture_model = _capture_model

331 +

332 + cgu.profile_cudagraph_memory(runner)

333 +

334 + assert pools_seen == [THROWAWAY_POOL]

335 + assert runner.speculator.cudagraph_manager is None

336 + # The real global pool is restored afterwards.

337 + assert _FakePlatform._global_graph_pool == GLOBAL_POOL

vllm/v1/worker/gpu/cudagraph_utils.py

+75 -53

@@ -723,62 +723,84 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:

723 gc.collect()

724 torch.accelerator.empty_cache()

725

726 - with set_current_vllm_config(runner.vllm_config):

727 - _init_minimal_kv_cache_for_profiling(runner)

726 + # Run the whole profiling phase against a throwaway CUDA graph pool by

727 + # pointing the global graph pool singleton at it: objects that bind the

728 + # pool lazily during profiling (speculator cudagraph managers, breakable

729 + # runners created mid-capture) then land on the throwaway pool too. Pools

730 + # bound before profiling (piecewise wrappers) are swapped explicitly in

731 + # the inner block. Profiling graphs captured into the persistent global

732 + # pool and then discarded would drop its use_count to 0, tripping the c10

733 + # allocator's create_or_incref_pool assert when the real capture reuses

734 + # that pool ("use_count > 0 INTERNAL ASSERT FAILED").

735 + platform_cls = type(current_platform)

736 + saved_global_pool = platform_cls._global_graph_pool

737 + throwaway_pool = current_platform.graph_pool_handle()

738 + platform_cls._global_graph_pool = throwaway_pool

739

729 - manager = runner.cudagraph_manager

730 - assert manager is not None

731 -

732 - # Don't count profiling captures; the real capture_model() runs later.

733 - saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured

734 - saved_capture_triggers = compilation_counter.num_gpu_runner_capture_triggers

735 - all_wrappers: list[Any] = []

736 - original_pools: dict[int, Any] = {}

740 try:

738 - if not manager.needs_capture():

739 - return 0

740 - # Capture all profiling graphs into a throwaway pool so their memory

741 - # is reclaimed on teardown rather than retained by the persistent

742 - # global pool (which the real capture reuses). This must include the

743 - # piecewise wrappers, not just the FULL-graph manager pool: graphs

744 - # captured into the global pool and then discarded drop its use_count

745 - # to 0, and the real capture on the same pool trips the c10 allocator's

746 - # create_or_incref_pool assert ("use_count > 0 INTERNAL ASSERT FAILED").

747 - manager.pool = current_platform.graph_pool_handle()

748 - if manager.use_breakable_cg:

749 - # The breakable runner is otherwise created lazily during capture,

750 - # after the pool swap below, and would capture into the global

751 - # pool. Create it now so its pool gets swapped too.

752 - manager.init_breakable_cg_runner(runner.model)

753 - all_wrappers = list(CUDAGraphWrapper._all_instances) + list(

754 - BreakableCUDAGraphWrapper._all_instances

755 - )

756 - for wrapper in all_wrappers:

757 - original_pools[id(wrapper)] = wrapper.graph_pool

758 - wrapper.graph_pool = manager.pool

759 - manager._max_full_descs_to_capture = _FULL_GRAPH_PROFILING_SAMPLES

760 - mem_samples: list[int] = []

761 - manager._capture_mem_samples = mem_samples

762 -

763 - measured = int(runner.capture_model())

764 -

765 - # The measured delta covers PIECEWISE, encoder and speculator graphs

766 - # plus the sampled FULL graphs; swap the sampled FULL cost for the

767 - # extrapolated total. FULL and PIECEWISE share one pool here just as

768 - # they share the global pool at runtime, so the overlap is not

769 - # double-counted.

770 - num_full_graphs = len(manager._capture_descs.get(CUDAGraphMode.FULL, []))

771 - full_estimate = _extrapolate_full_graph_memory(mem_samples, num_full_graphs)

772 - return max(measured - sum(mem_samples) + full_estimate, 0)

741 + with set_current_vllm_config(runner.vllm_config):

742 + _init_minimal_kv_cache_for_profiling(runner)

743 +

744 + manager = runner.cudagraph_manager

745 + assert manager is not None

746 +

747 + # Don't count profiling captures; the real capture_model() runs later.

748 + saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured

749 + saved_capture_triggers = compilation_counter.num_gpu_runner_capture_triggers

750 + all_wrappers: list[Any] = []

751 + original_pools: dict[int, Any] = {}

752 + speculator = getattr(runner, "speculator", None)

753 + spec_managers: list[tuple[str, CudaGraphManager]] = []

754 + try:

755 + if not manager.needs_capture():

756 + return 0

757 + manager.pool = throwaway_pool

758 + if manager.use_breakable_cg:

759 + # Create the breakable runner before the wrapper pool swap so

760 + # its pool is covered as well.

761 + manager.init_breakable_cg_runner(runner.model)

762 + all_wrappers = list(CUDAGraphWrapper._all_instances) + list(

763 + BreakableCUDAGraphWrapper._all_instances

764 + )

765 + for wrapper in all_wrappers:

766 + original_pools[id(wrapper)] = wrapper.graph_pool

767 + wrapper.graph_pool = throwaway_pool

768 + if speculator is not None:

769 + spec_managers = [

770 + (name, value)

771 + for name, value in vars(speculator).items()

772 + if isinstance(value, CudaGraphManager)

773 + ]

774 + manager._max_full_descs_to_capture = _FULL_GRAPH_PROFILING_SAMPLES

775 + mem_samples: list[int] = []

776 + manager._capture_mem_samples = mem_samples

777 +

778 + measured = int(runner.capture_model())

779 +

780 + # The measured delta covers PIECEWISE, encoder and speculator graphs

781 + # plus the sampled FULL graphs; swap the sampled FULL cost for the

782 + # extrapolated total. FULL and PIECEWISE share one pool here just as

783 + # they share the global pool at runtime, so the overlap is not

784 + # double-counted.

785 + num_full_graphs = len(manager._capture_descs.get(CUDAGraphMode.FULL, []))

786 + full_estimate = _extrapolate_full_graph_memory(mem_samples, num_full_graphs)

787 + return max(measured - sum(mem_samples) + full_estimate, 0)

788 + finally:

789 + compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured

790 + compilation_counter.num_gpu_runner_capture_triggers = saved_capture_triggers

791 + CUDAGraphWrapper.clear_all_graphs()

792 + BreakableCUDAGraphWrapper.clear_all_graphs()

793 + for wrapper in all_wrappers:

794 + if id(wrapper) in original_pools:

795 + wrapper.graph_pool = original_pools[id(wrapper)]

796 + # Drop the speculator's cudagraph managers; the real

797 + # initialize_kv_cache re-creates them. Their profiling graphs

798 + # release the throwaway pool here rather than after the real init.

799 + for name, _ in spec_managers:

800 + setattr(speculator, name, None)

801 + _teardown_profiling_state(runner)

802 finally:

774 - compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured

775 - compilation_counter.num_gpu_runner_capture_triggers = saved_capture_triggers

776 - CUDAGraphWrapper.clear_all_graphs()

777 - BreakableCUDAGraphWrapper.clear_all_graphs()

778 - for wrapper in all_wrappers:

779 - if id(wrapper) in original_pools:

780 - wrapper.graph_pool = original_pools[id(wrapper)]

781 - _teardown_profiling_state(runner)

803 + platform_cls._global_graph_pool = saved_global_pool

804

805

806 def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> int: