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 3, 22:45 UTC
[Model Runner V2] Fix `seq_lens_cpu_upper_bound` (#42202)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
May 11, 2026 at 10:37 UTC
9af6a5ed75bbb095c3e7e78b807ca0fe29ef90bb
2 files changed
+18
-35
vllm/v1/worker/gpu/model_runner.py
+15
-31
@@ -712,15 +712,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
712
self.sampler.apply_staged_writes()
713
714
def update_requests(self, scheduler_output: SchedulerOutput) -> None:
715
- # Add new blocks for the existing requests.
715
+ # Add new blocks and update num_computed_tokens for the existing requests.
716
reqs = scheduler_output.scheduled_cached_reqs
717
- for req_new_block_ids, req_id in zip(reqs.new_block_ids, reqs.req_ids):
717
+ num_computed_tokens_np = self.req_states.num_computed_tokens_np
718
+ for req_id, num_computed_tokens, req_new_block_ids in zip(
719
+ reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids
720
+ ):
721
+ req_index = self.req_states.req_id_to_index[req_id]
722
+ num_computed_tokens_np[req_index] = num_computed_tokens
723
if req_new_block_ids is not None:
719
- req_index = self.req_states.req_id_to_index[req_id]
724
self.block_tables.append_block_ids(
725
req_index, req_new_block_ids, overwrite=False
726
)
727
728
+ # Update num_computed_prefill_tokens.
729
+ np.minimum(
730
+ self.req_states.num_computed_tokens_np,
731
+ self.req_states.prefill_len.np,
732
+ out=self.req_states.num_computed_prefill_tokens,
733
+ )
734
+
735
def prepare_inputs(
736
self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor
737
) -> InputBatch:
@@ -787,10 +798,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
798
async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
799
query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
800
query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]
790
- is_prefilling_np = (
791
- self.req_states.num_computed_prefill_tokens[idx_mapping_np]
792
- < self.req_states.prefill_len.np[idx_mapping_np]
793
- )
801
+ is_prefilling_np = self.req_states.is_prefilling(idx_mapping_np)
802
803
# Get prefill tokens if any.
804
if np.any(is_prefilling_np):
@@ -973,18 +981,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
981
self.req_states.total_len.gpu,
982
)
983
976
- # Update the number of computed prefill tokens.
977
- idx_mapping_np = input_batch.idx_mapping_np
978
- computed_prefill = self.req_states.num_computed_prefill_tokens
979
- computed_prefill[idx_mapping_np] += input_batch.num_scheduled_tokens
980
- np.minimum(
981
- computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
982
- )
983
- # Advance the CPU mirror optimistically (assume all scheduled accepted).
984
- self.req_states.num_computed_tokens_np[idx_mapping_np] += (
985
- input_batch.num_scheduled_tokens
986
- )
987
-
984
self.model_state.postprocess_state(input_batch, num_sampled)
985
986
@torch.inference_mode()
@@ -1357,18 +1353,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1353
input_batch.query_start_loc,
1354
)
1355
1360
- # Update the number of computed prefill tokens.
1361
- idx_mapping_np = input_batch.idx_mapping_np
1362
- computed_prefill = self.req_states.num_computed_prefill_tokens
1363
- computed_prefill[idx_mapping_np] += input_batch.num_scheduled_tokens
1364
- np.minimum(
1365
- computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
1366
- )
1367
- # Advance the CPU mirror optimistically (assume all scheduled accepted).
1368
- self.req_states.num_computed_tokens_np[idx_mapping_np] += (
1369
- input_batch.num_scheduled_tokens
1370
- )
1371
-
1356
def shutdown(self) -> None:
1357
"""Release GPU tensors (model weights, KV caches, workspace) so that
1358
memory is reclaimable when running in the same process."""
vllm/v1/worker/gpu/states.py
+3
-4
@@ -104,9 +104,8 @@ class RequestState:
104
self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
105
self.num_computed_tokens_np[req_idx] = num_computed_tokens
106
self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)
107
- self.num_computed_tokens_np[req_idx] = num_computed_tokens
107
109
- if num_computed_tokens > 0 and num_computed_tokens <= prefill_len:
108
+ if 0 < num_computed_tokens <= prefill_len:
109
# For PD disagg or resumed requests: set last_sampled to the last
110
# computed token so the first decode step gets the right input_id.
111
# For fresh prefill requests (num_computed_tokens == 0) the tensor
@@ -134,8 +133,8 @@ class RequestState:
133
self.free_indices.append(req_idx)
134
return True
135
137
- def any_prefills(self, idx_mapping_np: np.ndarray) -> bool:
138
- return np.any(
136
+ def is_prefilling(self, idx_mapping_np: np.ndarray) -> np.ndarray:
137
+ return (
138
self.num_computed_prefill_tokens[idx_mapping_np]
139
< self.prefill_len.np[idx_mapping_np]
140
)