A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 16:00 UTC
[Model Runner V2] Minor fix for cudagraph_utils (#29256)
Woosuk Kwon committed
Nov 22, 2025 at 20:12 UTC
20ee418adc279f29e76e7770c4f688c4fc070274
2 files changed
+6
-14
vllm/v1/worker/gpu/cudagraph_utils.py
+5
-14
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import gc
4
-from contextlib import contextmanager
3
+from unittest.mock import patch
4
5
import numpy as np
6
import torch
@@ -140,6 +139,7 @@ class CudaGraphManager:
139
attn_metadata,
140
self.vllm_config,
141
num_tokens=batch_size,
142
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
143
num_tokens_across_dp=num_tokens_across_dp,
144
):
145
hidden_states = model(
@@ -148,15 +148,16 @@ class CudaGraphManager:
148
)
149
if self.hidden_states is None:
150
self.hidden_states = torch.empty_like(hidden_states)
151
- torch.cuda.synchronize()
151
152
# Capture the graph.
153
graph = torch.cuda.CUDAGraph()
154
with (
155
+ patch("torch.cuda.empty_cache", lambda: None),
156
set_forward_context(
157
attn_metadata,
158
self.vllm_config,
159
num_tokens=batch_size,
160
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
161
num_tokens_across_dp=num_tokens_across_dp,
162
),
163
torch.cuda.graph(graph, self.pool),
@@ -183,7 +184,7 @@ class CudaGraphManager:
184
if is_global_first_rank():
185
sizes_to_capture = tqdm(sizes_to_capture, desc="Capturing CUDA graphs")
186
186
- with freeze_gc(), graph_capture(device=self.device):
187
+ with graph_capture(device=self.device):
188
for batch_size in sizes_to_capture:
189
self.capture_graph(
190
batch_size,
@@ -199,13 +200,3 @@ class CudaGraphManager:
200
self.graphs[batch_size].replay()
201
assert self.hidden_states is not None
202
return self.hidden_states[:batch_size]
202
-
203
-
204
-@contextmanager
205
-def freeze_gc():
206
- gc.collect()
207
- gc.freeze()
208
- try:
209
- yield
210
- finally:
211
- gc.unfreeze()
vllm/v1/worker/gpu/model_runner.py
+1
@@ -298,6 +298,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
298
return 0
299
300
start_time = time.perf_counter()
301
+ torch.cuda.empty_cache()
302
start_free_gpu_memory = torch.cuda.mem_get_info()[0]
303
304
with self.maybe_setup_dummy_loras(self.lora_config):