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 14, 05:30 UTC
[Core] Small simplification in `GPUModelRunner._update_states()` (#26508)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Oct 9, 2025 at 19:53 UTC
aafb99a4d4dbb3f0f8512ff21ab29a2a673e6c9d
1 file changed
+2
-7
vllm/v1/worker/gpu_model_runner.py
+2
-7
@@ -708,6 +708,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
708
# Update the cached states.
709
710
req_state.num_computed_tokens = num_computed_tokens
711
+ req_index = self.input_batch.req_id_to_index.get(req_id)
712
713
if not is_last_rank:
714
# When using PP, the scheduler sends the sampled tokens back,
@@ -728,19 +729,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
729
# Some output tokens were discarded due to a sync-KV-load
730
# failure. Align the cached state.
731
del req_state.output_token_ids[num_output_tokens:]
731
-
732
- req_index = self.input_batch.req_id_to_index.get(req_id)
732
if req_index is not None:
734
- old_end_idx = self.input_batch.num_tokens_no_spec[req_index]
733
end_idx = (
734
self.input_batch.num_prompt_tokens[req_index]
735
+ num_output_tokens
736
)
737
self.input_batch.num_tokens[req_index] = end_idx
738
self.input_batch.num_tokens_no_spec[req_index] = end_idx
741
- self.input_batch.is_token_ids[req_index, end_idx:old_end_idx] = (
742
- False
743
- )
739
740
# Update the block IDs.
741
if not resumed_from_preemption:
@@ -749,12 +744,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
744
for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
745
block_ids.extend(new_ids)
746
else:
747
+ assert req_index is None
748
assert new_block_ids is not None
749
# The request is resumed from preemption.
750
# Replace the existing block IDs with the new ones.
751
req_state.block_ids = new_block_ids
752
757
- req_index = self.input_batch.req_id_to_index.get(req_id)
753
if req_index is None:
754
# The request is not in the persistent batch.
755
# The request was either preempted and resumed later, or was not