vllm.v1.worker.gpu.cudagraph_utils

Classes:

• BatchExecutionDescriptor –

  Describes the shape of the batch and CG mode to run; this is used to make shape

• CreateForwardFn –

  Factory that prepares inputs (OUTSIDE the graph) and returns a

• CudaGraphManager –

  Captures and dispatches the CUDA graphs of one runner.

• ModelCudaGraphManager –

  CudaGraphManager with model-specific capture and hidden state management.

Functions:

• has_compiled_submodule –

  Whether any submodule is an active @support_torch_compile module.

• profile_cudagraph_memory –

  Estimate the GPU memory needed for CUDA graph capture.

BatchExecutionDescriptor dataclass

Describes the shape of the batch and CG mode to run; this is used to make shape matches between the capture and runtime.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

@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; the runner passes None for batches that are
    # not decode-graph eligible, even if their queries fit within the bound.
    max_query_len: int | None = None
    num_active_loras: int = 0
    # Number of microbatches the batch is split into (DBO). 1 means no splitting.
    num_ubatches: int = 1

CreateForwardFn

Bases: 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.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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]: ...

CudaGraphManager

Captures and dispatches the CUDA graphs of one runner.

Uniform decode batches of query length q take graphs of n requests, for each q in decode_query_lens and n in decode_request_sizes (default decode_cudagraph_capture_request_sizes) up to the requests q maps to. With varlen_decode (adaptive verification) they take token-keyed graphs of get_varlen_decode_cudagraph_token_sizes instead. Other batches, and all batches when one routine runs every batch, take token-keyed graphs of prefill_cudagraph_capture_token_sizes.

Methods:

• capture –

  Capture CUDA graphs.

• captured_token_counts –

  Sorted token counts with a captured graph, ignoring LoRA variants.

• dispatch –

  Find the graph to run a batch with, else an eager descriptor.

• release_graphs –

  Drop the captured graphs so a later capture() can refill them.

• run_fullgraph –

  Replay a captured FULL cudagraph.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

