[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>

tests/v1/cudagraph/test_breakable_cudagraph.py

@@ -8,7 +8,7 @@ from __future__ import annotations
 
 import threading
 from contextlib import nullcontext
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -536,3 +536,133 @@ def test_nested_decorated_op_runs_inline(cuda_capture_stream):
     # 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113
     assert torch.equal(x, torch.full((4,), 113.0, device="cuda"))
     assert inner_calls == 2  # replay invokes the outer's lambda again
+
+
+# ---------------------------------------------------------------------------
+# BreakableCUDAGraphWrapper: runtime_mode dispatch gating
+# ---------------------------------------------------------------------------
+
+
+def _mock_vllm_config():
+    config = MagicMock()
+    config.parallel_config.data_parallel_size = 1
+    config.parallel_config.use_sequence_parallel_moe = False
+    config.parallel_config.is_moe_model = False
+    return config
+
+
+def test_wrapper_falls_through_on_runtime_mode_mismatch():
+    """A PIECEWISE-scoped wrapper runs other modes eagerly, capturing
+    nothing, so an outer CUDAGraphWrapper can own them."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+    from vllm.forward_context import BatchDescriptor, set_forward_context
+
+    calls = []
+
+    def model(x):
+        calls.append(x)
+        return x + 1
+
+    config = _mock_vllm_config()
+    wrapper = BreakableCUDAGraphWrapper(
+        model, config, runtime_mode=CUDAGraphMode.PIECEWISE
+    )
+    desc = BatchDescriptor(num_tokens=4)
+
+    with set_forward_context(
+        None,
+        config,
+        cudagraph_runtime_mode=CUDAGraphMode.FULL,
+        batch_descriptor=desc,
+    ):
+        assert wrapper(1) == 2
+
+    assert calls == [1]
+    assert not wrapper.entries
+
+
+def test_wrapper_captures_on_runtime_mode_match():
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+    from vllm.forward_context import BatchDescriptor, set_forward_context
+
+    calls = []
+
+    def model(x):
+        calls.append(x)
+        return x + 1
+
+    config = _mock_vllm_config()
+    wrapper = BreakableCUDAGraphWrapper(
+        model, config, runtime_mode=CUDAGraphMode.PIECEWISE
+    )
+    desc = BatchDescriptor(num_tokens=4)
+
+    sentinel = object()
+    with (
+        patch.object(
+            BreakableCUDAGraphWrapper, "_capture", return_value=sentinel
+        ) as mock_capture,
+        set_forward_context(
+            None,
+            config,
+            cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+            batch_descriptor=desc,
+        ),
+    ):
+        assert wrapper(1) is sentinel
+
+    mock_capture.assert_called_once()
+    assert calls == []
+    assert desc in wrapper.entries
+
+
+def test_wrapper_default_matches_any_non_none_mode():
+    """runtime_mode=None matches any non-NONE mode (MRV2 relies on this)."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+    from vllm.forward_context import BatchDescriptor, set_forward_context
+
+    def model(x):
+        return x + 1
+
+    config = _mock_vllm_config()
+    wrapper = BreakableCUDAGraphWrapper(model, config)
+    desc = BatchDescriptor(num_tokens=4)
+
+    sentinel = object()
+    with (
+        patch.object(
+            BreakableCUDAGraphWrapper, "_capture", return_value=sentinel
+        ) as mock_capture,
+        set_forward_context(
+            None,
+            config,
+            cudagraph_runtime_mode=CUDAGraphMode.FULL,
+            batch_descriptor=desc,
+        ),
+    ):
+        assert wrapper(1) is sentinel
+
+    mock_capture.assert_called_once()
+    assert desc in wrapper.entries
+
+
+def test_unwrap_recurses_through_nested_wrappers():
+    """unwrap() reaches the raw model through nested wrappers."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.compilation.cuda_graph import CUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+
+    def model(x):
+        return x + 1
+
+    config = _mock_vllm_config()
+    inner = BreakableCUDAGraphWrapper(
+        model, config, runtime_mode=CUDAGraphMode.PIECEWISE
+    )
+    outer = CUDAGraphWrapper(inner, config, runtime_mode=CUDAGraphMode.FULL)
+
+    assert inner.unwrap() is model
+    assert outer.unwrap() is model

vllm/compilation/breakable_cudagraph.py

@@ -256,6 +256,10 @@ class BreakableCUDAGraphWrapper:
         * If runtime mode mismatch / NONE, run eagerly.
         * Otherwise, lazily capture per ``batch_descriptor`` and replay
           on subsequent invocations with the same descriptor.
+
+    When ``runtime_mode`` is given, only that mode is captured/replayed;
+    other non-NONE modes fall through, so the wrapper can be nested under
+    a ``CUDAGraphWrapper`` for another mode.
     """
 
     _all_instances: ClassVar[weakref.WeakSet[BreakableCUDAGraphWrapper]] = (
@@ -271,15 +275,14 @@ class BreakableCUDAGraphWrapper:
         self,
         runnable: Callable[..., Any],
         vllm_config: VllmConfig,
+        runtime_mode: CUDAGraphMode | None = None,
     ) -> None:
-        # Unlike the original CUDAGraphWrapper which strictly matches a
-        # single runtime_mode, this wrapper captures whatever the
-        # dispatcher emits (any non-NONE runtime_mode) -- breakable's
-        # capture is identical for prefill and decode, so there's nothing
-        # to dispatch on at the runtime_mode level. Entries are keyed by
-        # BatchDescriptor which already encodes batch shape / uniformity.
+        # runtime_mode=None matches any non-NONE dispatch mode; entries
+        # are keyed by BatchDescriptor, which already encodes batch
+        # shape / uniformity.
         self.runnable = runnable
         self.vllm_config = vllm_config
+        self.runtime_mode = runtime_mode
         self.compilation_config = vllm_config.compilation_config
         self.graph_pool = current_platform.get_global_graph_pool()
         self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
@@ -298,7 +301,9 @@ class BreakableCUDAGraphWrapper:
         raise AttributeError(key)
 
     def unwrap(self) -> Callable[..., Any]:
-        return self.runnable
+        # Recurse through nested wrappers to reach the original runnable.
+        runnable = self.runnable
+        return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable
 
     @property
     def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper:
@@ -317,11 +322,13 @@ class BreakableCUDAGraphWrapper:
         batch_descriptor = forward_context.batch_descriptor
         cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
 
-        # Capture whenever the dispatcher says "some cudagraph mode" --
-        # breakable produces the same artifact regardless of PIECEWISE
-        # vs FULL, so we match either. Entries are keyed by batch
-        # descriptor, which already encodes prefill/decode distinctions.
-        if cudagraph_runtime_mode == CUDAGraphMode.NONE:
+        # By default (runtime_mode=None) match any non-NONE mode; a set
+        # runtime_mode restricts capture so other modes can be handled by
+        # an outer wrapper (e.g. CUDAGraphWrapper for FULL).
+        if cudagraph_runtime_mode == CUDAGraphMode.NONE or (
+            self.runtime_mode is not None
+            and cudagraph_runtime_mode != self.runtime_mode
+        ):
             return self.runnable(*args, **kwargs)
 
         assert batch_descriptor is not None

vllm/compilation/cuda_graph.py

@@ -221,7 +221,9 @@ class CUDAGraphWrapper:
 
     def unwrap(self) -> Callable[..., Any]:
         # in case we need to access the original runnable.
-        return self.runnable
+        runnable = self.runnable
+        # Recurse through nested wrappers.
+        return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable
 
     @property
     def cudagraph_wrapper(self) -> "CUDAGraphWrapper":

vllm/v1/worker/gpu_model_runner.py

@@ -5461,13 +5461,22 @@ class GPUModelRunner(
             and cudagraph_mode != CUDAGraphMode.NONE
             and not self.parallel_config.use_ubatching
         ):
-            self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config)
+            # Scoped to PIECEWISE dispatch; FULL cudagraphs (below) are
+            # unaffected. PIECEWISE dispatch can also arise after wrapping
+            # (drafters under a FULL target mode, or a later FULL ->
+            # FULL_AND_PIECEWISE upgrade in _check_and_update_cudagraph_mode).
+            self.model = BreakableCUDAGraphWrapper(
+                self.model, self.vllm_config, runtime_mode=CUDAGraphMode.PIECEWISE
+            )
             drafter = getattr(self, "drafter", None)
             if drafter is not None and hasattr(drafter, "model"):
                 drafter.model = BreakableCUDAGraphWrapper(
-                    drafter.model, self.vllm_config
+                    drafter.model,
+                    self.vllm_config,
+                    runtime_mode=CUDAGraphMode.PIECEWISE,
                 )
-        elif (
+
+        if (
             cudagraph_mode.has_full_cudagraphs()
             and not self.parallel_config.use_ubatching
         ):