modify tpu model runner

vllm-project · Feb 12, 2025 · d08e04a · d08e04a
1 parent 78c6be2
commit d08e04a
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/vllm/worker/tpu_model_runner.py b/vllm/worker/tpu_model_runner.py
@@ -839,12 +839,14 @@ def forward(
             if seq_len > 1:
                 # prefill
                 head_indicies *= num_blocks * block_size // MIN_PREFILL_SEQ_LEN
+                slot_mapping = slot_mapping.repeat(num_kv_heads, 1)
+                slot_mapping = slot_mapping + head_indicies.view(-1, 1)
             else:
                 # decoding
                 head_indicies *= num_blocks * block_size
-            slot_mapping = slot_mapping.repeat_interleave(num_kv_heads).view(
-                -1, num_kv_heads)
-            slot_mapping = slot_mapping + head_indicies.view(1, -1)
+                slot_mapping = slot_mapping.repeat_interleave(num_kv_heads).view(
+                    -1, num_kv_heads)
+                slot_mapping = slot_mapping + head_indicies.view(1, -1)
             slot_mapping = slot_mapping.flatten()
             attn_metadata.slot_mapping = slot_mapping