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 10, 00:15 UTC
[Model Runner V2] Refactor `update_states` (#32562)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 18, 2026 at 17:32 UTC
9a1f16da1e423ede2c2f52a9850cbfbb39cefe96
1 file changed
+24
-17
vllm/v1/worker/gpu/model_runner.py
+24
-17
@@ -425,7 +425,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
425
self._dummy_run(self.max_num_tokens, skip_attn=False)
426
torch.cuda.synchronize()
427
428
- def update_states(self, scheduler_output: SchedulerOutput) -> None:
428
+ def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
429
if scheduler_output.preempted_req_ids is not None:
430
for req_id in scheduler_output.preempted_req_ids:
431
self.req_states.remove_request(req_id)
@@ -436,11 +436,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
436
if self.supports_mm_inputs:
437
self.encoder_runner.remove_request(req_id)
438
439
+ def free_states(self, scheduler_output: SchedulerOutput) -> None:
440
if self.supports_mm_inputs:
441
for mm_hash in scheduler_output.free_encoder_mm_hashes:
442
self.encoder_runner.free_encoder_cache(mm_hash)
443
443
- # Add new requests.
444
+ def add_requests(self, scheduler_output: SchedulerOutput) -> None:
445
for new_req_data in scheduler_output.scheduled_new_reqs:
446
assert new_req_data.prompt_token_ids is not None
447
assert new_req_data.prefill_token_ids is not None
@@ -476,6 +477,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
477
req_index, prompt_len, new_req_data.sampling_params
478
)
479
480
+ if scheduler_output.scheduled_new_reqs:
481
+ self.req_states.apply_staged_writes()
482
+ self.sampler.apply_staged_writes(
483
+ self.req_states.prefill_token_ids.gpu,
484
+ self.req_states.prefill_len.np,
485
+ self.req_states.prompt_len,
486
+ )
487
+ if self.uses_mrope:
488
+ self.mrope_states.apply_staged_writes()
489
+
490
+ def update_requests(self, scheduler_output: SchedulerOutput) -> None:
491
# Add new blocks for the existing requests.
492
cached_reqs = scheduler_output.scheduled_cached_reqs
493
for i, req_id in enumerate(cached_reqs.req_ids):
@@ -486,16 +498,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
498
req_index, req_new_block_ids, overwrite=False
499
)
500
489
- self.req_states.apply_staged_writes()
490
- self.block_tables.apply_staged_writes()
491
- self.sampler.apply_staged_writes(
492
- self.req_states.prefill_token_ids.gpu,
493
- self.req_states.prefill_len.np,
494
- self.req_states.prompt_len,
495
- )
496
- if self.uses_mrope:
497
- self.mrope_states.apply_staged_writes()
498
-
501
def prepare_inputs(
502
self,
503
scheduler_output: SchedulerOutput,
@@ -951,15 +953,20 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
953
dummy_run: bool = False,
954
) -> ModelRunnerOutput | None:
955
assert intermediate_tensors is None
954
- if scheduler_output.total_num_scheduled_tokens == 0 and not dummy_run:
955
- # No need to run the model.
956
- self.update_states(scheduler_output)
957
- return EMPTY_MODEL_RUNNER_OUTPUT
956
+ if not dummy_run:
957
+ # Update the request states.
958
+ self.finish_requests(scheduler_output)
959
+ self.free_states(scheduler_output)
960
+ self.add_requests(scheduler_output)
961
+ self.update_requests(scheduler_output)
962
+ self.block_tables.apply_staged_writes()
963
+ if scheduler_output.total_num_scheduled_tokens == 0:
964
+ # No need to run the model.
965
+ return EMPTY_MODEL_RUNNER_OUTPUT
966
967
cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp = (
968
self.get_cudagraph_and_dp_padding(scheduler_output)
969
)
962
- self.update_states(scheduler_output)
970
if num_tokens_after_padding == 0:
971
# All DP ranks have zero tokens to run.
972
return EMPTY_MODEL_RUNNER_OUTPUT