Pass token num cuda graph exec

Signed-off-by: Diego-Castan <[email protected]>

Author: diegocastanibm

vllm/v1/worker/gpu_model_runner.py

@@ -2516,7 +2516,7 @@ def profile_run(self) -> None:
         self.encoder_cache.clear()
         gc.collect()
 
-    def capture_model(self) -> None:
+    def capture_model(self, specific_token_num: Optional[int]) -> None:
         if not self.use_cuda_graph:
             logger.warning(
                 "Skipping CUDA graph capture. To turn on CUDA graph capture, "
@@ -2550,16 +2550,15 @@ def freeze_gc():
         with freeze_gc(), graph_capture(device=self.device):
             full_cg = self.full_cuda_graph
             # Only rank 0 should print progress bar during capture
-            compilation_cases = reversed(self.cudagraph_batch_sizes)
-            if is_global_first_rank():
+            compilation_cases = [specific_token_num] if specific_token_num else reversed(self.cudagraph_batch_sizes)
+
+            if is_global_first_rank() and specific_token_num is None:
                 compilation_cases = tqdm(
                     list(compilation_cases),
                     disable=not self.load_config.use_tqdm_on_load,
                     desc="Capturing CUDA graph shapes")
             for num_tokens in compilation_cases:
                 # We skip EPLB here since we don't want to record dummy metrics
-                logger.info("DIEGO: compilation for number of tokens %d",
-                    num_tokens)
                 for _ in range(
                         self.compilation_config.cudagraph_num_of_warmups):
                     self._dummy_run(num_tokens,

vllm/v1/worker/gpu_worker.py

@@ -64,6 +64,9 @@ def __init__(
 
         # Buffers saved before sleep
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
+        
+        # executed cuda graph
+        self._token_compiled_cudagraphs: set[int] = set() 
 
         # Torch profiler. Enabled and configured through env vars:
         # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
@@ -310,8 +313,8 @@ def compile_or_warm_up_model(self) -> None:
         for size in sorted(warmup_sizes, reverse=True):
             logger.info("Compile and warming up model for size %d", size)
             self.model_runner._dummy_run(size, skip_eplb=True)
-        if not self.model_config.enforce_eager:
-            self.model_runner.capture_model()
+        # if not self.model_config.enforce_eager:
+        #     self.model_runner.capture_model()
 
         # Warm up sampler and preallocate memory buffer for logits and other
         # sampling related tensors of max possible shape to avoid memory
@@ -355,6 +358,12 @@ def execute_model(
                 get_pp_group().recv_tensor_dict(
                     all_gather_group=get_tp_group()))
 
+        # Adding capture model in execution time
+        if scheduler_output.total_num_scheduled_tokens not in self._token_compiled_cudagraphs:
+            logger.info("DIEGO: CUDAgraph in execution time for %d input tokens", scheduler_output.total_num_scheduled_tokens)
+            self._token_compiled_cudagraphs.add(scheduler_output.total_num_scheduled_tokens)
+            self.model_runner.capture_model(scheduler_output.total_num_scheduled_tokens)
+        
         output = self.model_runner.execute_model(scheduler_output,
                                                  intermediate_tensors)