| 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 logging |
| 29 | import threading |
| 30 | import weakref |
| 31 | from collections.abc import Callable |
| 32 | from typing import Any, ClassVar, TypeVar |
| 33 | |
| 34 | import torch |
| 35 | |
| 36 | import vllm.envs as envs |
| 37 | from vllm.compilation.monitor import validate_cudagraph_capturing_enabled |
| 38 | from vllm.config import CUDAGraphMode, VllmConfig |
| 39 | from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id |
| 40 | from vllm.forward_context import ( |
| 41 | BatchDescriptor, |
| 42 | get_forward_context, |
| 43 | is_forward_context_available, |
| 44 | ) |
| 45 | from vllm.logger import init_logger |
| 46 | from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation |
| 47 | from vllm.model_executor.offloader.base import get_offloader |
| 48 | from vllm.platforms import current_platform |
| 49 | from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors |
| 50 | |
| 51 | logger = init_logger(__name__) |
| 52 | |
| 53 | |
| 54 | def is_breakable_cudagraph_enabled() -> bool: |
| 55 | return bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH) |
| 56 | |
| 57 | |
| 58 | F = TypeVar("F", bound=Callable[..., Any]) |
| 59 | |
| 60 | |
| 61 | def _weak_ref_capture_arg(arg: Any) -> Any: |
| 62 | if isinstance(arg, QuantizedActivation): |
| 63 | return arg.weak_ref() |
| 64 | return weak_ref_tensor(arg) |
| 65 | |
| 66 | |
| 67 | def eager_break_during_capture(fn: F) -> F: |
| 68 | """Decorator that turns a custom-op Python kernel into a "break point" |
| 69 | for the breakable cudagraph capture. |
| 70 | |
| 71 | When the decorated function is invoked outside of a |
| 72 | :class:`BreakableCUDAGraphCapture` context, it executes normally. |
| 73 | |
| 74 | When invoked inside a capture context, it ends the current cudagraph |
| 75 | segment, runs the function eagerly on the capture stream, records the |
| 76 | callable for replay, and starts a fresh segment. |
| 77 | |
| 78 | **In-place output buffer required.** Decorated ops must write into a |
| 79 | caller-provided output tensor; a fresh tensor returned by ``fn`` would |
| 80 | change address each replay and break downstream graph segments. |
| 81 | |
| 82 | **Decorator order matters.** Apply as the *outermost* decorator if |
| 83 | there are other decorators that introduce host-side side effects |
| 84 | around the call -- the canonical example is |
| 85 | ``@maybe_transfer_kv_layer`` for PD-disaggregation, whose |
| 86 | ``wait_for_layer_load`` and ``save_kv_layer`` calls must run in the |
| 87 | eager segment, not inside the captured cudagraph. Putting |
| 88 | ``@eager_break_during_capture`` *inside* such a decorator would |
| 89 | record those side effects into the graph and hang on replay. |
| 90 | |
| 91 | The correct order is:: |
| 92 | |
| 93 | @eager_break_during_capture # outermost |
| 94 | @maybe_transfer_kv_layer |
| 95 | def unified_attention_with_output(...): |
| 96 | ... |
| 97 | """ |
| 98 | if not is_breakable_cudagraph_enabled(): |
| 99 | return fn |
| 100 | |
| 101 | @functools.wraps(fn) |
| 102 | def wrapper(*args: Any, **kwargs: Any) -> Any: |
| 103 | capture = BreakableCUDAGraphCapture.current() |
| 104 | if capture is None: |
| 105 | return fn(*args, **kwargs) |
| 106 | if not capture._capturing: |
| 107 | return fn(*args, **kwargs) |
| 108 | if is_forward_context_available(): |
| 109 | mode = get_forward_context().cudagraph_runtime_mode |
| 110 | if mode == CUDAGraphMode.FULL: |
| 111 | return fn(*args, **kwargs) |
| 112 | |
| 113 | # Weak-ref args: strong refs in the replay lambda pin cudagraph-pool |
| 114 | # slots across batch descriptors. cudagraph owns the slot, so the |
| 115 | # weak_ref is safe to deref on replay. |
| 116 | weak_args = tuple(_weak_ref_capture_arg(a) for a in args) |
| 117 | weak_kwargs = {k: _weak_ref_capture_arg(v) for k, v in kwargs.items()} |
| 118 | return capture.add_eager(lambda: fn(*weak_args, **weak_kwargs)) |
| 119 | |
| 120 | return wrapper # type: ignore[return-value] |
| 121 | |
| 122 | |
| 123 | # --------------------------------------------------------------------------- |
| 124 | # Capture context |
| 125 | # --------------------------------------------------------------------------- |
| 126 | |
| 127 | |
| 128 | class BreakableCUDAGraphCapture: |
| 129 | """Stream-capture context that supports eager breaks via :meth:`add_eager`. |
| 130 | |
| 131 | Usage:: |
| 132 | |
| 133 | cap = BreakableCUDAGraphCapture(pool=...) |
| 134 | with cap: |
| 135 | output = model(*static_inputs) |
| 136 | # Later, after copying new inputs into the static buffers: |
| 137 | cap.replay() |
| 138 | # Output tensors live at the same addresses as during capture. |
| 139 | |
| 140 | Thread-local: only one capture may be active per thread. |
| 141 | """ |
| 142 | |
| 143 | _tls = threading.local() |
| 144 | |
| 145 | @classmethod |
| 146 | def current(cls) -> BreakableCUDAGraphCapture | None: |
| 147 | return getattr(cls._tls, "active", None) |
| 148 | |
| 149 | @classmethod |
| 150 | def is_active(cls) -> bool: |
| 151 | return cls.current() is not None |
| 152 | |
| 153 | def __init__(self, pool: Any | None = None) -> None: |
| 154 | self.pool = pool |
| 155 | self.segments: list[Callable[[], Any]] = [] |
| 156 | self._num_graphs: int = 0 |
| 157 | self._num_eager_breaks: int = 0 |
| 158 | self._current_graph: torch.cuda.CUDAGraph | None = None |
| 159 | self._capturing: bool = False |
| 160 | |
| 161 | # --- context manager protocol ---------------------------------------- |
| 162 | |
| 163 | def __enter__(self) -> BreakableCUDAGraphCapture: |
| 164 | if getattr(BreakableCUDAGraphCapture._tls, "active", None) is not None: |
| 165 | raise RuntimeError("Nested BreakableCUDAGraphCapture is not supported.") |
| 166 | BreakableCUDAGraphCapture._tls.active = self |
| 167 | self._begin_segment() |
| 168 | return self |
| 169 | |
| 170 | def __exit__(self, exc_type, exc, tb) -> None: |
| 171 | try: |
| 172 | self._end_segment() |
| 173 | finally: |
| 174 | BreakableCUDAGraphCapture._tls.active = None |
| 175 | |
| 176 | # --- segment management ---------------------------------------------- |
| 177 | |
| 178 | def _begin_segment(self) -> None: |
| 179 | assert not self._capturing |
| 180 | g = torch.cuda.CUDAGraph() |
| 181 | if self.pool is not None: |
| 182 | g.capture_begin(pool=self.pool) |
| 183 | else: |
| 184 | g.capture_begin() |
| 185 | self._current_graph = g |
| 186 | self._capturing = True |
| 187 | |
| 188 | def _end_segment(self) -> None: |
| 189 | if not self._capturing: |
| 190 | return |
| 191 | assert self._current_graph is not None |
| 192 | self._current_graph.capture_end() |
| 193 | self.segments.append(self._current_graph.replay) |
| 194 | self._num_graphs += 1 |
| 195 | self._current_graph = None |
| 196 | self._capturing = False |
| 197 | |
| 198 | def add_eager(self, fn: Callable[[], Any]) -> Any: |
| 199 | """End the current capture segment, run ``fn`` eagerly on the |
| 200 | capture stream, record ``fn`` for replay, and start a new segment. |
| 201 | |
| 202 | Returns whatever ``fn`` returned during this (capture-time) call. |
| 203 | Replay does not return values; callers should propagate any |
| 204 | downstream dependencies via static output buffers. |
| 205 | """ |
| 206 | self._end_segment() |
| 207 | result = fn() |
| 208 | self.segments.append(fn) |
| 209 | self._num_eager_breaks += 1 |
| 210 | self._begin_segment() |
| 211 | return result |
| 212 | |
| 213 | # --- replay ---------------------------------------------------------- |
| 214 | |
| 215 | def replay(self) -> None: |
| 216 | for r in self.segments: |
| 217 | r() |
| 218 | |
| 219 | # --- introspection --------------------------------------------------- |
| 220 | |
| 221 | @property |
| 222 | def num_graphs(self) -> int: |
| 223 | return self._num_graphs |
| 224 | |
| 225 | @property |
| 226 | def num_eager_breaks(self) -> int: |
| 227 | return self._num_eager_breaks |
| 228 | |
| 229 | def __repr__(self) -> str: |
| 230 | return ( |
| 231 | f"BreakableCUDAGraphCapture(graphs={self.num_graphs}, " |
| 232 | f"eager_breaks={self.num_eager_breaks})" |
| 233 | ) |
| 234 | |
| 235 | |
| 236 | # --------------------------------------------------------------------------- |
| 237 | # Wrapper that mirrors CUDAGraphWrapper's interface |
| 238 | # --------------------------------------------------------------------------- |
| 239 | |
| 240 | |
| 241 | @dataclasses.dataclass |
| 242 | class _BreakableEntry: |
| 243 | batch_descriptor: BatchDescriptor |
| 244 | capture: BreakableCUDAGraphCapture | None = None |
| 245 | output: Any = None |
| 246 | input_addresses: list[int] | None = None |
| 247 | |
| 248 | |
| 249 | class BreakableCUDAGraphWrapper: |
| 250 | """Drop-in replacement for :class:`CUDAGraphWrapper` that uses |
| 251 | :class:`BreakableCUDAGraphCapture` instead of a single monolithic |
| 252 | ``torch.cuda.graph()`` capture. |
| 253 | |
| 254 | Same dispatch contract as ``CUDAGraphWrapper``: |
| 255 | * If no ``forward_context`` is available, run the underlying |
| 256 | callable eagerly. |
| 257 | * If runtime mode mismatch / NONE, run eagerly. |
| 258 | * Otherwise, lazily capture per ``batch_descriptor`` and replay |
| 259 | on subsequent invocations with the same descriptor. |
| 260 | |
| 261 | When ``runtime_mode`` is given, only that mode is captured/replayed; |
| 262 | other non-NONE modes fall through, so the wrapper can be nested under |
| 263 | a ``CUDAGraphWrapper`` for another mode. |
| 264 | """ |
| 265 | |
| 266 | _all_instances: ClassVar[weakref.WeakSet[BreakableCUDAGraphWrapper]] = ( |
| 267 | weakref.WeakSet() |
| 268 | ) |
| 269 | |
| 270 | @classmethod |
| 271 | def clear_all_graphs(cls) -> None: |
| 272 | for instance in list(cls._all_instances): |
| 273 | instance.clear_graphs() |
| 274 | |
| 275 | def __init__( |
| 276 | self, |
| 277 | runnable: Callable[..., Any], |
| 278 | vllm_config: VllmConfig, |
| 279 | runtime_mode: CUDAGraphMode | None = None, |
| 280 | ) -> None: |
| 281 | # runtime_mode=None matches any non-NONE dispatch mode; entries |
| 282 | # are keyed by BatchDescriptor, which already encodes batch |
| 283 | # shape / uniformity. |
| 284 | self.runnable = runnable |
| 285 | self.vllm_config = vllm_config |
| 286 | self.runtime_mode = runtime_mode |
| 287 | self.compilation_config = vllm_config.compilation_config |
| 288 | self.graph_pool = current_platform.get_global_graph_pool() |
| 289 | self.is_debugging_mode = logger.isEnabledFor(logging.DEBUG) |
| 290 | |
| 291 | self.entries: dict[BatchDescriptor, _BreakableEntry] = {} |
| 292 | BreakableCUDAGraphWrapper._all_instances.add(self) |
| 293 | |
| 294 | logger.info_once("Breakable CUDA graph enabled") |
| 295 | |
| 296 | # --- vllm-style attribute forwarding --------------------------------- |
| 297 | |
| 298 | def __getattr__(self, key: str) -> Any: |
| 299 | runnable = self.__dict__.get("runnable") |
| 300 | if runnable is not None and hasattr(runnable, key): |
| 301 | return getattr(runnable, key) |
| 302 | raise AttributeError(key) |
| 303 | |
| 304 | def unwrap(self) -> Callable[..., Any]: |
| 305 | # Recurse through nested wrappers to reach the original runnable. |
| 306 | runnable = self.runnable |
| 307 | return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable |
| 308 | |
| 309 | @property |
| 310 | def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper: |
| 311 | return self |
| 312 | |
| 313 | def clear_graphs(self) -> None: |
| 314 | self.entries.clear() |
| 315 | |
| 316 | # --- dispatch -------------------------------------------------------- |
| 317 | |
| 318 | def __call__(self, *args: Any, **kwargs: Any) -> Any: |
| 319 | if not is_forward_context_available(): |
| 320 | return self.runnable(*args, **kwargs) |
| 321 | |
| 322 | forward_context = get_forward_context() |
| 323 | batch_descriptor = forward_context.batch_descriptor |
| 324 | cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode |
| 325 | |
| 326 | # By default (runtime_mode=None) match any non-NONE mode; a set |
| 327 | # runtime_mode restricts capture so other modes can be handled by |
| 328 | # an outer wrapper (e.g. CUDAGraphWrapper for FULL). |
| 329 | if cudagraph_runtime_mode == CUDAGraphMode.NONE or ( |
| 330 | self.runtime_mode is not None |
| 331 | and cudagraph_runtime_mode != self.runtime_mode |
| 332 | ): |
| 333 | return self.runnable(*args, **kwargs) |
| 334 | |
| 335 | assert batch_descriptor is not None |
| 336 | entry = self.entries.get(batch_descriptor) |
| 337 | if entry is None: |
| 338 | entry = _BreakableEntry(batch_descriptor=batch_descriptor) |
| 339 | self.entries[batch_descriptor] = entry |
| 340 | |
| 341 | if entry.capture is None: |
| 342 | return self._capture(entry, args, kwargs) |
| 343 | return self._replay(entry, args, kwargs) |
| 344 | |
| 345 | # --- capture / replay paths ----------------------------------------- |
| 346 | |
| 347 | @staticmethod |
| 348 | def _collect_tensor_addresses( |
| 349 | args: tuple[Any, ...], kwargs: dict[str, Any] |
| 350 | ) -> list[int]: |
| 351 | """Flatten tensor data_ptrs from positional and keyword args in a |
| 352 | stable order (positionals first, then kwargs in insertion order). |
| 353 | |
| 354 | Used for the DEBUG-mode address-stability check; covers both call |
| 355 | styles since vLLM models are typically invoked with kwargs. |
| 356 | """ |
| 357 | addrs = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)] |
| 358 | addrs.extend( |
| 359 | v.data_ptr() for v in kwargs.values() if isinstance(v, torch.Tensor) |
| 360 | ) |
| 361 | return addrs |
| 362 | |
| 363 | def _capture( |
| 364 | self, |
| 365 | entry: _BreakableEntry, |
| 366 | args: tuple[Any, ...], |
| 367 | kwargs: dict[str, Any], |
| 368 | ) -> Any: |
| 369 | validate_cudagraph_capturing_enabled() |
| 370 | |
| 371 | entry.input_addresses = self._collect_tensor_addresses(args, kwargs) |
| 372 | |
| 373 | if self.graph_pool is not None: |
| 374 | set_graph_pool_id(self.graph_pool) |
| 375 | else: |
| 376 | set_graph_pool_id(current_platform.graph_pool_handle()) |
| 377 | |
| 378 | # Match torch.cuda.graph()'s pre-capture cleanup, which we bypass. |
| 379 | # Skip it when gc is disabled: bulk capture runs under |
| 380 | # freeze_gc_for_cudagraph_capture, which already did this cleanup, |
| 381 | # and repeating it per descriptor dominates capture time. |
| 382 | if gc.isenabled(): |
| 383 | gc.collect() |
| 384 | torch.accelerator.empty_cache() |
| 385 | # Sync the offloader's copy stream before capture so any in-flight |
| 386 | # pre-capture prefetches are complete and don't leak into the graph. |
| 387 | get_offloader().sync_prev_onload() |
| 388 | |
| 389 | capture = BreakableCUDAGraphCapture(pool=self.graph_pool) |
| 390 | with capture: |
| 391 | output = self.runnable(*args, **kwargs) |
| 392 | # Join the offloader's copy stream while we still hold the last |
| 393 | # segment open, so the join is captured into the graph (otherwise |
| 394 | # we get an "unjoined stream" error on subsequent forwards). |
| 395 | get_offloader().join_after_forward() |
| 396 | # Convert output to a weak ref *inside* the capture context so the |
| 397 | # strong ref is dropped before the last segment closes, letting |
| 398 | # the cudagraph pool reclaim/reuse that memory immediately for |
| 399 | # the next batch descriptor's capture. |
| 400 | output = weak_ref_tensors(output) |
| 401 | |
| 402 | entry.capture = capture |
| 403 | entry.output = weak_ref_tensors(output) |
| 404 | |
| 405 | logger.debug( |
| 406 | "Captured breakable cudagraph for %s: %r", |
| 407 | entry.batch_descriptor, |
| 408 | capture, |
| 409 | ) |
| 410 | # Return the (already-weak) output from the captured run so the |
| 411 | # caller of model(...) gets a tensor pointing at the cudagraph pool's memory |
| 412 | return output |
| 413 | |
| 414 | def _replay( |
| 415 | self, |
| 416 | entry: _BreakableEntry, |
| 417 | args: tuple[Any, ...], |
| 418 | kwargs: dict[str, Any], |
| 419 | ) -> Any: |
| 420 | if self.is_debugging_mode and entry.input_addresses is not None: |
| 421 | new_addresses = self._collect_tensor_addresses(args, kwargs) |
| 422 | assert new_addresses == entry.input_addresses, ( |
| 423 | "Input tensor addresses changed between capture and replay " |
| 424 | f"for {entry.batch_descriptor}. Expected " |
| 425 | f"{entry.input_addresses}, got {new_addresses}." |
| 426 | ) |
| 427 | # Sync the offloader's copy stream before replay so any external |
| 428 | # dependencies from pre-capture prefetches are satisfied. |
| 429 | get_offloader().sync_prev_onload() |
| 430 | assert entry.capture is not None |
| 431 | entry.capture.replay() |
| 432 | return entry.output |