class CudaGraphManager:
    """Captures and dispatches the CUDA graphs of one runner.

    Uniform decode batches of query length `q` take graphs of `n` requests, for
    each `q` in `decode_query_lens` and `n` in `decode_request_sizes` (default
    `decode_cudagraph_capture_request_sizes`) up to the requests `q` maps to.
    With `varlen_decode` (adaptive verification) they take token-keyed graphs of
    `get_varlen_decode_cudagraph_token_sizes` instead. Other batches, and all
    batches when one routine runs every batch, take token-keyed graphs of
    `prefill_cudagraph_capture_token_sizes`.
    """

    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        cudagraph_mode: CUDAGraphMode,
        decode_query_lens: Mapping[int, int],
        lora_capture_cases: list[int] | None = None,
        varlen_decode: bool = False,
        ubatch_runner: "UBatchRunner | None" = None,
        decode_request_sizes: Sequence[int] | None = None,
    ):
        self.vllm_config = vllm_config
        self.decode_request_sizes = decode_request_sizes
        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
        # Uniform decode query length -> the most requests a batch of it has.
        self.decode_max_reqs = {
            q: min(max_reqs, self.max_num_reqs)
            for q, max_reqs in sorted(decode_query_lens.items())
        }
        # Widest decode query; varlen decode graphs take 1..this many per request.
        self.decode_query_len = max(self.decode_max_reqs)
        self.varlen_decode = varlen_decode
        # DBO supports FULL CUDA graphs only.
        self.ubatch_runner = ubatch_runner

        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
        # The stream of the last capture; graphs replayed inside these ones (in
        # their eager breaks) can capture on it to share the pool's free blocks.
        self.capture_stream: torch.cuda.Stream | None = 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

        # (num_tokens, num_active_loras) -> token-keyed descriptors to try.
        self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}
        # (query_len, num_active_loras, num_ubatches) -> uniform decode
        # descriptors indexed by request count: entry n is the smallest graph
        # that takes n requests.
        self._decode_descs: dict[
            tuple[int, 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(self.compilation_config)
            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 _maybe_ubatch_twin(
        self, desc: BatchExecutionDescriptor
    ) -> BatchExecutionDescriptor | None:
        """Return a microbatched capture candidate when eligible.

        Uniform query lengths preserve the captured request split. Use the DP
        dispatch thresholds so all ranks generate the same candidates.
        """
        if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL:
            return None
        if desc.num_reqs is None:
            return None
        uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs)
        if remainder or desc.uniform_token_count not in (None, uniform_token_count):
            return None
        parallel_config = self.vllm_config.parallel_config
        num_ubatches = get_num_ubatches(parallel_config)
        if desc.num_tokens < num_ubatches:
            return None
        if not check_ubatch_thresholds(
            parallel_config, desc.num_tokens, uniform_decode=True
        ):
            return None
        return replace(
            desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count
        )

    def _init_candidates(self) -> None:
        """Build the descriptors to capture and the dispatch tables."""
        if self.cudagraph_mode == CUDAGraphMode.NONE:
            return

        compilation_config = self.compilation_config
        prefill_sizes = sorted(
            compilation_config.prefill_cudagraph_capture_token_sizes or []
        )
        decode_sizes = self.decode_request_sizes
        if decode_sizes is None:
            decode_sizes = compilation_config.decode_cudagraph_capture_request_sizes
        decode_sizes = decode_sizes or []
        max_decode_tokens = (
            self.max_num_reqs * self.vllm_config.uniform_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()

        decode_descs: list[BatchExecutionDescriptor] = []
        token_keyed_descs: list[BatchExecutionDescriptor] = []
        # With one routine for every batch, decode batches take token-keyed graphs.
        capture_decode = separate_decode_routine and decode_mode != CUDAGraphMode.NONE
        if capture_decode and self.varlen_decode:
            # Varlen decode graphs take any mix of 1..decode_query_len tokens per
            # request, worst case 1 token per request (or max_num_reqs).
            varlen_sizes = self.vllm_config.get_varlen_decode_cudagraph_token_sizes(
                self.decode_query_len
            )
            for num_tokens, num_active_loras in product(
                varlen_sizes, self.lora_capture_cases
            ):
                token_keyed_descs.append(
                    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,
                    )
                )
        elif capture_decode:
            # Attention backends size CUDA graph buffers and pick graph-compatible
            # paths by the recorded max (e.g. FA3 max_num_splits, FlashInfer's
            # CUDA graph decode wrappers), so a wider graph would compute wrongly.
            # TODO(#59537): when multi-token drafters extend past the target's
            # cap, record their shapes in the recorded max too.
            decode_shapes = decode_cudagraph_capture_shapes(
                decode_sizes,
                self.decode_max_reqs,
                decode_cudagraph_capture_max_tokens(compilation_config),
            )
            for (n, q), num_active_loras in product(
                decode_shapes, self.lora_capture_cases
            ):
                desc = BatchExecutionDescriptor(
                    cg_mode=decode_mode,
                    num_tokens=n * q,
                    num_reqs=n,
                    uniform_token_count=q,
                    num_active_loras=num_active_loras,
                )
                decode_descs.append(desc)
                ubatch_desc = self._maybe_ubatch_twin(desc)
                if ubatch_desc is not None:
                    decode_descs.append(ubatch_desc)

        for num_tokens, num_active_loras in product(
            prefill_sizes, self.lora_capture_cases
        ):
            # recoverSSM cannot capture a dummy query wider than its workspace.
            if mixed_mode and (
                not self.vllm_config.cache_config.use_kda_recoverssm
                or num_tokens <= max_decode_tokens
            ):
                # 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,
                )
                token_keyed_descs.append(desc)

                ubatch_desc = self._maybe_ubatch_twin(desc)
                if ubatch_desc is not None:
                    token_keyed_descs.append(ubatch_desc)

        descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = (
            defaultdict(list)
        )
        for desc in decode_descs + token_keyed_descs:
            descs_by_mode[desc.cg_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

        # A uniform decode batch looks up its request count among the graphs of
        # its own query length, so one query length never shadows another.
        decode_groups: defaultdict[
            tuple[int, int, int], list[BatchExecutionDescriptor]
        ] = defaultdict(list)
        for desc in decode_descs:
            assert desc.uniform_token_count is not None
            decode_groups[
                (desc.uniform_token_count, desc.num_active_loras, desc.num_ubatches)
            ].append(desc)
        for decode_key, decode_group in decode_groups.items():
            table: list[BatchExecutionDescriptor] = []
            for desc in sorted(decode_group, key=lambda d: d.num_tokens):
                assert desc.num_reqs is not None
                table.extend([desc] * (desc.num_reqs + 1 - len(table)))
            self._decode_descs[decode_key] = table

        for mode in (CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE):
            mode_descs = sorted(
                (d for d in token_keyed_descs if d.cg_mode == mode),
                key=lambda d: d.num_tokens,
            )
            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
                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

    @property
    def dp_size(self) -> int:
        # Not cached: elastic EP rewrites parallel_config in place on scale.
        return self.vllm_config.parallel_config.data_parallel_size

    def release_graphs(self) -> None:
        """Drop the captured graphs so a later capture() can refill them.

        Elastic EP reallocates the MoE workspace when it grows, which leaves
        every captured graph holding a stale data pointer. `_capture_descs` is
        kept, so `needs_capture()` still reports the work to redo.
        """
        self.graphs.clear()
        self._graphs_captured = False
        if self.breakable_cg_runner is not None:
            BreakableCUDAGraphWrapper.clear_all_graphs()

    def needs_capture(self) -> bool:
        return len(self._capture_descs) > 0

    def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream:
        """Capture on the stream used by the microbatch threads."""
        if desc.num_ubatches > 1:
            assert self.ubatch_runner is not None
            return self.ubatch_runner.capture_stream
        return current_stream()

    @torch.inference_mode()
    def capture(
        self,
        create_forward_fn: CreateForwardFn,
        progress_bar_desc: str = "Capturing CUDA graphs",
        stream: torch.cuda.Stream | None = None,
    ) -> 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.
            progress_bar_desc: Description shown on the capture progress bar.
            stream: Capture on this stream instead of a new one.

        """
        with graph_capture(self.device, stream) as context, ExitStack() as stack:
            self.capture_stream = context.stream
            if self.ubatch_runner is not None:
                # Join parked threads on failure to avoid blocking later captures.
                stack.callback(self.ubatch_runner.abort_pending_run)
            # Capture in order: PIECEWISE first, then FULL. FULL graphs no wider
            # than the widest PIECEWISE graph fit in its 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._capture_mem_samples is not None:
                            torch.accelerator.synchronize()
                            free_before = torch.accelerator.get_memory_info()[0]
                        with (
                            capture_pool(self.pool, self.vllm_config) as pool,
                            torch.cuda.graph(
                                graph, pool, stream=self._capture_stream(desc)
                            ),
                        ):
                            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,
        num_ubatches: int = 1,
    ) -> BatchExecutionDescriptor:
        """Find the graph to run a batch with, else an eager descriptor.

        A uniform decode batch takes the smallest decode graph of its query
        length that covers `num_tokens`, which may be padded (e.g. to the DP
        maximum), if that graph holds its requests. Other batches, and decode
        batches no decode graph takes, take the smallest compatible token-keyed
        graph.
        """
        effective_loras = self._resolve_effective_loras(num_active_loras)
        if self._graphs_captured and num_tokens > 0:
            candidates = self._candidates.get((num_tokens, effective_loras), [])
            if uniform_token_count is not None:
                decode_descs = self._decode_descs.get(
                    (uniform_token_count, effective_loras, num_ubatches), []
                )
                n = cdiv(num_tokens, uniform_token_count)
                if n < len(decode_descs):
                    candidates = [decode_descs[n], *candidates]
            for desc in candidates:
                if _is_compatible(
                    desc,
                    num_reqs,
                    num_tokens,
                    uniform_token_count,
                    effective_loras,
                    max_query_len,
                    num_ubatches,
                ):
                    return desc
        return BatchExecutionDescriptor(
            cg_mode=CUDAGraphMode.NONE,
            num_tokens=num_tokens,
            num_reqs=num_reqs,
            num_active_loras=effective_loras,
            num_ubatches=num_ubatches,
        )

    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, key_fn: Callable[..., Hashable] | None = None
    ) -> None:
        if self.breakable_cg_runner is None:
            self.breakable_cg_runner = BreakableCUDAGraphWrapper(
                model, self.vllm_config, key_fn=key_fn
            )

    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)

_build_lora_dispatch_map()

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.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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]

_capture_stream(desc)

Capture on the stream used by the microbatch threads.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream:
    """Capture on the stream used by the microbatch threads."""
    if desc.num_ubatches > 1:
        assert self.ubatch_runner is not None
        return self.ubatch_runner.capture_stream
    return current_stream()

_init_candidates()

Build the descriptors to capture and the dispatch tables.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

def _init_candidates(self) -> None:
    """Build the descriptors to capture and the dispatch tables."""
    if self.cudagraph_mode == CUDAGraphMode.NONE:
        return

    compilation_config = self.compilation_config
    prefill_sizes = sorted(
        compilation_config.prefill_cudagraph_capture_token_sizes or []
    )
    decode_sizes = self.decode_request_sizes
    if decode_sizes is None:
        decode_sizes = compilation_config.decode_cudagraph_capture_request_sizes
    decode_sizes = decode_sizes or []
    max_decode_tokens = (
        self.max_num_reqs * self.vllm_config.uniform_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()

    decode_descs: list[BatchExecutionDescriptor] = []
    token_keyed_descs: list[BatchExecutionDescriptor] = []
    # With one routine for every batch, decode batches take token-keyed graphs.
    capture_decode = separate_decode_routine and decode_mode != CUDAGraphMode.NONE
    if capture_decode and self.varlen_decode:
        # Varlen decode graphs take any mix of 1..decode_query_len tokens per
        # request, worst case 1 token per request (or max_num_reqs).
        varlen_sizes = self.vllm_config.get_varlen_decode_cudagraph_token_sizes(
            self.decode_query_len
        )
        for num_tokens, num_active_loras in product(
            varlen_sizes, self.lora_capture_cases
        ):
            token_keyed_descs.append(
                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,
                )
            )
    elif capture_decode:
        # Attention backends size CUDA graph buffers and pick graph-compatible
        # paths by the recorded max (e.g. FA3 max_num_splits, FlashInfer's
        # CUDA graph decode wrappers), so a wider graph would compute wrongly.
        # TODO(#59537): when multi-token drafters extend past the target's
        # cap, record their shapes in the recorded max too.
        decode_shapes = decode_cudagraph_capture_shapes(
            decode_sizes,
            self.decode_max_reqs,
            decode_cudagraph_capture_max_tokens(compilation_config),
        )
        for (n, q), num_active_loras in product(
            decode_shapes, self.lora_capture_cases
        ):
            desc = BatchExecutionDescriptor(
                cg_mode=decode_mode,
                num_tokens=n * q,
                num_reqs=n,
                uniform_token_count=q,
                num_active_loras=num_active_loras,
            )
            decode_descs.append(desc)
            ubatch_desc = self._maybe_ubatch_twin(desc)
            if ubatch_desc is not None:
                decode_descs.append(ubatch_desc)

    for num_tokens, num_active_loras in product(
        prefill_sizes, self.lora_capture_cases
    ):
        # recoverSSM cannot capture a dummy query wider than its workspace.
        if mixed_mode and (
            not self.vllm_config.cache_config.use_kda_recoverssm
            or num_tokens <= max_decode_tokens
        ):
            # 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,
            )
            token_keyed_descs.append(desc)

            ubatch_desc = self._maybe_ubatch_twin(desc)
            if ubatch_desc is not None:
                token_keyed_descs.append(ubatch_desc)

    descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = (
        defaultdict(list)
    )
    for desc in decode_descs + token_keyed_descs:
        descs_by_mode[desc.cg_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

    # A uniform decode batch looks up its request count among the graphs of
    # its own query length, so one query length never shadows another.
    decode_groups: defaultdict[
        tuple[int, int, int], list[BatchExecutionDescriptor]
    ] = defaultdict(list)
    for desc in decode_descs:
        assert desc.uniform_token_count is not None
        decode_groups[
            (desc.uniform_token_count, desc.num_active_loras, desc.num_ubatches)
        ].append(desc)
    for decode_key, decode_group in decode_groups.items():
        table: list[BatchExecutionDescriptor] = []
        for desc in sorted(decode_group, key=lambda d: d.num_tokens):
            assert desc.num_reqs is not None
            table.extend([desc] * (desc.num_reqs + 1 - len(table)))
        self._decode_descs[decode_key] = table

    for mode in (CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE):
        mode_descs = sorted(
            (d for d in token_keyed_descs if d.cg_mode == mode),
            key=lambda d: d.num_tokens,
        )
        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
            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

_maybe_ubatch_twin(desc)

Return a microbatched capture candidate when eligible.

Uniform query lengths preserve the captured request split. Use the DP dispatch thresholds so all ranks generate the same candidates.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

def _maybe_ubatch_twin(
    self, desc: BatchExecutionDescriptor
) -> BatchExecutionDescriptor | None:
    """Return a microbatched capture candidate when eligible.

    Uniform query lengths preserve the captured request split. Use the DP
    dispatch thresholds so all ranks generate the same candidates.
    """
    if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL:
        return None
    if desc.num_reqs is None:
        return None
    uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs)
    if remainder or desc.uniform_token_count not in (None, uniform_token_count):
        return None
    parallel_config = self.vllm_config.parallel_config
    num_ubatches = get_num_ubatches(parallel_config)
    if desc.num_tokens < num_ubatches:
        return None
    if not check_ubatch_thresholds(
        parallel_config, desc.num_tokens, uniform_decode=True
    ):
        return None
    return replace(
        desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count
    )

_resolve_effective_loras(num_active_loras)

Map an actual active-LoRA count to its captured graph case.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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)

capture(create_forward_fn, progress_bar_desc='Capturing CUDA graphs', stream=None)

Capture CUDA graphs.

Parameters:

• create_forward_fn

  (CreateForwardFn) –

  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.

• progress_bar_desc

  (str, default: 'Capturing CUDA graphs' ) –

  Description shown on the capture progress bar.

• stream

  (Stream | None, default: None ) –

  Capture on this stream instead of a new one.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

@torch.inference_mode()
def capture(
    self,
    create_forward_fn: CreateForwardFn,
    progress_bar_desc: str = "Capturing CUDA graphs",
    stream: torch.cuda.Stream | None = None,
) -> 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.
        progress_bar_desc: Description shown on the capture progress bar.
        stream: Capture on this stream instead of a new one.

    """
    with graph_capture(self.device, stream) as context, ExitStack() as stack:
        self.capture_stream = context.stream
        if self.ubatch_runner is not None:
            # Join parked threads on failure to avoid blocking later captures.
            stack.callback(self.ubatch_runner.abort_pending_run)
        # Capture in order: PIECEWISE first, then FULL. FULL graphs no wider
        # than the widest PIECEWISE graph fit in its 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._capture_mem_samples is not None:
                        torch.accelerator.synchronize()
                        free_before = torch.accelerator.get_memory_info()[0]
                    with (
                        capture_pool(self.pool, self.vllm_config) as pool,
                        torch.cuda.graph(
                            graph, pool, stream=self._capture_stream(desc)
                        ),
                    ):
                        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

captured_token_counts()

Sorted token counts with a captured graph, ignoring LoRA variants.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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}
    )

