Skip to content

GitLab

Unverified Commit af35d3a3 authored by Nicolò Lucchesi's avatar Nicolò Lucchesi Committed by GitHub
Browse files

[TPU][V1][Bugfix] Fix chunked prefill with padding (#15037) 


Signed-off-by: default avatarNickLucche <nlucches@redhat.com>
parent 3b457143
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
vllm/v1/worker/tpu_model_runner.py
View file @ af35d3a3
...@@ -410,6 +410,9 @@ class TPUModelRunner: ...@@ -410,6 +410,9 @@ class TPUModelRunner:
# Do the padding and copy the tensors to the TPU. # Do the padding and copy the tensors to the TPU.
padded_total_num_scheduled_tokens = _get_padded_token_len( padded_total_num_scheduled_tokens = _get_padded_token_len(
total_num_scheduled_tokens) total_num_scheduled_tokens)
# Zero out to avoid spurious values from prev iteration (last cp chunk)
self.input_ids_cpu[
total_num_scheduled_tokens:padded_total_num_scheduled_tokens] = 0
self.input_ids = self.input_ids_cpu[: self.input_ids = self.input_ids_cpu[:
padded_total_num_scheduled_tokens].to( padded_total_num_scheduled_tokens].to(
self.device) self.device)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment