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, 16:00 UTC
[V1][TPU] Do not compile sampling more than needed (#15883)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Apr 3, 2025 at 03:36 UTC
bd7599d34ae9aa46994f8e233e2293f6f2d13019
1 file changed
+3
-1
vllm/v1/worker/tpu_model_runner.py
+3
-1
@@ -862,7 +862,9 @@ class TPUModelRunner:
862
out = self.model.sample_from_hidden(dummy_hidden,
863
sampling_meta)
864
out = out.cpu()
865
- if num_reqs_to_sample >= self.max_num_reqs:
865
+ # Requests can't be more than tokens. But do compile for the
866
+ # next bigger value in case num_tokens uses bucketed padding.
867
+ if num_reqs_to_sample >= min(num_tokens, self.max_num_reqs):
868
break
869
# Make sure to compile the `max_num_reqs` upper-limit case
870
num_reqs_to_sample = _get_padded_num_reqs_with_upper_limit(