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, 11:00 UTC
[Bugfix][TPU] Use np array when updating cache slot_mapping (#17971)
Signed-off-by: Siyuan Liu <lsiyuan@google.com>
Siyuan Liu committed
May 11, 2025 at 21:58 UTC
430783018cbfad69c6ff3a52479bf7b556b65247
1 file changed
+1
-1
vllm/v1/worker/tpu_model_runner.py
+1
-1
@@ -531,7 +531,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
531
np.add(block_numbers * self.block_size,
532
block_offsets,
533
out=self.input_batch.block_table.
534
- slot_mapping_cpu[:total_num_scheduled_tokens])
534
+ slot_mapping_np[:total_num_scheduled_tokens])
535
536
# Prepare the attention metadata.
537
self.query_start_loc_np[0] = 0