xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-29 08:17:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Linkun <github@lkchen.net>
This commit is contained in:
lkchen 2025-07-11 23:04:45 -07:00 committed by GitHub
parent 147afb448b
commit f56d2996ca
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/v1/worker/gpu_model_runner.py
View File

@ -2270,8 +2270,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
# Only rank 0 should print progress bar during capture # Only rank 0 should print progress bar during capture
compilation_cases = reversed(self.cudagraph_batch_sizes) compilation_cases = reversed(self.cudagraph_batch_sizes)
if is_global_first_rank(): if is_global_first_rank():
compilation_cases = tqdm(list(compilation_cases), compilation_cases = tqdm(
desc="Capturing CUDA graph shapes") list(compilation_cases),
disable=not self.load_config.use_tqdm_on_load,
desc="Capturing CUDA graph shapes")
for num_tokens in compilation_cases: for num_tokens in compilation_cases:
# We skip EPLB here since we don't want to record dummy metrics # We skip EPLB here since we don't want to record dummy metrics
for _ in range( for _ in range(

1
vllm/worker/model_runner.py
View File

@ -1587,6 +1587,7 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
if get_tensor_model_parallel_rank() == 0: if get_tensor_model_parallel_rank() == 0:
compilation_cases = tqdm( compilation_cases = tqdm(
list(compilation_cases), list(compilation_cases),
disable=not self.load_config.use_tqdm_on_load,
desc="Capturing CUDA graph shapes") desc="Capturing CUDA graph shapes")
for batch_size, use_inputs_embeds in compilation_cases: for batch_size, use_inputs_embeds in compilation_cases:
attn_metadata = ( attn_metadata = (