dispatch(num_reqs, num_tokens, uniform_token_count, num_active_loras, max_query_len=None, num_ubatches=1)

Find the graph to run a batch with, else an eager descriptor.

A uniform decode batch takes the smallest decode graph of its query length that covers num_tokens, which may be padded (e.g. to the DP maximum), if that graph holds its requests. Other batches, and decode batches no decode graph takes, take the smallest compatible token-keyed graph.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

def dispatch(
    self,
    num_reqs: int,
    num_tokens: int,
    uniform_token_count: int | None,
    num_active_loras: int,
    max_query_len: int | None = None,
    num_ubatches: int = 1,
) -> BatchExecutionDescriptor:
    """Find the graph to run a batch with, else an eager descriptor.

    A uniform decode batch takes the smallest decode graph of its query
    length that covers `num_tokens`, which may be padded (e.g. to the DP
    maximum), if that graph holds its requests. Other batches, and decode
    batches no decode graph takes, take the smallest compatible token-keyed
    graph.
    """
    effective_loras = self._resolve_effective_loras(num_active_loras)
    if self._graphs_captured and num_tokens > 0:
        candidates = self._candidates.get((num_tokens, effective_loras), [])
        if uniform_token_count is not None:
            decode_descs = self._decode_descs.get(
                (uniform_token_count, effective_loras, num_ubatches), []
            )
            n = cdiv(num_tokens, uniform_token_count)
            if n < len(decode_descs):
                candidates = [decode_descs[n], *candidates]
        for desc in candidates:
            if _is_compatible(
                desc,
                num_reqs,
                num_tokens,
                uniform_token_count,
                effective_loras,
                max_query_len,
                num_ubatches,
            ):
                return desc
    return BatchExecutionDescriptor(
        cg_mode=CUDAGraphMode.NONE,
        num_tokens=num_tokens,
        num_reqs=num_reqs,
        num_active_loras=effective_loras,
        num_ubatches=num_ubatches,
    )

