| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import dataclasses |
| 5 | import logging |
| 6 | import weakref |
| 7 | from collections import Counter |
| 8 | from collections.abc import Callable |
| 9 | from contextlib import ExitStack |
| 10 | from typing import Any, ClassVar |
| 11 | from unittest.mock import patch |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.compilation.counter import compilation_counter |
| 16 | from vllm.compilation.monitor import validate_cudagraph_capturing_enabled |
| 17 | from vllm.config import CUDAGraphMode, VllmConfig |
| 18 | from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id |
| 19 | from vllm.forward_context import ( |
| 20 | BatchDescriptor, |
| 21 | get_forward_context, |
| 22 | is_forward_context_available, |
| 23 | ) |
| 24 | from vllm.logger import init_logger |
| 25 | from vllm.model_executor.offloader.base import get_offloader |
| 26 | from vllm.platforms import current_platform |
| 27 | from vllm.utils.torch_utils import current_stream, weak_ref_tensors |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | |
| 32 | @dataclasses.dataclass(frozen=True) |
| 33 | class CUDAGraphStat: |
| 34 | num_unpadded_tokens: int |
| 35 | num_padded_tokens: int |
| 36 | num_paddings: int |
| 37 | runtime_mode: str |
| 38 | |
| 39 | |
| 40 | class CUDAGraphLogging: |
| 41 | """Aggregate and log cudagraph metrics.""" |
| 42 | |
| 43 | COLUMN_HEADERS = [ |
| 44 | "Unpadded Tokens", |
| 45 | "Padded Tokens", |
| 46 | "Num Paddings", |
| 47 | "Runtime Mode", |
| 48 | "Count", |
| 49 | ] |
| 50 | |
| 51 | def __init__( |
| 52 | self, cg_mode: CUDAGraphMode, cg_capture_sizes: list[int] | None |
| 53 | ) -> None: |
| 54 | self.reset() |
| 55 | self.cg_mode = str(cg_mode) |
| 56 | self.cg_capture_sizes = str(cg_capture_sizes or []) |
| 57 | |
| 58 | self.settings_header = ( |
| 59 | "**CUDAGraph Config Settings:**\n\n" |
| 60 | f"- Mode: {self.cg_mode}\n" |
| 61 | f"- Capture sizes: {self.cg_capture_sizes}\n\n" |
| 62 | "**CUDAGraph Stats:**\n\n" |
| 63 | ) |
| 64 | |
| 65 | def reset(self) -> None: |
| 66 | self.stats: list[CUDAGraphStat] = [] |
| 67 | |
| 68 | def observe(self, cudagraph_stat: CUDAGraphStat) -> None: |
| 69 | self.stats.append(cudagraph_stat) |
| 70 | |
| 71 | def generate_metric_table(self) -> str: |
| 72 | stats_counts = Counter(self.stats) |
| 73 | |
| 74 | # Convert stats to rows of strings, in descending order of observed frequencies |
| 75 | rows = [] |
| 76 | for stat, count in sorted( |
| 77 | stats_counts.items(), key=lambda item: item[1], reverse=True |
| 78 | ): |
| 79 | rows.append( |
| 80 | [ |
| 81 | str(stat.num_unpadded_tokens), |
| 82 | str(stat.num_padded_tokens), |
| 83 | str(stat.num_paddings), |
| 84 | stat.runtime_mode, |
| 85 | str(count), |
| 86 | ] |
| 87 | ) |
| 88 | |
| 89 | # Calculate column widths (max of header and data) |
| 90 | col_widths = [] |
| 91 | for i, header_text in enumerate(self.COLUMN_HEADERS): |
| 92 | max_width = len(header_text) |
| 93 | for row in rows: |
| 94 | max_width = max(max_width, len(row[i])) |
| 95 | col_widths.append(max_width) |
| 96 | |
| 97 | table_header_list = [ |
| 98 | h.ljust(w) for h, w in zip(self.COLUMN_HEADERS, col_widths) |
| 99 | ] |
| 100 | table_header = "| " + " | ".join(table_header_list) + " |\n" |
| 101 | |
| 102 | table_separator = "|" + "|".join("-" * (w + 2) for w in col_widths) + "|\n" |
| 103 | |
| 104 | # Create data rows with proper alignment |
| 105 | data_rows = [] |
| 106 | for row in rows: |
| 107 | formatted_row = [ |
| 108 | str(val).ljust(width) for val, width in zip(row, col_widths) |
| 109 | ] |
| 110 | data_rows.append("| " + " | ".join(formatted_row) + " |") |
| 111 | |
| 112 | return ( |
| 113 | self.settings_header |
| 114 | + table_header |
| 115 | + table_separator |
| 116 | + "\n".join(data_rows) |
| 117 | + "\n" |
| 118 | ) |
| 119 | |
| 120 | def log(self, log_fn: Callable[..., Any] = logger.info) -> None: |
| 121 | if not self.stats: |
| 122 | return |
| 123 | log_fn(self.generate_metric_table()) |
| 124 | self.reset() |
| 125 | |
| 126 | |
| 127 | @dataclasses.dataclass |
| 128 | class CUDAGraphEntry: |
| 129 | batch_descriptor: BatchDescriptor |
| 130 | cudagraph: torch.cuda.CUDAGraph | None = None |
| 131 | output: Any | None = None |
| 132 | |
| 133 | # for cudagraph debugging, track the input addresses |
| 134 | # during capture, and check if they are the same during replay |
| 135 | input_addresses: list[int] | None = None |
| 136 | |
| 137 | |
| 138 | @dataclasses.dataclass |
| 139 | class CUDAGraphOptions: |
| 140 | debug_log_enable: bool = True |
| 141 | gc_disable: bool = False |
| 142 | weak_ref_output: bool = True |
| 143 | |
| 144 | |
| 145 | class CUDAGraphWrapper: |
| 146 | """Wraps a runnable to add CUDA graph capturing and replaying ability. And |
| 147 | provide attribute access to the underlying `runnable` via `__getattr__`. |
| 148 | |
| 149 | The workflow of this wrapper in the cudagraph dispatching is as follows: |
| 150 | 1. At initialization, a runtime mode is assigned to the wrapper (FULL or |
| 151 | PIECEWISE). |
| 152 | 2. At runtime, the wrapper receives a runtime_mode and a |
| 153 | batch_descriptor(key) from the forward context and blindly trust them |
| 154 | for cudagraph dispatching. |
| 155 | 3. If runtime_mode is NONE or runtime_mode does not match the mode of the |
| 156 | wrapper, just call the runnable directly. |
| 157 | 4. Otherwise, i.e., the runtime_mode matches the mode of the wrapper, |
| 158 | the wrapper will perform cudagraph capture(if key does not exist, create |
| 159 | a new entry and cache it) or replay (if key exists in the cache). |
| 160 | |
| 161 | Note: CUDAGraphWrapper does not store persistent buffers or copy any |
| 162 | runtime inputs into that buffers for replay. We assume implementing them |
| 163 | is done outside of the wrapper. That is because we do not make any |
| 164 | assumption on the dynamic shape (batch size) of the runtime inputs, as a |
| 165 | trade-off for staying orthogonal to compilation logic. Nevertheless, |
| 166 | tracing and checking the input addresses to be consistent during replay is |
| 167 | guaranteed when vLLM debug logging is enabled. |
| 168 | """ |
| 169 | |
| 170 | _all_instances: ClassVar[weakref.WeakSet["CUDAGraphWrapper"]] = weakref.WeakSet() |
| 171 | |
| 172 | @classmethod |
| 173 | def clear_all_graphs(cls) -> None: |
| 174 | """Clear captured graphs from all CUDAGraphWrapper instances.""" |
| 175 | for instance in list(cls._all_instances): |
| 176 | instance.clear_graphs() |
| 177 | |
| 178 | def __init__( |
| 179 | self, |
| 180 | runnable: Callable[..., Any], |
| 181 | vllm_config: VllmConfig, |
| 182 | runtime_mode: CUDAGraphMode, |
| 183 | cudagraph_options: CUDAGraphOptions | None = None, |
| 184 | ) -> None: |
| 185 | self.runnable = runnable |
| 186 | self.vllm_config = vllm_config |
| 187 | self.runtime_mode = runtime_mode |
| 188 | self.compilation_config = vllm_config.compilation_config |
| 189 | |
| 190 | self.first_run_finished = False |
| 191 | self.is_debugging_mode = logger.isEnabledFor(logging.DEBUG) |
| 192 | self._runnable_str = str(runnable) if self.is_debugging_mode else None |
| 193 | |
| 194 | # assert runtime_mode is not NONE(no cudagraph), otherwise, we don't |
| 195 | # need to initialize a CUDAGraphWrapper. |
| 196 | assert self.runtime_mode != CUDAGraphMode.NONE |
| 197 | # TODO: in the future, if we want to use multiple |
| 198 | # streams, it might not be safe to share a global pool. |
| 199 | # only investigate this when we use multiple streams |
| 200 | self.graph_pool = current_platform.get_global_graph_pool() |
| 201 | |
| 202 | if cudagraph_options is None: |
| 203 | cudagraph_options = CUDAGraphOptions() |
| 204 | self.cudagraph_options = cudagraph_options |
| 205 | # the entries for different batch descriptors that we need to capture |
| 206 | # cudagraphs for. |
| 207 | self.concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry] = {} |
| 208 | |
| 209 | CUDAGraphWrapper._all_instances.add(self) |
| 210 | |
| 211 | def __getattr__(self, key: str) -> Any: |
| 212 | # allow accessing the attributes of the runnable. |
| 213 | if hasattr(self.runnable, key): |
| 214 | return getattr(self.runnable, key) |
| 215 | if self.is_debugging_mode: |
| 216 | raise AttributeError( |
| 217 | f"Attribute {key} not exists in the runnable of " |
| 218 | f"cudagraph wrapper: {self._runnable_str}" |
| 219 | ) |
| 220 | raise AttributeError |
| 221 | |
| 222 | def unwrap(self) -> Callable[..., Any]: |
| 223 | # in case we need to access the original 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": |
| 230 | return self |
| 231 | |
| 232 | def clear_graphs(self) -> None: |
| 233 | self.concrete_cudagraph_entries.clear() |
| 234 | |
| 235 | def __call__(self, *args: Any, **kwargs: Any) -> Any | None: |
| 236 | if not is_forward_context_available(): |
| 237 | # No forward context means we are outside the normal |
| 238 | # inference path (e.g. a vision encoder forward pass). |
| 239 | # Just run the underlying function without cudagraphs. |
| 240 | return self.runnable(*args, **kwargs) |
| 241 | |
| 242 | forward_context = get_forward_context() |
| 243 | batch_descriptor = forward_context.batch_descriptor |
| 244 | cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode |
| 245 | |
| 246 | if ( |
| 247 | cudagraph_runtime_mode == CUDAGraphMode.NONE |
| 248 | or cudagraph_runtime_mode != self.runtime_mode |
| 249 | ): |
| 250 | # CUDAGraphMode.NONE could mean the profile run, a warmup run, or |
| 251 | # running without cudagraphs. |
| 252 | # We do not trigger capture/replay if the runtime mode is not |
| 253 | # matches. This enables properly dispatching to the correct |
| 254 | # CUDAGraphWrapper when nesting multiple instances with different |
| 255 | # runtime modes. |
| 256 | return self.runnable(*args, **kwargs) |
| 257 | |
| 258 | assert batch_descriptor is not None |
| 259 | if batch_descriptor not in self.concrete_cudagraph_entries: |
| 260 | # create a new entry for this batch descriptor |
| 261 | self.concrete_cudagraph_entries[batch_descriptor] = CUDAGraphEntry( |
| 262 | batch_descriptor=batch_descriptor |
| 263 | ) |
| 264 | |
| 265 | entry = self.concrete_cudagraph_entries[batch_descriptor] |
| 266 | |
| 267 | if entry.cudagraph is None: |
| 268 | if self.cudagraph_options.debug_log_enable: |
| 269 | # Since we capture cudagraph for many different shapes and |
| 270 | # capturing is fast, we don't need to log it for every |
| 271 | # shape. E.g. we only log it for the first subgraph in |
| 272 | # piecewise mode. |
| 273 | logger.debug( |
| 274 | "Capturing a cudagraph on (%s,%s)", |
| 275 | self.runtime_mode.name, |
| 276 | entry.batch_descriptor, |
| 277 | ) |
| 278 | # validate that cudagraph capturing is legal at this point. |
| 279 | validate_cudagraph_capturing_enabled() |
| 280 | |
| 281 | input_addresses = [ |
| 282 | x.data_ptr() for x in args if isinstance(x, torch.Tensor) |
| 283 | ] |
| 284 | entry.input_addresses = input_addresses |
| 285 | cudagraph = torch.cuda.CUDAGraph() |
| 286 | |
| 287 | with ExitStack() as stack: |
| 288 | if self.cudagraph_options.gc_disable: |
| 289 | # during every model forward for piecewise cudagraph |
| 290 | # mode, we will capture many pieces of cudagraphs |
| 291 | # (roughly one per layer). running gc again and again |
| 292 | # across layers will make the cudagraph capture very slow. |
| 293 | # therefore, we only run gc for the first graph, |
| 294 | # and disable gc for the rest of the graphs. |
| 295 | stack.enter_context( |
| 296 | patch("gc.collect", lambda *args, **kwargs: None) |
| 297 | ) |
| 298 | stack.enter_context( |
| 299 | patch( |
| 300 | "torch.accelerator.empty_cache", |
| 301 | lambda *args, **kwargs: None, |
| 302 | ) |
| 303 | ) |
| 304 | |
| 305 | if self.graph_pool is not None: |
| 306 | set_graph_pool_id(self.graph_pool) |
| 307 | else: |
| 308 | set_graph_pool_id(current_platform.graph_pool_handle()) |
| 309 | |
| 310 | # Sync offloader's copy stream before capture. |
| 311 | # Ensure any pre-capture prefetches from offloader are complete. |
| 312 | get_offloader().sync_prev_onload() |
| 313 | |
| 314 | # mind-exploding: carefully manage the reference and memory. |
| 315 | with torch.cuda.graph( |
| 316 | cudagraph, |
| 317 | pool=self.graph_pool, |
| 318 | stream=current_stream(), |
| 319 | ): |
| 320 | # `output` is managed by pytorch's cudagraph pool |
| 321 | output = self.runnable(*args, **kwargs) |
| 322 | # Join offloader's copy stream after forward to avoid |
| 323 | # unjoined stream error. The last layer's start_prefetch |
| 324 | # forks copy_stream, but wait_prefetch only happens in |
| 325 | # the next forward pass. |
| 326 | get_offloader().join_after_forward() |
| 327 | if self.cudagraph_options.weak_ref_output: |
| 328 | # by converting it to weak ref, |
| 329 | # the original `output` will immediately be released |
| 330 | # to save memory. It is only safe to do this for |
| 331 | # the last graph in piecewise cuadgraph mode, because |
| 332 | # the output of the last graph will not be used by |
| 333 | # any other cuda graph. |
| 334 | output = weak_ref_tensors(output) |
| 335 | |
| 336 | # here we always use weak ref for the output |
| 337 | # to save memory |
| 338 | entry.output = weak_ref_tensors(output) |
| 339 | entry.cudagraph = cudagraph |
| 340 | |
| 341 | compilation_counter.num_cudagraph_captured += 1 |
| 342 | |
| 343 | # important: we need to return the output, rather than |
| 344 | # the weak ref of the output, so that pytorch can correctly |
| 345 | # manage the memory during cuda graph capture |
| 346 | return output |
| 347 | |
| 348 | if self.is_debugging_mode: |
| 349 | # check if the input addresses are the same |
| 350 | new_input_addresses = [ |
| 351 | x.data_ptr() for x in args if isinstance(x, torch.Tensor) |
| 352 | ] |
| 353 | assert new_input_addresses == entry.input_addresses, ( |
| 354 | f"Input addresses for cudagraphs are different " |
| 355 | f"during replay. Expected {entry.input_addresses}, " |
| 356 | f"got {new_input_addresses}" |
| 357 | ) |
| 358 | |
| 359 | # Sync offloader before replay - ensures any external dependencies |
| 360 | # from pre-capture prefetches are satisfied. |
| 361 | get_offloader().sync_prev_onload() |
| 362 | entry.cudagraph.replay() |
| 363 | return entry.output |