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, 18:30 UTC
[TPU][V1][Bugfix] Fix chunked prefill with padding (#15037)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Mar 18, 2025 at 15:34 UTC
af35d3a3ccb502dcadacb5e24c0d3c4bc8e8657c
1 file changed
+3
vllm/v1/worker/tpu_model_runner.py
+3
@@ -410,6 +410,9 @@ class TPUModelRunner:
410
# Do the padding and copy the tensors to the TPU.
411
padded_total_num_scheduled_tokens = _get_padded_token_len(
412
total_num_scheduled_tokens)
413
+ # Zero out to avoid spurious values from prev iteration (last cp chunk)
414
+ self.input_ids_cpu[
415
+ total_num_scheduled_tokens:padded_total_num_scheduled_tokens] = 0
416
self.input_ids = self.input_ids_cpu[:
417
padded_total_num_scheduled_tokens].to(
418
self.device)