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, 22:45 UTC
[Model Runner V2] Enable CUDA graph for Eagle3 (#35040)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 21, 2026 at 21:42 UTC
2cbf9656ce6013f7b531bc5a0909d03b88c14862
2 files changed
+44
-9
vllm/v1/worker/gpu/cudagraph_utils.py
+39
-7
@@ -25,10 +25,17 @@ from vllm.v1.worker.utils import AttentionGroup
25
26
27
class CudaGraphManager:
28
- def __init__(self, vllm_config: VllmConfig, uses_mrope: bool, device: torch.device):
28
+ def __init__(
29
+ self,
30
+ vllm_config: VllmConfig,
31
+ uses_mrope: bool,
32
+ use_aux_hidden_state_outputs: bool,
33
+ device: torch.device,
34
+ ):
35
self.vllm_config = vllm_config
36
self.scheduler_config = vllm_config.scheduler_config
37
self.uses_mrope = uses_mrope
38
+ self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
39
self.device = device
40
41
self.max_model_len = vllm_config.model_config.max_model_len
@@ -63,6 +70,7 @@ class CudaGraphManager:
70
if self.cudagraph_mode != CUDAGraphMode.NONE:
71
self.pool = torch.cuda.graph_pool_handle()
72
self.hidden_states: torch.Tensor | None = None
73
+ self.aux_hidden_states: list[torch.Tensor] = []
74
75
def needs_capture(self) -> bool:
76
return len(self.cudagraph_sizes) > 0
@@ -134,13 +142,22 @@ class CudaGraphManager:
142
num_tokens_across_dp=num_tokens_across_dp,
143
slot_mapping=slot_mappings,
144
):
137
- hidden_states = model(
145
+ model_output = model(
146
input_ids=input_ids,
147
positions=positions,
148
inputs_embeds=inputs_embeds,
149
)
142
- if self.hidden_states is None:
143
- self.hidden_states = torch.empty_like(hidden_states)
150
+ if self.use_aux_hidden_state_outputs:
151
+ hidden_states, aux_hidden_states = model_output
152
+ else:
153
+ hidden_states = model_output
154
+ aux_hidden_states = None
155
+
156
+ # Allocate output buffers if not already done.
157
+ if self.hidden_states is None:
158
+ self.hidden_states = torch.empty_like(hidden_states)
159
+ if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:
160
+ self.aux_hidden_states = [torch.empty_like(x) for x in aux_hidden_states]
161
162
capture_fn(
163
num_tokens=num_tokens,
@@ -183,13 +200,23 @@ class CudaGraphManager:
200
),
201
torch.cuda.graph(graph, self.pool),
202
):
186
- hidden_states = model(
203
+ model_output = model(
204
input_ids=input_ids,
205
positions=positions,
206
inputs_embeds=inputs_embeds,
207
)
208
+ if self.use_aux_hidden_state_outputs:
209
+ hidden_states, aux_hidden_states = model_output
210
+ else:
211
+ hidden_states = model_output
212
+ aux_hidden_states = None
213
+
214
+ # Copy outputs to the output buffers.
215
assert self.hidden_states is not None
216
self.hidden_states[:num_tokens] = hidden_states
217
+ if self.use_aux_hidden_state_outputs:
218
+ for i, aux_hidden in enumerate(aux_hidden_states):
219
+ self.aux_hidden_states[i][:num_tokens] = aux_hidden
220
self.graphs[num_tokens] = graph
221
222
def _capture_piecewise_graph(
@@ -298,11 +325,16 @@ class CudaGraphManager:
325
cudagraph_size = None
326
return cudagraph_mode, cudagraph_size
327
301
- def run_fullgraph(self, num_tokens: int) -> torch.Tensor:
328
+ def run_fullgraph(
329
+ self, num_tokens: int
330
+ ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
331
assert num_tokens in self.graphs, f"No cudagraph for {num_tokens} tokens"
332
self.graphs[num_tokens].replay()
333
assert self.hidden_states is not None
305
- return self.hidden_states[:num_tokens]
334
+ hidden_states = self.hidden_states[:num_tokens]
335
+ if not self.use_aux_hidden_state_outputs:
336
+ return hidden_states
337
+ return hidden_states, [x[:num_tokens] for x in self.aux_hidden_states]
338
339
340
def get_cudagraph_sizes(
vllm/v1/worker/gpu/model_runner.py
+5
-2
@@ -197,7 +197,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
197
198
# CUDA graphs.
199
self.cudagraph_manager = CudaGraphManager(
200
- self.vllm_config, self.uses_mrope, self.device
200
+ self.vllm_config,
201
+ self.uses_mrope,
202
+ self.use_aux_hidden_state_outputs,
203
+ self.device,
204
)
205
# Structured outputs worker.
206
self.structured_outputs_worker = StructuredOutputsWorker(
@@ -1044,7 +1047,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1047
aux_hidden_states,
1048
input_batch,
1049
kv_connector_output,
1047
- )
1050
+ ) # type: ignore
1051
return None
1052
1053
@torch.inference_mode()