[MRV1] Scope breakable cudagraphs to the piecewise path only (#56312)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Nick Hill committed
Sep 11, 2026 at 12:55 UTC
cc5dd0a857cde0ec224c63fee40276269e32f278
4 files changed
+165
-17
tests/v1/cudagraph/test_breakable_cudagraph.py
+131
-1
@@ -8,7 +8,7 @@ from __future__ import annotations
8
9
import threading
10
from contextlib import nullcontext
11
-from unittest.mock import patch
11
+from unittest.mock import MagicMock, patch
12
13
import pytest
14
import torch
@@ -536,3 +536,133 @@ def test_nested_decorated_op_runs_inline(cuda_capture_stream):
536
# 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113
537
assert torch.equal(x, torch.full((4,), 113.0, device="cuda"))
538
assert inner_calls == 2 # replay invokes the outer's lambda again
539
+
540
+
541
+# ---------------------------------------------------------------------------
542
+# BreakableCUDAGraphWrapper: runtime_mode dispatch gating
543
+# ---------------------------------------------------------------------------
544
+
545
+
546
+def _mock_vllm_config():
547
+ config = MagicMock()
548
+ config.parallel_config.data_parallel_size = 1
549
+ config.parallel_config.use_sequence_parallel_moe = False
550
+ config.parallel_config.is_moe_model = False
551
+ return config
552
+
553
+
554
+def test_wrapper_falls_through_on_runtime_mode_mismatch():
555
+ """A PIECEWISE-scoped wrapper runs other modes eagerly, capturing
556
+ nothing, so an outer CUDAGraphWrapper can own them."""
557
+ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
558
+ from vllm.config import CUDAGraphMode
559
+ from vllm.forward_context import BatchDescriptor, set_forward_context
560
+
561
+ calls = []
562
+
563
+ def model(x):
564
+ calls.append(x)
565
+ return x + 1
566
+
567
+ config = _mock_vllm_config()
568
+ wrapper = BreakableCUDAGraphWrapper(
569
+ model, config, runtime_mode=CUDAGraphMode.PIECEWISE
570
+ )
571
+ desc = BatchDescriptor(num_tokens=4)
572
+
573
+ with set_forward_context(
574
+ None,
575
+ config,
576
+ cudagraph_runtime_mode=CUDAGraphMode.FULL,
577
+ batch_descriptor=desc,
578
+ ):
579
+ assert wrapper(1) == 2
580
+
581
+ assert calls == [1]
582
+ assert not wrapper.entries
583
+
584
+
585
+def test_wrapper_captures_on_runtime_mode_match():
586
+ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
587
+ from vllm.config import CUDAGraphMode
588
+ from vllm.forward_context import BatchDescriptor, set_forward_context
589
+
590
+ calls = []
591
+
592
+ def model(x):
593
+ calls.append(x)
594
+ return x + 1
595
+
596
+ config = _mock_vllm_config()
597
+ wrapper = BreakableCUDAGraphWrapper(
598
+ model, config, runtime_mode=CUDAGraphMode.PIECEWISE
599
+ )
600
+ desc = BatchDescriptor(num_tokens=4)
601
+
602
+ sentinel = object()
603
+ with (
604
+ patch.object(
605
+ BreakableCUDAGraphWrapper, "_capture", return_value=sentinel
606
+ ) as mock_capture,
607
+ set_forward_context(
608
+ None,
609
+ config,
610
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
611
+ batch_descriptor=desc,
612
+ ),
613
+ ):
614
+ assert wrapper(1) is sentinel
615
+
616
+ mock_capture.assert_called_once()
617
+ assert calls == []
618
+ assert desc in wrapper.entries
619
+
620
+
621
+def test_wrapper_default_matches_any_non_none_mode():
622
+ """runtime_mode=None matches any non-NONE mode (MRV2 relies on this)."""
623
+ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
624
+ from vllm.config import CUDAGraphMode
625
+ from vllm.forward_context import BatchDescriptor, set_forward_context
626
+
627
+ def model(x):
628
+ return x + 1
629
+
630
+ config = _mock_vllm_config()
631
+ wrapper = BreakableCUDAGraphWrapper(model, config)
632
+ desc = BatchDescriptor(num_tokens=4)
633
+
634
+ sentinel = object()
635
+ with (
636
+ patch.object(
637
+ BreakableCUDAGraphWrapper, "_capture", return_value=sentinel
638
+ ) as mock_capture,
639
+ set_forward_context(
640
+ None,
641
+ config,
642
+ cudagraph_runtime_mode=CUDAGraphMode.FULL,
643
+ batch_descriptor=desc,
644
+ ),
645
+ ):
646
+ assert wrapper(1) is sentinel
647
+
648
+ mock_capture.assert_called_once()
649
+ assert desc in wrapper.entries
650
+
651
+
652
+def test_unwrap_recurses_through_nested_wrappers():
653
+ """unwrap() reaches the raw model through nested wrappers."""
654
+ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
655
+ from vllm.compilation.cuda_graph import CUDAGraphWrapper
656
+ from vllm.config import CUDAGraphMode
657
+
658
+ def model(x):
659
+ return x + 1
660
+
661
+ config = _mock_vllm_config()
662
+ inner = BreakableCUDAGraphWrapper(
663
+ model, config, runtime_mode=CUDAGraphMode.PIECEWISE
664
+ )
665
+ outer = CUDAGraphWrapper(inner, config, runtime_mode=CUDAGraphMode.FULL)
666
+
667
+ assert inner.unwrap() is model
668
+ assert outer.unwrap() is model
vllm/compilation/breakable_cudagraph.py
+19
-12
@@ -256,6 +256,10 @@ class BreakableCUDAGraphWrapper:
256
* If runtime mode mismatch / NONE, run eagerly.
257
* Otherwise, lazily capture per ``batch_descriptor`` and replay
258
on subsequent invocations with the same descriptor.
259
+
260
+ When ``runtime_mode`` is given, only that mode is captured/replayed;
261
+ other non-NONE modes fall through, so the wrapper can be nested under
262
+ a ``CUDAGraphWrapper`` for another mode.
263
"""
264
265
_all_instances: ClassVar[weakref.WeakSet[BreakableCUDAGraphWrapper]] = (
@@ -271,15 +275,14 @@ class BreakableCUDAGraphWrapper:
275
self,
276
runnable: Callable[..., Any],
277
vllm_config: VllmConfig,
278
+ runtime_mode: CUDAGraphMode | None = None,
279
) -> None:
275
- # Unlike the original CUDAGraphWrapper which strictly matches a
276
- # single runtime_mode, this wrapper captures whatever the
277
- # dispatcher emits (any non-NONE runtime_mode) -- breakable's
278
- # capture is identical for prefill and decode, so there's nothing
279
- # to dispatch on at the runtime_mode level. Entries are keyed by
280
- # BatchDescriptor which already encodes batch shape / uniformity.
280
+ # runtime_mode=None matches any non-NONE dispatch mode; entries
281
+ # are keyed by BatchDescriptor, which already encodes batch
282
+ # shape / uniformity.
283
self.runnable = runnable
284
self.vllm_config = vllm_config
285
+ self.runtime_mode = runtime_mode
286
self.compilation_config = vllm_config.compilation_config
287
self.graph_pool = current_platform.get_global_graph_pool()
288
self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
@@ -298,7 +301,9 @@ class BreakableCUDAGraphWrapper:
301
raise AttributeError(key)
302
303
def unwrap(self) -> Callable[..., Any]:
301
- return self.runnable
304
+ # Recurse through nested wrappers to reach the original runnable.
305
+ runnable = self.runnable
306
+ return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable
307
308
@property
309
def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper:
@@ -317,11 +322,13 @@ class BreakableCUDAGraphWrapper:
322
batch_descriptor = forward_context.batch_descriptor
323
cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
324
320
- # Capture whenever the dispatcher says "some cudagraph mode" --
321
- # breakable produces the same artifact regardless of PIECEWISE
322
- # vs FULL, so we match either. Entries are keyed by batch
323
- # descriptor, which already encodes prefill/decode distinctions.
324
- if cudagraph_runtime_mode == CUDAGraphMode.NONE:
325
+ # By default (runtime_mode=None) match any non-NONE mode; a set
326
+ # runtime_mode restricts capture so other modes can be handled by
327
+ # an outer wrapper (e.g. CUDAGraphWrapper for FULL).
328
+ if cudagraph_runtime_mode == CUDAGraphMode.NONE or (
329
+ self.runtime_mode is not None
330
+ and cudagraph_runtime_mode != self.runtime_mode
331
+ ):
332
return self.runnable(*args, **kwargs)
333
334
assert batch_descriptor is not None
vllm/compilation/cuda_graph.py
+3
-1
@@ -221,7 +221,9 @@ class CUDAGraphWrapper:
221
222
def unwrap(self) -> Callable[..., Any]:
223
# in case we need to access the original runnable.
224
- return self.runnable
224
+ runnable = self.runnable
225
+ # Recurse through nested wrappers.
226
+ return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable
227
228
@property
229
def cudagraph_wrapper(self) -> "CUDAGraphWrapper":
vllm/v1/worker/gpu_model_runner.py
+12
-3
@@ -5461,13 +5461,22 @@ class GPUModelRunner(
5461
and cudagraph_mode != CUDAGraphMode.NONE
5462
and not self.parallel_config.use_ubatching
5463
):
5464
- self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config)
5464
+ # Scoped to PIECEWISE dispatch; FULL cudagraphs (below) are
5465
+ # unaffected. PIECEWISE dispatch can also arise after wrapping
5466
+ # (drafters under a FULL target mode, or a later FULL ->
5467
+ # FULL_AND_PIECEWISE upgrade in _check_and_update_cudagraph_mode).
5468
+ self.model = BreakableCUDAGraphWrapper(
5469
+ self.model, self.vllm_config, runtime_mode=CUDAGraphMode.PIECEWISE
5470
+ )
5471
drafter = getattr(self, "drafter", None)
5472
if drafter is not None and hasattr(drafter, "model"):
5473
drafter.model = BreakableCUDAGraphWrapper(
5468
- drafter.model, self.vllm_config
5474
+ drafter.model,
5475
+ self.vllm_config,
5476
+ runtime_mode=CUDAGraphMode.PIECEWISE,
5477
)
5470
- elif (
5478
+
5479
+ if (
5480
cudagraph_mode.has_full_cudagraphs()
5481
and not self.parallel_config.use_ubatching
5482
):