917 lines
41 KiB
Python
917 lines
41 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import gc
|
|
import itertools
|
|
from collections import defaultdict
|
|
from collections.abc import Callable, Iterable
|
|
from dataclasses import dataclass
|
|
from itertools import groupby, product
|
|
from typing import TYPE_CHECKING, 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.compilation.cuda_graph import CUDAGraphWrapper
|
|
from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
|
|
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.utils.torch_utils import current_stream
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig
|
|
from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
|
|
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, clear_layer_kv_caches
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
|
|
from vllm.v1.worker.gpu.pcp_manager import PCPManager
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class AttentionState(NamedTuple):
|
|
attn_metadata: dict[str, Any] | None
|
|
slot_mappings: dict[str, torch.Tensor]
|
|
|
|
|
|
@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
|
|
# Upper bound on per-request query length. Varlen decode graphs leave
|
|
# uniform_token_count unset, so this is what keeps a prefill batch out of one.
|
|
max_query_len: int | None = None
|
|
num_active_loras: int = 0
|
|
|
|
|
|
class CreateForwardFn(Protocol):
|
|
"""Factory that prepares inputs (OUTSIDE the graph) and returns a
|
|
forward_fn. Called with warmup=True for the warmup pass and warmup=False
|
|
for the captured pass."""
|
|
|
|
def __call__(
|
|
self,
|
|
desc: BatchExecutionDescriptor,
|
|
warmup: bool,
|
|
) -> Callable[[CUDAGraphMode], None]: ...
|
|
|
|
|
|
def _is_compatible(
|
|
desc: BatchExecutionDescriptor,
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
uniform_token_count: int | None,
|
|
num_active_loras: int,
|
|
max_query_len: int | None,
|
|
) -> bool:
|
|
# desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count
|
|
# desc.num_reqs=None means no request padding needed (PIECEWISE)
|
|
# desc.max_query_len=None means the graph does not constrain query length; a
|
|
# caller that does not track max_query_len must not match one that does
|
|
return (
|
|
(
|
|
desc.uniform_token_count is None
|
|
or desc.uniform_token_count == uniform_token_count
|
|
)
|
|
and (
|
|
desc.max_query_len is None
|
|
or (max_query_len is not None and desc.max_query_len >= max_query_len)
|
|
)
|
|
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 has_compiled_submodule(model: nn.Module) -> bool:
|
|
"""Whether any submodule is an active @support_torch_compile module."""
|
|
return any(
|
|
isinstance(m, TorchCompileWithNoGuardsWrapper)
|
|
and not getattr(m, "do_not_compile", True)
|
|
for m in model.modules()
|
|
)
|
|
|
|
|
|
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,
|
|
varlen_decode: bool = False,
|
|
):
|
|
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.varlen_decode = varlen_decode
|
|
|
|
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
|
|
|
|
# Profiling hooks, set only by profile_cudagraph_memory() below: cap
|
|
# FULL-mode capture at the N largest descriptors and record each
|
|
# captured FULL graph's memory delta for extrapolation.
|
|
self._max_full_descs_to_capture: int | None = None
|
|
self._capture_mem_samples: list[int] | None = None
|
|
|
|
self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}
|
|
self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
|
|
|
|
# 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
|
|
|
|
self._init_candidates()
|
|
|
|
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_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()
|
|
):
|
|
# 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
|
|
)
|
|
dense_schedule = build_dynamic_sd_schedule_lookup(
|
|
speculative_config.num_speculative_tokens_per_batch_size,
|
|
vllm_max_batch_size=self.max_num_reqs,
|
|
vllm_num_speculative_tokens=self.vllm_config.num_speculative_tokens,
|
|
)
|
|
decode_query_lens = sorted(
|
|
{
|
|
num_spec + num_new_sampled_tokens_per_step
|
|
for num_spec in dense_schedule[1:]
|
|
}
|
|
)
|
|
else:
|
|
decode_query_lens = [self.decode_query_len]
|
|
|
|
capture_varlen_decode = (
|
|
separate_decode_routine and bool(decode_mode) and self.varlen_decode
|
|
)
|
|
for num_tokens, num_active_loras in product(
|
|
capture_sizes, self.lora_capture_cases
|
|
):
|
|
# Varlen decode graphs take any mix of 1..decode_query_len tokens per
|
|
# request, worst case 1 token per request (or max_num_reqs)
|
|
if capture_varlen_decode and num_tokens <= max_decode_tokens:
|
|
desc = BatchExecutionDescriptor(
|
|
cg_mode=decode_mode,
|
|
num_tokens=num_tokens,
|
|
num_reqs=min(num_tokens, self.max_num_reqs),
|
|
max_query_len=self.decode_query_len,
|
|
num_active_loras=num_active_loras,
|
|
)
|
|
descs_by_mode[decode_mode].append(desc)
|
|
# Capture uniform decode specfifc graphs if required
|
|
# (i.e. separate decode routine)
|
|
elif separate_decode_routine and decode_mode and not self.varlen_decode:
|
|
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)
|
|
|
|
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.
|
|
# For breakable PW graphs, break-point kernels read the real batch
|
|
# from the forward context; in-graph kernels handle the token padding
|
|
# themselves from the padded slot_mapping (rows with slot == -1).
|
|
num_reqs = None
|
|
if mixed_mode == CUDAGraphMode.FULL:
|
|
num_reqs = min(num_tokens, self.max_num_reqs)
|
|
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)
|
|
|
|
for mode, descs in descs_by_mode.items():
|
|
descs.sort(key=lambda d: d.num_tokens, reverse=True)
|
|
self._capture_descs[mode] = descs
|
|
|
|
for mode in (CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE):
|
|
mode_descs = tuple(reversed(descs_by_mode.get(mode, [])))
|
|
for num_active_loras in self.lora_capture_cases:
|
|
lora_descs = [
|
|
d for d in mode_descs if d.num_active_loras == num_active_loras
|
|
]
|
|
current_range_start = 0
|
|
# Dynamic speculative decoding can produce multiple graphs with the same
|
|
# num_tokens. Group them so each graph covers the same candidate range.
|
|
for num_tokens, group in groupby(lora_descs, lambda d: d.num_tokens):
|
|
matching = list(group)
|
|
for i in range(current_range_start, num_tokens + 1):
|
|
key = (i, num_active_loras)
|
|
self._candidates.setdefault(key, []).extend(matching)
|
|
current_range_start = num_tokens + 1
|
|
|
|
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",
|
|
) -> None:
|
|
"""Capture CUDA graphs.
|
|
|
|
Args:
|
|
create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and
|
|
returns a forward_fn. For FULL and breakable PIECEWISE modes,
|
|
it is invoked once with warmup=True and again with warmup=False
|
|
because attention backends may mutate or lazily initialize
|
|
metadata during warmup.
|
|
"""
|
|
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 (
|
|
mode == CUDAGraphMode.FULL
|
|
and self._max_full_descs_to_capture is not None
|
|
):
|
|
# Profiling only: capture a sample of the largest FULL
|
|
# graphs; the total cost is extrapolated from their
|
|
# per-graph memory deltas.
|
|
descs = descs[: self._max_full_descs_to_capture]
|
|
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 = 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
|
|
and not self.use_breakable_cg
|
|
):
|
|
forward_fn(CUDAGraphMode.PIECEWISE)
|
|
else:
|
|
# Capture with fresh attention state.
|
|
forward_fn = create_forward_fn(desc, warmup=False)
|
|
if desc.cg_mode == CUDAGraphMode.PIECEWISE:
|
|
forward_fn(CUDAGraphMode.PIECEWISE)
|
|
continue
|
|
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()
|
|
if self.pool is not None:
|
|
set_graph_pool_id(self.pool)
|
|
else:
|
|
set_graph_pool_id(current_platform.graph_pool_handle())
|
|
if self._capture_mem_samples is not None:
|
|
torch.accelerator.synchronize()
|
|
free_before = torch.accelerator.get_memory_info()[0]
|
|
with torch.cuda.graph(
|
|
graph, self.pool, stream=current_stream()
|
|
):
|
|
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()
|
|
if self._capture_mem_samples is not None:
|
|
torch.accelerator.synchronize()
|
|
free_after = torch.accelerator.get_memory_info()[0]
|
|
self._capture_mem_samples.append(free_before - free_after)
|
|
self.graphs[desc] = graph
|
|
compilation_counter.num_cudagraph_captured += 1
|
|
self._graphs_captured = True
|
|
|
|
def captured_token_counts(self) -> list[int]:
|
|
"""Sorted token counts with a captured graph, ignoring LoRA variants."""
|
|
return sorted(
|
|
{desc.num_tokens for desc in self.graphs if desc.num_active_loras == 0}
|
|
)
|
|
|
|
def dispatch(
|
|
self,
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
uniform_token_count: int | None,
|
|
num_active_loras: int,
|
|
max_query_len: int | None = None,
|
|
) -> 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,
|
|
max_query_len,
|
|
):
|
|
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,
|
|
varlen_decode: bool = False,
|
|
):
|
|
super().__init__(
|
|
vllm_config,
|
|
device,
|
|
cudagraph_mode,
|
|
decode_query_len,
|
|
lora_capture_cases=lora_capture_cases,
|
|
varlen_decode=varlen_decode,
|
|
)
|
|
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,
|
|
pcp_manager: "PCPManager | None" = None,
|
|
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",
|
|
) -> None:
|
|
"""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)
|
|
|
|
if self.cudagraph_mode.has_piecewise_cudagraphs() and not (
|
|
self.use_breakable_cg or has_compiled_submodule(model)
|
|
):
|
|
raise RuntimeError(
|
|
f"{type(model).__name__}: piecewise CUDA graphs "
|
|
f"(cudagraph_mode={self.cudagraph_mode.name}) unavailable, "
|
|
"model is not torch-compiled and breakable CUDA graph is off. "
|
|
"Set VLLM_USE_BREAKABLE_CUDAGRAPH=1 or cudagraph_mode=NONE/FULL."
|
|
)
|
|
|
|
def create_forward_fn(
|
|
desc: BatchExecutionDescriptor,
|
|
warmup: bool,
|
|
) -> Callable[[CUDAGraphMode], None]:
|
|
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,
|
|
full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
|
|
max_query_len=desc.max_query_len,
|
|
pcp_manager=pcp_manager,
|
|
)
|
|
|
|
# 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:
|
|
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
|
|
|
|
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,
|
|
full_cudagraph: bool,
|
|
max_query_len: int | None = None,
|
|
pcp_manager: "PCPManager | None" = None,
|
|
) -> AttentionState:
|
|
input_batch = InputBatch.make_dummy(
|
|
num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
|
|
)
|
|
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
|
|
slot_mapping_provider: BlockTables | PCPManager = block_tables
|
|
if pcp_manager is not None:
|
|
slot_mapping_provider = pcp_manager
|
|
slot_mappings = slot_mapping_provider.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]
|
|
|
|
# NOTE(woosuk): Attention metadata is required not just by standard attention
|
|
# kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
|
|
# DSV4 compressor), which maintain their own states and require special metadata
|
|
# such as block tables.
|
|
# During CUDA graph capture:
|
|
# - For FULL CUDA graphs: We set for_capture=True so that both attention and
|
|
# attention-like ops produce capturable metadata compatible with CUDA graphs.
|
|
# - For PIECEWISE CUDA graphs: We still build attention metadata, but set
|
|
# for_capture=False. This is because:
|
|
# * Attention-like ops (such as sconv or DSV4 compressor) may not be used as
|
|
# breakpoints in PIECEWISE CUDA graphs, so we must generate their attention
|
|
# metadata so they can execute and be captured during graph capture.
|
|
# * Standard attention ops that are treated as breakpoints will be executed
|
|
# eagerly at capture time (not included in the graph itself), and for these,
|
|
# setting for_capture=False is essential. Some attention backends
|
|
# (like linear attention) cannot generate capturable metadata for prefill,
|
|
# so for_capture=False ensures they execute without issue.
|
|
# * We assume that attention-like operations intended for capture will still
|
|
# produce capturable metadata, even when for_capture=False. While this
|
|
# assumption is brittle, it currently works in practice.
|
|
# In summary: We always generate attention metadata for both FULL and PIECEWISE
|
|
# CUDA graphs, setting for_capture=True for FULL graphs, and for_capture=False
|
|
# for PIECEWISE graphs, to ensure correct execution and capture.
|
|
attn_metadata = model_state.prepare_attn(
|
|
input_batch,
|
|
CUDAGraphMode.NONE,
|
|
input_block_tables,
|
|
slot_mappings,
|
|
attn_groups,
|
|
kv_cache_config,
|
|
for_capture=full_cudagraph,
|
|
)
|
|
return AttentionState(attn_metadata, slot_mappings_by_layer)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CUDA graph memory profiling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Number of FULL graphs captured during profiling; the total FULL capture
|
|
# cost is extrapolated from this sample to avoid a second full capture.
|
|
_FULL_GRAPH_PROFILING_SAMPLES = 2
|
|
# Floor for the extrapolated per-graph cost (driver overhead per graph).
|
|
_MIN_PER_GRAPH_BYTES = 1 << 20
|
|
|
|
|
|
@torch.inference_mode()
|
|
def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
|
|
"""Estimate the GPU memory needed for CUDA graph capture.
|
|
|
|
Called during memory profiling, *before* the real KV cache is allocated,
|
|
so that ``Worker.determine_available_memory`` can reserve headroom for
|
|
graph capture. Bootstraps a minimal KV cache, runs ``capture_model()``
|
|
once, then releases everything so the real init/capture path starts clean.
|
|
|
|
FULL graphs bake in KV cache pointers, so only the largest few are
|
|
captured (into a throwaway pool) and their total cost is extrapolated.
|
|
PIECEWISE, encoder and speculator graphs are measured in full. All
|
|
profiling captures are discarded afterwards: replaying graphs recorded
|
|
against the throwaway profiling state is unsafe (e.g. inductor graph
|
|
partition reclaims the storages of earlier cudagraph recordings once the
|
|
real capture records new ones, leading to use-after-free crashes).
|
|
"""
|
|
if runner.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:
|
|
return 0
|
|
|
|
gc.collect()
|
|
torch.accelerator.empty_cache()
|
|
|
|
# Run the whole profiling phase against a throwaway CUDA graph pool by
|
|
# pointing the global graph pool singleton at it: objects that bind the
|
|
# pool lazily during profiling (speculator cudagraph managers, breakable
|
|
# runners created mid-capture) then land on the throwaway pool too. Pools
|
|
# bound before profiling (piecewise wrappers) are swapped explicitly in
|
|
# the inner block. Profiling graphs captured into the persistent global
|
|
# pool and then discarded would drop its use_count to 0, tripping the c10
|
|
# allocator's create_or_incref_pool assert when the real capture reuses
|
|
# that pool ("use_count > 0 INTERNAL ASSERT FAILED").
|
|
platform_cls = type(current_platform)
|
|
saved_global_pool = platform_cls._global_graph_pool
|
|
throwaway_pool = current_platform.graph_pool_handle()
|
|
platform_cls._global_graph_pool = throwaway_pool
|
|
|
|
try:
|
|
with set_current_vllm_config(runner.vllm_config):
|
|
_init_minimal_kv_cache_for_profiling(runner)
|
|
|
|
manager = runner.cudagraph_manager
|
|
assert manager is not None
|
|
|
|
# Don't count profiling captures; the real capture_model() runs later.
|
|
saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured
|
|
saved_capture_triggers = compilation_counter.num_gpu_runner_capture_triggers
|
|
all_wrappers: list[Any] = []
|
|
original_pools: dict[int, Any] = {}
|
|
speculator = getattr(runner, "speculator", None)
|
|
spec_manager_names: list[str] = []
|
|
try:
|
|
if not manager.needs_capture():
|
|
return 0
|
|
manager.pool = throwaway_pool
|
|
if manager.use_breakable_cg:
|
|
# Create the breakable runner before the wrapper pool swap so
|
|
# its pool is covered as well.
|
|
manager.init_breakable_cg_runner(runner.model)
|
|
all_wrappers = list(CUDAGraphWrapper._all_instances) + list(
|
|
BreakableCUDAGraphWrapper._all_instances
|
|
)
|
|
for wrapper in all_wrappers:
|
|
original_pools[id(wrapper)] = wrapper.graph_pool
|
|
wrapper.graph_pool = throwaway_pool
|
|
if speculator is not None:
|
|
spec_manager_names = [
|
|
name
|
|
for name, value in vars(speculator).items()
|
|
if isinstance(value, CudaGraphManager)
|
|
]
|
|
manager._max_full_descs_to_capture = _FULL_GRAPH_PROFILING_SAMPLES
|
|
mem_samples: list[int] = []
|
|
manager._capture_mem_samples = mem_samples
|
|
|
|
measured = int(runner.capture_model())
|
|
|
|
# The measured delta covers PIECEWISE, encoder and speculator graphs
|
|
# plus the sampled FULL graphs; swap the sampled FULL cost for the
|
|
# extrapolated total. FULL and PIECEWISE share one pool here just as
|
|
# they share the global pool at runtime, so the overlap is not
|
|
# double-counted.
|
|
num_full_graphs = len(manager._capture_descs.get(CUDAGraphMode.FULL, []))
|
|
full_estimate = _extrapolate_full_graph_memory(mem_samples, num_full_graphs)
|
|
return max(measured - sum(mem_samples) + full_estimate, 0)
|
|
finally:
|
|
compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured
|
|
compilation_counter.num_gpu_runner_capture_triggers = saved_capture_triggers
|
|
CUDAGraphWrapper.clear_all_graphs()
|
|
BreakableCUDAGraphWrapper.clear_all_graphs()
|
|
for wrapper in all_wrappers:
|
|
if id(wrapper) in original_pools:
|
|
wrapper.graph_pool = original_pools[id(wrapper)]
|
|
# Drop the speculator's cudagraph managers; the real
|
|
# initialize_kv_cache re-creates them. Their profiling graphs
|
|
# release the throwaway pool here rather than after the real init.
|
|
for name in spec_manager_names:
|
|
setattr(speculator, name, None)
|
|
# Drop local references before teardown detaches the runner's
|
|
# manager and flushes the allocator.
|
|
del manager
|
|
_teardown_profiling_state(runner)
|
|
finally:
|
|
platform_cls._global_graph_pool = saved_global_pool
|
|
|
|
|
|
def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> int:
|
|
"""Extrapolate the total FULL capture cost from samples of the largest
|
|
graphs. The first capture allocates the pool baseline; later graphs mostly
|
|
reuse it, so the second sample is taken as the per-graph cost."""
|
|
if not mem_samples:
|
|
return 0
|
|
first_capture = mem_samples[0]
|
|
per_graph = max(mem_samples[1], _MIN_PER_GRAPH_BYTES) if len(mem_samples) > 1 else 0
|
|
return first_capture + (total_graphs - 1) * per_graph
|
|
|
|
|
|
def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None:
|
|
"""Allocate the smallest KV cache that still lets every graph be captured."""
|
|
from vllm.v1.core.kv_cache_utils import (
|
|
get_kv_cache_config_from_groups,
|
|
get_kv_cache_groups,
|
|
)
|
|
|
|
kv_cache_spec = runner.get_kv_cache_spec()
|
|
kv_cache_groups = get_kv_cache_groups(runner.vllm_config, kv_cache_spec)
|
|
# At least one block per sequence is required to capture the graphs.
|
|
min_blocks = (
|
|
min(runner.max_num_reqs, runner.compilation_config.max_cudagraph_capture_size)
|
|
or 1
|
|
)
|
|
saved_override = runner.cache_config.num_gpu_blocks_override
|
|
runner.cache_config.num_gpu_blocks_override = min_blocks
|
|
try:
|
|
minimal_config = get_kv_cache_config_from_groups(
|
|
runner.vllm_config, kv_cache_groups, available_memory=0
|
|
)
|
|
finally:
|
|
runner.cache_config.num_gpu_blocks_override = saved_override
|
|
|
|
runner.initialize_kv_cache(minimal_config, is_profiling=True)
|
|
runner.cache_config.num_gpu_blocks = minimal_config.num_blocks
|
|
|
|
|
|
def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
|
|
"""Release the profiling KV cache and captured graphs while keeping model
|
|
weights, so the real ``initialize_kv_cache`` starts from a clean slate."""
|
|
torch.accelerator.synchronize()
|
|
if hasattr(runner.model_state, "_mamba_ctx"):
|
|
runner.model_state._mamba_ctx = None
|
|
# Invalidate the align-mode Mamba group metadata cached from the
|
|
# profiling KVCacheConfig: the real (e.g. PP-projected) config may
|
|
# place Mamba layers into a different group layout, so it must be
|
|
# re-derived from the real config.
|
|
if hasattr(runner.model_state, "_mamba_group_ids"):
|
|
runner.model_state._mamba_group_ids = []
|
|
if hasattr(runner.model_state, "_mamba_spec"):
|
|
runner.model_state._mamba_spec = None
|
|
if hasattr(runner, "kv_caches"):
|
|
runner.kv_caches.clear()
|
|
if hasattr(runner, "attn_groups"):
|
|
runner.attn_groups.clear()
|
|
if hasattr(runner, "kv_cache_config"):
|
|
del runner.kv_cache_config
|
|
# Dropping the manager releases the profiling graphs and throwaway pool.
|
|
runner.cudagraph_manager = None
|
|
# Release encoder graphs captured during profiling; the real
|
|
# capture_model() re-captures them.
|
|
if runner.model_state.supports_mm_inputs:
|
|
runner.model_state.encoder_runner.clear()
|
|
# Detach profiling KV tensors held by attention layers. The layers live
|
|
# in the static forward context for compiled models.
|
|
layers: Iterable[Any] = runner.compilation_config.static_forward_context.values()
|
|
if (model := getattr(runner, "model", None)) is not None:
|
|
layers = itertools.chain(layers, model.modules())
|
|
clear_layer_kv_caches(layers)
|
|
runner.cache_config.num_gpu_blocks = None
|
|
runner.maybe_remove_all_loras(runner.lora_config)
|
|
gc.collect()
|
|
torch.accelerator.empty_cache()
|