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 7, 00:15 UTC
[Model Runner V2] Use NamedTuple for `execute_model_state` (#35930)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 9, 2026 at 11:17 UTC
10a5f4d53d0dc7390802ad99bf5d27b2423094e9
1 file changed
+38
-34
vllm/v1/worker/gpu/model_runner.py
+38
-34
@@ -21,6 +21,7 @@ import functools
21
import gc
22
import time
23
from copy import deepcopy
24
+from typing import Any, NamedTuple
25
26
import numpy as np
27
import torch
@@ -44,7 +45,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
45
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
46
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
47
from vllm.v1.kv_cache_interface import KVCacheConfig
47
-from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
48
+from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
49
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
50
from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
51
from vllm.v1.worker.gpu.attn_utils import (
@@ -213,7 +214,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
214
self.pooling_runner: PoolingRunner | None = None
215
216
# For transferring state from execute_model to subsequent sample_tokens call.
216
- self.execute_model_state: tuple | None = None
217
+ self.execute_model_state: ExecuteModelState | None = None
218
219
def update_max_model_len(self, max_model_len: int) -> None:
220
self.max_model_len = max_model_len
@@ -375,16 +376,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
376
return None, None
377
378
assert self.execute_model_state is not None
378
- (
379
- input_batch,
380
- model_inputs,
381
- attn_metadata,
382
- slot_mappings_by_layer,
383
- hidden_states,
384
- aux_hidden_states,
385
- kv_connector_output,
386
- num_tokens_across_dp,
387
- ) = self.execute_model_state
379
+ input_batch = self.execute_model_state.input_batch
380
+ attn_metadata = self.execute_model_state.attn_metadata
381
+ slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
382
+ hidden_states = self.execute_model_state.hidden_states
383
+ aux_hidden_states = self.execute_model_state.aux_hidden_states
384
+ num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp
385
self.execute_model_state = None
386
387
# dummy run the eagle speculator's propose to ensure DP/EP sync.
@@ -989,15 +986,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
986
aux_hidden_states = None
987
988
kv_connector_output = self.kv_connector.post_forward(scheduler_output)
992
- self.execute_model_state = (
993
- input_batch,
994
- model_inputs,
995
- attn_metadata,
996
- slot_mappings_by_layer,
997
- hidden_states,
998
- aux_hidden_states,
999
- kv_connector_output,
1000
- num_tokens_across_dp,
989
+ self.execute_model_state = ExecuteModelState(
990
+ input_batch=input_batch,
991
+ attn_metadata=attn_metadata,
992
+ slot_mappings_by_layer=slot_mappings_by_layer,
993
+ hidden_states=hidden_states,
994
+ aux_hidden_states=aux_hidden_states,
995
+ kv_connector_output=kv_connector_output,
996
+ num_tokens_across_dp=num_tokens_across_dp,
997
)
998
999
if not self.is_last_pp_rank:
@@ -1016,16 +1012,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1012
if self.execute_model_state is None:
1013
# The prior execute_model call must have failed.
1014
return None
1019
- (
1020
- input_batch,
1021
- model_inputs,
1022
- attn_metadata,
1023
- slot_mappings_by_layer,
1024
- hidden_states,
1025
- aux_hidden_states,
1026
- kv_connector_output,
1027
- num_tokens_across_dp,
1028
- ) = self.execute_model_state
1015
+
1016
+ input_batch = self.execute_model_state.input_batch
1017
+ attn_metadata = self.execute_model_state.attn_metadata
1018
+ slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
1019
+ hidden_states = self.execute_model_state.hidden_states
1020
+ aux_hidden_states = self.execute_model_state.aux_hidden_states
1021
+ kv_connector_output = self.execute_model_state.kv_connector_output
1022
+ num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp
1023
self.execute_model_state = None
1024
1025
if not self.is_last_pp_rank:
@@ -1116,9 +1110,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1110
# The prior execute_model call must have failed.
1111
return None
1112
1119
- input_batch, _, _, _, hidden_states, _, kv_connector_output, _ = (
1120
- self.execute_model_state
1121
- )
1113
+ input_batch = self.execute_model_state.input_batch
1114
+ hidden_states = self.execute_model_state.hidden_states
1115
+ kv_connector_output = self.execute_model_state.kv_connector_output
1116
self.execute_model_state = None
1117
1118
if not self.is_last_pp_rank:
@@ -1164,3 +1158,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1158
np.minimum(
1159
computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
1160
)
1161
+
1162
+
1163
+class ExecuteModelState(NamedTuple):
1164
+ input_batch: InputBatch
1165
+ attn_metadata: dict[str, Any] | None
1166
+ slot_mappings_by_layer: dict[str, torch.Tensor] | None
1167
+ hidden_states: torch.Tensor | IntermediateTensors
1168
+ aux_hidden_states: list[torch.Tensor] | None
1169
+ kv_connector_output: KVConnectorOutput | None
1170
+ num_tokens_across_dp: torch.Tensor | None