From cc5dd0a857cde0ec224c63fee40276269e32f278 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 11 Sep 2026 12:55:34 -0700 Subject: [PATCH] [MRV1] Scope breakable cudagraphs to the piecewise path only (#56312) Signed-off-by: Nick Hill Co-authored-by: Andreas Karatzas --- .../v1/cudagraph/test_breakable_cudagraph.py | 132 +++++++++++++++++- vllm/compilation/breakable_cudagraph.py | 31 ++-- vllm/compilation/cuda_graph.py | 4 +- vllm/v1/worker/gpu_model_runner.py | 15 +- 4 files changed, 165 insertions(+), 17 deletions(-) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index ded83d800d..0e10e4a51f 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/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 diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 8e25d05c28..ca6aa17137 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/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 diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index b63d861997..5577160883 100644 --- a/vllm/compilation/cuda_graph.py +++ b/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": diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index a14d00bd58..4ff50f82b5 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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 ):