[Core][MRV2] Freeze gc during V2 CG capture; skip per-descriptor cleanup (#54646)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Aug 31, 2026 at 19:37 UTC
c28feab98919739ae6d2041c31c94a3b11718590
4 files changed
+78
-69
vllm/compilation/breakable_cudagraph.py
+7
-9
@@ -365,15 +365,13 @@ class BreakableCUDAGraphWrapper:
365
else:
366
set_graph_pool_id(current_platform.graph_pool_handle())
367
368
- # Match torch.cuda.graph()'s pre-capture cleanup once per descriptor.
369
- # We drive capture_begin/end directly and bypass torch.cuda.graph(),
370
- # so its built-in gc + empty_cache never fire. Run them here once
371
- # per _capture call -- NOT inside _begin_segment, since this capture
372
- # session may issue many begin/end pairs (one per layer's break),
373
- # and repeated gc would tank capture time the way it did for the
374
- # pre-`gc_disable` piecewise path.
375
- gc.collect()
376
- torch.accelerator.empty_cache()
368
+ # Match torch.cuda.graph()'s pre-capture cleanup, which we bypass.
369
+ # Skip it when gc is disabled: bulk capture runs under
370
+ # freeze_gc_for_cudagraph_capture, which already did this cleanup,
371
+ # and repeating it per descriptor dominates capture time.
372
+ if gc.isenabled():
373
+ gc.collect()
374
+ torch.accelerator.empty_cache()
375
# Sync the offloader's copy stream before capture so any in-flight
376
# pre-capture prefetches are complete and don't leak into the graph.
377
get_offloader().sync_prev_onload()
vllm/utils/gc_utils.py
+29
-1
@@ -4,7 +4,7 @@ import gc
4
import json
5
import time
6
from collections import Counter
7
-from contextlib import suppress
7
+from contextlib import contextmanager, suppress
8
from typing import Any
9
10
import vllm.envs as envs
@@ -93,6 +93,34 @@ class GCDebugger:
93
)
94
95
96
+@contextmanager
97
+def freeze_gc_for_cudagraph_capture():
98
+ """Freeze and disable gc for the duration of bulk CUDA graph capture.
99
+
100
+ A gc cycle during stream capture can invalidate the captured graph, e.g.
101
+ a finalized Triton kernel unloads its module. Opt out with
102
+ VLLM_ENABLE_CUDAGRAPH_GC=1.
103
+ """
104
+ gc_was_enabled = gc.isenabled()
105
+ gc.collect()
106
+ should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC
107
+ if should_freeze:
108
+ gc.freeze()
109
+ gc.disable()
110
+ try:
111
+ yield
112
+ finally:
113
+ if should_freeze:
114
+ try:
115
+ gc.unfreeze()
116
+ gc.collect()
117
+ finally:
118
+ if gc_was_enabled:
119
+ gc.enable()
120
+ else:
121
+ gc.disable()
122
+
123
+
124
def freeze_gc_heap() -> None:
125
"""
126
Freeze all objects tracked by the garbage collector. It should be invoked
vllm/v1/worker/gpu/model_runner.py
+40
-36
@@ -61,6 +61,7 @@ from vllm.multimodal.encoder_budget import (
61
)
62
from vllm.sequence import IntermediateTensors
63
from vllm.tasks import SupportedTask
64
+from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
65
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
66
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
67
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -916,44 +917,47 @@ class GPUModelRunner(LoRAModelRunnerMixin):
917
compilation_counter.num_gpu_runner_capture_triggers += 1
918
919
start_time = time.perf_counter()
919
- gc.collect()
920
- torch.accelerator.empty_cache()
921
- start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
922
-
923
- with self.maybe_setup_dummy_loras(self.lora_config):
924
- if capture_encoder:
925
- self.model_state.encoder_runner.capture()
926
-
927
- if capture_decoder:
928
- input_buffers = self.input_buffers
929
- if self.pcp_manager is not None:
930
- input_buffers = self.pcp_manager.input_buffers
931
- self.cudagraph_manager.capture(
932
- self.model,
933
- self.model_state,
934
- input_buffers,
935
- self.intermediate_tensors,
936
- self.block_tables,
937
- self.attn_groups,
938
- self.kv_cache_config,
939
- pcp_manager=self.pcp_manager,
940
- has_lora=self.lora_config is not None,
941
- use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
942
- lora_capture_hook=create_lora_capture_hook(self.lora_config, self),
943
- )
944
- if self.speculator is not None:
945
- with use_workspace_lane(self._draft_workspace_lane):
946
- self.speculator.capture()
947
- if self.adaptive_verification is not None:
948
- with self.step_timing.collect() as timings:
949
- for batch in self.adaptive_verification.batches_to_profile(
950
- self.cudagraph_manager.captured_token_counts()
951
- ):
952
- self._dummy_run(**batch)
953
- self.adaptive_verification.set_initial_cost_curves(timings)
920
+ with freeze_gc_for_cudagraph_capture():
921
+ torch.accelerator.empty_cache()
922
+ start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
923
+
924
+ with self.maybe_setup_dummy_loras(self.lora_config):
925
+ if capture_encoder:
926
+ self.model_state.encoder_runner.capture()
927
+
928
+ if capture_decoder:
929
+ input_buffers = self.input_buffers
930
+ if self.pcp_manager is not None:
931
+ input_buffers = self.pcp_manager.input_buffers
932
+ self.cudagraph_manager.capture(
933
+ self.model,
934
+ self.model_state,
935
+ input_buffers,
936
+ self.intermediate_tensors,
937
+ self.block_tables,
938
+ self.attn_groups,
939
+ self.kv_cache_config,
940
+ pcp_manager=self.pcp_manager,
941
+ has_lora=self.lora_config is not None,
942
+ use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
943
+ lora_capture_hook=create_lora_capture_hook(
944
+ self.lora_config, self
945
+ ),
946
+ )
947
+ if self.speculator is not None:
948
+ with use_workspace_lane(self._draft_workspace_lane):
949
+ self.speculator.capture()
950
+ if self.adaptive_verification is not None:
951
+ with self.step_timing.collect() as timings:
952
+ for batch in self.adaptive_verification.batches_to_profile(
953
+ self.cudagraph_manager.captured_token_counts()
954
+ ):
955
+ self._dummy_run(**batch)
956
+ self.adaptive_verification.set_initial_cost_curves(timings)
957
+
958
+ end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
959
960
end_time = time.perf_counter()
956
- end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
961
elapsed_time = end_time - start_time
962
cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
963
# This usually takes 5~20 seconds.
vllm/v1/worker/gpu_model_runner.py
+2
-23
@@ -124,6 +124,7 @@ from vllm.sequence import IntermediateTensors
124
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
125
from vllm.tracing import instrument
126
from vllm.utils import length_from_prompt_token_ids_or_embeds
127
+from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
128
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
129
from vllm.utils.math_utils import cdiv, round_up
130
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
@@ -6614,29 +6615,7 @@ class GPUModelRunner(
6615
6616
logger.debug("Initialized minimal KV cache for CUDA graph profiling")
6617
6617
- @staticmethod
6618
- @contextmanager
6619
- def _freeze_gc():
6620
- gc_was_enabled = gc.isenabled()
6621
- gc.collect()
6622
- should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC
6623
- if should_freeze:
6624
- gc.freeze()
6625
- # A Triton kernel finalized during stream capture unloads its
6626
- # module and invalidates the captured graph.
6627
- gc.disable()
6628
- try:
6629
- yield
6630
- finally:
6631
- if should_freeze:
6632
- try:
6633
- gc.unfreeze()
6634
- gc.collect()
6635
- finally:
6636
- if gc_was_enabled:
6637
- gc.enable()
6638
- else:
6639
- gc.disable()
6618
+ _freeze_gc = staticmethod(freeze_gc_for_cudagraph_capture)
6619
6620
def shutdown(self) -> None:
6621
"""Release GPU tensors (model weights, KV caches, workspace) so that