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, 04:45 UTC
[Frontend] Add optional token-level progress bar to `LLM.beam_search` (#19301)
Signed-off-by: Ruosen Li <rxl190028@utdallas.edu> Signed-off-by: Aaron Pham <contact@aarnphm.xyz> Signed-off-by: Ubuntu <ubuntu@ip-172-31-71-179.ec2.internal> Co-authored-by: 22quinn <33176974+22quinn@users.noreply.github.com>
NekoMimiUnagi committed
Jun 19, 2025 at 02:21 UTC
466166dcfdc40f85f4043a94b9a53099af4a0850
1 file changed
+14
-1
vllm/entrypoints/llm.py
+14
-1
@@ -552,6 +552,7 @@ class LLM:
552
prompts: list[Union[TokensPrompt, TextPrompt]],
553
params: BeamSearchParams,
554
lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
555
+ use_tqdm: bool = False,
556
) -> list[BeamSearchOutput]:
557
"""
558
Generate sequences using beam search.
@@ -561,6 +562,7 @@ class LLM:
562
of token IDs.
563
params: The beam search parameters.
564
lora_request: LoRA request to use for generation, if any.
565
+ use_tqdm: Whether to use tqdm to display the progress bar.
566
"""
567
# TODO: how does beam search work together with length penalty,
568
# frequency, penalty, and stopping criteria, etc.?
@@ -623,7 +625,18 @@ class LLM:
625
**mm_kwargs,
626
), )
627
626
- for _ in range(max_tokens):
628
+ token_iter = range(max_tokens)
629
+ if use_tqdm:
630
+ token_iter = tqdm(token_iter,
631
+ desc="Beam search",
632
+ unit="token",
633
+ unit_scale=False)
634
+ logger.warning(
635
+ "The progress bar shows the upper bound on token steps and "
636
+ "may finish early due to stopping conditions. It does not "
637
+ "reflect instance-level progress.")
638
+
639
+ for _ in token_iter:
640
all_beams: list[BeamSearchSequence] = list(
641
sum((instance.beams for instance in instances), []))
642
pos = [0] + list(