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 2, 08:30 UTC
[MRV2] Support breakable CUDA graph (#44050)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 30, 2026 at 09:40 UTC
27fa5aa3b952a6108de127423397e50364a95fcb
3 files changed
+59
-5
vllm/v1/worker/gpu/cudagraph_utils.py
+35
-3
@@ -9,6 +9,10 @@ import torch
9
import torch.nn as nn
10
from tqdm import tqdm
11
12
+from vllm.compilation.breakable_cudagraph import (
13
+ BreakableCUDAGraphWrapper,
14
+ is_breakable_cudagraph_enabled,
15
+)
16
from vllm.compilation.counter import compilation_counter
17
from vllm.config import VllmConfig
18
from vllm.config.compilation import CUDAGraphMode
@@ -116,6 +120,13 @@ class CudaGraphManager:
120
)
121
self._init_candidates()
122
123
+ # Breakable CUDA graph (PW CUDA graph without torch.compile)
124
+ self.use_breakable_cg = (
125
+ is_breakable_cudagraph_enabled()
126
+ and self.cudagraph_mode.has_piecewise_cudagraphs()
127
+ )
128
+ self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None
129
+
130
def _init_candidates(self) -> None:
131
"""Build priority-ordered candidate lists for each token count."""
132
capture_sizes = self.compilation_config.cudagraph_capture_sizes
@@ -283,6 +294,20 @@ class CudaGraphManager:
294
get_offloader().sync_prev_onload()
295
self.graphs[desc].replay()
296
297
+ def init_breakable_cg_runner(self, model: nn.Module) -> None:
298
+ if self.breakable_cg_runner is None:
299
+ self.breakable_cg_runner = BreakableCUDAGraphWrapper(
300
+ model, self.vllm_config
301
+ )
302
+ self.breakable_cg_runner.graph_pool = self.pool
303
+
304
+ def run_pw_graph(self, model: nn.Module, model_inputs: dict[str, Any]) -> Any:
305
+ if not self.use_breakable_cg:
306
+ # Default: Use torch-compiled piecewise cudagraph.
307
+ return model(**model_inputs)
308
+ assert self.breakable_cg_runner is not None
309
+ return self.breakable_cg_runner(**model_inputs)
310
+
311
312
class ModelCudaGraphManager(CudaGraphManager):
313
"""CudaGraphManager with model-specific capture and hidden state management."""
@@ -316,6 +341,8 @@ class ModelCudaGraphManager(CudaGraphManager):
341
) -> dict[BatchExecutionDescriptor, CapturedAttentionState]:
342
"""Capture CUDA graphs for model forward pass."""
343
self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
344
+ if self.use_breakable_cg:
345
+ self.init_breakable_cg_runner(model)
346
347
def create_forward_fn(
348
desc: BatchExecutionDescriptor,
@@ -370,11 +397,16 @@ class ModelCudaGraphManager(CudaGraphManager):
397
slot_mapping=slot_mappings,
398
batch_descriptor=batch_descriptor,
399
):
373
- model_output = model(**model_inputs)
400
+ if cg_mode == CUDAGraphMode.PIECEWISE:
401
+ # PIECEWISE graph (compiled PW or breakable, chosen inside
402
+ # run_pw_graph).
403
+ model_output = self.run_pw_graph(model, model_inputs)
404
+ else:
405
+ model_output = model(**model_inputs)
406
407
if cg_mode == CUDAGraphMode.PIECEWISE:
376
- # PW CUDA graph internally handles the model outputs.
377
- # No need to keep track of the hidden states.
408
+ # PW CUDA graph (compiled or breakable) internally handles the
409
+ # model outputs. No need to keep track of the hidden states.
410
return None
411
412
if self.is_last_pp_rank:
vllm/v1/worker/gpu/model_runner.py
+11
-1
@@ -1189,7 +1189,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1189
skip_compiled=skip_compiled,
1190
):
1191
self.kv_connector.pre_forward(scheduler_output)
1192
- model_output = self.model(**model_inputs)
1192
+ if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
1193
+ # Run the PIECEWISE graph (compiled PW cudagraph or breakable
1194
+ # cudagraph, chosen inside run_pw_graph). cg_mode is only
1195
+ # PIECEWISE after the cudagraph manager exists.
1196
+ assert self.cudagraph_manager is not None
1197
+ model_output = self.cudagraph_manager.run_pw_graph(
1198
+ self.model, model_inputs
1199
+ )
1200
+ else:
1201
+ # Eager (NONE): call the raw model directly.
1202
+ model_output = self.model(**model_inputs)
1203
1204
if self.is_last_pp_rank:
1205
if self.use_aux_hidden_state_outputs:
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+13
-1
@@ -215,12 +215,22 @@ class EagleSpeculator:
215
)
216
inputs_embeds = self.inputs_embeds[:num_tokens]
217
218
- ret_hidden_states = self.model(
218
+ model_inputs = dict(
219
input_ids=self.input_buffers.input_ids[:num_tokens],
220
positions=self.input_buffers.positions[:num_tokens],
221
hidden_states=self.hidden_states[:num_tokens],
222
inputs_embeds=inputs_embeds,
223
)
224
+ if cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE:
225
+ # Draft prefill with PIECEWISE cudagraph (compiled PW or breakable),
226
+ # chosen inside run_pw_graph.
227
+ assert self.prefill_cudagraph_manager is not None
228
+ ret_hidden_states = self.prefill_cudagraph_manager.run_pw_graph(
229
+ self.model, model_inputs
230
+ )
231
+ else:
232
+ # Eager (NONE): call the raw model directly.
233
+ ret_hidden_states = self.model(**model_inputs)
234
if self.method == "mtp":
235
last_hidden_states = ret_hidden_states
236
hidden_states = ret_hidden_states
@@ -431,6 +441,8 @@ class EagleSpeculator:
441
# For PIECEWISE, only the model's compiled regions are captured
442
# and the rest (compute_logits, gumbel_sample) runs eagerly.
443
assert self.prefill_cudagraph_manager is not None
444
+ if self.prefill_cudagraph_manager.use_breakable_cg:
445
+ self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
446
self.prefill_cudagraph_manager.capture(
447
self.prefill,
448
attn_states,