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 15, 06:00 UTC
[UX] Add Feedback During CUDAGraph Capture (#19501)
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Robert Shaw committed
Jun 11, 2025 at 14:09 UTC
97a9465bbca114c223da604f5e25da90f9f16bc1
1 file changed
+4
-1
vllm/v1/worker/gpu_model_runner.py
+4
-1
@@ -12,6 +12,7 @@ import numpy as np
12
import torch
13
import torch.distributed
14
import torch.nn as nn
15
+from tqdm import tqdm
16
17
import vllm.envs as envs
18
from vllm.attention import AttentionType, get_attn_backend
@@ -2034,7 +2035,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2035
# can reuse the memory pool allocated for the large shapes.
2036
with graph_capture(device=self.device):
2037
skip_attn = not self.vllm_config.compilation_config.full_cuda_graph
2037
- for num_tokens in reversed(self.cudagraph_batch_sizes):
2038
+ for num_tokens in tqdm(reversed(self.cudagraph_batch_sizes),
2039
+ desc="Capturing CUDA graphs",
2040
+ total=len(self.cudagraph_batch_sizes)):
2041
for _ in range(self.vllm_config.compilation_config.
2042
cudagraph_num_of_warmups):
2043
self._dummy_run(num_tokens, skip_attn=skip_attn)