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, 16:30 UTC
[Model Runner V2] Fix error-handling (#35063)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 26, 2026 at 11:00 UTC
b6d5a17298548e77cf5af456e029e5beb26b253c
1 file changed
+9
-3
vllm/v1/worker/gpu/model_runner.py
+9
-3
@@ -227,6 +227,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
227
# KV Connector if configured.
228
self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
229
230
+ # For transferring state from execute_model to subsequent sample_tokens call.
231
+ self.execute_model_state: tuple | None = None
232
+
233
def update_max_model_len(self, max_model_len: int) -> None:
234
self.max_model_len = max_model_len
235
self.req_states.max_model_len = max_model_len
@@ -388,6 +391,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
391
392
assert self.execute_model_state is not None
393
hidden_states, _, input_batch, _ = self.execute_model_state
394
+ self.execute_model_state = None
395
assert hidden_states is not None # Last PP rank always has hidden_states
396
sample_hidden_states = hidden_states[input_batch.logits_indices]
397
return hidden_states, sample_hidden_states
@@ -1036,18 +1040,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1040
aux_hidden_states,
1041
input_batch,
1042
kv_connector_output,
1039
- ) # type: ignore
1043
+ )
1044
return None
1045
1046
@torch.inference_mode()
1047
def sample_tokens(
1048
self, grammar_output: GrammarOutput | None
1049
) -> AsyncOutput | ModelRunnerOutput | None:
1046
- assert self.execute_model_state is not None
1050
+ if self.execute_model_state is None:
1051
+ # The prior execute_model call must have failed.
1052
+ return None
1053
hidden_states, aux_hidden_states, input_batch, kv_connector_output = (
1054
self.execute_model_state
1055
)
1050
- self.execute_model_state = None # type: ignore
1056
+ self.execute_model_state = None
1057
1058
if not self.is_last_pp_rank:
1059
# Non-last PP rank: hidden_states is None because this rank produced