release_graphs()

Drop the captured graphs so a later capture() can refill them.

Elastic EP reallocates the MoE workspace when it grows, which leaves every captured graph holding a stale data pointer. _capture_descs is kept, so needs_capture() still reports the work to redo.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

def release_graphs(self) -> None:
    """Drop the captured graphs so a later capture() can refill them.

    Elastic EP reallocates the MoE workspace when it grows, which leaves
    every captured graph holding a stale data pointer. `_capture_descs` is
    kept, so `needs_capture()` still reports the work to redo.
    """
    self.graphs.clear()
    self._graphs_captured = False
    if self.breakable_cg_runner is not None:
        BreakableCUDAGraphWrapper.clear_all_graphs()

run_fullgraph(desc)

Replay a captured FULL cudagraph.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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()

ModelCudaGraphManager

Bases: CudaGraphManager

CudaGraphManager with model-specific capture and hidden state management.

Methods:

• capture –

  Capture CUDA graphs for model forward pass.

• run_fullgraph –

  Replay a captured FULL cudagraph and return hidden states.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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_lens: Mapping[int, int],
        lora_capture_cases: list[int] | None = None,
        varlen_decode: bool = False,
        ubatch_runner: "UBatchRunner | None" = None,
    ):
        super().__init__(
            vllm_config,
            device,
            cudagraph_mode,
            decode_query_lens,
            lora_capture_cases=lora_capture_cases,
            varlen_decode=varlen_decode,
            ubatch_runner=ubatch_runner,
        )
        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."
            )

        # The first output stored (of the widest PIECEWISE graph) may be narrower
        # than the widest decode graph, so size the buffers for the widest.
        max_tokens = max(
            (d.num_tokens for descs in self._capture_descs.values() for d in descs),
            default=0,
        )

        def persistent_like(x: torch.Tensor) -> torch.Tensor:
            return x.new_empty((max(max_tokens, x.shape[0]), *x.shape[1:]))

        def store_capture_output(num_tokens: int, model_output: Any) -> None:
            """Copy outputs to persistent buffers, allocating on first use."""
            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 = persistent_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 = [
                        persistent_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(
                        {
                            k: persistent_like(v)
                            for k, v in intermediate_tensors.tensors.items()
                        }
                    )
                for k, v in intermediate_tensors.tensors.items():
                    self.intermediate_tensors[k][:num_tokens] = v

        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],
                "intermediate_tensors": None,
                **model_state.prepare_dummy_inputs(num_reqs, num_tokens),
            }
            if not self.is_first_pp_rank:
                # Update for non-first PP ranks.
                if not requires_raw_input_tokens(model):
                    model_inputs["input_ids"] = None
                model_inputs["inputs_embeds"] = None
                assert intermediate_tensors is not None
                model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens]

            if desc.num_ubatches > 1:
                # Prepare and park threads before capture; finish runs inside it.
                assert self.ubatch_runner is not None
                input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
                input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
                set_dcp_capture_context(
                    input_batch,
                    model_state,
                    block_tables,
                    input_block_tables,
                    kv_cache_config,
                )
                ubatch_state = self.ubatch_runner.prepare(
                    input_batch,
                    input_block_tables,
                    block_tables.get_dummy_slot_mappings(num_tokens),
                    desc,
                    for_capture=True,
                )
                # Capture with dummy rows marked as padding.
                input_buffers.is_padding.fill_(True)
                finish = self.ubatch_runner.begin_capturable_run(
                    model, model_inputs, ubatch_state, for_capture=True
                )

                def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None:
                    assert cg_mode != CUDAGraphMode.PIECEWISE, (
                        "DBO does not support PIECEWISE cudagraphs"
                    )
                    store_capture_output(num_tokens, finish())

                return ubatch_forward_fn

            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 or desc.uniform_token_count,
                pcp_manager=pcp_manager,
            )

            # Capture with dummy rows marked as padding.
            input_buffers.is_padding.fill_(True)
            is_padding = input_buffers.is_padding[:num_tokens]
            if pcp_manager is not None:
                is_padding = pcp_manager.moe_is_padding(is_padding)

            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=is_padding,
                ):
                    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

                store_capture_output(num_tokens, model_output)

            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]

