xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-16 04:45:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
Nicolò Lucchesi 2025-03-18 15:34:45 +01:00 committed by GitHub
parent 3b457143d2
commit af35d3a3cc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/v1/worker/tpu_model_runner.py
View File

@ -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)