[Misc] Respect no_use_tqdm_on_load flag while capturing CUDA graph (#20834)

Signed-off-by: Linkun <github@lkchen.net>

Author: lk-chen

vllm/v1/worker/gpu_model_runner.py

@@ -2270,8 +2270,10 @@ def capture_model(self) -> None:
             # Only rank 0 should print progress bar during capture
             compilation_cases = reversed(self.cudagraph_batch_sizes)
             if is_global_first_rank():
-                compilation_cases = tqdm(list(compilation_cases),
-                                         desc="Capturing CUDA graph shapes")
+                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
                 for _ in range(

vllm/worker/model_runner.py

@@ -1587,6 +1587,7 @@ def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:
                 if get_tensor_model_parallel_rank() == 0:
                     compilation_cases = tqdm(
                         list(compilation_cases),
+                        disable=not self.load_config.use_tqdm_on_load,
                         desc="Capturing CUDA graph shapes")
                 for batch_size, use_inputs_embeds in compilation_cases:
                     attn_metadata = (