[Model Runner V2] Optimize CUDA graph capture time (vllm-project#29275)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Author: WoosukKwon

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -106,7 +106,10 @@ def capture_graph(
         input_buffers.query_start_loc.np[: batch_size + 1] = np.arange(batch_size + 1)
         input_buffers.query_start_loc.np[batch_size:] = batch_size
         input_buffers.query_start_loc.copy_to_gpu()
-        input_buffers.seq_lens[:batch_size] = self.max_model_len
+        # HACK(woosuk): To optimize warmup time, we use 1 (instead of max_model_len)
+        # for seq_lens. This leads to a mismatch between seq_lens (GPU) and
+        # seq_lens_np (CPU), which might cause issues in some attention backends.
+        input_buffers.seq_lens[:batch_size] = 1
         input_buffers.seq_lens[batch_size:] = 0
 
         input_block_tables = [x[:batch_size] for x in block_tables.input_block_tables]

vllm/v1/worker/gpu/model_runner.py

@@ -313,6 +313,7 @@ def capture_model(self) -> int:
             return 0
 
         start_time = time.perf_counter()
+        gc.collect()
         torch.cuda.empty_cache()
         start_free_gpu_memory = torch.cuda.mem_get_info()[0]