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] Add model states [1/N] (#35350)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 26, 2026 at 11:20 UTC
c66aa48e993b74c46f83261654827b7349b2208c
4 files changed
+119
-101
vllm/v1/worker/gpu/cudagraph_utils.py
+23
-36
@@ -22,6 +22,7 @@ from vllm.v1.worker.gpu.attn_utils import (
22
from vllm.v1.worker.gpu.block_table import BlockTables
23
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
24
from vllm.v1.worker.gpu.input_batch import InputBuffers
25
+from vllm.v1.worker.gpu.model_states import ModelState
26
from vllm.v1.worker.utils import AttentionGroup
27
28
@@ -29,13 +30,11 @@ class CudaGraphManager:
30
def __init__(
31
self,
32
vllm_config: VllmConfig,
32
- uses_mrope: bool,
33
use_aux_hidden_state_outputs: bool,
34
device: torch.device,
35
):
36
self.vllm_config = vllm_config
37
self.scheduler_config = vllm_config.scheduler_config
38
- self.uses_mrope = uses_mrope
38
self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
39
self.device = device
40
@@ -88,8 +87,8 @@ class CudaGraphManager:
87
num_tokens: int,
88
capture_cg_mode: CUDAGraphMode,
89
model: nn.Module,
90
+ model_state: ModelState,
91
input_buffers: InputBuffers,
92
- mrope_positions: torch.Tensor | None,
92
inputs_embeds: torch.Tensor | None,
93
block_tables: BlockTables,
94
attn_groups: list[list[AttentionGroup]],
@@ -113,13 +112,18 @@ class CudaGraphManager:
112
)
113
else:
114
num_reqs = min(num_tokens, self.max_num_reqs)
116
- input_ids = input_buffers.input_ids[:num_tokens]
117
- positions = input_buffers.positions[:num_tokens]
118
- if self.uses_mrope:
119
- assert mrope_positions is not None
120
- positions = mrope_positions[:, :num_tokens]
121
- if inputs_embeds is not None:
122
- inputs_embeds = inputs_embeds[:num_tokens]
115
+
116
+ model_inputs = {
117
+ "input_ids": input_buffers.input_ids[:num_tokens],
118
+ "positions": input_buffers.positions[:num_tokens],
119
+ "inputs_embeds": (
120
+ inputs_embeds[:num_tokens] if inputs_embeds is not None else None
121
+ ),
122
+ # NOTE: Values returned by `prepare_dummy_inputs` will override the
123
+ # default values above.
124
+ **model_state.prepare_dummy_inputs(num_reqs, num_tokens),
125
+ }
126
+
127
attn_metadata, slot_mappings = prepare_inputs_to_capture(
128
num_reqs,
129
num_tokens,
@@ -143,11 +147,7 @@ class CudaGraphManager:
147
num_tokens_across_dp=num_tokens_across_dp,
148
slot_mapping=slot_mappings,
149
):
146
- model_output = model(
147
- input_ids=input_ids,
148
- positions=positions,
149
- inputs_embeds=inputs_embeds,
150
- )
150
+ model_output = model(**model_inputs)
151
if self.use_aux_hidden_state_outputs:
152
hidden_states, aux_hidden_states = model_output
153
else:
@@ -164,9 +164,7 @@ class CudaGraphManager:
164
num_tokens=num_tokens,
165
num_reqs=num_reqs,
166
model=model,
167
- input_ids=input_ids,
168
- positions=positions,
169
- inputs_embeds=inputs_embeds,
167
+ model_inputs=model_inputs,
168
num_tokens_across_dp=num_tokens_across_dp,
169
attn_metadata=attn_metadata,
170
slot_mappings=slot_mappings,
@@ -178,9 +176,7 @@ class CudaGraphManager:
176
num_tokens: int,
177
num_reqs: int,
178
model: nn.Module,
181
- input_ids: torch.Tensor,
182
- positions: torch.Tensor,
183
- inputs_embeds: torch.Tensor | None,
179
+ model_inputs: dict[str, torch.Tensor | None],
180
num_tokens_across_dp: torch.Tensor,
181
attn_metadata: dict[str, Any] | None,
182
slot_mappings: dict[str, torch.Tensor] | None,
@@ -206,11 +202,8 @@ class CudaGraphManager:
202
),
203
torch.cuda.graph(graph, self.pool),
204
):
209
- model_output = model(
210
- input_ids=input_ids,
211
- positions=positions,
212
- inputs_embeds=inputs_embeds,
213
- )
205
+ model_output = model(**model_inputs)
206
+
207
# Join offloader's copy stream after forward to avoid unjoined
208
# stream error. The last layer's start_prefetch forks copy_stream,
209
# but wait_prefetch only happens in the next forward pass.
@@ -235,9 +228,7 @@ class CudaGraphManager:
228
num_tokens: int,
229
num_reqs: int,
230
model: nn.Module,
238
- input_ids: torch.Tensor,
239
- positions: torch.Tensor,
240
- inputs_embeds: torch.Tensor | None,
231
+ model_inputs: dict[str, torch.Tensor | None],
232
num_tokens_across_dp: torch.Tensor,
233
attn_metadata: dict[str, Any] | None,
234
slot_mappings: dict[str, torch.Tensor] | None,
@@ -256,18 +247,14 @@ class CudaGraphManager:
247
batch_descriptor=batch_descriptor,
248
slot_mapping=slot_mappings,
249
):
259
- model(
260
- input_ids=input_ids,
261
- positions=positions,
262
- inputs_embeds=inputs_embeds,
263
- )
250
+ model(**model_inputs)
251
252
@torch.inference_mode()
253
def capture(
254
self,
255
model: nn.Module,
256
+ model_state: ModelState,
257
input_buffers: InputBuffers,
270
- mrope_positions: torch.Tensor | None,
258
inputs_embeds: torch.Tensor | None,
259
block_tables: BlockTables,
260
attn_groups: list[list[AttentionGroup]],
@@ -278,8 +265,8 @@ class CudaGraphManager:
265
device=self.device,
266
capture_fn=self.capture_graph,
267
model=model,
268
+ model_state=model_state,
269
input_buffers=input_buffers,
282
- mrope_positions=mrope_positions,
270
inputs_embeds=inputs_embeds,
271
block_tables=block_tables,
272
attn_groups=attn_groups,
vllm/v1/worker/gpu/input_batch.py
-3
@@ -65,8 +65,6 @@ class InputBatch:
65
input_ids: torch.Tensor
66
# [num_tokens_after_padding]
67
positions: torch.Tensor
68
- # [3, num_tokens_after_padding]
69
- mrope_positions: torch.Tensor | None
68
# [num_tokens_after_padding, hidden_size]
69
inputs_embeds: torch.Tensor | None
70
@@ -143,7 +141,6 @@ class InputBatch:
141
seq_lens=seq_lens,
142
input_ids=input_ids,
143
positions=positions,
146
- mrope_positions=None,
144
inputs_embeds=None,
145
attn_metadata=None, # type: ignore
146
slot_mappings=None, # type: ignore
vllm/v1/worker/gpu/model_runner.py
+22
-62
@@ -77,7 +77,7 @@ from vllm.v1.worker.gpu.kv_connector import (
77
)
78
from vllm.v1.worker.gpu.lora_utils import LoraState
79
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
80
-from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
80
+from vllm.v1.worker.gpu.model_states import ModelState
81
from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
82
from vllm.v1.worker.gpu.sample.output import SamplerOutput
83
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
@@ -140,14 +140,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
140
dtype=self.dtype,
141
device=self.device,
142
)
143
- self.uses_mrope = self.model_config.uses_mrope
144
- if self.uses_mrope:
145
- self.mrope_states = MRopeState(
146
- max_num_reqs=self.max_num_reqs,
147
- max_num_tokens=self.max_num_tokens,
148
- max_model_len=self.max_model_len,
149
- device=self.device,
150
- )
143
144
self.use_async_scheduling = self.scheduler_config.async_scheduling
145
self.output_copy_stream = torch.cuda.Stream(self.device)
@@ -212,7 +204,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
204
# CUDA graphs.
205
self.cudagraph_manager = CudaGraphManager(
206
self.vllm_config,
215
- self.uses_mrope,
207
self.use_aux_hidden_state_outputs,
208
self.device,
209
)
@@ -271,6 +262,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
262
if self.speculator is not None:
263
prepare_communication_buffer_for_model(self.speculator)
264
265
+ # Initialize the components that require the model.
266
+ self.model_state = ModelState(self.vllm_config, self.model, self.device)
267
+
268
def get_model(self) -> nn.Module:
269
return self.model
270
@@ -481,16 +475,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
475
start_free_gpu_memory = torch.cuda.mem_get_info()[0]
476
477
with self.maybe_setup_dummy_loras(self.lora_config):
484
- mrope_positions = None
485
- if self.uses_mrope:
486
- mrope_positions = self.mrope_states.mrope_positions
478
inputs_embeds = None
479
if self.supports_mm_inputs:
480
inputs_embeds = self.encoder_runner.inputs_embeds
481
self.cudagraph_manager.capture(
482
model=self.model,
483
+ model_state=self.model_state,
484
input_buffers=self.input_buffers,
493
- mrope_positions=mrope_positions,
485
inputs_embeds=inputs_embeds,
486
block_tables=self.block_tables,
487
attn_groups=self.attn_groups,
@@ -554,14 +545,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
545
if self.supports_mm_inputs:
546
self.encoder_runner.add_request(req_id, new_req_data.mm_features)
547
557
- # Pre-compute M-RoPE positions for prefill.
558
- if self.uses_mrope:
559
- self.mrope_states.init_prefill_mrope_positions(
560
- req_index,
561
- self.model, # type: ignore
562
- new_req_data.prefill_token_ids,
563
- mm_features=new_req_data.mm_features,
564
- )
548
+ self.model_state.add_request(req_index, new_req_data)
549
550
self.block_tables.append_block_ids(
551
req_index, new_req_data.block_ids, overwrite=True
@@ -577,8 +561,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
561
if scheduler_output.scheduled_new_reqs:
562
self.req_states.apply_staged_writes()
563
self.sampler.apply_staged_writes()
580
- if self.uses_mrope:
581
- self.mrope_states.apply_staged_writes()
564
+ self.model_state.apply_staged_writes()
565
566
def update_requests(self, scheduler_output: SchedulerOutput) -> None:
567
# Add new blocks for the existing requests.
@@ -692,15 +675,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
675
)
676
dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs]
677
695
- # Prepare M-RoPE positions.
696
- if self.uses_mrope:
697
- self.mrope_states.prepare_mrope_positions(
698
- idx_mapping,
699
- query_start_loc,
700
- self.req_states.prefill_len.gpu,
701
- self.req_states.num_computed_tokens.gpu,
702
- )
703
-
678
# Some input token ids are directly read from the last sampled tokens
679
# and draft tokens. Also, get the logits indices to sample tokens from.
680
logits_indices = combine_sampled_and_draft_tokens(
@@ -744,10 +718,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
718
719
input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
720
positions = self.input_buffers.positions[:num_tokens_after_padding]
747
- mrope_positions = None
748
- if self.uses_mrope:
749
- mrope_positions = self.mrope_states.mrope_positions
750
- mrope_positions = mrope_positions[:, :num_tokens_after_padding]
721
return InputBatch(
722
req_ids=req_ids,
723
num_reqs=num_reqs,
@@ -764,7 +734,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
734
seq_lens=seq_lens,
735
input_ids=input_ids,
736
positions=positions,
767
- mrope_positions=mrope_positions,
737
inputs_embeds=None,
738
attn_metadata=attn_metadata,
739
slot_mappings=slot_mappings_by_layer,
@@ -959,14 +928,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):
928
input_buffers=self.input_buffers,
929
device=self.device,
930
)
962
- if self.uses_mrope:
963
- input_batch.mrope_positions = self.mrope_states.mrope_positions[
964
- :, :num_tokens_after_padding
965
- ]
931
if not skip_attn_for_dummy_run:
932
self.prepare_dummy_attn_metadata(input_batch)
933
# FIXME(woosuk): Fix warmup for LoRA.
934
935
+ model_inputs = {
936
+ "input_ids": input_batch.input_ids,
937
+ "positions": input_batch.positions,
938
+ "inputs_embeds": input_batch.inputs_embeds,
939
+ # NOTE: Values returned by `prepare_inputs` will override the default
940
+ # values above.
941
+ **self.model_state.prepare_inputs(input_batch, self.req_states),
942
+ }
943
+ if not self.is_first_pp_rank:
944
+ # Update for non-first PP ranks.
945
+ model_inputs["input_ids"] = None
946
+ model_inputs["inputs_embeds"] = None
947
+ model_inputs["intermediate_tensors"] = intermediate_tensors
948
+
949
# Run model.
950
if cudagraph_runtime_mode == CUDAGraphMode.FULL:
951
# Use explicit cudagraph replay for FULL mode.
@@ -983,20 +962,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
962
aux_hidden_states = None
963
else:
964
# For piecewise and eager mode, just call model().
986
- positions = input_batch.positions
987
- if self.uses_mrope:
988
- assert input_batch.mrope_positions is not None
989
- positions = input_batch.mrope_positions
990
-
991
- if self.is_first_pp_rank:
992
- input_ids = input_batch.input_ids
993
- inputs_embeds = input_batch.inputs_embeds
994
- assert intermediate_tensors is None
995
- else:
996
- input_ids = None
997
- inputs_embeds = None
998
- assert intermediate_tensors is not None
999
-
965
batch_descriptor = BatchDescriptor(
966
num_tokens=input_batch.num_tokens_after_padding,
967
has_lora=self.lora_config is not None,
@@ -1012,12 +977,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
977
slot_mapping=input_batch.slot_mappings,
978
):
979
self.kv_connector.pre_forward(scheduler_output)
1015
- model_output = self.model(
1016
- input_ids=input_ids,
1017
- positions=positions,
1018
- inputs_embeds=inputs_embeds,
1019
- intermediate_tensors=intermediate_tensors,
1020
- )
980
+ model_output = self.model(**model_inputs)
981
if self.use_aux_hidden_state_outputs:
982
hidden_states, aux_hidden_states = model_output
983
else:
vllm/v1/worker/gpu/model_states.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+import torch.nn as nn
5
+
6
+from vllm.config import VllmConfig
7
+from vllm.v1.core.sched.output import NewRequestData
8
+from vllm.v1.worker.gpu.input_batch import InputBatch
9
+from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
10
+from vllm.v1.worker.gpu.states import RequestState
11
+
12
+
13
+class ModelState:
14
+ def __init__(self, vllm_config: VllmConfig, model: nn.Module, device: torch.device):
15
+ self.vllm_config = vllm_config
16
+ self.model_config = vllm_config.model_config
17
+ self.scheduler_config = vllm_config.scheduler_config
18
+ self.model = model
19
+ self.device = device
20
+
21
+ self.max_model_len = self.model_config.max_model_len
22
+ self.max_num_reqs = self.scheduler_config.max_num_seqs
23
+ self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
24
+
25
+ self.uses_mrope = self.model_config.uses_mrope
26
+ if self.uses_mrope:
27
+ self.mrope_state = MRopeState(
28
+ max_num_reqs=self.max_num_reqs,
29
+ max_num_tokens=self.max_num_tokens,
30
+ max_model_len=self.max_model_len,
31
+ device=self.device,
32
+ )
33
+
34
+ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
35
+ if self.uses_mrope:
36
+ # Pre-compute M-RoPE positions for prefill.
37
+ assert new_req_data.prefill_token_ids is not None
38
+ self.mrope_state.init_prefill_mrope_positions(
39
+ req_index,
40
+ self.model, # type: ignore
41
+ new_req_data.prefill_token_ids,
42
+ mm_features=new_req_data.mm_features,
43
+ )
44
+
45
+ def apply_staged_writes(self) -> None:
46
+ if self.uses_mrope:
47
+ self.mrope_state.apply_staged_writes()
48
+
49
+ def prepare_inputs(
50
+ self, input_batch: InputBatch, req_states: RequestState
51
+ ) -> dict[str, torch.Tensor | None]:
52
+ if not self.uses_mrope:
53
+ # Common case (1D positions).
54
+ return {}
55
+
56
+ # Prepare M-RoPE positions.
57
+ self.mrope_state.prepare_mrope_positions(
58
+ input_batch.idx_mapping,
59
+ input_batch.query_start_loc,
60
+ req_states.prefill_len.gpu,
61
+ req_states.num_computed_tokens.gpu,
62
+ )
63
+ mrope_positions = self.mrope_state.mrope_positions[
64
+ :, : input_batch.num_tokens_after_padding
65
+ ]
66
+ return {"positions": mrope_positions}
67
+
68
+ def prepare_dummy_inputs(
69
+ self, num_reqs: int, num_tokens: int
70
+ ) -> dict[str, torch.Tensor | None]:
71
+ if not self.uses_mrope:
72
+ return {}
73
+ mrope_positions = self.mrope_state.mrope_positions[:, :num_tokens]
74
+ return {"positions": mrope_positions}