capture(model, model_state, input_buffers, intermediate_tensors, block_tables, attn_groups, kv_cache_config, pcp_manager=None, has_lora=False, use_aux_hidden_state_outputs=False, lora_capture_hook=None, progress_bar_desc='Capturing CUDA graphs')

Capture CUDA graphs for model forward pass.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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."
        )

    # The first output stored (of the widest PIECEWISE graph) may be narrower
    # than the widest decode graph, so size the buffers for the widest.
    max_tokens = max(
        (d.num_tokens for descs in self._capture_descs.values() for d in descs),
        default=0,
    )

    def persistent_like(x: torch.Tensor) -> torch.Tensor:
        return x.new_empty((max(max_tokens, x.shape[0]), *x.shape[1:]))

    def store_capture_output(num_tokens: int, model_output: Any) -> None:
        """Copy outputs to persistent buffers, allocating on first use."""
        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 = persistent_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 = [
                    persistent_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(
                    {
                        k: persistent_like(v)
                        for k, v in intermediate_tensors.tensors.items()
                    }
                )
            for k, v in intermediate_tensors.tensors.items():
                self.intermediate_tensors[k][:num_tokens] = v

    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],
            "intermediate_tensors": None,
            **model_state.prepare_dummy_inputs(num_reqs, num_tokens),
        }
        if not self.is_first_pp_rank:
            # Update for non-first PP ranks.
            if not requires_raw_input_tokens(model):
                model_inputs["input_ids"] = None
            model_inputs["inputs_embeds"] = None
            assert intermediate_tensors is not None
            model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens]

        if desc.num_ubatches > 1:
            # Prepare and park threads before capture; finish runs inside it.
            assert self.ubatch_runner is not None
            input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
            input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
            set_dcp_capture_context(
                input_batch,
                model_state,
                block_tables,
                input_block_tables,
                kv_cache_config,
            )
            ubatch_state = self.ubatch_runner.prepare(
                input_batch,
                input_block_tables,
                block_tables.get_dummy_slot_mappings(num_tokens),
                desc,
                for_capture=True,
            )
            # Capture with dummy rows marked as padding.
            input_buffers.is_padding.fill_(True)
            finish = self.ubatch_runner.begin_capturable_run(
                model, model_inputs, ubatch_state, for_capture=True
            )

            def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None:
                assert cg_mode != CUDAGraphMode.PIECEWISE, (
                    "DBO does not support PIECEWISE cudagraphs"
                )
                store_capture_output(num_tokens, finish())

            return ubatch_forward_fn

        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 or desc.uniform_token_count,
            pcp_manager=pcp_manager,
        )

        # Capture with dummy rows marked as padding.
        input_buffers.is_padding.fill_(True)
        is_padding = input_buffers.is_padding[:num_tokens]
        if pcp_manager is not None:
            is_padding = pcp_manager.moe_is_padding(is_padding)

        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=is_padding,
            ):
                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

            store_capture_output(num_tokens, model_output)

        return forward_fn

    super().capture(create_forward_fn, progress_bar_desc)

