@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8a56da3845
Fork 0

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 3, 13:00 UTC

[Experimental] Breakable CUDA graph (#42304)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

Jiangyun Zhu committed May 16, 2026 at 22:04 UTC 8a56da3845270837424ef4b7ee83ca97a7883025
11 files changed +845 -5
.buildkite/test_areas/cuda.yaml
+2 -1
@@ -27,4 +27,5 @@ steps:
27 - vllm/compilation
28 commands:
29 - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
30 - - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
\ No newline at end of file
30 + - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
31 + - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
\ No newline at end of file
tests/v1/cudagraph/test_breakable_cudagraph.py new
+367
@@ -0,0 +1,367 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for the breakable cudagraph primitives.
5 +"""
6 +
7 +from __future__ import annotations
8 +
9 +import os
10 +import threading
11 +
12 +import pytest
13 +import torch
14 +
15 +os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
16 +
17 +
18 +@pytest.fixture(autouse=True)
19 +def _reset_breakable_tls():
20 + """Defensively clear thread-local capture state between tests so a
21 + failure in one test can't leak "nested capture" errors into the next."""
22 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
23 +
24 + BreakableCUDAGraphCapture._tls.active = None
25 + yield
26 + BreakableCUDAGraphCapture._tls.active = None
27 +
28 +
29 +@pytest.fixture
30 +def cuda_capture_stream():
31 + """A non-default CUDA stream suitable for cudagraph capture.
32 +
33 + ``CUDAGraph.capture_begin`` refuses to capture from the default
34 + stream, so all capture-using tests need to run under
35 + ``torch.cuda.stream(...)`` for a separate stream.
36 + """
37 + if not torch.cuda.is_available():
38 + pytest.skip("CUDA required")
39 + stream = torch.cuda.Stream()
40 + with torch.cuda.stream(stream):
41 + yield stream
42 + torch.cuda.current_stream().wait_stream(stream)
43 +
44 +
45 +# ---------------------------------------------------------------------------
46 +# eager_break_during_capture: outside capture
47 +# ---------------------------------------------------------------------------
48 +
49 +
50 +def test_decorator_passthrough_outside_capture():
51 + from vllm.compilation.breakable_cudagraph import eager_break_during_capture
52 +
53 + calls = []
54 +
55 + @eager_break_during_capture
56 + def f(x):
57 + calls.append(x)
58 + return x * 2
59 +
60 + assert f(3) == 6
61 + assert calls == [3]
62 +
63 +
64 +# ---------------------------------------------------------------------------
65 +# BreakableCUDAGraphCapture: thread-local + nested rejection
66 +# ---------------------------------------------------------------------------
67 +
68 +
69 +def test_current_is_none_when_inactive():
70 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
71 +
72 + assert BreakableCUDAGraphCapture.current() is None
73 + assert BreakableCUDAGraphCapture.is_active() is False
74 +
75 +
76 +def test_thread_local_active_during_context(cuda_capture_stream):
77 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
78 +
79 + cap = BreakableCUDAGraphCapture()
80 + with cap:
81 + assert BreakableCUDAGraphCapture.current() is cap
82 + assert BreakableCUDAGraphCapture.is_active() is True
83 + assert BreakableCUDAGraphCapture.current() is None
84 +
85 +
86 +def test_nested_capture_raises(cuda_capture_stream):
87 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
88 +
89 + outer = BreakableCUDAGraphCapture()
90 + inner = BreakableCUDAGraphCapture()
91 + with outer, pytest.raises(RuntimeError, match="Nested.*not supported"), inner:
92 + pass
93 +
94 +
95 +def test_active_state_isolated_across_threads(cuda_capture_stream):
96 + """Verify the thread-local 'active capture' slot is per-thread.
97 +
98 + We don't run concurrent captures here -- CUDA only supports one
99 + in-flight capture per stream and we keep tests cheap. We just check
100 + that the worker thread sees its own slot as None while the main
101 + thread has a capture active.
102 + """
103 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
104 +
105 + worker_view: dict[str, BreakableCUDAGraphCapture | None] = {}
106 +
107 + def worker():
108 + worker_view["state"] = BreakableCUDAGraphCapture.current()
109 +
110 + main_cap = BreakableCUDAGraphCapture()
111 + with main_cap:
112 + # Main thread has a live capture.
113 + assert BreakableCUDAGraphCapture.current() is main_cap
114 + t = threading.Thread(target=worker)
115 + t.start()
116 + t.join()
117 +
118 + # Worker thread saw None -- thread-local separation works.
119 + assert worker_view["state"] is None
120 + # Main thread's slot is cleared on exit.
121 + assert BreakableCUDAGraphCapture.current() is None
122 +
123 +
124 +# ---------------------------------------------------------------------------
125 +# Segment list construction
126 +# ---------------------------------------------------------------------------
127 +
128 +
129 +def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream):
130 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
131 +
132 + x = torch.zeros(4, device="cuda")
133 + cap = BreakableCUDAGraphCapture()
134 + with cap:
135 + x.add_(1.0)
136 + assert len(cap.segments) == 1
137 + assert cap.num_graphs == 1
138 + assert cap.num_eager_breaks == 0
139 +
140 +
141 +def test_add_eager_creates_alternating_graph_eager_graph(cuda_capture_stream):
142 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
143 +
144 + x = torch.zeros(4, device="cuda")
145 + counter = {"eager_calls": 0}
146 +
147 + def eager_step():
148 + counter["eager_calls"] += 1
149 + x.add_(10.0)
150 +
151 + cap = BreakableCUDAGraphCapture()
152 + with cap:
153 + x.add_(1.0)
154 + cap.add_eager(eager_step)
155 + x.add_(1.0)
156 + cap.add_eager(eager_step)
157 + x.add_(1.0)
158 + # 3 graph segments + 2 eager segments, interleaved as G E G E G.
159 + assert len(cap.segments) == 5
160 + assert cap.num_graphs == 3
161 + assert cap.num_eager_breaks == 2
162 + # Eager fn is stored as-is in the segment list, so we can confirm
163 + # the alternation pattern by identity check.
164 + assert cap.segments[1] is eager_step
165 + assert cap.segments[3] is eager_step
166 + assert counter["eager_calls"] == 2 # only the in-capture invocation
167 +
168 +
169 +# ---------------------------------------------------------------------------
170 +# Capture vs eager numerical equivalence
171 +# ---------------------------------------------------------------------------
172 +
173 +
174 +def test_capture_replay_matches_eager_simple(cuda_capture_stream):
175 + """Verify that replay reproduces the same end-state as a single eager
176 + forward, with an eager break in the middle.
177 +
178 + Note: during capture, the *captured* kernels are recorded but NOT
179 + executed (that's CUDA-graph semantics). Only the eager segments
180 + actually mutate state at capture time. So we check correctness after
181 + ``replay()``, not after ``with cap:`` exits.
182 + """
183 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
184 +
185 + x = torch.zeros(8, device="cuda")
186 + log: list[str] = []
187 +
188 + def eager_break_op():
189 + x.mul_(2.0)
190 + log.append("eager")
191 +
192 + cap = BreakableCUDAGraphCapture()
193 + with cap:
194 + x.add_(1.0) # recorded into graph[0]
195 + cap.add_eager(eager_break_op) # runs eagerly: x *= 2
196 + x.add_(5.0) # recorded into graph[1]
197 +
198 + # Capture-time: graph kernels were recorded only; eager segment ran
199 + # once on x == 0, leaving x == 0.
200 + torch.accelerator.synchronize()
201 + assert torch.equal(x, torch.zeros(8, device="cuda"))
202 + assert log == ["eager"]
203 +
204 + # Replay with a fresh input: 10 -> 11 -> 22 -> 27.
205 + x.fill_(10.0)
206 + cap.replay()
207 + torch.accelerator.synchronize()
208 + assert torch.equal(x, torch.full((8,), 27.0, device="cuda"))
209 + assert log == ["eager", "eager"]
210 +
211 + # Replay again with another input: 100 -> 101 -> 202 -> 207.
212 + x.fill_(100.0)
213 + cap.replay()
214 + torch.accelerator.synchronize()
215 + assert torch.equal(x, torch.full((8,), 207.0, device="cuda"))
216 + assert log == ["eager", "eager", "eager"]
217 +
218 +
219 +def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream):
220 + """Verify @eager_break_during_capture correctly routes through
221 + add_eager when inside a capture context, and runs straight through
222 + when there's no active capture."""
223 + from vllm.compilation.breakable_cudagraph import (
224 + BreakableCUDAGraphCapture,
225 + eager_break_during_capture,
226 + )
227 +
228 + @eager_break_during_capture
229 + def attention_like(t: torch.Tensor) -> None:
230 + # In-place double; stands in for "real" attention work.
231 + t.mul_(2.0)
232 +
233 + x = torch.zeros(4, device="cuda")
234 +
235 + # Outside capture: decorator should just call through.
236 + x.fill_(3.0)
237 + attention_like(x)
238 + torch.accelerator.synchronize()
239 + assert torch.equal(x, torch.full((4,), 6.0, device="cuda"))
240 +
241 + # Inside capture: decorator should split the graph. Only the eager
242 + # segment actually mutates state during capture.
243 + x.fill_(0.0)
244 + cap = BreakableCUDAGraphCapture()
245 + with cap:
246 + x.add_(5.0) # recorded
247 + attention_like(x) # eager: x *= 2 (on x == 0, no-op)
248 + x.add_(1.0) # recorded
249 + torch.accelerator.synchronize()
250 + assert torch.equal(x, torch.zeros(4, device="cuda"))
251 + # 2 graph segments + 1 eager segment, ordered G E G; the arithmetic
252 + # equivalence check below verifies the ordering.
253 + assert len(cap.segments) == 3
254 + assert cap.num_graphs == 2
255 + assert cap.num_eager_breaks == 1
256 +
257 + # Replay: 2 -> 7 -> 14 -> 15.
258 + x.fill_(2.0)
259 + cap.replay()
260 + torch.accelerator.synchronize()
261 + assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))
262 +
263 +
264 +# ---------------------------------------------------------------------------
265 +# Replay ordering
266 +# ---------------------------------------------------------------------------
267 +
268 +
269 +def test_replay_invokes_eager_segments_in_order(cuda_capture_stream):
270 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
271 +
272 + log: list[str] = []
273 + x = torch.zeros(1, device="cuda")
274 +
275 + def make_eager(name):
276 + def step():
277 + log.append(name)
278 + x.add_(1.0)
279 +
280 + return step
281 +
282 + cap = BreakableCUDAGraphCapture()
283 + with cap:
284 + x.add_(1.0)
285 + cap.add_eager(make_eager("A"))
286 + x.add_(1.0)
287 + cap.add_eager(make_eager("B"))
288 + x.add_(1.0)
289 + cap.add_eager(make_eager("C"))
290 + x.add_(1.0)
291 +
292 + # Capture-time invocation order
293 + assert log == ["A", "B", "C"]
294 +
295 + log.clear()
296 + cap.replay()
297 + torch.accelerator.synchronize()
298 + assert log == ["A", "B", "C"]
299 +
300 +
301 +# ---------------------------------------------------------------------------
302 +# Capture cleanup releases thread-local even if body raises
303 +# ---------------------------------------------------------------------------
304 +
305 +
306 +def test_exception_in_body_clears_active(cuda_capture_stream):
307 + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
308 +
309 + cap = BreakableCUDAGraphCapture()
310 + with pytest.raises(RuntimeError, match="boom"), cap:
311 + raise RuntimeError("boom")
312 +
313 + # active must be reset even after an exception inside the body
314 + assert BreakableCUDAGraphCapture.current() is None
315 +
316 +
317 +# ---------------------------------------------------------------------------
318 +# Nested decorated ops: inner op must not trigger a recursive eager break
319 +# ---------------------------------------------------------------------------
320 +
321 +
322 +def test_nested_decorated_op_runs_inline(cuda_capture_stream):
323 + """A decorated op invoked from inside another decorated op's eager
324 + body must execute inline -- starting a second eager break mid-flight
325 + corrupts the segment state and explodes ``_begin_segment``'s assert.
326 +
327 + This mirrors the deepseek_v4_attention case where the outer attention
328 + op's impl internally dispatches sparse_attn_indexer (also decorated).
329 + """
330 + from vllm.compilation.breakable_cudagraph import (
331 + BreakableCUDAGraphCapture,
332 + eager_break_during_capture,
333 + )
334 +
335 + x = torch.zeros(4, device="cuda")
336 + inner_calls = 0
337 +
338 + @eager_break_during_capture
339 + def inner_op(t: torch.Tensor) -> None:
340 + nonlocal inner_calls
341 + inner_calls += 1
342 + t.add_(1.0)
343 +
344 + @eager_break_during_capture
345 + def outer_op(t: torch.Tensor) -> None:
346 + # outer body calls another decorated op -- this is the case that
347 + # used to assert in _begin_segment.
348 + inner_op(t)
349 + t.add_(10.0)
350 +
351 + cap = BreakableCUDAGraphCapture()
352 + with cap:
353 + x.add_(2.0) # recorded in graph[0]
354 + outer_op(x) # one eager break, inner runs inline
355 + x.add_(100.0) # recorded in graph[1]
356 +
357 + # Exactly one eager break (the outer); inner must NOT add a second.
358 + assert cap.num_graphs == 2
359 + assert cap.num_eager_breaks == 1
360 + assert inner_calls == 1 # only the capture-time invocation
361 +
362 + x.fill_(0.0)
363 + cap.replay()
364 + torch.accelerator.synchronize()
365 + # 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113
366 + assert torch.equal(x, torch.full((4,), 113.0, device="cuda"))
367 + assert inner_calls == 2 # replay invokes the outer's lambda again
vllm/compilation/breakable_cudagraph.py new
+424
@@ -0,0 +1,424 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Breakable CUDA graph capture/replay.
4 +
5 +This is an alternative to :class:`CUDAGraphWrapper` that replaces vLLM's
6 +torch.compile-based FX graph splitting with runtime stream-capture
7 +breaks.
8 +
9 +The idea (inspired by sgl-project/sglang#19102): instead of pre-splitting
10 +the model into many pieces at attention boundaries, a
11 +single capture context drives the whole forward and intercepts
12 +attention / kv-cache custom ops at the dispatcher to end the current
13 +stream capture, run the op eagerly, and resume capture.
14 +
15 +The captured artifact is a list of zero-arg callables -- the bound
16 +``CUDAGraph.replay`` for graph segments, or the user fn for eager
17 +segments -- replayed in order at inference time.
18 +
19 +Eager segments must operate on the same static buffers used during
20 +capture so subsequent graph segments read the same memory addresses.
21 +"""
22 +
23 +from __future__ import annotations
24 +
25 +import dataclasses
26 +import functools
27 +import gc
28 +import threading
29 +import weakref
30 +from collections.abc import Callable
31 +from typing import Any, ClassVar, TypeVar
32 +
33 +import torch
34 +
35 +import vllm.envs as envs
36 +from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
37 +from vllm.config import CUDAGraphMode, VllmConfig
38 +from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
39 +from vllm.forward_context import (
40 + BatchDescriptor,
41 + get_forward_context,
42 + is_forward_context_available,
43 +)
44 +from vllm.logger import init_logger
45 +from vllm.model_executor.offloader.base import get_offloader
46 +from vllm.platforms import current_platform
47 +from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors
48 +
49 +logger = init_logger(__name__)
50 +
51 +
52 +def is_breakable_cudagraph_enabled() -> bool:
53 + return bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH)
54 +
55 +
56 +F = TypeVar("F", bound=Callable[..., Any])
57 +
58 +
59 +def eager_break_during_capture(fn: F) -> F:
60 + """Decorator that turns a custom-op Python kernel into a "break point"
61 + for the breakable cudagraph capture.
62 +
63 + When the decorated function is invoked outside of a
64 + :class:`BreakableCUDAGraphCapture` context, it executes normally.
65 +
66 + When invoked inside a capture context, it ends the current cudagraph
67 + segment, runs the function eagerly on the capture stream, records the
68 + callable for replay, and starts a fresh segment.
69 +
70 + **In-place output buffer required.** Decorated ops must write into a
71 + caller-provided output tensor; a fresh tensor returned by ``fn`` would
72 + change address each replay and break downstream graph segments.
73 +
74 + **Decorator order matters.** Apply as the *outermost* decorator if
75 + there are other decorators that introduce host-side side effects
76 + around the call -- the canonical example is
77 + ``@maybe_transfer_kv_layer`` for PD-disaggregation, whose
78 + ``wait_for_layer_load`` and ``save_kv_layer`` calls must run in the
79 + eager segment, not inside the captured cudagraph. Putting
80 + ``@eager_break_during_capture`` *inside* such a decorator would
81 + record those side effects into the graph and hang on replay.
82 +
83 + The correct order is::
84 +
85 + @eager_break_during_capture # outermost
86 + @maybe_transfer_kv_layer
87 + def unified_attention_with_output(...):
88 + ...
89 + """
90 + if not is_breakable_cudagraph_enabled():
91 + return fn
92 +
93 + @functools.wraps(fn)
94 + def wrapper(*args: Any, **kwargs: Any) -> Any:
95 + capture = BreakableCUDAGraphCapture.current()
96 + if capture is None:
97 + return fn(*args, **kwargs)
98 + if not capture._capturing:
99 + return fn(*args, **kwargs)
100 + if is_forward_context_available():
101 + mode = get_forward_context().cudagraph_runtime_mode
102 + if mode == CUDAGraphMode.FULL:
103 + return fn(*args, **kwargs)
104 +
105 + # Weak-ref args: strong refs in the replay lambda pin cudagraph-pool
106 + # slots across batch descriptors. cudagraph owns the slot, so the
107 + # weak_ref is safe to deref on replay.
108 + weak_args = tuple(
109 + weak_ref_tensor(a) if isinstance(a, torch.Tensor) else a for a in args
110 + )
111 + weak_kwargs = {
112 + k: weak_ref_tensor(v) if isinstance(v, torch.Tensor) else v
113 + for k, v in kwargs.items()
114 + }
115 + return capture.add_eager(lambda: fn(*weak_args, **weak_kwargs))
116 +
117 + return wrapper # type: ignore[return-value]
118 +
119 +
120 +# ---------------------------------------------------------------------------
121 +# Capture context
122 +# ---------------------------------------------------------------------------
123 +
124 +
125 +class BreakableCUDAGraphCapture:
126 + """Stream-capture context that supports eager breaks via :meth:`add_eager`.
127 +
128 + Usage::
129 +
130 + cap = BreakableCUDAGraphCapture(pool=...)
131 + with cap:
132 + output = model(*static_inputs)
133 + # Later, after copying new inputs into the static buffers:
134 + cap.replay()
135 + # Output tensors live at the same addresses as during capture.
136 +
137 + Thread-local: only one capture may be active per thread.
138 + """
139 +
140 + _tls = threading.local()
141 +
142 + @classmethod
143 + def current(cls) -> BreakableCUDAGraphCapture | None:
144 + return getattr(cls._tls, "active", None)
145 +
146 + @classmethod
147 + def is_active(cls) -> bool:
148 + return cls.current() is not None
149 +
150 + def __init__(self, pool: Any | None = None) -> None:
151 + self.pool = pool
152 + self.segments: list[Callable[[], Any]] = []
153 + self._num_graphs: int = 0
154 + self._num_eager_breaks: int = 0
155 + self._current_graph: torch.cuda.CUDAGraph | None = None
156 + self._capturing: bool = False
157 +
158 + # --- context manager protocol ----------------------------------------
159 +
160 + def __enter__(self) -> BreakableCUDAGraphCapture:
161 + if getattr(BreakableCUDAGraphCapture._tls, "active", None) is not None:
162 + raise RuntimeError("Nested BreakableCUDAGraphCapture is not supported.")
163 + BreakableCUDAGraphCapture._tls.active = self
164 + self._begin_segment()
165 + return self
166 +
167 + def __exit__(self, exc_type, exc, tb) -> None:
168 + try:
169 + self._end_segment()
170 + finally:
171 + BreakableCUDAGraphCapture._tls.active = None
172 +
173 + # --- segment management ----------------------------------------------
174 +
175 + def _begin_segment(self) -> None:
176 + assert not self._capturing
177 + g = torch.cuda.CUDAGraph()
178 + if self.pool is not None:
179 + g.capture_begin(pool=self.pool)
180 + else:
181 + g.capture_begin()
182 + self._current_graph = g
183 + self._capturing = True
184 +
185 + def _end_segment(self) -> None:
186 + if not self._capturing:
187 + return
188 + assert self._current_graph is not None
189 + self._current_graph.capture_end()
190 + self.segments.append(self._current_graph.replay)
191 + self._num_graphs += 1
192 + self._current_graph = None
193 + self._capturing = False
194 +
195 + def add_eager(self, fn: Callable[[], Any]) -> Any:
196 + """End the current capture segment, run ``fn`` eagerly on the
197 + capture stream, record ``fn`` for replay, and start a new segment.
198 +
199 + Returns whatever ``fn`` returned during this (capture-time) call.
200 + Replay does not return values; callers should propagate any
201 + downstream dependencies via static output buffers.
202 + """
203 + self._end_segment()
204 + result = fn()
205 + self.segments.append(fn)
206 + self._num_eager_breaks += 1
207 + self._begin_segment()
208 + return result
209 +
210 + # --- replay ----------------------------------------------------------
211 +
212 + def replay(self) -> None:
213 + for r in self.segments:
214 + r()
215 +
216 + # --- introspection ---------------------------------------------------
217 +
218 + @property
219 + def num_graphs(self) -> int:
220 + return self._num_graphs
221 +
222 + @property
223 + def num_eager_breaks(self) -> int:
224 + return self._num_eager_breaks
225 +
226 + def __repr__(self) -> str:
227 + return (
228 + f"BreakableCUDAGraphCapture(graphs={self.num_graphs}, "
229 + f"eager_breaks={self.num_eager_breaks})"
230 + )
231 +
232 +
233 +# ---------------------------------------------------------------------------
234 +# Wrapper that mirrors CUDAGraphWrapper's interface
235 +# ---------------------------------------------------------------------------
236 +
237 +
238 +@dataclasses.dataclass
239 +class _BreakableEntry:
240 + batch_descriptor: BatchDescriptor
241 + capture: BreakableCUDAGraphCapture | None = None
242 + output: Any = None
243 + input_addresses: list[int] | None = None
244 +
245 +
246 +class BreakableCUDAGraphWrapper:
247 + """Drop-in replacement for :class:`CUDAGraphWrapper` that uses
248 + :class:`BreakableCUDAGraphCapture` instead of a single monolithic
249 + ``torch.cuda.graph()`` capture.
250 +
251 + Same dispatch contract as ``CUDAGraphWrapper``:
252 + * If no ``forward_context`` is available, run the underlying
253 + callable eagerly.
254 + * If runtime mode mismatch / NONE, run eagerly.
255 + * Otherwise, lazily capture per ``batch_descriptor`` and replay
256 + on subsequent invocations with the same descriptor.
257 + """
258 +
259 + _all_instances: ClassVar[weakref.WeakSet[BreakableCUDAGraphWrapper]] = (
260 + weakref.WeakSet()
261 + )
262 +
263 + @classmethod
264 + def clear_all_graphs(cls) -> None:
265 + for instance in list(cls._all_instances):
266 + instance.clear_graphs()
267 +
268 + def __init__(
269 + self,
270 + runnable: Callable[..., Any],
271 + vllm_config: VllmConfig,
272 + ) -> None:
273 + # Unlike the original CUDAGraphWrapper which strictly matches a
274 + # single runtime_mode, this wrapper captures whatever the
275 + # dispatcher emits (any non-NONE runtime_mode) -- breakable's
276 + # capture is identical for prefill and decode, so there's nothing
277 + # to dispatch on at the runtime_mode level. Entries are keyed by
278 + # BatchDescriptor which already encodes batch shape / uniformity.
279 + self.runnable = runnable
280 + self.vllm_config = vllm_config
281 + self.compilation_config = vllm_config.compilation_config
282 + self.graph_pool = current_platform.get_global_graph_pool()
283 + self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
284 +
285 + self.entries: dict[BatchDescriptor, _BreakableEntry] = {}
286 + BreakableCUDAGraphWrapper._all_instances.add(self)
287 +
288 + logger.info_once("Breakable CUDA graph enabled")
289 +
290 + # --- vllm-style attribute forwarding ---------------------------------
291 +
292 + def __getattr__(self, key: str) -> Any:
293 + runnable = self.__dict__.get("runnable")
294 + if runnable is not None and hasattr(runnable, key):
295 + return getattr(runnable, key)
296 + raise AttributeError(key)
297 +
298 + def unwrap(self) -> Callable[..., Any]:
299 + return self.runnable
300 +
301 + @property
302 + def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper:
303 + return self
304 +
305 + def clear_graphs(self) -> None:
306 + self.entries.clear()
307 +
308 + # --- dispatch --------------------------------------------------------
309 +
310 + def __call__(self, *args: Any, **kwargs: Any) -> Any:
311 + if not is_forward_context_available():
312 + return self.runnable(*args, **kwargs)
313 +
314 + forward_context = get_forward_context()
315 + batch_descriptor = forward_context.batch_descriptor
316 + cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
317 +
318 + # Capture whenever the dispatcher says "some cudagraph mode" --
319 + # breakable produces the same artifact regardless of PIECEWISE
320 + # vs FULL, so we match either. Entries are keyed by batch
321 + # descriptor, which already encodes prefill/decode distinctions.
322 + if cudagraph_runtime_mode == CUDAGraphMode.NONE:
323 + return self.runnable(*args, **kwargs)
324 +
325 + assert batch_descriptor is not None
326 + entry = self.entries.get(batch_descriptor)
327 + if entry is None:
328 + entry = _BreakableEntry(batch_descriptor=batch_descriptor)
329 + self.entries[batch_descriptor] = entry
330 +
331 + if entry.capture is None:
332 + return self._capture(entry, args, kwargs)
333 + return self._replay(entry, args, kwargs)
334 +
335 + # --- capture / replay paths -----------------------------------------
336 +
337 + @staticmethod
338 + def _collect_tensor_addresses(
339 + args: tuple[Any, ...], kwargs: dict[str, Any]
340 + ) -> list[int]:
341 + """Flatten tensor data_ptrs from positional and keyword args in a
342 + stable order (positionals first, then kwargs in insertion order).
343 +
344 + Used for the DEBUG-mode address-stability check; covers both call
345 + styles since vLLM models are typically invoked with kwargs.
346 + """
347 + addrs = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]
348 + addrs.extend(
349 + v.data_ptr() for v in kwargs.values() if isinstance(v, torch.Tensor)
350 + )
351 + return addrs
352 +
353 + def _capture(
354 + self,
355 + entry: _BreakableEntry,
356 + args: tuple[Any, ...],
357 + kwargs: dict[str, Any],
358 + ) -> Any:
359 + validate_cudagraph_capturing_enabled()
360 +
361 + entry.input_addresses = self._collect_tensor_addresses(args, kwargs)
362 +
363 + if self.graph_pool is not None:
364 + set_graph_pool_id(self.graph_pool)
365 + else:
366 + set_graph_pool_id(current_platform.graph_pool_handle())
367 +
368 + # Match torch.cuda.graph()'s pre-capture cleanup once per descriptor.
369 + # We drive capture_begin/end directly and bypass torch.cuda.graph(),
370 + # so its built-in gc + empty_cache never fire. Run them here once
371 + # per _capture call -- NOT inside _begin_segment, since this capture
372 + # session may issue many begin/end pairs (one per layer's break),
373 + # and repeated gc would tank capture time the way it did for the
374 + # pre-`gc_disable` piecewise path.
375 + gc.collect()
376 + torch.accelerator.empty_cache()
377 + # Sync the offloader's copy stream before capture so any in-flight
378 + # pre-capture prefetches are complete and don't leak into the graph.
379 + get_offloader().sync_prev_onload()
380 +
381 + capture = BreakableCUDAGraphCapture(pool=self.graph_pool)
382 + with capture:
383 + output = self.runnable(*args, **kwargs)
384 + # Join the offloader's copy stream while we still hold the last
385 + # segment open, so the join is captured into the graph (otherwise
386 + # we get an "unjoined stream" error on subsequent forwards).
387 + get_offloader().join_after_forward()
388 + # Convert output to a weak ref *inside* the capture context so the
389 + # strong ref is dropped before the last segment closes, letting
390 + # the cudagraph pool reclaim/reuse that memory immediately for
391 + # the next batch descriptor's capture.
392 + output = weak_ref_tensors(output)
393 +
394 + entry.capture = capture
395 + entry.output = weak_ref_tensors(output)
396 +
397 + logger.debug(
398 + "Captured breakable cudagraph for %s: %r",
399 + entry.batch_descriptor,
400 + capture,
401 + )
402 + # Return the (already-weak) output from the captured run so the
403 + # caller of model(...) gets a tensor pointing at the cudagraph pool's memory
404 + return output
405 +
406 + def _replay(
407 + self,
408 + entry: _BreakableEntry,
409 + args: tuple[Any, ...],
410 + kwargs: dict[str, Any],
411 + ) -> Any:
412 + if self.is_debugging_mode and entry.input_addresses is not None:
413 + new_addresses = self._collect_tensor_addresses(args, kwargs)
414 + assert new_addresses == entry.input_addresses, (
415 + "Input tensor addresses changed between capture and replay "
416 + f"for {entry.batch_descriptor}. Expected "
417 + f"{entry.input_addresses}, got {new_addresses}."
418 + )
419 + # Sync the offloader's copy stream before replay so any external
420 + # dependencies from pre-capture prefetches are satisfied.
421 + get_offloader().sync_prev_onload()
422 + assert entry.capture is not None
423 + entry.capture.replay()
424 + return entry.output
vllm/config/vllm.py
+12 -1
@@ -1035,6 +1035,13 @@ class VllmConfig:
1035 )
1036 self.compilation_config.mode = CompilationMode.NONE
1037
1038 + if envs.VLLM_USE_BREAKABLE_CUDAGRAPH:
1039 + logger.warning_once(
1040 + "VLLM_USE_BREAKABLE_CUDAGRAPH is set, disabling vLLM's "
1041 + "torch.compile pipeline. Equivalent to -cc.mode=none."
1042 + )
1043 + self.compilation_config.mode = CompilationMode.NONE
1044 +
1045 if self.compilation_config.backend == "eager" or (
1046 self.compilation_config.mode is not None
1047 and self.compilation_config.mode != CompilationMode.VLLM_COMPILE
@@ -1102,6 +1109,7 @@ class VllmConfig:
1109 if (
1110 self.compilation_config.cudagraph_mode.requires_piecewise_compilation()
1111 and self.compilation_config.mode != CompilationMode.VLLM_COMPILE
1112 + and not envs.VLLM_USE_BREAKABLE_CUDAGRAPH
1113 ):
1114 logger.info(
1115 "Cudagraph mode %s is not compatible with compilation mode %s."
@@ -1335,7 +1343,10 @@ class VllmConfig:
1343 )
1344
1345 if self.compilation_config.cudagraph_mode.requires_piecewise_compilation():
1338 - assert self.compilation_config.mode == CompilationMode.VLLM_COMPILE, (
1346 + assert (
1347 + self.compilation_config.mode == CompilationMode.VLLM_COMPILE
1348 + or envs.VLLM_USE_BREAKABLE_CUDAGRAPH
1349 + ), (
1350 "Compilation mode should be CompilationMode.VLLM_COMPILE "
1351 "when cudagraph_mode piecewise cudagraphs is used, "
1352 f"cudagraph_mode={self.compilation_config.cudagraph_mode}"
vllm/envs.py
+5
@@ -145,6 +145,7 @@ if TYPE_CHECKING:
145 VLLM_DP_SIZE: int = 1
146 VLLM_USE_STANDALONE_COMPILE: bool = True
147 VLLM_ENABLE_PREGRAD_PASSES: bool = True
148 + VLLM_USE_BREAKABLE_CUDAGRAPH: bool = False
149 VLLM_DP_MASTER_IP: str = ""
150 VLLM_DP_MASTER_PORT: int = 0
151 VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
@@ -637,6 +638,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
638 "VLLM_ENABLE_PREGRAD_PASSES": lambda: (
639 os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "1") == "1"
640 ),
641 + # Experimental: breakable cudagraph does not rely on torch.compile
642 + "VLLM_USE_BREAKABLE_CUDAGRAPH": lambda: (
643 + os.environ.get("VLLM_USE_BREAKABLE_CUDAGRAPH", "0") == "1"
644 + ),
645 # Debug pattern matching inside custom passes.
646 # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3').
647 "VLLM_PATTERN_MATCH_DEBUG": lambda: os.environ.get(
vllm/model_executor/layers/attention/mla_attention.py
+2
@@ -200,6 +200,7 @@ from tqdm import tqdm
200 import vllm.envs as envs
201 from vllm import _custom_ops as ops
202 from vllm._aiter_ops import rocm_aiter_ops
203 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
204 from vllm.config import (
205 CacheConfig,
206 ModelConfig,
@@ -1036,6 +1037,7 @@ direct_register_custom_op(
1037 )
1038
1039
1040 +@eager_break_during_capture
1041 @maybe_transfer_kv_layer
1042 def unified_mla_attention_with_output(
1043 q: torch.Tensor,
vllm/model_executor/layers/deepseek_v4_attention.py
+2
@@ -14,6 +14,7 @@ import torch.nn.functional as F
14 from transformers import DeepseekV2Config, DeepseekV3Config
15
16 import vllm.envs as envs
17 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
18 from vllm.model_executor.layers.linear import (
19 ReplicatedLinear,
20 )
@@ -553,6 +554,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
554 )
555
556
557 +@eager_break_during_capture
558 def deepseek_v4_attention(
559 hidden_states: torch.Tensor,
560 positions: torch.Tensor,
vllm/model_executor/layers/sparse_attn_indexer.py
+2
@@ -6,6 +6,7 @@ import torch
6
7 import vllm.envs as envs
8 from vllm._aiter_ops import rocm_aiter_ops
9 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
10 from vllm.forward_context import get_forward_context
11 from vllm.logger import init_logger
12 from vllm.model_executor.custom_op import CustomOp
@@ -81,6 +82,7 @@ def kv_cache_as_quant_view(
82 return kv_cache.unsqueeze(-2)
83
84
85 +@eager_break_during_capture
86 def sparse_attn_indexer(
87 hidden_states: torch.Tensor,
88 k_cache_prefix: LayerNameType,
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+2
@@ -8,6 +8,7 @@ from importlib.util import find_spec
8 import torch
9 import torch.nn.functional as F
10
11 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
12 from vllm.forward_context import get_forward_context
13 from vllm.platforms import current_platform
14 from vllm.triton_utils import tl, triton
@@ -836,6 +837,7 @@ def rocm_aiter_sparse_attn_indexer_native(
837 return topk_indices_buffer
838
839
840 +@eager_break_during_capture
841 def rocm_aiter_sparse_attn_indexer(
842 hidden_states: torch.Tensor,
843 k_cache_prefix: LayerNameType,
vllm/v1/cudagraph_dispatcher.py
+5
@@ -46,9 +46,14 @@ class CudagraphDispatcher:
46 CUDAGraphMode.FULL: set(),
47 }
48
49 + from vllm.compilation.breakable_cudagraph import (
50 + is_breakable_cudagraph_enabled,
51 + )
52 +
53 assert (
54 not self.compilation_config.cudagraph_mode.requires_piecewise_compilation()
55 or self.compilation_config.is_attention_compiled_piecewise()
56 + or is_breakable_cudagraph_enabled()
57 ), (
58 "Compilation mode should be CompilationMode.VLLM_COMPILE when "
59 "cudagraph_mode piecewise cudagraphs is used, "
vllm/v1/worker/gpu_model_runner.py
+22 -3
@@ -21,6 +21,10 @@ import torch.nn as nn
21 from tqdm import tqdm
22
23 import vllm.envs as envs
24 +from vllm.compilation.breakable_cudagraph import (
25 + BreakableCUDAGraphWrapper,
26 + is_breakable_cudagraph_enabled,
27 +)
28 from vllm.compilation.counter import compilation_counter
29 from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper
30 from vllm.compilation.monitor import set_cudagraph_capturing_enabled
@@ -3127,7 +3131,9 @@ class GPUModelRunner(
3131 def get_model(self) -> nn.Module:
3132 if not hasattr(self, "model"):
3133 raise ValueError("Cannot get model before model has been initialized")
3130 - if isinstance(self.model, (CUDAGraphWrapper, UBatchWrapper)):
3134 + if isinstance(
3135 + self.model, (CUDAGraphWrapper, UBatchWrapper, BreakableCUDAGraphWrapper)
3136 + ):
3137 # get raw model out of the cudagraph wrapper.
3138 return self.model.unwrap()
3139 return self.model
@@ -5130,6 +5136,12 @@ class GPUModelRunner(
5136 cudagraph_mode = self.compilation_config.cudagraph_mode
5137 assert cudagraph_mode is not None
5138 if (
5139 + is_breakable_cudagraph_enabled()
5140 + and cudagraph_mode != CUDAGraphMode.NONE
5141 + and not self.parallel_config.use_ubatching
5142 + ):
5143 + self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config)
5144 + elif (
5145 cudagraph_mode.has_full_cudagraphs()
5146 and not self.parallel_config.use_ubatching
5147 ):
@@ -6206,7 +6218,10 @@ class GPUModelRunner(
6218 # Use a temporary pool for profiling to avoid fragmentation in the main pool.
6219 profiling_pool = current_platform.graph_pool_handle()
6220 original_pools: dict[int, Any] = {}
6209 - for instance in list(CUDAGraphWrapper._all_instances):
6221 + all_wrappers = list(CUDAGraphWrapper._all_instances) + list(
6222 + BreakableCUDAGraphWrapper._all_instances
6223 + )
6224 + for instance in all_wrappers:
6225 original_pools[id(instance)] = instance.graph_pool
6226 instance.graph_pool = profiling_pool
6227
@@ -6257,7 +6272,11 @@ class GPUModelRunner(
6272
6273 set_cudagraph_capturing_enabled(False)
6274 CUDAGraphWrapper.clear_all_graphs()
6260 - for instance in list(CUDAGraphWrapper._all_instances):
6275 + BreakableCUDAGraphWrapper.clear_all_graphs()
6276 + all_wrappers = list(CUDAGraphWrapper._all_instances) + list(
6277 + BreakableCUDAGraphWrapper._all_instances
6278 + )
6279 + for instance in all_wrappers:
6280 if id(instance) in original_pools:
6281 instance.graph_pool = original_pools[id(instance)]
6282 for key_set in self.cudagraph_dispatcher.cudagraph_keys.values():