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, 07:15 UTC
[BugFix] Fix tpu_model_runner block_id concatenation (#19228)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Jun 6, 2025 at 16:28 UTC
46ecc579733f13a555ca42da76c1234c586271eb
3 files changed
+11
-4
tests/v1/tpu/worker/test_tpu_model_runner.py
+1
-1
@@ -226,7 +226,7 @@ def test_update_states_request_resumed(model_runner):
226
req_id=req_id,
227
resumed_from_preemption=False,
228
new_token_ids=[],
229
- new_block_ids=[],
229
+ new_block_ids=[[]],
230
num_computed_tokens=0,
231
)
232
vllm/v1/worker/gpu_model_runner.py
+5
-2
@@ -460,8 +460,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
460
# Update the block IDs.
461
if not req_data.resumed_from_preemption:
462
# Append the new blocks to the existing block IDs.
463
- for i in range(len(self.kv_cache_config.kv_cache_groups)):
464
- req_state.block_ids[i].extend(req_data.new_block_ids[i])
463
+ for block_ids, new_block_ids in zip( # type: ignore[call-overload]
464
+ req_state.block_ids,
465
+ req_data.new_block_ids,
466
+ strict=True):
467
+ block_ids.extend(new_block_ids)
468
else:
469
# The request is resumed from preemption.
470
# Replace the existing block IDs with the new ones.
vllm/v1/worker/tpu_model_runner.py
+5
-1
@@ -413,7 +413,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
413
req_state.num_computed_tokens = req_data.num_computed_tokens
414
if not req_data.resumed_from_preemption:
415
# Append the new blocks to the existing block IDs.
416
- req_state.block_ids.extend(req_data.new_block_ids)
416
+ for block_ids, new_block_ids in zip( # type: ignore[call-overload]
417
+ req_state.block_ids,
418
+ req_data.new_block_ids,
419
+ strict=True):
420
+ block_ids.extend(new_block_ids)
421
else:
422
# The request is resumed from preemption.
423
# Replace the existing block IDs with the new ones.