youngkingdom/vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[BugFix] Fix tpu_model_runner block_id concatenation (#19228)

Browse Source 
Signed-off-by: Nick Hill <nhill@redhat.com>
This commit is contained in:
Nick Hill
2025-06-06 16:28:17 -07:00
committed by GitHub
parent b6a3a9f76d
commit 46ecc57973
3 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/v1/tpu/worker/test_tpu_model_runner.py
View File

@ -226,7 +226,7 @@ def test_update_states_request_resumed(model_runner):
req_id=req_id,
resumed_from_preemption=False,
new_token_ids=[],
new_block_ids=[],
new_block_ids=[[]],
num_computed_tokens=0,
)

7
vllm/v1/worker/gpu_model_runner.py
View File

@ -460,8 +460,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
# Update the block IDs.
if not req_data.resumed_from_preemption:
# Append the new blocks to the existing block IDs.
for i in range(len(self.kv_cache_config.kv_cache_groups)):
req_state.block_ids[i].extend(req_data.new_block_ids[i])
for block_ids, new_block_ids in zip( # type: ignore[call-overload]
req_state.block_ids,
req_data.new_block_ids,
strict=True):
block_ids.extend(new_block_ids)
else:
# The request is resumed from preemption.
# Replace the existing block IDs with the new ones.

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

@ -413,7 +413,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
req_state.num_computed_tokens = req_data.num_computed_tokens
if not req_data.resumed_from_preemption:
# Append the new blocks to the existing block IDs.
req_state.block_ids.extend(req_data.new_block_ids)
for block_ids, new_block_ids in zip( # type: ignore[call-overload]
req_state.block_ids,
req_data.new_block_ids,
strict=True):
block_ids.extend(new_block_ids)
else:
# The request is resumed from preemption.
# Replace the existing block IDs with the new ones.