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] Optimize CUDA graph capture time (#29275)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 23, 2025 at 11:15 UTC
62d54ba46db25b95de2d21e46f4b57b5502ed747
2 files changed
+5
-1
vllm/v1/worker/gpu/cudagraph_utils.py
+4
-1
@@ -106,7 +106,10 @@ class CudaGraphManager:
106
input_buffers.query_start_loc.np[: batch_size + 1] = np.arange(batch_size + 1)
107
input_buffers.query_start_loc.np[batch_size:] = batch_size
108
input_buffers.query_start_loc.copy_to_gpu()
109
- input_buffers.seq_lens[:batch_size] = self.max_model_len
109
+ # HACK(woosuk): To optimize warmup time, we use 1 (instead of max_model_len)
110
+ # for seq_lens. This leads to a mismatch between seq_lens (GPU) and
111
+ # seq_lens_np (CPU), which might cause issues in some attention backends.
112
+ input_buffers.seq_lens[:batch_size] = 1
113
input_buffers.seq_lens[batch_size:] = 0
114
115
input_block_tables = [x[:batch_size] for x in block_tables.input_block_tables]
vllm/v1/worker/gpu/model_runner.py
+1
@@ -313,6 +313,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
313
return 0
314
315
start_time = time.perf_counter()
316
+ gc.collect()
317
torch.cuda.empty_cache()
318
start_free_gpu_memory = torch.cuda.mem_get_info()[0]
319