run_fullgraph(desc)

Replay a captured FULL cudagraph and return hidden states.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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]

_extrapolate_full_graph_memory(mem_samples, total_graphs)

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.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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

_init_minimal_kv_cache_for_profiling(runner)

Allocate the smallest KV cache that still lets every graph be captured.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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_config_builder import get_kv_cache_config_builder
    from vllm.v1.core.kv_cache_planning import get_kv_cache_groups_from_builder

    kv_cache_spec = runner.get_kv_cache_spec()
    builder = get_kv_cache_config_builder(runner.vllm_config)
    kv_cache_groups = get_kv_cache_groups_from_builder(
        builder, 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
    )
    minimal_config = builder.get_kv_cache_config_from_groups(
        runner.vllm_config, kv_cache_groups, num_blocks=min_blocks
    )

    runner.initialize_kv_cache(minimal_config, is_profiling=True)
    runner.cache_config.num_gpu_blocks = minimal_config.num_blocks

_teardown_profiling_state(runner)

Release the profiling KV cache and captured graphs while keeping model weights, so the real initialize_kv_cache starts from a clean slate.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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)
    unbind_hisparse_caches(runner.compilation_config.static_forward_context)
    runner.cache_config.num_gpu_blocks = None
    runner.maybe_remove_all_loras(runner.lora_config)
    gc.collect()
    torch.accelerator.empty_cache()

has_compiled_submodule(model)

Whether any submodule is an active @support_torch_compile module.

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

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()
    )

profile_cudagraph_memory(runner)

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).

Source code in vllm/v1/worker/gpu/cudagraph_utils.py

@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
        and not runner.needs_cudagraph_capture()
    ):
        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 runner.needs_cudagraph_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

            with capture_outside_cumem_pool():
                measured = int(runner.capture_model(profile_only=True))

            # 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