# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import defaultdict from collections.abc import Callable from dataclasses import dataclass from itertools import product from typing import Any, NamedTuple, Protocol import torch import torch.nn as nn from tqdm import tqdm from vllm.compilation.breakable_cudagraph import ( BreakableCUDAGraphWrapper, is_breakable_cudagraph_enabled, ) from vllm.compilation.counter import compilation_counter from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import ( get_pp_group, graph_capture, is_global_first_rank, ) from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import round_up from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) class AttentionState(NamedTuple): attn_metadata: dict[str, Any] | None slot_mappings: dict[str, torch.Tensor] class AttentionStatePair(NamedTuple): warmup: AttentionState captured: AttentionState @dataclass(frozen=True) class BatchExecutionDescriptor: """Describes the shape of the batch and CG mode to run; this is used to make shape matches between the capture and runtime.""" cg_mode: CUDAGraphMode num_tokens: int num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs) uniform_token_count: int | None = None num_active_loras: int = 0 class CreateForwardFn(Protocol): """Factory that prepares inputs (OUTSIDE the graph) and returns a tuple of (forward_fn, attn_state). Called with warmup=True for the warmup pass and warmup=False for the captured pass.""" def __call__( self, desc: BatchExecutionDescriptor, warmup: bool, ) -> tuple[Callable[[CUDAGraphMode], None], AttentionState]: ... def _is_compatible( desc: BatchExecutionDescriptor, num_reqs: int, num_tokens: int, uniform_token_count: int | None, num_active_loras: int, ) -> bool: # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count # desc.num_reqs=None means no request padding needed (PIECEWISE) return ( ( desc.uniform_token_count is None or desc.uniform_token_count == uniform_token_count ) and (desc.num_reqs is None or desc.num_reqs >= num_reqs) and desc.num_tokens >= num_tokens and desc.num_active_loras == num_active_loras ) def get_uniform_token_count( num_reqs: int, num_tokens: int, max_query_len: int, ) -> int | None: """ Return the uniform token count if batch is uniform, else None. A batch is uniform if all requests have the same number of tokens. """ if (max_query_len == num_tokens // num_reqs) and ( num_tokens == max_query_len * num_reqs ): return max_query_len return None class CudaGraphManager: def __init__( self, vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, decode_query_len: int, lora_capture_cases: list[int] | None = None, ): self.vllm_config = vllm_config self.device = device self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs self.compilation_config = vllm_config.compilation_config assert self.compilation_config is not None self.cudagraph_mode = cudagraph_mode self.decode_query_len = decode_query_len self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size self.is_first_pp_rank = get_pp_group().is_first_rank self.is_last_pp_rank = get_pp_group().is_last_rank self.lora_capture_cases = lora_capture_cases or [0] # Precompute actual num_active_loras -> captured case mapping so that # dispatch() is a plain dict lookup instead of a per-call bisect. self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map() self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {} self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None self._graphs_captured = False self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {} self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} self._init_candidates() # Breakable CUDA graph (PW CUDA graph without torch.compile) self.use_breakable_cg = ( is_breakable_cudagraph_enabled() and self.cudagraph_mode.has_piecewise_cudagraphs() ) self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]: """Precompute actual num_active_loras -> effective captured case. Mirrors the num_tokens candidate expansion in ``_init_candidates``: every possible active-LoRA count is mapped ahead of time to the smallest captured case that can serve it, so ``dispatch`` is a plain dict lookup instead of a per-call bisect. """ captured_with_lora = sorted(c for c in self.lora_capture_cases if c > 0) if not captured_with_lora: return {}, 0 dispatch_map: dict[int, int] = {} case_idx = 0 for n in range(1, captured_with_lora[-1] + 1): while captured_with_lora[case_idx] < n: case_idx += 1 dispatch_map[n] = captured_with_lora[case_idx] return dispatch_map, captured_with_lora[-1] def _resolve_effective_loras(self, num_active_loras: int) -> int: """Map an actual active-LoRA count to its captured graph case.""" if num_active_loras <= 0 or not self._lora_dispatch_map: return num_active_loras # Counts above the largest captured case clamp to it. return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case) def _init_candidates(self) -> None: """Build priority-ordered candidate lists for each token count.""" capture_sizes = self.compilation_config.cudagraph_capture_sizes if not (self.cudagraph_mode and capture_sizes): return capture_sizes = sorted(capture_sizes) max_decode_tokens = self.max_num_reqs * self.decode_query_len decode_mode = self.cudagraph_mode.decode_mode() mixed_mode = self.cudagraph_mode.mixed_mode() separate_decode_routine = self.cudagraph_mode.separate_routine() max_cg_capture_size = self.compilation_config.max_cudagraph_capture_size descs_by_token_lora: dict[tuple[int, int], list[BatchExecutionDescriptor]] = ( defaultdict(list) ) descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = ( defaultdict(list) ) # When using Dynamic SD, num_speculative_tokens is the max number of # draft tokens. The scheduler might use a smaller number so we need # to capture graphs for all possible values during decode. speculative_config = self.vllm_config.speculative_config if ( speculative_config and speculative_config.uses_dynamic_speculative_decoding() ): num_spec_per_batch_size = ( speculative_config.num_speculative_tokens_per_batch_size ) # uses_dynamic_speculative_decoding() guarantees this is set. assert num_spec_per_batch_size is not None # decode_query_len = num_speculative_steps + num_new_sampled_tokens # _per_step. Recover num_new_sampled_tokens_per_step # from the values the manager already has. num_new_sampled_tokens_per_step = ( self.decode_query_len - self.vllm_config.num_speculative_tokens ) # Each entry is (range_start, range_end, num_speculative_tokens). decode_query_lens = [ x[2] + num_new_sampled_tokens_per_step for x in num_spec_per_batch_size ] else: decode_query_lens = [self.decode_query_len] for num_tokens, num_active_loras in product( capture_sizes, self.lora_capture_cases ): # Capture uniform decode specfifc graphs if required # (i.e. separate decode routine) if separate_decode_routine and decode_mode: for decode_query_len in decode_query_lens: rounded_num_tokens = round_up(num_tokens, decode_query_len) rounded_num_reqs = rounded_num_tokens // decode_query_len if ( rounded_num_tokens > max_decode_tokens or rounded_num_tokens > max_cg_capture_size or rounded_num_reqs > self.max_num_reqs ): continue desc = BatchExecutionDescriptor( cg_mode=decode_mode, num_tokens=rounded_num_tokens, num_reqs=rounded_num_reqs, uniform_token_count=decode_query_len, num_active_loras=num_active_loras, ) # avoid duplicate graphs if desc not in descs_by_mode[decode_mode]: descs_by_mode[decode_mode].append(desc) descs_by_token_lora[ (rounded_num_tokens, num_active_loras) ].append(desc) if mixed_mode: # for PIECEWISE graphs there is no limit on requests when replaying # i.e. no request padding is needed # so we leave it as None num_reqs = ( min(num_tokens, self.max_num_reqs) if mixed_mode == CUDAGraphMode.FULL else None ) desc = BatchExecutionDescriptor( cg_mode=mixed_mode, num_tokens=num_tokens, num_reqs=num_reqs, num_active_loras=num_active_loras, ) descs_by_mode[mixed_mode].append(desc) descs_by_token_lora[(num_tokens, num_active_loras)].append(desc) if not descs_by_token_lora: return all_token_counts = sorted({k[0] for k in descs_by_token_lora}) current_range_start = 0 for token_cg_size in all_token_counts: for i in range(current_range_start, token_cg_size + 1): for num_active_loras in self.lora_capture_cases: staging_key = (token_cg_size, num_active_loras) if staging_key in descs_by_token_lora: self._candidates[(i, num_active_loras)] = descs_by_token_lora[ staging_key ] current_range_start = token_cg_size + 1 for mode, descs in descs_by_mode.items(): descs.sort(key=lambda d: d.num_tokens, reverse=True) self._capture_descs[mode] = descs def needs_capture(self) -> bool: return len(self._capture_descs) > 0 @torch.inference_mode() def capture( self, create_forward_fn: CreateForwardFn, progress_bar_desc: str = "Capturing CUDA graphs", ) -> dict[BatchExecutionDescriptor, AttentionStatePair]: """Capture CUDA graphs. Args: create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and returns a tuple of (forward_fn, attn_state). For FULL cudagraph mode, it is invoked once with warmup=True for the warmup pass, and again with warmup=False for the captured pass. For attention backends that perform lazy metadata initialization (e.g. FlashMLA), FULL cudagraph capture requires distinct metadatas for warmup and capture. """ attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {} with graph_capture(device=self.device): # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger # activations so FULL activations should fit in already allocated # buffers in the graph pool. for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]: if mode not in self._capture_descs: continue descs = self._capture_descs[mode] if is_global_first_rank(): descs = tqdm(descs, desc=f"{progress_bar_desc} ({mode.name})") for desc in descs: # Prepare inputs and get forward function forward_fn, warmup_attn_state = create_forward_fn(desc, warmup=True) # Warmup forward_fn(CUDAGraphMode.NONE) # Capture logger.debug( "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc ) if desc.cg_mode == CUDAGraphMode.PIECEWISE: attn_states[desc] = AttentionStatePair( warmup_attn_state, warmup_attn_state ) forward_fn(CUDAGraphMode.PIECEWISE) else: # Capture with fresh attention state. forward_fn, capture_attn_state = create_forward_fn( desc, warmup=False ) attn_states[desc] = AttentionStatePair( warmup_attn_state, capture_attn_state ) assert desc not in self.graphs, ( f"Graph already captured for {desc}" ) graph = torch.cuda.CUDAGraph() # Sync offloader's copy stream before capture. # Ensure any pre-capture prefetches from offloader are complete. get_offloader().sync_prev_onload() with torch.cuda.graph(graph, self.pool): forward_fn(CUDAGraphMode.NONE) # Join offloader's copy stream after forward to avoid # unjoined stream error. The last layer's start_prefetch # forks copy_stream, but wait_prefetch only happens in # the next forward pass. get_offloader().join_after_forward() self.graphs[desc] = graph compilation_counter.num_cudagraph_captured += 1 self._graphs_captured = True return attn_states def dispatch( self, num_reqs: int, num_tokens: int, uniform_token_count: int | None, num_active_loras: int, ) -> BatchExecutionDescriptor: """Find matching cudagraph descriptor from priority-ordered candidates.""" effective_loras = self._resolve_effective_loras(num_active_loras) key = (num_tokens, effective_loras) if self._graphs_captured and num_tokens > 0 and key in self._candidates: for desc in self._candidates[key]: if _is_compatible( desc, num_reqs, num_tokens, uniform_token_count, effective_loras, ): return desc return BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens, num_reqs=num_reqs, num_active_loras=effective_loras, ) def run_fullgraph(self, desc: BatchExecutionDescriptor): """Replay a captured FULL cudagraph.""" assert desc.cg_mode == CUDAGraphMode.FULL, ( f"Expected FULL mode, got {desc.cg_mode}" ) assert desc in self.graphs, f"No cudagraph for {desc}" # Sync offloader before replay - needed when transitioning from # eager/piecewise to full cudagraph (e.g., prefill → decode). # The previous eager iteration's start_prefetch may have queued # H2D copies on copy_stream that the graph's captured events # cannot see. Without this, replay could overwrite static buffers # while those copies are still in flight. get_offloader().sync_prev_onload() self.graphs[desc].replay() def init_breakable_cg_runner(self, model: nn.Module) -> None: if self.breakable_cg_runner is None: self.breakable_cg_runner = BreakableCUDAGraphWrapper( model, self.vllm_config ) def run_pw_graph(self, model: nn.Module, model_inputs: dict[str, Any]) -> Any: if not self.use_breakable_cg: # Default: Use torch-compiled piecewise cudagraph. return model(**model_inputs) assert self.breakable_cg_runner is not None return self.breakable_cg_runner(**model_inputs) class ModelCudaGraphManager(CudaGraphManager): """CudaGraphManager with model-specific capture and hidden state management.""" def __init__( self, vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, decode_query_len: int, lora_capture_cases: list[int] | None = None, ): super().__init__( vllm_config, device, cudagraph_mode, decode_query_len, lora_capture_cases=lora_capture_cases, ) self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] self.use_aux_hidden_state_outputs = False self.intermediate_tensors: IntermediateTensors | None = None def capture( self, model: nn.Module, model_state: ModelState, input_buffers: InputBuffers, intermediate_tensors: IntermediateTensors | None, block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, has_lora: bool = False, use_aux_hidden_state_outputs: bool = False, lora_capture_hook: Callable[[int, int, int], None] | None = None, progress_bar_desc: str = "Capturing CUDA graphs", ) -> dict[BatchExecutionDescriptor, AttentionStatePair]: """Capture CUDA graphs for model forward pass.""" self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs if self.use_breakable_cg: self.init_breakable_cg_runner(model) def create_forward_fn( desc: BatchExecutionDescriptor, warmup: bool, ) -> tuple[ Callable[[CUDAGraphMode], None], AttentionState, ]: num_tokens = desc.num_tokens num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs) # Set LoRA state before capture so kernels see correct adapters. if lora_capture_hook is not None: lora_capture_hook(desc.num_active_loras, num_reqs, num_tokens) num_tokens_across_dp = ( torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu") if self.dp_size > 1 else None ) model_inputs = { "input_ids": input_buffers.input_ids[:num_tokens], "positions": input_buffers.positions[:num_tokens], **model_state.prepare_dummy_inputs(num_reqs, num_tokens), } if not self.is_first_pp_rank: # Update for non-first PP ranks. model_inputs["input_ids"] = None model_inputs["inputs_embeds"] = None assert intermediate_tensors is not None model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] attn_metadata, slot_mappings = prepare_inputs_to_capture( num_reqs, num_tokens, model_state, input_buffers, block_tables, attn_groups, kv_cache_config, skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), ) # Capture with dummy rows marked as padding. input_buffers.is_padding.fill_(True) def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None if cg_mode == CUDAGraphMode.PIECEWISE: assert attn_metadata is None batch_descriptor = BatchDescriptor( num_tokens=num_tokens, has_lora=has_lora, num_active_loras=desc.num_active_loras, ) with set_forward_context( attn_metadata, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=cg_mode, num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, is_padding=input_buffers.is_padding[:num_tokens], ): if cg_mode == CUDAGraphMode.PIECEWISE: # PIECEWISE graph (compiled PW or breakable, chosen inside # run_pw_graph). model_output = self.run_pw_graph(model, model_inputs) else: model_output = model(**model_inputs) if cg_mode == CUDAGraphMode.PIECEWISE: # PW CUDA graph (compiled or breakable) internally handles the # model outputs. No need to keep track of the hidden states. return None if self.is_last_pp_rank: # Last PP rank (common case). if self.use_aux_hidden_state_outputs: hidden_states, aux_hidden_states = model_output else: hidden_states = model_output aux_hidden_states = [] if self.hidden_states is None: self.hidden_states = torch.empty_like(hidden_states) self.hidden_states[:num_tokens] = hidden_states if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: self.aux_hidden_states = [ torch.empty_like(x) for x in aux_hidden_states ] for i, aux in enumerate(aux_hidden_states): self.aux_hidden_states[i][:num_tokens] = aux else: # Non-last PP rank. assert isinstance(model_output, IntermediateTensors) intermediate_tensors = model_output if self.intermediate_tensors is None: self.intermediate_tensors = IntermediateTensors.empty_like( intermediate_tensors ) for k, v in intermediate_tensors.tensors.items(): self.intermediate_tensors[k][:num_tokens] = v return forward_fn, AttentionState(attn_metadata, slot_mappings) return super().capture(create_forward_fn, progress_bar_desc) def run_fullgraph( self, desc: BatchExecutionDescriptor ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors: """Replay a captured FULL cudagraph and return hidden states.""" super().run_fullgraph(desc) if not self.is_last_pp_rank: assert self.intermediate_tensors is not None return self.intermediate_tensors[: desc.num_tokens] assert self.hidden_states is not None hidden_states = self.hidden_states[: desc.num_tokens] if not self.use_aux_hidden_state_outputs: return hidden_states return hidden_states, [x[: desc.num_tokens] for x in self.aux_hidden_states] def prepare_inputs_to_capture( num_reqs: int, num_tokens: int, model_state: ModelState, input_buffers: InputBuffers, block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, skip_attn: bool = False, ) -> AttentionState: input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) input_block_tables = block_tables.get_dummy_block_tables(num_reqs) slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens) slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, kv_cache_config ) # HACK(woosuk): Special handling for DCP. if block_tables.cp_size > 1: prepare_dcp_local_seq_lens( input_buffers.dcp_local_seq_lens, input_batch.seq_lens, num_reqs, block_tables.cp_size, block_tables.cp_rank, block_tables.cp_interleave, ) input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs] attn_metadata = None if not skip_attn: attn_metadata = model_state.prepare_attn( input_batch, CUDAGraphMode.NONE, input_block_tables, slot_mappings, attn_groups, kv_cache_config, for_capture=True, ) return AttentionState(attn_metadata, slot_mappings_by_layer)