| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import gc |
| 4 | import itertools |
| 5 | from collections import defaultdict |
| 6 | from collections.abc import Callable, Iterable |
| 7 | from contextlib import ExitStack |
| 8 | from dataclasses import dataclass, replace |
| 9 | from itertools import groupby, product |
| 10 | from typing import TYPE_CHECKING, Any, NamedTuple, Protocol |
| 11 | |
| 12 | import torch |
| 13 | import torch.nn as nn |
| 14 | from tqdm import tqdm |
| 15 | |
| 16 | from vllm.compilation.breakable_cudagraph import ( |
| 17 | BreakableCUDAGraphWrapper, |
| 18 | is_breakable_cudagraph_enabled, |
| 19 | ) |
| 20 | from vllm.compilation.counter import compilation_counter |
| 21 | from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper |
| 22 | from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper |
| 23 | from vllm.config import VllmConfig, set_current_vllm_config |
| 24 | from vllm.config.compilation import CUDAGraphMode |
| 25 | from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id |
| 26 | from vllm.distributed.parallel_state import ( |
| 27 | get_pp_group, |
| 28 | graph_capture, |
| 29 | is_global_first_rank, |
| 30 | ) |
| 31 | from vllm.forward_context import BatchDescriptor, set_forward_context |
| 32 | from vllm.logger import init_logger |
| 33 | from vllm.model_executor.offloader.base import get_offloader |
| 34 | from vllm.platforms import current_platform |
| 35 | from vllm.sequence import IntermediateTensors |
| 36 | from vllm.utils.math_utils import round_up |
| 37 | from vllm.utils.torch_utils import current_stream |
| 38 | from vllm.v1.hisparse.binding import release_hisparse_profiling_cache |
| 39 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 40 | from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup |
| 41 | from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer |
| 42 | from vllm.v1.worker.gpu.block_table import BlockTables |
| 43 | from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens |
| 44 | from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers |
| 45 | from vllm.v1.worker.gpu.model_states.interface import ModelState |
| 46 | from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches |
| 47 | from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches |
| 48 | |
| 49 | if TYPE_CHECKING: |
| 50 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner |
| 51 | from vllm.v1.worker.gpu.pcp_manager import PCPManager |
| 52 | from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner |
| 53 | |
| 54 | logger = init_logger(__name__) |
| 55 | |
| 56 | |
| 57 | class AttentionState(NamedTuple): |
| 58 | attn_metadata: dict[str, Any] | None |
| 59 | slot_mappings: dict[str, torch.Tensor] |
| 60 | |
| 61 | |
| 62 | @dataclass(frozen=True) |
| 63 | class BatchExecutionDescriptor: |
| 64 | """Describes the shape of the batch and CG mode to run; this is used to make shape |
| 65 | matches between the capture and runtime.""" |
| 66 | |
| 67 | cg_mode: CUDAGraphMode |
| 68 | num_tokens: int |
| 69 | num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs) |
| 70 | uniform_token_count: int | None = None |
| 71 | # Upper bound on per-request query length. Varlen decode graphs leave |
| 72 | # uniform_token_count unset, so this is what keeps a prefill batch out of one: |
| 73 | # the runner passes None for any batch with a prefill. |
| 74 | max_query_len: int | None = None |
| 75 | num_active_loras: int = 0 |
| 76 | # Number of microbatches the batch is split into (DBO). 1 means no splitting. |
| 77 | num_ubatches: int = 1 |
| 78 | |
| 79 | |
| 80 | def make_cudagraph_stats( |
| 81 | batch_desc: BatchExecutionDescriptor, num_tokens: int |
| 82 | ) -> CUDAGraphStat: |
| 83 | return CUDAGraphStat( |
| 84 | num_unpadded_tokens=num_tokens, |
| 85 | num_padded_tokens=batch_desc.num_tokens, |
| 86 | num_paddings=batch_desc.num_tokens - num_tokens, |
| 87 | runtime_mode=str(batch_desc.cg_mode), |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | class CreateForwardFn(Protocol): |
| 92 | """Factory that prepares inputs (OUTSIDE the graph) and returns a |
| 93 | forward_fn. Called with warmup=True for the warmup pass and warmup=False |
| 94 | for the captured pass.""" |
| 95 | |
| 96 | def __call__( |
| 97 | self, |
| 98 | desc: BatchExecutionDescriptor, |
| 99 | warmup: bool, |
| 100 | ) -> Callable[[CUDAGraphMode], None]: ... |
| 101 | |
| 102 | |
| 103 | def _is_compatible( |
| 104 | desc: BatchExecutionDescriptor, |
| 105 | num_reqs: int, |
| 106 | num_tokens: int, |
| 107 | uniform_token_count: int | None, |
| 108 | num_active_loras: int, |
| 109 | max_query_len: int | None, |
| 110 | num_ubatches: int, |
| 111 | ) -> bool: |
| 112 | # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count |
| 113 | # desc.num_reqs=None means no request padding needed (PIECEWISE) |
| 114 | # desc.max_query_len=None means the graph does not constrain query length; a |
| 115 | # caller that does not track max_query_len must not match one that does |
| 116 | # A graph captured for N microbatches can only serve a batch split N ways. |
| 117 | return ( |
| 118 | ( |
| 119 | desc.uniform_token_count is None |
| 120 | or desc.uniform_token_count == uniform_token_count |
| 121 | ) |
| 122 | and ( |
| 123 | desc.max_query_len is None |
| 124 | or (max_query_len is not None and desc.max_query_len >= max_query_len) |
| 125 | ) |
| 126 | and (desc.num_reqs is None or desc.num_reqs >= num_reqs) |
| 127 | and desc.num_tokens >= num_tokens |
| 128 | and desc.num_active_loras == num_active_loras |
| 129 | and desc.num_ubatches == num_ubatches |
| 130 | ) |
| 131 | |
| 132 | |
| 133 | def has_compiled_submodule(model: nn.Module) -> bool: |
| 134 | """Whether any submodule is an active @support_torch_compile module.""" |
| 135 | return any( |
| 136 | isinstance(m, TorchCompileWithNoGuardsWrapper) |
| 137 | and not getattr(m, "do_not_compile", True) |
| 138 | for m in model.modules() |
| 139 | ) |
| 140 | |
| 141 | |
| 142 | class CudaGraphManager: |
| 143 | def __init__( |
| 144 | self, |
| 145 | vllm_config: VllmConfig, |
| 146 | device: torch.device, |
| 147 | cudagraph_mode: CUDAGraphMode, |
| 148 | decode_query_len: int, |
| 149 | lora_capture_cases: list[int] | None = None, |
| 150 | varlen_decode: bool = False, |
| 151 | ubatch_runner: "UBatchRunner | None" = None, |
| 152 | ): |
| 153 | self.vllm_config = vllm_config |
| 154 | self.device = device |
| 155 | self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs |
| 156 | self.compilation_config = vllm_config.compilation_config |
| 157 | assert self.compilation_config is not None |
| 158 | self.cudagraph_mode = cudagraph_mode |
| 159 | self.decode_query_len = decode_query_len |
| 160 | self.varlen_decode = varlen_decode |
| 161 | # DBO supports FULL CUDA graphs only. |
| 162 | self.ubatch_runner = ubatch_runner |
| 163 | |
| 164 | self.tp_size = vllm_config.parallel_config.tensor_parallel_size |
| 165 | self.is_first_pp_rank = get_pp_group().is_first_rank |
| 166 | self.is_last_pp_rank = get_pp_group().is_last_rank |
| 167 | self.lora_capture_cases = lora_capture_cases or [0] |
| 168 | # Precompute actual num_active_loras -> captured case mapping so that |
| 169 | # dispatch() is a plain dict lookup instead of a per-call bisect. |
| 170 | self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map() |
| 171 | |
| 172 | self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {} |
| 173 | self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None |
| 174 | |
| 175 | self._graphs_captured = False |
| 176 | |
| 177 | # Profiling hooks, set only by profile_cudagraph_memory() below: cap |
| 178 | # FULL-mode capture at the N largest descriptors and record each |
| 179 | # captured FULL graph's memory delta for extrapolation. |
| 180 | self._max_full_descs_to_capture: int | None = None |
| 181 | self._capture_mem_samples: list[int] | None = None |
| 182 | |
| 183 | self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {} |
| 184 | self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} |
| 185 | |
| 186 | # Breakable CUDA graph (PW CUDA graph without torch.compile) |
| 187 | self.use_breakable_cg = ( |
| 188 | is_breakable_cudagraph_enabled() |
| 189 | and self.cudagraph_mode.has_piecewise_cudagraphs() |
| 190 | ) |
| 191 | self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None |
| 192 | |
| 193 | self._init_candidates() |
| 194 | |
| 195 | def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]: |
| 196 | """Precompute actual num_active_loras -> effective captured case. |
| 197 | |
| 198 | Mirrors the num_tokens candidate expansion in ``_init_candidates``: |
| 199 | every possible active-LoRA count is mapped ahead of time to the |
| 200 | smallest captured case that can serve it, so ``dispatch`` is a plain |
| 201 | dict lookup instead of a per-call bisect. |
| 202 | """ |
| 203 | captured_with_lora = sorted(c for c in self.lora_capture_cases if c > 0) |
| 204 | if not captured_with_lora: |
| 205 | return {}, 0 |
| 206 | dispatch_map: dict[int, int] = {} |
| 207 | case_idx = 0 |
| 208 | for n in range(1, captured_with_lora[-1] + 1): |
| 209 | while captured_with_lora[case_idx] < n: |
| 210 | case_idx += 1 |
| 211 | dispatch_map[n] = captured_with_lora[case_idx] |
| 212 | return dispatch_map, captured_with_lora[-1] |
| 213 | |
| 214 | def _resolve_effective_loras(self, num_active_loras: int) -> int: |
| 215 | """Map an actual active-LoRA count to its captured graph case.""" |
| 216 | if num_active_loras <= 0 or not self._lora_dispatch_map: |
| 217 | return num_active_loras |
| 218 | # Counts above the largest captured case clamp to it. |
| 219 | return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case) |
| 220 | |
| 221 | def _maybe_ubatch_twin( |
| 222 | self, desc: BatchExecutionDescriptor |
| 223 | ) -> BatchExecutionDescriptor | None: |
| 224 | """Return a microbatched capture candidate when eligible. |
| 225 | |
| 226 | Uniform query lengths preserve the captured request split. Use the DP |
| 227 | dispatch thresholds so all ranks generate the same candidates. |
| 228 | """ |
| 229 | if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL: |
| 230 | return None |
| 231 | if desc.num_reqs is None: |
| 232 | return None |
| 233 | uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs) |
| 234 | if remainder or desc.uniform_token_count not in (None, uniform_token_count): |
| 235 | return None |
| 236 | parallel_config = self.vllm_config.parallel_config |
| 237 | num_ubatches = get_num_ubatches(parallel_config) |
| 238 | if desc.num_tokens < num_ubatches: |
| 239 | return None |
| 240 | if not check_ubatch_thresholds( |
| 241 | parallel_config, desc.num_tokens, uniform_decode=True |
| 242 | ): |
| 243 | return None |
| 244 | return replace( |
| 245 | desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count |
| 246 | ) |
| 247 | |
| 248 | def _init_candidates(self) -> None: |
| 249 | """Build priority-ordered candidate lists for each token count.""" |
| 250 | capture_sizes = self.compilation_config.cudagraph_capture_sizes |
| 251 | if not (self.cudagraph_mode and capture_sizes): |
| 252 | return |
| 253 | |
| 254 | capture_sizes = sorted(capture_sizes) |
| 255 | max_decode_tokens = self.max_num_reqs * self.decode_query_len |
| 256 | decode_mode = self.cudagraph_mode.decode_mode() |
| 257 | mixed_mode = self.cudagraph_mode.mixed_mode() |
| 258 | separate_decode_routine = self.cudagraph_mode.separate_routine() |
| 259 | max_cg_capture_size = self.compilation_config.max_cudagraph_capture_size |
| 260 | |
| 261 | descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = ( |
| 262 | defaultdict(list) |
| 263 | ) |
| 264 | |
| 265 | # When using Dynamic SD, num_speculative_tokens is the max number of |
| 266 | # draft tokens. The scheduler might use a smaller number so we need |
| 267 | # to capture graphs for all possible values during decode. |
| 268 | speculative_config = self.vllm_config.speculative_config |
| 269 | if ( |
| 270 | speculative_config |
| 271 | and speculative_config.uses_dynamic_speculative_decoding() |
| 272 | ): |
| 273 | # decode_query_len = num_speculative_steps + num_new_sampled_tokens |
| 274 | # _per_step. Recover num_new_sampled_tokens_per_step |
| 275 | # from the values the manager already has. |
| 276 | num_new_sampled_tokens_per_step = ( |
| 277 | self.decode_query_len - self.vllm_config.num_speculative_tokens |
| 278 | ) |
| 279 | dense_schedule = build_dynamic_sd_schedule_lookup( |
| 280 | speculative_config.num_speculative_tokens_per_batch_size, |
| 281 | vllm_max_batch_size=self.max_num_reqs, |
| 282 | vllm_num_speculative_tokens=self.vllm_config.num_speculative_tokens, |
| 283 | ) |
| 284 | decode_query_lens = sorted( |
| 285 | { |
| 286 | num_spec + num_new_sampled_tokens_per_step |
| 287 | for num_spec in dense_schedule[1:] |
| 288 | } |
| 289 | ) |
| 290 | else: |
| 291 | decode_query_lens = [self.decode_query_len] |
| 292 | |
| 293 | capture_varlen_decode = ( |
| 294 | separate_decode_routine and bool(decode_mode) and self.varlen_decode |
| 295 | ) |
| 296 | for num_tokens, num_active_loras in product( |
| 297 | capture_sizes, self.lora_capture_cases |
| 298 | ): |
| 299 | # Varlen decode graphs take any mix of 1..decode_query_len tokens per |
| 300 | # request, worst case 1 token per request (or max_num_reqs) |
| 301 | if capture_varlen_decode and num_tokens <= max_decode_tokens: |
| 302 | desc = BatchExecutionDescriptor( |
| 303 | cg_mode=decode_mode, |
| 304 | num_tokens=num_tokens, |
| 305 | num_reqs=min(num_tokens, self.max_num_reqs), |
| 306 | max_query_len=self.decode_query_len, |
| 307 | num_active_loras=num_active_loras, |
| 308 | ) |
| 309 | descs_by_mode[decode_mode].append(desc) |
| 310 | # Capture uniform decode specfifc graphs if required |
| 311 | # (i.e. separate decode routine) |
| 312 | elif separate_decode_routine and decode_mode and not self.varlen_decode: |
| 313 | for decode_query_len in decode_query_lens: |
| 314 | rounded_num_tokens = round_up(num_tokens, decode_query_len) |
| 315 | rounded_num_reqs = rounded_num_tokens // decode_query_len |
| 316 | |
| 317 | if ( |
| 318 | rounded_num_tokens > max_decode_tokens |
| 319 | or rounded_num_tokens > max_cg_capture_size |
| 320 | or rounded_num_reqs > self.max_num_reqs |
| 321 | ): |
| 322 | continue |
| 323 | |
| 324 | desc = BatchExecutionDescriptor( |
| 325 | cg_mode=decode_mode, |
| 326 | num_tokens=rounded_num_tokens, |
| 327 | num_reqs=rounded_num_reqs, |
| 328 | uniform_token_count=decode_query_len, |
| 329 | num_active_loras=num_active_loras, |
| 330 | ) |
| 331 | |
| 332 | # avoid duplicate graphs |
| 333 | if desc not in descs_by_mode[decode_mode]: |
| 334 | descs_by_mode[decode_mode].append(desc) |
| 335 | |
| 336 | ubatch_desc = self._maybe_ubatch_twin(desc) |
| 337 | if ubatch_desc is not None and ( |
| 338 | ubatch_desc not in descs_by_mode[decode_mode] |
| 339 | ): |
| 340 | descs_by_mode[decode_mode].append(ubatch_desc) |
| 341 | |
| 342 | # recoverSSM cannot capture a dummy query wider than its workspace. |
| 343 | if mixed_mode and ( |
| 344 | not self.vllm_config.cache_config.use_kda_recoverssm |
| 345 | or num_tokens <= max_decode_tokens |
| 346 | ): |
| 347 | # for PIECEWISE graphs there is no limit on requests when replaying |
| 348 | # i.e. no request padding is needed, so we leave it as None. |
| 349 | # For breakable PW graphs, break-point kernels read the real batch |
| 350 | # from the forward context; in-graph kernels handle the token padding |
| 351 | # themselves from the padded slot_mapping (rows with slot == -1). |
| 352 | num_reqs = None |
| 353 | if mixed_mode == CUDAGraphMode.FULL: |
| 354 | num_reqs = min(num_tokens, self.max_num_reqs) |
| 355 | desc = BatchExecutionDescriptor( |
| 356 | cg_mode=mixed_mode, |
| 357 | num_tokens=num_tokens, |
| 358 | num_reqs=num_reqs, |
| 359 | num_active_loras=num_active_loras, |
| 360 | ) |
| 361 | descs_by_mode[mixed_mode].append(desc) |
| 362 | |
| 363 | ubatch_desc = self._maybe_ubatch_twin(desc) |
| 364 | if ubatch_desc is not None: |
| 365 | descs_by_mode[mixed_mode].append(ubatch_desc) |
| 366 | |
| 367 | for mode, descs in descs_by_mode.items(): |
| 368 | descs.sort(key=lambda d: d.num_tokens, reverse=True) |
| 369 | self._capture_descs[mode] = descs |
| 370 | |
| 371 | for mode in (CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE): |
| 372 | mode_descs = tuple(reversed(descs_by_mode.get(mode, []))) |
| 373 | for num_active_loras in self.lora_capture_cases: |
| 374 | lora_descs = [ |
| 375 | d for d in mode_descs if d.num_active_loras == num_active_loras |
| 376 | ] |
| 377 | current_range_start = 0 |
| 378 | # Dynamic speculative decoding can produce multiple graphs with the same |
| 379 | # num_tokens. Group them so each graph covers the same candidate range. |
| 380 | for num_tokens, group in groupby(lora_descs, lambda d: d.num_tokens): |
| 381 | matching = list(group) |
| 382 | for i in range(current_range_start, num_tokens + 1): |
| 383 | key = (i, num_active_loras) |
| 384 | self._candidates.setdefault(key, []).extend(matching) |
| 385 | current_range_start = num_tokens + 1 |
| 386 | |
| 387 | @property |
| 388 | def dp_size(self) -> int: |
| 389 | # Not cached: elastic EP rewrites parallel_config in place on scale. |
| 390 | return self.vllm_config.parallel_config.data_parallel_size |
| 391 | |
| 392 | def release_graphs(self) -> None: |
| 393 | """Drop the captured graphs so a later capture() can refill them. |
| 394 | |
| 395 | Elastic EP reallocates the MoE workspace when it grows, which leaves |
| 396 | every captured graph holding a stale data pointer. `_capture_descs` is |
| 397 | kept, so `needs_capture()` still reports the work to redo. |
| 398 | """ |
| 399 | self.graphs.clear() |
| 400 | self._graphs_captured = False |
| 401 | if self.breakable_cg_runner is not None: |
| 402 | BreakableCUDAGraphWrapper.clear_all_graphs() |
| 403 | |
| 404 | def needs_capture(self) -> bool: |
| 405 | return len(self._capture_descs) > 0 |
| 406 | |
| 407 | def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream: |
| 408 | """Capture on the stream used by the microbatch threads.""" |
| 409 | if desc.num_ubatches > 1: |
| 410 | assert self.ubatch_runner is not None |
| 411 | return self.ubatch_runner.capture_stream |
| 412 | return current_stream() |
| 413 | |
| 414 | @torch.inference_mode() |
| 415 | def capture( |
| 416 | self, |
| 417 | create_forward_fn: CreateForwardFn, |
| 418 | progress_bar_desc: str = "Capturing CUDA graphs", |
| 419 | ) -> None: |
| 420 | """Capture CUDA graphs. |
| 421 | |
| 422 | Args: |
| 423 | create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and |
| 424 | returns a forward_fn. For FULL and breakable PIECEWISE modes, |
| 425 | it is invoked once with warmup=True and again with warmup=False |
| 426 | because attention backends may mutate or lazily initialize |
| 427 | metadata during warmup. |
| 428 | progress_bar_desc: Description shown on the capture progress bar. |
| 429 | |
| 430 | """ |
| 431 | with graph_capture(device=self.device), ExitStack() as stack: |
| 432 | if self.ubatch_runner is not None: |
| 433 | # Join parked threads on failure to avoid blocking later captures. |
| 434 | stack.callback(self.ubatch_runner.abort_pending_run) |
| 435 | # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger |
| 436 | # activations so FULL activations should fit in already allocated |
| 437 | # buffers in the graph pool. |
| 438 | for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]: |
| 439 | if mode not in self._capture_descs: |
| 440 | continue |
| 441 | |
| 442 | descs = self._capture_descs[mode] |
| 443 | if ( |
| 444 | mode == CUDAGraphMode.FULL |
| 445 | and self._max_full_descs_to_capture is not None |
| 446 | ): |
| 447 | # Profiling only: capture a sample of the largest FULL |
| 448 | # graphs; the total cost is extrapolated from their |
| 449 | # per-graph memory deltas. |
| 450 | descs = descs[: self._max_full_descs_to_capture] |
| 451 | if is_global_first_rank(): |
| 452 | descs = tqdm(descs, desc=f"{progress_bar_desc} ({mode.name})") |
| 453 | for desc in descs: |
| 454 | # Prepare inputs and get forward function |
| 455 | forward_fn = create_forward_fn(desc, warmup=True) |
| 456 | |
| 457 | # Warmup |
| 458 | forward_fn(CUDAGraphMode.NONE) |
| 459 | |
| 460 | # Capture |
| 461 | logger.debug( |
| 462 | "CG Capture: mode=%s, batch_desc=%s", |
| 463 | desc.cg_mode.name, |
| 464 | desc, |
| 465 | ) |
| 466 | if ( |
| 467 | desc.cg_mode == CUDAGraphMode.PIECEWISE |
| 468 | and not self.use_breakable_cg |
| 469 | ): |
| 470 | forward_fn(CUDAGraphMode.PIECEWISE) |
| 471 | else: |
| 472 | # Capture with fresh attention state. |
| 473 | forward_fn = create_forward_fn(desc, warmup=False) |
| 474 | if desc.cg_mode == CUDAGraphMode.PIECEWISE: |
| 475 | forward_fn(CUDAGraphMode.PIECEWISE) |
| 476 | continue |
| 477 | assert desc not in self.graphs, ( |
| 478 | f"Graph already captured for {desc}" |
| 479 | ) |
| 480 | graph = torch.cuda.CUDAGraph() |
| 481 | # Sync offloader's copy stream before capture. |
| 482 | # Ensure any pre-capture prefetches from offloader are complete. |
| 483 | get_offloader().sync_prev_onload() |
| 484 | if self.pool is not None: |
| 485 | set_graph_pool_id(self.pool) |
| 486 | else: |
| 487 | set_graph_pool_id(current_platform.graph_pool_handle()) |
| 488 | if self._capture_mem_samples is not None: |
| 489 | torch.accelerator.synchronize() |
| 490 | free_before = torch.accelerator.get_memory_info()[0] |
| 491 | with torch.cuda.graph( |
| 492 | graph, self.pool, stream=self._capture_stream(desc) |
| 493 | ): |
| 494 | forward_fn(CUDAGraphMode.NONE) |
| 495 | # Join offloader's copy stream after forward to avoid |
| 496 | # unjoined stream error. The last layer's start_prefetch |
| 497 | # forks copy_stream, but wait_prefetch only happens in |
| 498 | # the next forward pass. |
| 499 | get_offloader().join_after_forward() |
| 500 | if self._capture_mem_samples is not None: |
| 501 | torch.accelerator.synchronize() |
| 502 | free_after = torch.accelerator.get_memory_info()[0] |
| 503 | self._capture_mem_samples.append(free_before - free_after) |
| 504 | self.graphs[desc] = graph |
| 505 | compilation_counter.num_cudagraph_captured += 1 |
| 506 | |
| 507 | self._graphs_captured = True |
| 508 | |
| 509 | def captured_token_counts(self) -> list[int]: |
| 510 | """Sorted token counts with a captured graph, ignoring LoRA variants.""" |
| 511 | return sorted( |
| 512 | {desc.num_tokens for desc in self.graphs if desc.num_active_loras == 0} |
| 513 | ) |
| 514 | |
| 515 | def dispatch( |
| 516 | self, |
| 517 | num_reqs: int, |
| 518 | num_tokens: int, |
| 519 | uniform_token_count: int | None, |
| 520 | num_active_loras: int, |
| 521 | max_query_len: int | None = None, |
| 522 | num_ubatches: int = 1, |
| 523 | ) -> BatchExecutionDescriptor: |
| 524 | """Find matching cudagraph descriptor from priority-ordered candidates.""" |
| 525 | effective_loras = self._resolve_effective_loras(num_active_loras) |
| 526 | key = (num_tokens, effective_loras) |
| 527 | if self._graphs_captured and num_tokens > 0 and key in self._candidates: |
| 528 | for desc in self._candidates[key]: |
| 529 | if _is_compatible( |
| 530 | desc, |
| 531 | num_reqs, |
| 532 | num_tokens, |
| 533 | uniform_token_count, |
| 534 | effective_loras, |
| 535 | max_query_len, |
| 536 | num_ubatches, |
| 537 | ): |
| 538 | return desc |
| 539 | return BatchExecutionDescriptor( |
| 540 | cg_mode=CUDAGraphMode.NONE, |
| 541 | num_tokens=num_tokens, |
| 542 | num_reqs=num_reqs, |
| 543 | num_active_loras=effective_loras, |
| 544 | num_ubatches=num_ubatches, |
| 545 | ) |
| 546 | |
| 547 | def run_fullgraph(self, desc: BatchExecutionDescriptor): |
| 548 | """Replay a captured FULL cudagraph.""" |
| 549 | assert desc.cg_mode == CUDAGraphMode.FULL, ( |
| 550 | f"Expected FULL mode, got {desc.cg_mode}" |
| 551 | ) |
| 552 | assert desc in self.graphs, f"No cudagraph for {desc}" |
| 553 | # Sync offloader before replay - needed when transitioning from |
| 554 | # eager/piecewise to full cudagraph (e.g., prefill → decode). |
| 555 | # The previous eager iteration's start_prefetch may have queued |
| 556 | # H2D copies on copy_stream that the graph's captured events |
| 557 | # cannot see. Without this, replay could overwrite static buffers |
| 558 | # while those copies are still in flight. |
| 559 | get_offloader().sync_prev_onload() |
| 560 | self.graphs[desc].replay() |
| 561 | |
| 562 | def init_breakable_cg_runner(self, model: nn.Module) -> None: |
| 563 | if self.breakable_cg_runner is None: |
| 564 | self.breakable_cg_runner = BreakableCUDAGraphWrapper( |
| 565 | model, self.vllm_config |
| 566 | ) |
| 567 | |
| 568 | def run_pw_graph(self, model: nn.Module, model_inputs: dict[str, Any]) -> Any: |
| 569 | if not self.use_breakable_cg: |
| 570 | # Default: Use torch-compiled piecewise cudagraph. |
| 571 | return model(**model_inputs) |
| 572 | assert self.breakable_cg_runner is not None |
| 573 | return self.breakable_cg_runner(**model_inputs) |
| 574 | |
| 575 | |
| 576 | class ModelCudaGraphManager(CudaGraphManager): |
| 577 | """CudaGraphManager with model-specific capture and hidden state management.""" |
| 578 | |
| 579 | def __init__( |
| 580 | self, |
| 581 | vllm_config: VllmConfig, |
| 582 | device: torch.device, |
| 583 | cudagraph_mode: CUDAGraphMode, |
| 584 | decode_query_len: int, |
| 585 | lora_capture_cases: list[int] | None = None, |
| 586 | varlen_decode: bool = False, |
| 587 | ubatch_runner: "UBatchRunner | None" = None, |
| 588 | ): |
| 589 | super().__init__( |
| 590 | vllm_config, |
| 591 | device, |
| 592 | cudagraph_mode, |
| 593 | decode_query_len, |
| 594 | lora_capture_cases=lora_capture_cases, |
| 595 | varlen_decode=varlen_decode, |
| 596 | ubatch_runner=ubatch_runner, |
| 597 | ) |
| 598 | self.hidden_states: torch.Tensor | None = None |
| 599 | self.aux_hidden_states: list[torch.Tensor] = [] |
| 600 | self.use_aux_hidden_state_outputs = False |
| 601 | self.intermediate_tensors: IntermediateTensors | None = None |
| 602 | |
| 603 | def capture( |
| 604 | self, |
| 605 | model: nn.Module, |
| 606 | model_state: ModelState, |
| 607 | input_buffers: InputBuffers, |
| 608 | intermediate_tensors: IntermediateTensors | None, |
| 609 | block_tables: BlockTables, |
| 610 | attn_groups: list[list[AttentionGroup]], |
| 611 | kv_cache_config: KVCacheConfig, |
| 612 | pcp_manager: "PCPManager | None" = None, |
| 613 | has_lora: bool = False, |
| 614 | use_aux_hidden_state_outputs: bool = False, |
| 615 | lora_capture_hook: Callable[[int, int, int], None] | None = None, |
| 616 | progress_bar_desc: str = "Capturing CUDA graphs", |
| 617 | ) -> None: |
| 618 | """Capture CUDA graphs for model forward pass.""" |
| 619 | self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs |
| 620 | if self.use_breakable_cg: |
| 621 | self.init_breakable_cg_runner(model) |
| 622 | |
| 623 | if self.cudagraph_mode.has_piecewise_cudagraphs() and not ( |
| 624 | self.use_breakable_cg or has_compiled_submodule(model) |
| 625 | ): |
| 626 | raise RuntimeError( |
| 627 | f"{type(model).__name__}: piecewise CUDA graphs " |
| 628 | f"(cudagraph_mode={self.cudagraph_mode.name}) unavailable, " |
| 629 | "model is not torch-compiled and breakable CUDA graph is off. " |
| 630 | "Set VLLM_USE_BREAKABLE_CUDAGRAPH=1 or cudagraph_mode=NONE/FULL." |
| 631 | ) |
| 632 | |
| 633 | def store_capture_output(num_tokens: int, model_output: Any) -> None: |
| 634 | """Copy outputs to persistent buffers, allocating on first use.""" |
| 635 | if self.is_last_pp_rank: |
| 636 | # Last PP rank (common case). |
| 637 | if self.use_aux_hidden_state_outputs: |
| 638 | hidden_states, aux_hidden_states = model_output |
| 639 | else: |
| 640 | hidden_states = model_output |
| 641 | aux_hidden_states = [] |
| 642 | if self.hidden_states is None: |
| 643 | self.hidden_states = torch.empty_like(hidden_states) |
| 644 | self.hidden_states[:num_tokens] = hidden_states |
| 645 | if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: |
| 646 | self.aux_hidden_states = [ |
| 647 | torch.empty_like(x) for x in aux_hidden_states |
| 648 | ] |
| 649 | for i, aux in enumerate(aux_hidden_states): |
| 650 | self.aux_hidden_states[i][:num_tokens] = aux |
| 651 | else: |
| 652 | # Non-last PP rank. |
| 653 | assert isinstance(model_output, IntermediateTensors) |
| 654 | intermediate_tensors = model_output |
| 655 | if self.intermediate_tensors is None: |
| 656 | self.intermediate_tensors = IntermediateTensors.empty_like( |
| 657 | intermediate_tensors |
| 658 | ) |
| 659 | for k, v in intermediate_tensors.tensors.items(): |
| 660 | self.intermediate_tensors[k][:num_tokens] = v |
| 661 | |
| 662 | def create_forward_fn( |
| 663 | desc: BatchExecutionDescriptor, warmup: bool |
| 664 | ) -> Callable[[CUDAGraphMode], None]: |
| 665 | num_tokens = desc.num_tokens |
| 666 | num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs) |
| 667 | |
| 668 | # Set LoRA state before capture so kernels see correct adapters. |
| 669 | if lora_capture_hook is not None: |
| 670 | lora_capture_hook(desc.num_active_loras, num_reqs, num_tokens) |
| 671 | |
| 672 | num_tokens_across_dp = ( |
| 673 | torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu") |
| 674 | if self.dp_size > 1 |
| 675 | else None |
| 676 | ) |
| 677 | |
| 678 | model_inputs = { |
| 679 | "input_ids": input_buffers.input_ids[:num_tokens], |
| 680 | "positions": input_buffers.positions[:num_tokens], |
| 681 | "intermediate_tensors": None, |
| 682 | **model_state.prepare_dummy_inputs(num_reqs, num_tokens), |
| 683 | } |
| 684 | if not self.is_first_pp_rank: |
| 685 | # Update for non-first PP ranks. |
| 686 | model_inputs["input_ids"] = None |
| 687 | model_inputs["inputs_embeds"] = None |
| 688 | assert intermediate_tensors is not None |
| 689 | model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] |
| 690 | |
| 691 | if desc.num_ubatches > 1: |
| 692 | # Prepare and park threads before capture; finish runs inside it. |
| 693 | assert self.ubatch_runner is not None |
| 694 | ubatch_state = self.ubatch_runner.prepare( |
| 695 | InputBatch.make_dummy(num_reqs, num_tokens, input_buffers), |
| 696 | block_tables.get_dummy_block_tables(num_reqs), |
| 697 | block_tables.get_dummy_slot_mappings(num_tokens), |
| 698 | cg_mode=CUDAGraphMode.FULL, |
| 699 | for_capture=True, |
| 700 | ) |
| 701 | # Capture with dummy rows marked as padding. |
| 702 | input_buffers.is_padding.fill_(True) |
| 703 | finish = self.ubatch_runner.begin_capturable_run( |
| 704 | model, model_inputs, ubatch_state, for_capture=True |
| 705 | ) |
| 706 | |
| 707 | def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None: |
| 708 | assert cg_mode != CUDAGraphMode.PIECEWISE, ( |
| 709 | "DBO does not support PIECEWISE cudagraphs" |
| 710 | ) |
| 711 | store_capture_output(num_tokens, finish()) |
| 712 | |
| 713 | return ubatch_forward_fn |
| 714 | |
| 715 | attn_metadata, slot_mappings = prepare_inputs_to_capture( |
| 716 | num_reqs, |
| 717 | num_tokens, |
| 718 | model_state, |
| 719 | input_buffers, |
| 720 | block_tables, |
| 721 | attn_groups, |
| 722 | kv_cache_config, |
| 723 | full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, |
| 724 | max_query_len=desc.max_query_len or desc.uniform_token_count, |
| 725 | pcp_manager=pcp_manager, |
| 726 | ) |
| 727 | |
| 728 | # Capture with dummy rows marked as padding. |
| 729 | input_buffers.is_padding.fill_(True) |
| 730 | |
| 731 | def forward_fn(cg_mode: CUDAGraphMode) -> None: |
| 732 | batch_descriptor = None |
| 733 | if cg_mode == CUDAGraphMode.PIECEWISE: |
| 734 | batch_descriptor = BatchDescriptor( |
| 735 | num_tokens=num_tokens, |
| 736 | has_lora=has_lora, |
| 737 | num_active_loras=desc.num_active_loras, |
| 738 | ) |
| 739 | with set_forward_context( |
| 740 | attn_metadata, |
| 741 | self.vllm_config, |
| 742 | num_tokens=num_tokens, |
| 743 | cudagraph_runtime_mode=cg_mode, |
| 744 | num_tokens_across_dp=num_tokens_across_dp, |
| 745 | slot_mapping=slot_mappings, |
| 746 | batch_descriptor=batch_descriptor, |
| 747 | is_padding=input_buffers.is_padding[:num_tokens], |
| 748 | ): |
| 749 | if cg_mode == CUDAGraphMode.PIECEWISE: |
| 750 | # PIECEWISE graph (compiled PW or breakable, chosen inside |
| 751 | # run_pw_graph). |
| 752 | model_output = self.run_pw_graph(model, model_inputs) |
| 753 | else: |
| 754 | model_output = model(**model_inputs) |
| 755 | |
| 756 | if cg_mode == CUDAGraphMode.PIECEWISE: |
| 757 | # PW CUDA graph (compiled or breakable) internally handles the |
| 758 | # model outputs. No need to keep track of the hidden states. |
| 759 | return None |
| 760 | |
| 761 | store_capture_output(num_tokens, model_output) |
| 762 | |
| 763 | return forward_fn |
| 764 | |
| 765 | super().capture(create_forward_fn, progress_bar_desc) |
| 766 | |
| 767 | def run_fullgraph( |
| 768 | self, desc: BatchExecutionDescriptor |
| 769 | ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors: |
| 770 | """Replay a captured FULL cudagraph and return hidden states.""" |
| 771 | super().run_fullgraph(desc) |
| 772 | if not self.is_last_pp_rank: |
| 773 | assert self.intermediate_tensors is not None |
| 774 | return self.intermediate_tensors[: desc.num_tokens] |
| 775 | |
| 776 | assert self.hidden_states is not None |
| 777 | hidden_states = self.hidden_states[: desc.num_tokens] |
| 778 | if not self.use_aux_hidden_state_outputs: |
| 779 | return hidden_states |
| 780 | return hidden_states, [x[: desc.num_tokens] for x in self.aux_hidden_states] |
| 781 | |
| 782 | |
| 783 | def prepare_inputs_to_capture( |
| 784 | num_reqs: int, |
| 785 | num_tokens: int, |
| 786 | model_state: ModelState, |
| 787 | input_buffers: InputBuffers, |
| 788 | block_tables: BlockTables, |
| 789 | attn_groups: list[list[AttentionGroup]], |
| 790 | kv_cache_config: KVCacheConfig, |
| 791 | full_cudagraph: bool, |
| 792 | max_query_len: int | None = None, |
| 793 | pcp_manager: "PCPManager | None" = None, |
| 794 | ) -> AttentionState: |
| 795 | if full_cudagraph and max_query_len is None: |
| 796 | # Mixed graphs can replay a single prefill spanning the entire batch, |
| 797 | # even when the dummy batch distributes one token to each request. |
| 798 | max_query_len = num_tokens |
| 799 | input_batch = InputBatch.make_dummy( |
| 800 | num_reqs, num_tokens, input_buffers, max_query_len=max_query_len |
| 801 | ) |
| 802 | if pcp_manager is not None: |
| 803 | input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) |
| 804 | |
| 805 | block_table_provider = pcp_manager or block_tables |
| 806 | input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) |
| 807 | slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens) |
| 808 | slot_mappings_by_layer = build_slot_mappings_by_layer( |
| 809 | slot_mappings, kv_cache_config |
| 810 | ) |
| 811 | |
| 812 | if block_tables.cp_size > 1: |
| 813 | input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens( |
| 814 | input_buffers.dcp_local_seq_lens, |
| 815 | input_batch.seq_lens, |
| 816 | input_batch.num_reqs, |
| 817 | block_tables.cp_size, |
| 818 | block_tables.cp_rank, |
| 819 | block_tables.cp_interleave, |
| 820 | num_reqs_padded=input_batch.num_reqs_after_padding, |
| 821 | ) |
| 822 | |
| 823 | # NOTE(woosuk): Attention metadata is required not just by standard attention |
| 824 | # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv, |
| 825 | # DSV4 compressor), which maintain their own states and require special metadata |
| 826 | # such as block tables. |
| 827 | # During CUDA graph capture: |
| 828 | # - For FULL CUDA graphs: We set for_capture=True so that both attention and |
| 829 | # attention-like ops produce capturable metadata compatible with CUDA graphs. |
| 830 | # - For PIECEWISE CUDA graphs: We still build attention metadata, but set |
| 831 | # for_capture=False. This is because: |
| 832 | # * Attention-like ops (such as sconv or DSV4 compressor) may not be used as |
| 833 | # breakpoints in PIECEWISE CUDA graphs, so we must generate their attention |
| 834 | # metadata so they can execute and be captured during graph capture. |
| 835 | # * Standard attention ops that are treated as breakpoints will be executed |
| 836 | # eagerly at capture time (not included in the graph itself), and for these, |
| 837 | # setting for_capture=False is essential. Some attention backends |
| 838 | # (like linear attention) cannot generate capturable metadata for prefill, |
| 839 | # so for_capture=False ensures they execute without issue. |
| 840 | # * We assume that attention-like operations intended for capture will still |
| 841 | # produce capturable metadata, even when for_capture=False. While this |
| 842 | # assumption is brittle, it currently works in practice. |
| 843 | # In summary: We always generate attention metadata for both FULL and PIECEWISE |
| 844 | # CUDA graphs, setting for_capture=True for FULL graphs, and for_capture=False |
| 845 | # for PIECEWISE graphs, to ensure correct execution and capture. |
| 846 | attn_metadata = model_state.prepare_attn( |
| 847 | input_batch, |
| 848 | CUDAGraphMode.NONE, |
| 849 | input_block_tables, |
| 850 | slot_mappings, |
| 851 | attn_groups, |
| 852 | kv_cache_config, |
| 853 | for_capture=full_cudagraph, |
| 854 | ) |
| 855 | return AttentionState(attn_metadata, slot_mappings_by_layer) |
| 856 | |
| 857 | |
| 858 | # --------------------------------------------------------------------------- |
| 859 | # CUDA graph memory profiling |
| 860 | # --------------------------------------------------------------------------- |
| 861 | |
| 862 | # Number of FULL graphs captured during profiling; the total FULL capture |
| 863 | # cost is extrapolated from this sample to avoid a second full capture. |
| 864 | _FULL_GRAPH_PROFILING_SAMPLES = 2 |
| 865 | # Floor for the extrapolated per-graph cost (driver overhead per graph). |
| 866 | _MIN_PER_GRAPH_BYTES = 1 << 20 |
| 867 | |
| 868 | |
| 869 | @torch.inference_mode() |
| 870 | def profile_cudagraph_memory(runner: "GPUModelRunner") -> int: |
| 871 | """Estimate the GPU memory needed for CUDA graph capture. |
| 872 | |
| 873 | Called during memory profiling, *before* the real KV cache is allocated, |
| 874 | so that ``Worker.determine_available_memory`` can reserve headroom for |
| 875 | graph capture. Bootstraps a minimal KV cache, runs ``capture_model()`` |
| 876 | once, then releases everything so the real init/capture path starts clean. |
| 877 | |
| 878 | FULL graphs bake in KV cache pointers, so only the largest few are |
| 879 | captured (into a throwaway pool) and their total cost is extrapolated. |
| 880 | PIECEWISE, encoder and speculator graphs are measured in full. All |
| 881 | profiling captures are discarded afterwards: replaying graphs recorded |
| 882 | against the throwaway profiling state is unsafe (e.g. inductor graph |
| 883 | partition reclaims the storages of earlier cudagraph recordings once the |
| 884 | real capture records new ones, leading to use-after-free crashes). |
| 885 | """ |
| 886 | if runner.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: |
| 887 | return 0 |
| 888 | |
| 889 | gc.collect() |
| 890 | torch.accelerator.empty_cache() |
| 891 | |
| 892 | # Run the whole profiling phase against a throwaway CUDA graph pool by |
| 893 | # pointing the global graph pool singleton at it: objects that bind the |
| 894 | # pool lazily during profiling (speculator cudagraph managers, breakable |
| 895 | # runners created mid-capture) then land on the throwaway pool too. Pools |
| 896 | # bound before profiling (piecewise wrappers) are swapped explicitly in |
| 897 | # the inner block. Profiling graphs captured into the persistent global |
| 898 | # pool and then discarded would drop its use_count to 0, tripping the c10 |
| 899 | # allocator's create_or_incref_pool assert when the real capture reuses |
| 900 | # that pool ("use_count > 0 INTERNAL ASSERT FAILED"). |
| 901 | platform_cls = type(current_platform) |
| 902 | saved_global_pool = platform_cls._global_graph_pool |
| 903 | throwaway_pool = current_platform.graph_pool_handle() |
| 904 | platform_cls._global_graph_pool = throwaway_pool |
| 905 | |
| 906 | try: |
| 907 | with set_current_vllm_config(runner.vllm_config): |
| 908 | _init_minimal_kv_cache_for_profiling(runner) |
| 909 | |
| 910 | manager = runner.cudagraph_manager |
| 911 | assert manager is not None |
| 912 | |
| 913 | # Don't count profiling captures; the real capture_model() runs later. |
| 914 | saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured |
| 915 | saved_capture_triggers = compilation_counter.num_gpu_runner_capture_triggers |
| 916 | all_wrappers: list[Any] = [] |
| 917 | original_pools: dict[int, Any] = {} |
| 918 | speculator = getattr(runner, "speculator", None) |
| 919 | spec_manager_names: list[str] = [] |
| 920 | try: |
| 921 | if not manager.needs_capture(): |
| 922 | return 0 |
| 923 | manager.pool = throwaway_pool |
| 924 | if manager.use_breakable_cg: |
| 925 | # Create the breakable runner before the wrapper pool swap so |
| 926 | # its pool is covered as well. |
| 927 | manager.init_breakable_cg_runner(runner.model) |
| 928 | all_wrappers = list(CUDAGraphWrapper._all_instances) + list( |
| 929 | BreakableCUDAGraphWrapper._all_instances |
| 930 | ) |
| 931 | for wrapper in all_wrappers: |
| 932 | original_pools[id(wrapper)] = wrapper.graph_pool |
| 933 | wrapper.graph_pool = throwaway_pool |
| 934 | if speculator is not None: |
| 935 | spec_manager_names = [ |
| 936 | name |
| 937 | for name, value in vars(speculator).items() |
| 938 | if isinstance(value, CudaGraphManager) |
| 939 | ] |
| 940 | manager._max_full_descs_to_capture = _FULL_GRAPH_PROFILING_SAMPLES |
| 941 | mem_samples: list[int] = [] |
| 942 | manager._capture_mem_samples = mem_samples |
| 943 | |
| 944 | measured = int(runner.capture_model(profile_only=True)) |
| 945 | |
| 946 | # The measured delta covers PIECEWISE, encoder and speculator graphs |
| 947 | # plus the sampled FULL graphs; swap the sampled FULL cost for the |
| 948 | # extrapolated total. FULL and PIECEWISE share one pool here just as |
| 949 | # they share the global pool at runtime, so the overlap is not |
| 950 | # double-counted. |
| 951 | num_full_graphs = len(manager._capture_descs.get(CUDAGraphMode.FULL, [])) |
| 952 | full_estimate = _extrapolate_full_graph_memory(mem_samples, num_full_graphs) |
| 953 | return max(measured - sum(mem_samples) + full_estimate, 0) |
| 954 | finally: |
| 955 | compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured |
| 956 | compilation_counter.num_gpu_runner_capture_triggers = saved_capture_triggers |
| 957 | CUDAGraphWrapper.clear_all_graphs() |
| 958 | BreakableCUDAGraphWrapper.clear_all_graphs() |
| 959 | for wrapper in all_wrappers: |
| 960 | if id(wrapper) in original_pools: |
| 961 | wrapper.graph_pool = original_pools[id(wrapper)] |
| 962 | # Drop the speculator's cudagraph managers; the real |
| 963 | # initialize_kv_cache re-creates them. Their profiling graphs |
| 964 | # release the throwaway pool here rather than after the real init. |
| 965 | for name in spec_manager_names: |
| 966 | setattr(speculator, name, None) |
| 967 | # Drop local references before teardown detaches the runner's |
| 968 | # manager and flushes the allocator. |
| 969 | del manager |
| 970 | _teardown_profiling_state(runner) |
| 971 | finally: |
| 972 | platform_cls._global_graph_pool = saved_global_pool |
| 973 | |
| 974 | |
| 975 | def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> int: |
| 976 | """Extrapolate the total FULL capture cost from samples of the largest |
| 977 | graphs. The first capture allocates the pool baseline; later graphs mostly |
| 978 | reuse it, so the second sample is taken as the per-graph cost.""" |
| 979 | if not mem_samples: |
| 980 | return 0 |
| 981 | first_capture = mem_samples[0] |
| 982 | per_graph = max(mem_samples[1], _MIN_PER_GRAPH_BYTES) if len(mem_samples) > 1 else 0 |
| 983 | return first_capture + (total_graphs - 1) * per_graph |
| 984 | |
| 985 | |
| 986 | def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: |
| 987 | """Allocate the smallest KV cache that still lets every graph be captured.""" |
| 988 | from vllm.v1.core.kv_cache_utils import ( |
| 989 | get_kv_cache_config_from_groups, |
| 990 | get_kv_cache_groups, |
| 991 | ) |
| 992 | |
| 993 | kv_cache_spec = runner.get_kv_cache_spec() |
| 994 | kv_cache_groups = get_kv_cache_groups(runner.vllm_config, kv_cache_spec) |
| 995 | # At least one block per sequence is required to capture the graphs. |
| 996 | min_blocks = ( |
| 997 | min(runner.max_num_reqs, runner.compilation_config.max_cudagraph_capture_size) |
| 998 | or 1 |
| 999 | ) |
| 1000 | saved_override = runner.cache_config.num_gpu_blocks_override |
| 1001 | runner.cache_config.num_gpu_blocks_override = min_blocks |
| 1002 | try: |
| 1003 | minimal_config = get_kv_cache_config_from_groups( |
| 1004 | runner.vllm_config, kv_cache_groups, available_memory=0 |
| 1005 | ) |
| 1006 | finally: |
| 1007 | runner.cache_config.num_gpu_blocks_override = saved_override |
| 1008 | |
| 1009 | runner.initialize_kv_cache(minimal_config, is_profiling=True) |
| 1010 | runner.cache_config.num_gpu_blocks = minimal_config.num_blocks |
| 1011 | |
| 1012 | |
| 1013 | def _teardown_profiling_state(runner: "GPUModelRunner") -> None: |
| 1014 | """Release the profiling KV cache and captured graphs while keeping model |
| 1015 | weights, so the real ``initialize_kv_cache`` starts from a clean slate.""" |
| 1016 | torch.accelerator.synchronize() |
| 1017 | if hasattr(runner.model_state, "_mamba_ctx"): |
| 1018 | runner.model_state._mamba_ctx = None |
| 1019 | # Invalidate the align-mode Mamba group metadata cached from the |
| 1020 | # profiling KVCacheConfig: the real (e.g. PP-projected) config may |
| 1021 | # place Mamba layers into a different group layout, so it must be |
| 1022 | # re-derived from the real config. |
| 1023 | if hasattr(runner.model_state, "_mamba_group_ids"): |
| 1024 | runner.model_state._mamba_group_ids = [] |
| 1025 | if hasattr(runner.model_state, "_mamba_spec"): |
| 1026 | runner.model_state._mamba_spec = None |
| 1027 | if hasattr(runner, "kv_caches"): |
| 1028 | runner.kv_caches.clear() |
| 1029 | if hasattr(runner, "attn_groups"): |
| 1030 | runner.attn_groups.clear() |
| 1031 | if hasattr(runner, "kv_cache_config"): |
| 1032 | del runner.kv_cache_config |
| 1033 | # Dropping the manager releases the profiling graphs and throwaway pool. |
| 1034 | runner.cudagraph_manager = None |
| 1035 | # Release encoder graphs captured during profiling; the real |
| 1036 | # capture_model() re-captures them. |
| 1037 | if runner.model_state.supports_mm_inputs: |
| 1038 | runner.model_state.encoder_runner.clear() |
| 1039 | # Detach profiling KV tensors held by attention layers. The layers live |
| 1040 | # in the static forward context for compiled models. |
| 1041 | layers: Iterable[Any] = runner.compilation_config.static_forward_context.values() |
| 1042 | if (model := getattr(runner, "model", None)) is not None: |
| 1043 | layers = itertools.chain(layers, model.modules()) |
| 1044 | clear_layer_kv_caches(layers) |
| 1045 | release_hisparse_profiling_cache(runner.compilation_config.static_forward_context) |
| 1046 | runner.cache_config.num_gpu_blocks = None |
| 1047 | runner.maybe_remove_all_loras(runner.lora_config) |
| 1048 | gc.collect() |
| 1049 | torch.accelerator.empty_cache() |