vllm.v1.worker.gpu.cudagraph_utils ¶

@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
    # Number of microbatches the batch is split into (DBO). 1 means no splitting.
    num_ubatches: int = 1

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

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,
        ubatch_runner: "UBatchRunner | None" = None,
    ):
        self.vllm_config = vllm_config
        self.device = device
        self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
        self.compilation_config = vllm_config.compilation_config
        assert self.compilation_config is not None
        self.cudagraph_mode = cudagraph_mode
        self.decode_query_len = decode_query_len
        self.varlen_decode = varlen_decode
        # DBO supports FULL CUDA graphs only.
        self.ubatch_runner = ubatch_runner

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

                    ubatch_desc = self._maybe_ubatch_twin(desc)
                    if ubatch_desc is not None and (
                        ubatch_desc not in descs_by_mode[decode_mode]
                    ):
                        descs_by_mode[decode_mode].append(ubatch_desc)

            # 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,
                )
                descs_by_mode[mixed_mode].append(desc)

                ubatch_desc = self._maybe_ubatch_twin(desc)
                if ubatch_desc is not None:
                    descs_by_mode[mixed_mode].append(ubatch_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

    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",
    ) -> 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), ExitStack() as stack:
            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. 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=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 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,
                    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) -> 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,
        ubatch_runner: "UBatchRunner | None" = None,
    ):
        super().__init__(
            vllm_config,
            device,
            cudagraph_mode,
            decode_query_len,
            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."
            )

        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 = 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

        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]

            if desc.num_ubatches > 1:
                # Prepare and park threads before capture; finish runs inside it.
                assert self.ubatch_runner is not None
                ubatch_state = self.ubatch_runner.prepare(
                    InputBatch.make_dummy(num_reqs, num_tokens, input_buffers),
                    block_tables.get_dummy_block_tables(num_reqs),
                    block_tables.get_dummy_slot_mappings(num_tokens),
                    cg_mode=CUDAGraphMode.FULL,
                    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,
                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

                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]