vllm-omni/vllm_omni/worker/gpu_ar_model_runner.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""AR GPU Model Runner for vLLM-Omni.

Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
and also outputs sampled tokens.
"""

from __future__ import annotations

import gc
import threading
from collections.abc import Callable, Sequence
from contextlib import nullcontext
from copy import copy
from dataclasses import replace
from typing import Any, NamedTuple, cast

import numpy as np
import torch
from vllm.config import CUDAGraphMode
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
from vllm.forward_context import set_forward_context
from vllm.logger import init_logger
from vllm.sequence import IntermediateTensors
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
from vllm.v1.spec_decode.dflash import DFlashProposer
from vllm.v1.spec_decode.draft_model import DraftModelProposer
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
from vllm.v1.structured_output.utils import apply_grammar_bitmask
from vllm.v1.utils import record_function_or_nullcontext
from vllm.v1.worker.gpu_model_runner import (
    AsyncGPUModelRunnerOutput,
)
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
from vllm.v1.worker.utils import is_residual_scattered_for_sp

from vllm_omni.data_entry_keys import flatten_payload
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output
from vllm_omni.model_executor.duplex_sampling import DuplexSamplingRunnerMixin
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.utils.mm_outputs import (
    build_mm_cpu,
    partition_payload_list,
    snapshot_mm_payload,
)
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
from vllm_omni.worker.omni_connector_model_runner_mixin import (
    OmniConnectorModelRunnerMixin,
    needs_omni_connector,
)
from vllm_omni.worker.output.payload_build import build_omni_mm_payload
from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding, sanitize_min_tokens_stop_ids
from vllm_omni.worker.sparse_audio import resolve_sparse_mm_routing

logger = init_logger(__name__)


def _to_cpu_contiguous(tensor: torch.Tensor) -> torch.Tensor:
    tensor = tensor.detach()
    if tensor.device.type == "cpu":
        return tensor.contiguous()
    return tensor.to("cpu").contiguous()


def _clone_cuda_tensor_payload(value: Any, sources: list[torch.Tensor]) -> Any:
    """Clone CUDA tensors on the current stream before async CPU copies.

    The clone protects async Omni output snapshots from CUDA graph output
    buffers that may be reused by subsequent decode steps. CPU tensors are
    cloned synchronously because they are already host-owned snapshots.
    """
    if isinstance(value, torch.Tensor):
        if value.device.type == "cuda":
            cloned = value.detach().clone()
            sources.append(cloned)
            return cloned
        return value.detach().clone()
    if isinstance(value, dict):
        return {k: _clone_cuda_tensor_payload(v, sources) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        if (
            value
            and all(
                isinstance(v, torch.Tensor)
                and v.is_cuda
                and v.dtype == value[0].dtype
                and v.device == value[0].device
                and v.layout == torch.strided
                and v.is_contiguous()
                for v in value
            )
            # Keep frame/state packing bounded; large payloads use individual snapshots.
            and sum(v.numel() * v.element_size() for v in value) <= 1024 * 1024
        ):
            packed = torch.cat([v.detach().reshape(-1) for v in value])
            sources.append(packed)
            return _PackedTensorPayload(packed, [v.shape for v in value], isinstance(value, tuple))
        items = [_clone_cuda_tensor_payload(v, sources) for v in value]
        return tuple(items) if isinstance(value, tuple) else items
    return value


def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any:
    if isinstance(value, _PackedTensorPayload):
        flat = _copy_tensor_payload_to_cpu(value.tensor, pin_memory)
        items = []
        offset = 0
        for shape in value.shapes:
            count = shape.numel()
            items.append(flat[offset : offset + count].view(shape))
            offset += count
        return tuple(items) if value.is_tuple else items
    if isinstance(value, torch.Tensor):
        if value.device.type != "cuda":
            return value
        cpu = torch.empty_like(value, device="cpu", pin_memory=pin_memory)
        cpu.copy_(value, non_blocking=True)
        return cpu
    if isinstance(value, dict):
        return {k: _copy_tensor_payload_to_cpu(v, pin_memory) for k, v in value.items()}
    if isinstance(value, list):
        return [_copy_tensor_payload_to_cpu(v, pin_memory) for v in value]
    if isinstance(value, tuple):
        return tuple(_copy_tensor_payload_to_cpu(v, pin_memory) for v in value)
    return value


class _PackedTensorPayload(NamedTuple):
    tensor: torch.Tensor
    shapes: list[torch.Size]
    is_tuple: bool


class _AsyncCPUPayloadSnapshot:
    def __init__(
        self,
        payload: Any,
        ready_event: torch.cuda.Event | None,
        cuda_sources: list[torch.Tensor],
    ) -> None:
        self.payload = payload
        self._ready_event = ready_event
        self._cuda_sources = cuda_sources
        self._waited = False

    def wait(self) -> None:
        if self._waited:
            return
        if self._ready_event is not None:
            self._ready_event.synchronize()
        self._cuda_sources.clear()
        self._waited = True


def _snapshot_tensor_payload_to_cpu_async(
    value: Any,
    *,
    copy_stream: torch.cuda.Stream,
    pin_memory: bool,
) -> _AsyncCPUPayloadSnapshot:
    cuda_sources: list[torch.Tensor] = []
    cloned = _clone_cuda_tensor_payload(value, cuda_sources)
    if not cuda_sources:
        return _AsyncCPUPayloadSnapshot(cloned, None, cuda_sources)

    source_stream = torch.cuda.current_stream()
    ready_event = torch.cuda.Event()
    with torch.cuda.stream(copy_stream):
        copy_stream.wait_stream(source_stream)
        cpu_payload = _copy_tensor_payload_to_cpu(cloned, pin_memory)
        ready_event.record(copy_stream)
    return _AsyncCPUPayloadSnapshot(cpu_payload, ready_event, cuda_sources)


class _OmniOutputTensorSnapshot(NamedTuple):
    hidden_states: torch.Tensor
    staged_hidden_states_cpu: torch.Tensor | None
    multimodal_outputs: Any
    async_payload: _AsyncCPUPayloadSnapshot | None = None


class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
    def __init__(
        self,
        *,
        model_runner_output_builder: Callable[[], OmniModelRunnerOutput],
        cuda_device: torch.device | int | str | None = None,
        **kwargs: Any,
    ) -> None:
        sampled_token_ids = kwargs.pop("sampled_token_ids")
        logprobs_tensors = kwargs.pop("logprobs_tensors")
        invalid_req_indices = kwargs.pop("invalid_req_indices")
        async_output_copy_stream = kwargs.pop("async_output_copy_stream")
        vocab_size = kwargs.pop("vocab_size")
        routed_experts = kwargs.pop("routed_experts", None)
        num_nans = kwargs.pop("num_nans", None)
        # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault
        # for EP all2all fault tolerance (PR #43637). Omni doesn't use this
        # feature but must consume the kwarg to prevent TypeError from stray
        # kwargs and initialize the attribute so super().get_output() works.
        kwargs.pop("check_ep_fault", False)
        if kwargs:
            raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}")

        self._model_runner_output: OmniModelRunnerOutput | None = None
        self._invalid_req_indices = invalid_req_indices

        self.async_copy_ready_event = torch.Event()
        self._sampled_token_ids = sampled_token_ids
        self.vocab_size = vocab_size
        self._logprobs_tensors = logprobs_tensors
        self._routed_experts = routed_experts
        self._has_fault: torch.Tensor | None = None
        # Upstream b1e12d142d (PR #51304) added device-side NaN-in-logits
        # counts (num_nans) to AsyncGPUModelRunnerOutput. Omni keeps the
        # counts on the async copy stream and lets super().get_output()
        # populate num_nans_in_logits from the CPU copy.
        self._num_nans = num_nans

        default_stream = torch.cuda.current_stream()
        with torch.cuda.stream(async_output_copy_stream):
            async_output_copy_stream.wait_stream(default_stream)
            # Keep sampled-token feedback identical to upstream async
            # scheduling. This tensor drives the next decode step, so avoid
            # changing its host-copy allocation semantics while building Omni
            # output asynchronously.
            self.sampled_token_ids_cpu = self._sampled_token_ids.to("cpu", non_blocking=True)
            self._logprobs_tensors_cpu = self._logprobs_tensors.to_cpu_nonblocking() if self._logprobs_tensors else None
            self._routed_experts_cpu = (
                self._routed_experts.to_cpu_nonblocking() if self._routed_experts is not None else None
            )
            self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None
            self.async_copy_ready_event.record()

        self._model_runner_output_builder: Callable[[], OmniModelRunnerOutput] | None = model_runner_output_builder
        self._background_exception: BaseException | None = None
        self._background_thread: threading.Thread | None = None
        self._cuda_device = cuda_device
        self._background_thread = threading.Thread(
            target=self._build_output_in_background,
            daemon=True,
            name="omni-async-output-builder",
        )
        self._background_thread.start()

    def _build_model_runner_output_once(self) -> None:
        if self._model_runner_output is not None:
            return
        with record_function_or_nullcontext("omni_async_output:get_output/build_model_runner_output"):
            assert self._model_runner_output_builder is not None
            self._model_runner_output = self._model_runner_output_builder()
        self._model_runner_output_builder = None

    def _build_output_in_background(self) -> None:
        try:
            if self._cuda_device is not None:
                torch.cuda.set_device(self._cuda_device)
            self._build_model_runner_output_once()
        except BaseException as exc:  # noqa: BLE001 - re-raised by get_output().
            self._background_exception = exc

    def get_output(self) -> OmniModelRunnerOutput:
        background_thread = getattr(self, "_background_thread", None)
        if background_thread is not None:
            background_thread.join()
            self._background_thread = None
            background_exception = getattr(self, "_background_exception", None)
            if background_exception is not None:
                raise background_exception
        self._build_model_runner_output_once()
        # Upstream AsyncGPUModelRunnerOutput.get_output() accesses
        # self._has_fault for EP all2all fault tolerance (PR #43637).
        # Ensure the attribute exists even when __init__ was bypassed
        # (e.g. unit tests using object.__new__).
        if not hasattr(self, "_has_fault"):
            self._has_fault = None
        # Upstream b1e12d142d (PR #51304) also touches _num_nans/_num_nans_cpu
        # in get_output(). Guard them the same way for object.__new__ tests.
        if not hasattr(self, "_num_nans"):
            self._num_nans = None
        if not hasattr(self, "_num_nans_cpu"):
            self._num_nans_cpu = None
        with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"):
            return super().get_output()


def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]:
    """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``.

    Non-tensor scalars (int, float) are wrapped with ``torch.tensor()``.
    Values that cannot be safely converted are dropped with a warning.
    This enforces the tensor-only invariant required by the
    ``OmniEngineCoreOutput.multimodal_output`` wire field and msgspec
    serialization.
    """
    result: dict[str, torch.Tensor] = {}
    for key, val in payload.items():
        if isinstance(val, torch.Tensor):
            result[key] = val
        elif isinstance(val, (int, float, bool)):
            result[key] = torch.tensor(val)
        elif isinstance(val, (list, tuple)):
            try:
                result[key] = torch.tensor(val)
            except (ValueError, TypeError, RuntimeError):
                logger.warning(
                    "Dropping non-tensorizable multimodal output key '%s' (type=%s) from wire payload.",
                    key,
                    type(val).__name__,
                )
        else:
            logger.warning(
                "Dropping non-tensor multimodal output key '%s' (type=%s) from wire payload.",
                key,
                type(val).__name__,
            )
    return result


class ExecuteModelState(NamedTuple):
    scheduler_output: SchedulerOutput
    logits: torch.Tensor | None
    spec_decode_metadata: Any
    spec_decode_common_attn_metadata: Any
    hidden_states: torch.Tensor
    hidden_states_cpu: torch.Tensor | None
    sample_hidden_states: torch.Tensor
    aux_hidden_states: list[torch.Tensor] | None
    ec_connector_output: Any
    cudagraph_stats: Any
    # OMNI: multimodal_outputs field for omni-specific multimodal handling
    multimodal_outputs: Any
    # slot_mappings for attention/drafter (aligned with upstream v1 API)
    slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None


class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin):
    """Autoregressive GPU model runner that returns hidden states per request.

    Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and
    reuses Omni hooks for additional_information / multimodal outputs. This
    class only overrides sample_tokens to expose hidden states + multimodal
    outputs per request while keeping Async output semantics.
    """

    execute_model_state: ExecuteModelState | None
    kv_extracted_req_ids: list[str] | None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
        # each model stage has their own hidden size
        self.hidden_size = self.model_config.hf_text_config.hidden_size
        # The width of embeddings received by this stage. It can differ from
        # the stage's internal hidden size (e.g., Qwen2.5-Omni Talker receives
        # 3584-wide Thinker states and projects them to a 896-wide hidden state).
        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
        self.inputs_embeds = self._make_buffer(
            self.max_num_tokens,
            self.inputs_embeds_size,
            dtype=self.dtype,
            numpy=False,
        )
        # Initialize KV cache manager (preserve vllm_config fallback behavior)
        self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
        self._async_chunk = getattr(self.model_config, "async_chunk", False)
        # Connector init is derived from pipeline capabilities and active
        # producer/KV/connector configuration. A stage needs it when it is:
        # - a full-payload consumer, or
        # - a producer (custom_process_next_stage_input_func is set, meaning
        #   it sends inter-stage data to a downstream consumer via connector), or
        # - explicitly configured with a connector role via stage_connector_config.
        if needs_omni_connector(self.model_config):
            self.init_omni_connectors(
                model_config=self.model_config,
                kv_transfer_manager=self.kv_transfer_manager,
            )
        self._downstream_payload_cache: dict[str, bool] = {}
        self._init_duplex_sampling_state()

    def load_model(self, *args, **kwargs) -> None:
        super().load_model(*args, **kwargs)
        self._resolve_duplex_sampling_hook(force=True)

    def _make_buffer(self, *size, dtype, numpy=True):
        # Prevent ray from pinning the buffer due to large size
        from vllm_omni.distributed.ray_utils.utils import (
            calculate_total_bytes,
            maybe_disable_pin_memory_for_ray,
        )

        total_bytes = calculate_total_bytes(size, dtype)

        # Use the context manager to temporarily disable pinning if needed
        with maybe_disable_pin_memory_for_ray(self, total_bytes):
            return super()._make_buffer(*size, dtype=dtype, numpy=numpy)

    def _build_model_sampler_output_token_ids(self) -> list[list[int]]:
        """Build decoded-token history for custom model samplers.

        vLLM only populates sampling_metadata.output_token_ids when penalties or
        logits processors require it. CosyVoice3's custom RAS sampler also
        depends on this history, so we reconstruct it directly from the input
        batch for prefer_model_sampler models.
        """
        req_output_token_ids = getattr(self.input_batch, "req_output_token_ids", [])
        req_ids = list(getattr(self.input_batch, "req_ids", []))
        output_token_ids = [list(req_output_token_ids[idx] or []) for idx in range(len(req_ids))]

        sampled_token_ids_cpu = getattr(self.input_batch, "sampled_token_ids_cpu", None)
        async_copy_ready_event = getattr(self.input_batch, "async_copy_ready_event", None)
        prev_req_id_to_index = getattr(self.input_batch, "prev_req_id_to_index", None)
        if sampled_token_ids_cpu is None or not output_token_ids or prev_req_id_to_index is None:
            return output_token_ids

        sampled_token_ids: list[list[int]] | None = None
        for index, req_id in enumerate(req_ids):
            prev_index = prev_req_id_to_index.get(req_id)
            if prev_index is None:
                continue
            req_history = output_token_ids[index]
            if not req_history or req_history[-1] != -1:
                continue
            if sampled_token_ids is None:
                assert async_copy_ready_event is not None
                async_copy_ready_event.synchronize()
                sampled_token_ids = sampled_token_ids_cpu.tolist()
            new_ids = list(sampled_token_ids[prev_index])
            if not new_ids:
                continue
            num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)
            first_placeholder = req_history.index(-1)
            num_placeholders = len(req_history) - first_placeholder
            num_to_replace = min(num_sampled_ids, num_placeholders)
            req_history[first_placeholder : first_placeholder + num_to_replace] = new_ids[:num_to_replace]

        for index, req_history in enumerate(output_token_ids):
            if -1 in req_history:
                output_token_ids[index] = req_history[: req_history.index(-1)]

        return output_token_ids

    def _sampling_metadata_for_model_sampler(self, sampling_metadata):
        if getattr(self.model, "skips_model_sampler_output_token_history", False):
            return sampling_metadata
        output_token_ids = self._build_model_sampler_output_token_ids()
        if output_token_ids != sampling_metadata.output_token_ids:
            return replace(sampling_metadata, output_token_ids=output_token_ids)
        return sampling_metadata

    def _update_states(self, scheduler_output: SchedulerOutput) -> Callable | None:
        deferred_state_corrections_fn = super()._update_states(scheduler_output)
        self._update_duplex_sampling_states(scheduler_output)
        return deferred_state_corrections_fn

    def _request_final_stage_id(self, req_id: str) -> int | None:
        info = self.model_intermediate_buffer.get(req_id)
        if not isinstance(info, dict):
            req_state = self.requests.get(req_id)
            info = getattr(req_state, "additional_information_cpu", None)
        if not isinstance(info, dict):
            return None
        val = info.get("omni_final_stage_id")
        if val is None:
            return None
        try:
            return int(val)
        except (TypeError, ValueError):
            return None

    def _request_needs_downstream_stage_payload(self, req_id: str) -> bool:
        cached = self._downstream_payload_cache.get(req_id)
        if cached is not None:
            return cached
        final_stage_id = self._request_final_stage_id(req_id)
        if final_stage_id is None:
            # Conservative default while the marker is missing: keep the
            # payload, but do NOT memoize - the marker arrives via
            # `model_intermediate_buffer`, which may be unpopulated on the
            # first call (memoizing here pinned the request to True forever,
            # never refreshing once the marker landed).
            return True
        needs_payload = final_stage_id > 0
        self._downstream_payload_cache[req_id] = needs_payload
        return needs_payload

    def _resolve_pooler_payload_req_ids(self, req_ids_output_copy: list[str]) -> tuple[str, list[str]]:
        downstream_req_ids = [rid for rid in req_ids_output_copy if self._request_needs_downstream_stage_payload(rid)]
        engine_output_type = (self.vllm_config.model_config.engine_output_type or "").lower()
        if self._client_multimodal_output_keys():
            downstream_req_ids = req_ids_output_copy
        # Single-stage AR TTS models (e.g. VoxCPM2) finish on this stage but still
        # need multimodal payloads for final audio postprocess/output.
        elif engine_output_type == "audio" and not downstream_req_ids:
            downstream_req_ids = req_ids_output_copy
        return engine_output_type, downstream_req_ids

    def capture_model(self) -> int:
        result = super().capture_model()
        self._capture_talker_mtp_graphs()
        return result

    def shutdown(self) -> None:
        """Release omni-specific GPU resources before upstream shutdown.

        Order of operations (must match upstream's expectation):
          1. Unfreeze Python GC so model weights are collected immediately
             when self.model is set to None (upstream Worker.init_device
             calls gc.freeze() / freeze_gc_heap()).
          2. Destroy omni-specific CUDA graphs (talker MTP) so references to
             model parameters are released before self.model = None.
          3. Clear GPU-side buffers (input_ids, inputs_embeds) and per-request
             caches that may hold GPU tensor references.
          4. Call CUDAGraphWrapper.clear_all_graphs() unconditionally (not just
             on ROCm) to ensure all CUDA graphs including talker MTP are
             released before model weight teardown.
          5. Call BreakableCUDAGraphWrapper.clear_all_graphs() as well, to
             match the upstream ROCm-only pattern but also protect CUDA.
          6. Delegate to upstream GPUModelRunner.shutdown() which sets
             self.model = None, clears KV caches, resets workspace, etc.

        This prevents abrupt GPU memory release during EngineCore subprocess
        exit that can trigger GPU OOM signals when the parent process
        concurrently cleans up its own GPU state.
        """
        # 1. Unfreeze GC so model weights and GPU tensors are collected
        #    immediately when references are dropped (upstream Worker.shutdown
        #    also does this before any teardown).
        gc.unfreeze()

        # 2. Destroy talker MTP CUDA graph wrapper to release captured graphs.
        if hasattr(self, "talker_mtp") and self.talker_mtp is not None:
            self.talker_mtp = None
        self.has_talker_mtp = False

        # 3. Clear GPU-side buffers (small tensors, but every MiB helps).
        if hasattr(self, "input_ids") and self.input_ids is not None:
            self.input_ids = None
        if hasattr(self, "inputs_embeds") and self.inputs_embeds is not None:
            self.inputs_embeds = None

        # 4. Clear per-request caches that may hold GPU tensor references.
        if hasattr(self, "_downstream_payload_cache"):
            self._downstream_payload_cache.clear()
        if hasattr(self, "model_intermediate_buffer"):
            self.model_intermediate_buffer.clear()
        self._clear_duplex_sampling()

        # 5. Release all CUDA graphs unconditionally (upstream only does this
        #    on ROCm; on CUDA the graphs are only freed by Python GC during
        #    interpreter shutdown, which is too late to prevent memory spikes).
        from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
        from vllm.compilation.cuda_graph import CUDAGraphWrapper

        CUDAGraphWrapper.clear_all_graphs()
        BreakableCUDAGraphWrapper.clear_all_graphs()

        # 6. Delegate to upstream shutdown (model = None, KV caches, workspace).
        super().shutdown()

    def _capture_talker_mtp_graphs(self) -> None:
        from vllm.compilation.cuda_graph import CUDAGraphWrapper

        if not self.has_talker_mtp or not isinstance(self.talker_mtp, CUDAGraphWrapper):
            return

        from vllm.compilation.monitor import set_cudagraph_capturing_enabled
        from vllm.distributed.parallel_state import graph_capture

        capture_sizes = self.compilation_config.cudagraph_capture_sizes
        num_warmups = self.compilation_config.cudagraph_num_of_warmups
        capture_sizes = sorted(capture_sizes, reverse=True)
        logger.info("Capturing talker_mtp graphs for sizes %s", capture_sizes)

        set_cudagraph_capturing_enabled(True)
        try:
            with torch.inference_mode(), graph_capture(device=self.device):
                for bsz in capture_sizes:
                    _, batch_desc, _, _, _ = self._determine_batch_execution_and_padding(
                        num_tokens=bsz,
                        num_reqs=bsz,
                        num_scheduled_tokens_np=np.ones(bsz, dtype=np.int32),
                        max_num_scheduled_tokens=1,
                        use_cascade_attn=False,
                    )
                    n = batch_desc.num_tokens
                    ids = self.talker_mtp_input_ids.gpu[:n]
                    emb = self.talker_mtp_inputs_embeds.gpu[:n]
                    hid = self.last_talker_hidden.gpu[:n]
                    ts = self.text_step.gpu[:n]

                    for _ in range(num_warmups):
                        with set_forward_context(
                            None,
                            self.vllm_config,
                            cudagraph_runtime_mode=CUDAGraphMode.NONE,
                            batch_descriptor=batch_desc,
                        ):
                            self.talker_mtp(ids, emb, hid, ts)

                    with set_forward_context(
                        None,
                        self.vllm_config,
                        cudagraph_runtime_mode=CUDAGraphMode.FULL,
                        batch_descriptor=batch_desc,
                    ):
                        self.talker_mtp(ids, emb, hid, ts)
                    torch.accelerator.synchronize()

            logger.info("Captured talker_mtp graphs for %d sizes", len(capture_sizes))
        except RuntimeError as e:
            raise RuntimeError(
                f"talker_mtp graph capture failed for a model that declared talker_mtp_graph_safe=True: {e}"
            ) from e
        finally:
            set_cudagraph_capturing_enabled(False)

    def _model_needs_full_prefix_hidden_states(self) -> bool:
        """Opt-out hook for models whose postprocess only consumes the tail.

        When False, we skip both the per-step GPU->CPU hidden-state write into
        OmniTensorPrefixCache and the merged-tensor reconstruction on hits;
        postprocess receives the normal scheduled-token slice instead. Models
        that need the full cached_prefix + new_tail span (default) are not
        affected.
        """
        model = getattr(self, "model", None)
        return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))

    def _get_runner_assisted_full_attention_metadata_request(
        self,
        *,
        req_ids: Sequence[str],
        num_reqs: int,
        num_reqs_padded: int,
        num_scheduled_tokens_np: np.ndarray,
        num_computed_tokens_cpu: torch.Tensor | np.ndarray,
        max_num_scheduled_tokens: int,
    ) -> RunnerAssistedFullAttentionMetadataRequest | None:
        # Models without this hook keep the normal runner path. VoxCPM2 uses
        # it to ask the runner for padded FULL attention metadata while keeping
        # graph policy and lifecycle in the model layer.
        hook = getattr(self.model, "get_runner_assisted_full_attention_metadata_request", None)
        if not callable(hook):
            return None
        num_computed_tokens = num_computed_tokens_cpu
        if hasattr(num_computed_tokens, "numpy"):
            num_computed_tokens = num_computed_tokens.numpy()
        request = hook(
            req_ids=req_ids,
            num_reqs=num_reqs,
            num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
            num_computed_tokens=num_computed_tokens,
            max_num_scheduled_tokens=max_num_scheduled_tokens,
        )
        if request is None:
            return None
        if not isinstance(request, RunnerAssistedFullAttentionMetadataRequest):
            raise TypeError(
                "runner-assisted full attention metadata hook must return "
                "RunnerAssistedFullAttentionMetadataRequest or None, got "
                f"{type(request).__name__}"
            )
        padded_num_reqs = max(
            num_reqs_padded,
            min(int(request.num_reqs_padded), self.scheduler_config.max_num_seqs),
        )
        return RunnerAssistedFullAttentionMetadataRequest(
            num_reqs_padded=padded_num_reqs,
            for_cudagraph_capture=bool(request.for_cudagraph_capture),
        )

    def _refresh_runner_assisted_full_attention_metadata_buffers(
        self,
        *,
        num_reqs: int,
        num_reqs_padded: int,
        num_scheduled_tokens_np: np.ndarray,
    ) -> None:
        num_computed_tokens = self.input_batch.num_computed_tokens_cpu[:num_reqs]
        if hasattr(num_computed_tokens, "numpy"):
            num_computed_tokens = num_computed_tokens.numpy()
        np.add(
            num_computed_tokens,
            num_scheduled_tokens_np,
            out=self.optimistic_seq_lens_cpu[:num_reqs].numpy(),
        )
        self.optimistic_seq_lens_cpu[num_reqs:num_reqs_padded].fill_(0)
        self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True)

        cum_num_tokens = self._get_cumsum_and_arange(
            num_scheduled_tokens_np,
            self.query_pos.np,
        )
        self.query_start_loc.np[0] = 0
        self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
        self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1])
        self.query_start_loc.copy_to_gpu()

        self.input_batch.block_table.commit_block_table(num_reqs_padded)

    def _set_runner_assisted_full_attention_metadata_context(
        self,
        *,
        enabled: bool,
        num_reqs: int = 0,
    ) -> bool:
        hook = getattr(self.model, "set_runner_assisted_full_attention_metadata_context", None)
        if not callable(hook):
            return False
        hook(
            enabled=enabled,
            num_reqs=num_reqs,
        )
        return True

    def _deferred_prefix_cache_mm_keys(self) -> set[str]:
        """Model-declared multimodal keys whose prefix-cache writes are deferred."""
        model = getattr(self, "model", None)
        keys = getattr(model, "deferred_prefix_cache_mm_keys", ())
        return set(keys or ())

    def _maybe_update_prefix_cache(
        self,
        hidden_states: torch.Tensor,
        hidden_states_cpu: torch.Tensor | None,
        multimodal_outputs: dict,
        num_tokens_unpadded: int,
        num_tokens_padded: int,
    ):
        """If prefix caching is enabled and it's the last pipeline parallelism rank,
        retrieve the hidden states & multimodal outputs from the prefix cache based
        on our batch slot mappings.
        """
        # Cache hidden states if we've enabled hidden state prefix caching
        # unless this isn't the last pipeline parallelism rank.
        is_last_pp_rank = get_pp_group().is_last_rank
        if hidden_states_cpu is not None and not is_last_pp_rank:
            raise RuntimeError("hidden_states_cpu staging is only valid on the last pipeline parallel rank.")
        if self.omni_prefix_cache is not None and is_last_pp_rank:
            # If this happens, it generally means the model is not following the correct
            # interface yet and is therefore currently not compatible with prefix cache.
            hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
            # FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton
            # _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and
            # sync back to CPU to get the correctly computed slot mapping.
            slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu
            slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu()
            self.omni_prefix_cache.update_omni_tensor_prefix_cache(
                hidden_states=hs_for_cache,
                multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
                num_tokens_unpadded=num_tokens_unpadded,
                slot_mapping=slot_mapping_cpu,
                num_tokens_padded=num_tokens_padded,
                skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(),
                hidden_states_cpu=hidden_states_cpu,
            )

    def _maybe_get_combined_prefix_cache_tensors(
        self,
        hidden_states: torch.Tensor,
        hidden_states_cpu: torch.Tensor | None,
        multimodal_outputs: dict,
        num_scheduled_tokens: dict[str, int],
    ) -> tuple[dict[str, torch.Tensor] | None, dict | None]:
        """If prefix caching is enabled, extract the merged hidden states and multimodal outputs for
        all requests in the batch (including those that aren't a hit on Prefix cache).
        """
        # Prior to applying the post-processing func, extract
        # the prefix cached hidden states and multimodal states.
        combined_hidden_states, combined_multimodal_outputs = None, None
        is_last_pp_rank = get_pp_group().is_last_rank
        if hidden_states_cpu is not None and not is_last_pp_rank:
            raise RuntimeError("hidden_states_cpu staging is only valid on the last pipeline parallel rank.")
        if self.omni_prefix_cache is not None:
            if not is_last_pp_rank:
                raise RuntimeError("Omni prefix-cache tensor merge is only valid on the last pipeline parallel rank.")
            if (
                not self._model_needs_full_prefix_hidden_states()
                and not self.omni_prefix_cache.has_prefix_cached_new_req_ids()
            ):
                return None, None
            if self._model_needs_full_prefix_hidden_states():
                combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
                    query_start_loc=self.query_start_loc.cpu,
                    input_batch=self.input_batch,
                    hidden_states=hidden_states,
                    hidden_states_cpu=hidden_states_cpu,
                    num_scheduled_tokens=num_scheduled_tokens,
                )
            combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states(
                query_start_loc=self.query_start_loc.cpu,
                input_batch=self.input_batch,
                multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
                num_scheduled_tokens=num_scheduled_tokens,
            )
        return combined_hidden_states, combined_multimodal_outputs

    def _stage_deferred_prefix_cache_mm_outputs(
        self,
        *,
        scheduler_output: SchedulerOutput,
        multimodal_outputs: Any,
        query_start_loc_cpu: Any,
    ) -> None:
        if self.omni_prefix_cache is None:
            return

        deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys()
        if not deferred_mm_cache_keys:
            return

        self.omni_prefix_cache.stage_deferred_mm_outputs(
            query_start_loc=query_start_loc_cpu,
            input_batch=self.input_batch,
            multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
            num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
            deferred_mm_cache_keys=deferred_mm_cache_keys,
        )

    def _prepare_prefix_cache_pooler_payload_sources(
        self,
        *,
        hidden_states: torch.Tensor,
        staged_hidden_states_cpu: torch.Tensor | None,
        multimodal_outputs: Any,
        scheduler_output: SchedulerOutput,
        needs_scheduled_hidden_payload: bool,
    ) -> tuple[torch.Tensor | None, dict[str, torch.Tensor] | None, dict | None]:
        hidden_states_cpu = None
        if needs_scheduled_hidden_payload:
            if staged_hidden_states_cpu is None:
                raise RuntimeError("Prefix-cache hidden-state payload requires staged CPU hidden states.")
            hidden_states_cpu = staged_hidden_states_cpu

        combined_hidden_states, combined_multimodal_outputs = self._maybe_get_combined_prefix_cache_tensors(
            hidden_states,
            staged_hidden_states_cpu,
            multimodal_outputs,
            scheduler_output.num_scheduled_tokens,
        )
        return hidden_states_cpu, combined_hidden_states, combined_multimodal_outputs

    def _build_omni_pooler_payload(
        self,
        *,
        rid: str,
        idx: int,
        start: int,
        end: int,
        hidden_states_cpu: torch.Tensor | None,
        req_hidden_states_cpu: dict[str, torch.Tensor] | None,
        combined_hidden_states: dict[str, torch.Tensor] | None,
        combined_multimodal_outputs: dict | None,
        mm_cpu: dict[str, object] | None,
        audio_sparse_output: bool,
        sparse_mm_index: dict[str, int],
        hidden_seq_len: int,
        scheduled_seq_len: int,
    ) -> dict[str, object]:
        payload: dict[str, object] = {}
        if not audio_sparse_output:
            if req_hidden_states_cpu is not None and combined_hidden_states is None:
                req_hidden_states = req_hidden_states_cpu[rid]
            else:
                req_hidden_states = self._resolve_req_hidden_states(
                    hidden_states_cpu,
                    combined_hidden_states,
                    rid,
                    start,
                    end,
                )
            if req_hidden_states is not None:
                payload["hidden"] = req_hidden_states

        mm_payload = build_omni_mm_payload(
            combined_multimodal_outputs=combined_multimodal_outputs,
            mm_cpu=mm_cpu,
            rid=rid,
            idx=idx,
            start=start,
            end=end,
            audio_sparse_output=audio_sparse_output,
            sparse_mm_index=sparse_mm_index,
            hidden_seq_len=hidden_seq_len,
            scheduled_seq_len=scheduled_seq_len,
        )
        payload.update(mm_payload)
        return payload

    def _merge_model_kv_transfer_metadata(
        self,
        finished_reqs: dict[str, dict[str, Any]],
    ) -> dict[str, dict[str, Any]]:
        """Merge model-provided KV-transfer metadata copy-on-write.

        The outer mapping is always new; a per-request dict is shallow-copied
        (with a fresh merged ``custom_metadata``) only when the model supplies
        metadata for it, and passes through unchanged otherwise.
        `scheduler_output` (and its per-request dicts) can be shared with the
        engine-core process, so nothing here mutates them — the ngram block
        in `execute_model` `replace()`-copies for exactly the same reason. The
        merged mapping is only consumed by
        `kv_transfer_manager.handle_finished_requests_kv_transfer`.
        """
        merged: dict[str, dict[str, Any]] = {}
        for req_id, data in finished_reqs.items():
            try:
                # NOTE: seq_len is the same as num_computed_tokens_cpu in current
                # async scheduling, since both exclude async placeholders. We use
                # seq_len since we control it, just in case upstream async scheduler
                # semantics change in the future.
                model_meta = self.model.get_kv_transfer_metadata(
                    req_id,
                    num_computed_tokens=data.get("seq_len"),
                )
            except Exception:
                # A swallowed failure here used to ship the KV transfer WITHOUT
                # the model's custom metadata and surface much later as a
                # corrupt/incomplete transfer on the consumer stage. Fail at
                # the point of the defect instead.
                logger.exception("Failed to get custom KV-transfer metadata from model for %s", req_id)
                raise
            if model_meta:
                data = {
                    **data,
                    "custom_metadata": {**(data.get("custom_metadata") or {}), **model_meta},
                }
            merged[req_id] = data
        return merged

    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: SchedulerOutput,
        intermediate_tensors: IntermediateTensors | None = None,
    ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:
        if self.execute_model_state is not None:
            raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")

        if self.routed_experts_initialized:
            self.routed_experts_capturer.clear_buffer()

        if not getattr(self, "_warmup_state_cleared", False):
            self._warmup_state_cleared = True
            if hasattr(self.model, "_clear_warmup_state"):
                self.model._clear_warmup_state()

        # Async-write pipeline: apply any pending GPU->CPU prefix-cache writes
        # whose copy event has already fired. Non-blocking — entries whose D2H
        # is still in flight stay queued and will be picked up on the next
        # step's drain. This guarantees any downstream read of
        # ``omni_prefix_cache.hidden_states_cache`` /
        # ``omni_prefix_cache.mm_outputs_cache`` in this step sees the
        # state produced no later than the previous forward step.
        if self.omni_prefix_cache is not None:
            self.omni_prefix_cache.drain_ready_async_writes()

        # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
        finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
        if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
            finished_reqs = self._merge_model_kv_transfer_metadata(finished_reqs)
        self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
            finished_reqs=finished_reqs,
            kv_caches=self.kv_caches,
            block_size=self.cache_config.block_size,
            cache_dtype=str(self.cache_config.cache_dtype),
            request_id_resolver=self._resolve_global_request_id,
        )

        if hasattr(self, "_omni_connector"):
            for request in getattr(scheduler_output, "pending_input_registrations", []):
                self.register_chunk_recv(request)
            self.recv_full_payload_inputs(scheduler_output)
            if self._pending_full_payload_send:
                flush_ids = set(getattr(scheduler_output, "finished_req_ids", set()))
                flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests})
                if flush_ids:
                    self.flush_full_payload_outputs(flush_ids)

        if self.omni_prefix_cache is not None and scheduler_output.finished_req_ids:
            self.omni_prefix_cache.commit_deferred_mm_outputs(
                set(scheduler_output.finished_req_ids),
                self.input_batch,
            )

        if self.routed_experts_initialized:
            capturer = self.routed_experts_capturer
            if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
                capturer.finalize_pending_copy()

        # If ngram_gpu is used, we need to copy the scheduler_output to avoid
        # the modification has influence on the scheduler_output in engine core process.
        # The replace is much faster than deepcopy.
        if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
            num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy()
            spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy()
            scheduler_output = replace(
                scheduler_output,
                num_scheduled_tokens=num_scheduled_tokens_copy,
                scheduled_spec_decode_tokens=spec_decode_tokens_copy,
            )

        if has_kv_transfer_group():
            kv_connector_metadata = scheduler_output.kv_connector_metadata
            if kv_connector_metadata is not None:
                get_kv_transfer_group().handle_preemptions(kv_connector_metadata)

        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        with (
            record_function_or_nullcontext("gpu_model_runner: preprocess"),
            self.synchronize_input_prep(),
        ):
            # Update persistent batch states.
            deferred_state_corrections_fn = self._update_states(scheduler_output)

            # Notify model of finished requests for state cleanup
            if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
                self.model.on_requests_finished(scheduler_output.finished_req_ids)

            if has_ec_transfer() and not get_ec_transfer().is_consumer:
                with self.maybe_get_ec_connector_output(
                    scheduler_output,
                    encoder_cache=self.encoder_cache,
                ) as ec_connector_output:
                    self._execute_mm_encoder(scheduler_output)

                    kv_ids = self.kv_extracted_req_ids
                    self.kv_extracted_req_ids = None

                    output = make_empty_encoder_model_runner_output(scheduler_output)
                    if kv_ids:
                        output = copy(output)
                        output.kv_extracted_req_ids = kv_ids
                    return self.attach_omni_connector_output(output)

            # `<= 0`: upstream can schedule a negative span, which is truthy (#5196).
            if num_scheduled_tokens <= 0:
                if (
                    self.parallel_config.distributed_executor_backend == "external_launcher"
                    and self.parallel_config.data_parallel_size > 1
                ):
                    self._dummy_run(1)

                # Capture KV extraction results before early return;
                # sample_tokens() is skipped on this path so the IDs
                # would otherwise be silently overwritten next step.
                kv_ids = self.kv_extracted_req_ids
                self.kv_extracted_req_ids = None

                if not has_kv_transfer_group():
                    output = EMPTY_MODEL_RUNNER_OUTPUT
                else:
                    output = self.kv_connector_no_forward(scheduler_output, self.vllm_config)

                if kv_ids:
                    output = copy(output)
                    output.kv_extracted_req_ids = kv_ids

                return self.attach_omni_connector_output(output)

            if self.cache_config.kv_sharing_fast_prefill:
                assert not self.num_prompt_logprobs, (
                    "--kv-sharing-fast-prefill produces incorrect "
                    "logprobs for prompt tokens, tokens, please disable "
                    "it when the requests need prompt logprobs"
                )

            num_reqs = self.input_batch.num_reqs
            req_ids = self.input_batch.req_ids
            tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
            num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
            max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
            num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens

            logits_indices, spec_decode_metadata, max_num_sampled_tokens = self._prepare_inputs(
                scheduler_output, num_scheduled_tokens_np
            )

            cascade_attn_prefix_lens = None
            # Disable cascade attention when using microbatching (DBO)
            if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
                # Pre-compute cascade attention prefix lengths
                cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
                    num_scheduled_tokens_np,
                    self.input_batch.num_computed_tokens_cpu[:num_reqs],
                    scheduler_output.num_common_prefix_blocks,
                )

            (
                cudagraph_mode,
                batch_desc,
                should_ubatch,
                num_tokens_across_dp,
                cudagraph_stats,
            ) = self._determine_batch_execution_and_padding(
                num_tokens=num_tokens_unpadded,
                num_reqs=num_reqs,
                num_scheduled_tokens_np=num_scheduled_tokens_np,
                max_num_scheduled_tokens=max_num_scheduled_tokens,
                use_cascade_attn=cascade_attn_prefix_lens is not None,
                num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
            )
            num_tokens_padded = batch_desc.num_tokens
            num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
            num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu[:num_reqs]

            runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request(
                req_ids=req_ids[:num_reqs],
                num_reqs=num_reqs,
                num_reqs_padded=num_reqs_padded,
                num_scheduled_tokens_np=num_scheduled_tokens_np,
                num_computed_tokens_cpu=num_computed_tokens_cpu,
                max_num_scheduled_tokens=max_num_scheduled_tokens,
            )
            runner_assisted_full_attn_capture = False
            if runner_assisted_full_attn_request is not None:
                num_reqs_padded = runner_assisted_full_attn_request.num_reqs_padded
                runner_assisted_full_attn_capture = runner_assisted_full_attn_request.for_cudagraph_capture
                num_tokens_padded = max(
                    num_tokens_padded,
                    num_reqs_padded * max_num_scheduled_tokens,
                )
            runner_assisted_full_attn = runner_assisted_full_attn_request is not None

            ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
                should_ubatch,
                num_scheduled_tokens_np,
                num_tokens_padded,
                num_reqs_padded,
                self.parallel_config.num_ubatches,
            )
            pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn

            use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
            ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices

            # True if any attention backend handles KV cache update separately
            # from forward() (i.e., forward_includes_kv_cache_update=False). When true,
            # slot_mappings must use padded dimensions to match the key/value tensors.
            from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec

            has_separate_kv_update = not all(
                all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id])
                for id, spec in enumerate(self.kv_cache_config.kv_cache_groups)
                if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec)
            )

            slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
                num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded,
                num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs),
                num_tokens_unpadded=num_tokens_unpadded,
                ubatch_slices=ubatch_slices_padded,
            )
            if runner_assisted_full_attn:
                self._refresh_runner_assisted_full_attention_metadata_buffers(
                    num_reqs=num_reqs,
                    num_reqs_padded=num_reqs_padded,
                    num_scheduled_tokens_np=num_scheduled_tokens_np,
                )

            attn_num_reqs = num_reqs_padded if runner_assisted_full_attn else num_reqs
            attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
                num_tokens=num_tokens_unpadded,
                num_tokens_padded=num_tokens_padded if pad_attn else None,
                num_reqs=attn_num_reqs,
                num_reqs_padded=num_reqs_padded if pad_attn else None,
                max_query_len=max_num_scheduled_tokens,
                ubatch_slices=ubatch_slices_attn,
                logits_indices=logits_indices,
                max_num_sampled_tokens=max_num_sampled_tokens,
                use_spec_decode=use_spec_decode,
                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
                cascade_attn_prefix_lens=cascade_attn_prefix_lens,
                slot_mappings=slot_mappings_by_group,
                for_cudagraph_capture=runner_assisted_full_attn_capture,
            )
            self._maybe_attach_attention_metadata_extensions(
                attn_metadata=attn_metadata,
                num_reqs=num_reqs,
                num_reqs_padded=num_reqs_padded,
                max_query_len=max_num_scheduled_tokens,
                pad_attn=pad_attn,
                num_scheduled_tokens_np=num_scheduled_tokens_np,
                for_cudagraph_capture=runner_assisted_full_attn_capture,
            )

            (
                input_ids,
                inputs_embeds,
                positions,
                intermediate_tensors,
                model_kwargs,
                ec_connector_output,
            ) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)

        # Let the model adjust inputs before forward (e.g. restore input_ids
        # for multimodal position detection, fix decode position offsets).
        prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None)
        if callable(prepare_runner_inputs):
            input_ids, positions = prepare_runner_inputs(
                input_ids=input_ids,
                positions=positions,
                inputs_embeds=inputs_embeds,
                req_ids=req_ids[:num_reqs],
                num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs],
                num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
                input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
            )

        runner_assisted_context_enabled = False
        if runner_assisted_full_attn:
            runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context(
                enabled=True,
                num_reqs=num_reqs,
            )
            runner_assisted_full_attn = runner_assisted_context_enabled

        # Run the model.
        # Use persistent buffers for CUDA graphs.
        # When spec decode is enabled, defer connector finalization
        # (wait_for_save + clear metadata) until after draft model runs.
        defer_kv_connector_finalize = self.speculative_config is not None
        try:
            with (
                nullcontext(),
                set_forward_context(
                    attn_metadata,
                    self.vllm_config,
                    num_tokens=num_tokens_padded,
                    num_tokens_across_dp=num_tokens_across_dp,
                    cudagraph_runtime_mode=(CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode),
                    batch_descriptor=batch_desc,
                    ubatch_slices=ubatch_slices_padded,
                    slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
                ),
                record_function_or_nullcontext("gpu_model_runner: forward"),
                self.maybe_get_kv_connector_output(
                    scheduler_output,
                    defer_finalize=defer_kv_connector_finalize,
                ) as kv_connector_output,
            ):
                model_output = self._model_forward(
                    input_ids=input_ids,
                    positions=positions,
                    intermediate_tensors=intermediate_tensors,
                    inputs_embeds=inputs_embeds,
                    **model_kwargs,
                    sampling_metadata=self.input_batch.sampling_metadata,
                    logits_index=logits_indices,
                    sampler=self.sampler,
                )

                # [Omni] Map pending ropes metadata to req_ids.
                flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None)
                if callable(flush_pending_metadata):
                    flush_pending_metadata(req_ids[:num_reqs])

                # [Omni] Hand the model the batch's req_ids in logits order, for
                # models that gate logits per request. Only valid without spec
                # decode: there logits_indices carries several rows per request,
                # so row i no longer corresponds to req_ids[i].
                if spec_decode_metadata is None:
                    set_batch_req_ids = getattr(self.model, "set_batch_req_ids", None)
                    if callable(set_batch_req_ids):
                        set_batch_req_ids(req_ids[:num_reqs])
        finally:
            if runner_assisted_context_enabled:
                self._set_runner_assisted_full_attention_metadata_context(enabled=False)

        with record_function_or_nullcontext("gpu_model_runner: postprocess"):
            if self.use_aux_hidden_state_outputs:
                # True when EAGLE 3 is used.
                hidden_states, aux_hidden_states = model_output
            else:
                # Common case.
                hidden_states = model_output
                aux_hidden_states = None

            # A pooling stage skips mm extraction + the omni prefix cache (both assume an AR
            # [n_tok, hidden] output) and exits through the shared _pool() branch below.
            if self.is_pooling_model:
                multimodal_outputs = None
            else:
                hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
            hidden_states_cpu = None

            # Async-write pipeline (replaces the per-step blocking
            # ``.to("cpu")`` + ``aten::index_put_`` on pageable host memory).
            # Schedules non-blocking GPU->CPU copies on a dedicated stream;
            # the actual CPU scatter into ``hidden_states_cache`` /
            # ``mm_outputs_cache`` happens in ``drain_ready_async_writes``
            # at the top of subsequent execute_model() calls.
            if not self.is_pooling_model and self.omni_prefix_cache is not None and get_pp_group().is_last_rank:
                hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
                # Some models (e.g. qwen3-tts-talker) opt out of full-hidden-state
                # prefix caching but the downstream pooler payload path still
                # needs a CPU hidden-states view. Materialize it synchronously
                # in that case; the legacy behavior is preserved.
                if hs_for_cache is None and self._model_omni_pooler_payload_include_hidden():
                    hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous()
                slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu
                self.omni_prefix_cache.schedule_async_write(
                    hidden_states_gpu=hs_for_cache,
                    multimodal_outputs_gpu=(flatten_payload(multimodal_outputs) if multimodal_outputs else None),
                    slot_mapping_gpu=slot_mapping_gpu,
                    num_tokens_unpadded=num_tokens_unpadded,
                    num_tokens_padded=num_tokens_padded,
                    skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(),
                )

            if not self.broadcast_pp_output:
                # Common case.
                if not get_pp_group().is_last_rank:
                    # Return the intermediate tensors.
                    assert isinstance(hidden_states, IntermediateTensors)
                    self.kv_connector_output = kv_connector_output
                    return hidden_states

                if self.is_pooling_model:
                    # Return the pooling output.
                    return self._pool(
                        hidden_states,
                        num_scheduled_tokens,
                        num_scheduled_tokens_np,
                        kv_connector_output,
                    )

                sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)]
                # Try with sampling_metadata first; fall back to without for models that don't support it
                try:
                    logits = self.model.compute_logits(
                        sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
                    )
                except TypeError:
                    logits = self.model.compute_logits(sample_hidden_states)
            else:
                # Rare case.
                assert not self.is_pooling_model

                sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)]
                if not get_pp_group().is_last_rank:
                    all_gather_tensors = {
                        "residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded)
                    }
                    get_pp_group().send_tensor_dict(
                        hidden_states.tensors,
                        all_gather_group=get_tp_group(),
                        all_gather_tensors=all_gather_tensors,
                    )
                    logits = None
                else:
                    # Try with sampling_metadata first; fall back to without for models that don't support it
                    try:
                        logits = self.model.compute_logits(
                            sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
                        )
                    except TypeError:
                        logits = self.model.compute_logits(sample_hidden_states)

                model_output_broadcast_data: dict[str, Any] = {}
                if logits is not None:
                    model_output_broadcast_data["logits"] = logits.contiguous()

                broadcasted = get_pp_group().broadcast_tensor_dict(
                    model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
                )
                assert broadcasted is not None
                logits = broadcasted["logits"]

        self.execute_model_state = ExecuteModelState(
            scheduler_output,
            logits,
            spec_decode_metadata,
            spec_decode_common_attn_metadata,
            hidden_states,
            hidden_states_cpu,
            sample_hidden_states,
            aux_hidden_states,
            ec_connector_output,
            cudagraph_stats,
            multimodal_outputs,
            slot_mappings,  # OMNI: pass slot_mappings for drafter
        )
        self.kv_connector_output = kv_connector_output

        if deferred_state_corrections_fn:
            deferred_state_corrections_fn()

        if self._should_return_omni_routed_experts() and hasattr(self, "_positions_cpu"):
            self._omni_routed_experts_d2h(scheduler_output)

        return None

    def _sample(
        self,
        logits: torch.Tensor | None,
        spec_decode_metadata: Any,
    ):
        """Sample tokens, preferring the model's own sampler when declared.

        Model-interface contract for `prefer_model_sampler` models: a custom
        `model.sample(logits, sampling_metadata)` may return ``None`` to
        DECLINE the step, and the runner then falls back to the default
        sampler. Declarers rely on this deliberately for empty-logits or
        inapplicable-stage steps (e.g. cosyvoice3, glm_tts), so the
        fallthrough is load-bearing, not an error path; it is logged once per
        model for visibility.
        """
        sampling_metadata = self.input_batch.sampling_metadata
        if spec_decode_metadata is None:
            model_sample = getattr(self.model, "sample", None)
            self.input_batch.update_async_output_token_ids()
            if logits is not None and callable(model_sample) and getattr(self.model, "prefer_model_sampler", False):
                # Apply logit bias (min_tokens, allowed_token_ids) before
                # the custom model sampler — the standard GPU sampler does
                # this internally, but prefer_model_sampler bypasses it.
                if hasattr(self.sampler, "logit_bias_state"):
                    self.sampler.logit_bias_state.apply_logit_bias(
                        logits,
                        self.input_batch.expanded_idx_mapping,
                        self.input_batch.idx_mapping_np,
                        self.input_batch.positions[self.input_batch.logits_indices],
                    )
                prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata)
                self._apply_duplex_sampling(logits, prepared_sampling_metadata)
                sampler_output = model_sample(logits, prepared_sampling_metadata)
                if sampler_output is not None:
                    return sampler_output
                # Contract: None => fall back to the default sampler (see
                # docstring). If a custom sampler returns None by accident,
                # its tokens silently come from the default sampler - the
                # one-time log is the visibility hook for that case.
                logger.warning_once(
                    "prefer_model_sampler model %s returned None from sample(); falling back to the default sampler.",
                    type(self.model).__name__,
                )
            return self.sampler(
                logits=logits,
                sampling_metadata=sampling_metadata,
            )

        return super()._sample(logits, spec_decode_metadata)

    @staticmethod
    def _resolve_req_hidden_states(
        hidden_states_cpu: torch.Tensor | None,
        combined_hidden_states: dict[str, torch.Tensor] | None,
        rid: str,
        start: int,
        end: int,
    ) -> torch.Tensor | None:
        if combined_hidden_states is not None:
            # We always have all request IDs for prefix cache, even for
            # partial cache misses, so this should never happen.
            if rid not in combined_hidden_states:
                raise RuntimeError("Request IDs in the batch are missing from the merged states!")
            return combined_hidden_states[rid]
        # Prefix caching is disabled. hidden_states_cpu may legitimately be
        # None (e.g. sparse audio output or no scheduled hidden payload);
        # callers must omit the "hidden" key in that case.
        if hidden_states_cpu is None:
            return None
        return hidden_states_cpu[start:end]

    def _run_post_sample_talker_mtp(
        self,
        *,
        req_ids: list[str],
        valid_sampled_token_ids: list[list[int]],
        sampled_token_ids: torch.Tensor,
        invalid_req_indices: list[int],
        sample_hidden_states: torch.Tensor,
        multimodal_outputs: Any,
    ) -> Any:
        """Run an opt-in talker MTP with the token sampled in this step.

        Some frame-synchronous talkers need the current temporal hidden state
        and its newly sampled text token to produce the same frame's codec
        codes. A resumable ``max_tokens=1`` segment has no later decode
        preprocess in which to do that work, so those models expose
        ``post_sample_talker_mtp`` instead.
        """
        model = getattr(self, "model", None)
        hook = getattr(model, "post_sample_talker_mtp", None)
        if not callable(hook):
            return multimodal_outputs

        selected_indices: list[int] = []
        selected_req_ids: list[str] = []
        selected_req_infos: list[dict[str, Any]] = []
        duplex_indices: list[int] = []
        invalid_indices = set(invalid_req_indices)
        use_async_scheduling = bool(getattr(self, "use_async_scheduling", False))
        for idx, req_id in enumerate(req_ids):
            req_info = self.model_intermediate_buffer.get(req_id)
            duplex = req_info.get("duplex") if isinstance(req_info, dict) else None
            if not isinstance(duplex, dict) or duplex.get("data_plane") is not True:
                continue
            duplex_indices.append(idx)
            if use_async_scheduling:
                if idx in invalid_indices:
                    continue
            else:
                token_ids = valid_sampled_token_ids[idx] if idx < len(valid_sampled_token_ids) else []
                if not token_ids:
                    continue
                if len(token_ids) != 1:
                    raise ValueError(
                        "post_sample_talker_mtp requires exactly one sampled token "
                        f"per request, got {len(token_ids)} for {req_id}"
                    )
            selected_indices.append(idx)
            selected_req_ids.append(req_id)
            selected_req_infos.append(cast(dict[str, Any], req_info))

        if not selected_indices:
            return multimodal_outputs
        if sampled_token_ids.ndim != 2 or sampled_token_ids.shape[0] < len(req_ids):
            raise ValueError(
                "post_sample_talker_mtp sampled-token rows do not match requests: "
                f"shape={tuple(sampled_token_ids.shape)}, requests={len(req_ids)}"
            )
        if sampled_token_ids.shape[1] != 1:
            raise ValueError(
                "post_sample_talker_mtp does not support speculative token rows: "
                f"shape={tuple(sampled_token_ids.shape)}"
            )

        empty_audio = torch.empty(0, dtype=torch.long, device=sample_hidden_states.device)
        merged = dict(multimodal_outputs) if isinstance(multimodal_outputs, dict) else {}
        existing_codes = merged.get("codes")
        codes_payload = dict(existing_codes) if isinstance(existing_codes, dict) else {}
        existing_audio = codes_payload.get("audio")
        if isinstance(existing_audio, list) and len(existing_audio) == len(req_ids):
            per_request_audio = list(existing_audio)
        elif (
            isinstance(existing_audio, torch.Tensor)
            and existing_audio.ndim > 0
            and existing_audio.shape[0] == len(req_ids)
        ):
            per_request_audio = [existing_audio[idx : idx + 1] for idx in range(len(req_ids))]
        else:
            per_request_audio = [empty_audio for _ in req_ids]
        for idx in duplex_indices:
            per_request_audio[idx] = empty_audio

        indices = torch.tensor(
            selected_indices,
            dtype=torch.long,
            device=sample_hidden_states.device,
        )
        input_ids = sampled_token_ids.index_select(0, indices).reshape(-1).to(torch.long)
        codes = hook(
            input_ids=input_ids,
            hidden_states=sample_hidden_states.index_select(0, indices),
            req_ids=selected_req_ids,
            req_infos=selected_req_infos,
        )
        if not isinstance(codes, torch.Tensor) or codes.ndim != 2:
            raise TypeError(f"post_sample_talker_mtp must return a rank-2 Tensor, got {type(codes).__name__}")
        if codes.shape[0] != len(selected_req_ids):
            raise ValueError(
                f"post_sample_talker_mtp row count does not match requests: {codes.shape[0]} != {len(selected_req_ids)}"
            )
        for row, (idx, req_id) in enumerate(zip(selected_indices, selected_req_ids, strict=True)):
            request_codes = codes[row : row + 1]
            per_request_audio[idx] = request_codes
            self._update_intermediate_buffer(
                req_id,
                {"codes": {"audio": request_codes}},
            )

        # Empty rows are intentional: they suppress replay of the previous
        # frame while a chunked prefill has not sampled a token yet.
        codes_payload["audio"] = per_request_audio
        merged["codes"] = codes_payload
        return merged

    def _build_multimodal_outputs(
        self,
        per_req_payloads: Sequence[dict[str, object] | None] | None,
    ) -> list[dict[str, torch.Tensor] | None] | None:
        """Build per-request multimodal output payloads (dedicated channel).

        Reuses the per-request payloads assembled by the pooler-payload loop
        in sample_tokens() (which already handles prefix-cache merging,
        sparse audio output, and partial downstream batches) so the wire
        channel stays consistent with the full-payload accumulation path.
        Enforces the tensor-only invariant required by msgspec: scalars and
        lists are wrapped into tensors, and anything that cannot be safely
        converted is dropped.
        """
        if self.vllm_config.model_config.engine_output_type == "text" and not self._client_multimodal_output_keys():
            return None
        if per_req_payloads is None:
            return None
        wire_payloads: list[dict[str, torch.Tensor] | None] = []
        for payload in per_req_payloads:
            if not payload:
                wire_payloads.append(None)
            else:
                wire_payloads.append(_ensure_tensor_values(payload))
        if all(item is None for item in wire_payloads):
            return None
        return wire_payloads

    def _snapshot_query_start_loc_cpu(self) -> Any:
        query_start_loc_cpu = self.query_start_loc.cpu
        if callable(query_start_loc_cpu):
            query_start_loc_cpu = query_start_loc_cpu()
        if isinstance(query_start_loc_cpu, torch.Tensor):
            return query_start_loc_cpu.detach().cpu().clone()
        if isinstance(query_start_loc_cpu, np.ndarray):
            return query_start_loc_cpu.copy()
        if isinstance(query_start_loc_cpu, list):
            return list(query_start_loc_cpu)
        return query_start_loc_cpu

    @staticmethod
    def _snapshot_scheduler_output_for_async_omni_output(
        scheduler_output: SchedulerOutput,
    ) -> SchedulerOutput:
        updates: dict[str, Any] = {}
        for attr in ("num_scheduled_tokens", "scheduled_spec_decode_tokens"):
            val = getattr(scheduler_output, attr, None)
            if isinstance(val, dict):
                updates[attr] = val.copy()
            elif isinstance(val, list):
                updates[attr] = list(val)
        if not updates:
            return scheduler_output
        try:
            return replace(scheduler_output, **updates)
        except TypeError:
            return scheduler_output

    def _should_return_omni_routed_experts(self) -> bool:
        model_config = getattr(self, "model_config", None)
        if model_config is None:
            model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
        return bool(getattr(model_config, "enable_return_routed_experts", False)) and bool(
            getattr(self, "routed_experts_initialized", False)
        )

    @staticmethod
    def _model_omni_flag(model: Any, name: str, default: bool = False) -> bool:
        return bool(getattr(model, name, default)) if model is not None else default

    def _runner_model_omni_flag(self, name: str, default: bool = False) -> bool:
        return self._model_omni_flag(getattr(self, "model", None), name, default)

    def _client_multimodal_output_keys(self) -> tuple[str, ...]:
        raw = getattr(
            getattr(self, "model", None),
            "omni_client_multimodal_output_keys",
            (),
        )
        if not isinstance(raw, tuple) or any(not isinstance(key, str) or not key for key in raw):
            raise TypeError("omni_client_multimodal_output_keys must be a tuple of non-empty strings")
        if len(raw) != len(set(raw)):
            raise ValueError("omni_client_multimodal_output_keys must not contain duplicates")
        return raw

    def _model_omni_pooler_payload_include_hidden(self) -> bool:
        return self._runner_model_omni_flag("omni_pooler_payload_include_hidden", default=True)

    def _should_defer_full_payload_d2h(self) -> bool:
        """Keep opted-in full payloads on device until request completion."""
        return (
            getattr(self, "omni_prefix_cache", None) is None
            and self._runner_model_omni_flag("omni_payload_at_request_end")
            and self._should_accumulate_full_payload_output()
        )

    def _build_omni_step_outputs(
        self,
        pooler_inter: Sequence[dict[str, object] | None] | None,
        pooler_client: Sequence[dict[str, object] | None] | None,
        *,
        defer_full_payload_d2h: bool,
    ) -> tuple[
        list[dict[str, torch.Tensor] | None] | None,
        list[dict[str, torch.Tensor] | None] | None,
    ]:
        if defer_full_payload_d2h:
            # The connector-side full-payload accumulator owns the snapshots.
            # Nothing from an intermediate stage needs to cross the worker/core
            # boundary on each decode step.
            return None, None
        inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
        multimodal_outputs = (
            inter_stage_outputs if pooler_client is pooler_inter else self._build_multimodal_outputs(pooler_client)
        )
        return inter_stage_outputs, multimodal_outputs

    def _should_use_async_omni_output(self) -> bool:
        if not self.use_async_scheduling:
            return False
        if self.omni_prefix_cache is not None:
            return False
        if self.speculative_config is not None:
            return False

        model_config = getattr(self, "model_config", None)
        if model_config is None:
            model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
        if not bool(getattr(model_config, "async_chunk", False)):
            return False
        if bool(getattr(model_config, "enable_return_routed_experts", False)):
            return False

        model = getattr(self, "model", None)
        if not self._model_omni_flag(model, "use_async_omni_output"):
            return False
        if self._model_omni_flag(model, "has_postprocess") and not self._model_omni_flag(
            model, "eager_omni_postprocess_before_async_output"
        ):
            return False

        return True

    def _build_omni_async_snapshot_payload(
        self,
        *,
        hidden_states: torch.Tensor,
        staged_hidden_states_cpu: torch.Tensor | None,
        multimodal_outputs: Any,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {"multimodal_outputs": multimodal_outputs}
        if self._model_omni_pooler_payload_include_hidden():
            payload["hidden_states"] = hidden_states
            payload["staged_hidden_states_cpu"] = staged_hidden_states_cpu
        return payload

    def _snapshot_omni_output_tensors_for_async_output(
        self,
        *,
        use_async_omni_output: bool,
        hidden_states: torch.Tensor,
        staged_hidden_states_cpu: torch.Tensor | None,
        multimodal_outputs: Any,
    ) -> _OmniOutputTensorSnapshot:
        if not use_async_omni_output:
            return _OmniOutputTensorSnapshot(
                hidden_states=hidden_states,
                staged_hidden_states_cpu=staged_hidden_states_cpu,
                multimodal_outputs=multimodal_outputs,
            )

        with record_function_or_nullcontext("omni_async_output:snapshot_cpu_payload"):
            async_payload_snapshot = _snapshot_tensor_payload_to_cpu_async(
                self._build_omni_async_snapshot_payload(
                    hidden_states=hidden_states,
                    staged_hidden_states_cpu=staged_hidden_states_cpu,
                    multimodal_outputs=multimodal_outputs,
                ),
                copy_stream=self._get_or_create_omni_payload_copy_stream(),
                # NOTE: vLLM v0.24.0's GPUModelRunner no longer exposes a
                # ``self.pin_memory`` attribute (it uses a module-level
                # ``PIN_MEMORY`` constant instead), so the old
                # ``getattr(self, "pin_memory", False)`` silently fell back to
                # False. That allocated the async D2H snapshot destination in
                # *pageable* host memory, which turns ``copy_(non_blocking=True)``
                # into a fully synchronous, stream-stalling copy (~240 ms/step
                # on the 17.5k-token Thinker prefill). Resolve pinning from the
                # platform helper so the copy is a true async cudaMemcpyAsync.
                pin_memory=is_pin_memory_available(),
            )

        payload = async_payload_snapshot.payload
        hidden_states_snapshot = payload.get("hidden_states")
        if hidden_states_snapshot is None:
            # Models that omit hidden from the async snapshot only need
            # multimodal payloads (for example, talker codes.audio).
            hidden_states_snapshot = hidden_states[:0]

        return _OmniOutputTensorSnapshot(
            hidden_states=hidden_states_snapshot,
            staged_hidden_states_cpu=payload.get("staged_hidden_states_cpu"),
            multimodal_outputs=payload["multimodal_outputs"],
            async_payload=async_payload_snapshot,
        )

    def _maybe_run_eager_omni_postprocess_before_async_output(
        self,
        *,
        hidden_states: torch.Tensor,
        multimodal_outputs: Any,
        num_scheduled_tokens_np: np.ndarray,
        scheduler_output: SchedulerOutput,
        req_ids_output_copy: list[str],
        query_start_loc_cpu: Any,
    ) -> bool:
        """Apply model postprocess on live GPU tensors before async payload D2H."""
        model = getattr(self, "model", None)
        if not self._model_omni_flag(model, "has_postprocess"):
            return False
        if not self._model_omni_flag(model, "eager_omni_postprocess_before_async_output"):
            return False

        _, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
        if not downstream_req_ids:
            return False

        with record_function_or_nullcontext("omni_output_builder:eager_postprocess"):
            self._process_additional_information_updates(
                hidden_states,
                multimodal_outputs,
                num_scheduled_tokens_np,
                scheduler_output,
                None,
                None,
                req_ids_filter=set(downstream_req_ids),
                req_ids=req_ids_output_copy,
                query_start_loc_cpu=query_start_loc_cpu,
            )
        return True

    def _get_or_create_omni_payload_copy_stream(self) -> torch.cuda.Stream:
        stream = getattr(self, "_omni_payload_copy_stream", None)
        if stream is None:
            stream = torch.cuda.Stream()
            self._omni_payload_copy_stream = stream
        return stream

    def _build_omni_model_runner_output_from_snapshot(
        self,
        *,
        scheduler_output: SchedulerOutput,
        hidden_states: torch.Tensor,
        staged_hidden_states_cpu: torch.Tensor | None,
        multimodal_outputs: Any,
        req_ids_output_copy: list[str],
        req_id_to_index_output_copy: dict[str, int],
        valid_sampled_token_ids: list[list[int]],
        logprobs_lists: Any,
        prompt_logprobs_dict: dict[str, Any],
        num_nans_in_logits: Any,
        kv_connector_output: Any,
        ec_connector_output: Any,
        cudagraph_stats: Any,
        kv_extracted_req_ids: list[str] | None,
        num_scheduled_tokens_np: np.ndarray,
        query_start_loc_cpu: Any,
        postprocess_already_applied: bool = False,
    ) -> OmniModelRunnerOutput:
        combined_hidden_states = None
        combined_multimodal_outputs = None

        engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
        downstream_req_ids, sparse_mm_index, audio_sparse_output = resolve_sparse_mm_routing(
            engine_output_type=engine_output_type,
            req_ids_output_copy=req_ids_output_copy,
            downstream_req_ids=downstream_req_ids,
            multimodal_outputs=multimodal_outputs,
        )

        needs_pooler_payload = len(downstream_req_ids) > 0
        downstream_req_id_set = set(downstream_req_ids)
        defer_full_payload_d2h = needs_pooler_payload and self._should_defer_full_payload_d2h()
        hidden_states_cpu = None
        req_hidden_states_cpu: dict[str, torch.Tensor] | None = None
        include_hidden_payload = self._model_omni_pooler_payload_include_hidden()
        needs_scheduled_hidden_payload = (
            include_hidden_payload
            and needs_pooler_payload
            and (self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states())
        )
        self._stage_deferred_prefix_cache_mm_outputs(
            scheduler_output=scheduler_output,
            multimodal_outputs=multimodal_outputs,
            query_start_loc_cpu=query_start_loc_cpu,
        )

        if self.omni_prefix_cache is None and needs_scheduled_hidden_payload and not audio_sparse_output:
            num_valid_tokens = min(
                int(scheduler_output.total_num_scheduled_tokens),
                int(hidden_states.shape[0]),
            )
            if len(downstream_req_ids) == len(req_ids_output_copy):
                with record_function_or_nullcontext("omni_output_builder:hidden_d2h/scheduled"):
                    hidden_states_cpu = _to_cpu_contiguous(hidden_states[:num_valid_tokens])
            else:
                req_hidden_states_cpu = {}
                with record_function_or_nullcontext("omni_output_builder:hidden_d2h/per_request"):
                    for rid in downstream_req_ids:
                        idx = req_id_to_index_output_copy[rid]
                        start = int(query_start_loc_cpu[idx])
                        sched = int(num_scheduled_tokens_np[idx])
                        end = start + sched
                        req_hidden_states_cpu[rid] = _to_cpu_contiguous(hidden_states[start:end])

        # NOTE: pooler_output here is used only for the full-payload accumulation
        # path (accumulate_full_payload_output) and is NOT passed on the wire via
        # OmniModelRunnerOutput.pooler_output (which is set to None below).
        # The actual multimodal wire transport uses multimodal_outputs instead.
        pooler_output: list[dict[str, object]] | None = None
        if needs_pooler_payload:
            hidden_seq_len = int(hidden_states.shape[0])
            scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
            mm_cpu = None
            if self.omni_prefix_cache is not None:
                (
                    hidden_states_cpu,
                    combined_hidden_states,
                    combined_multimodal_outputs,
                ) = self._prepare_prefix_cache_pooler_payload_sources(
                    hidden_states=hidden_states,
                    staged_hidden_states_cpu=staged_hidden_states_cpu,
                    multimodal_outputs=multimodal_outputs,
                    scheduler_output=scheduler_output,
                    needs_scheduled_hidden_payload=needs_scheduled_hidden_payload,
                )
            if combined_multimodal_outputs is None:
                flat_mm = flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs
                if defer_full_payload_d2h:
                    with record_function_or_nullcontext("omni_output_builder:snapshot_mm_payload"):
                        mm_cpu = snapshot_mm_payload(flat_mm)
                else:
                    with record_function_or_nullcontext("omni_output_builder:build_mm_cpu"):
                        mm_cpu = build_mm_cpu(flat_mm)

            with record_function_or_nullcontext("omni_output_builder:process_additional_information"):
                if not postprocess_already_applied:
                    self._process_additional_information_updates(
                        hidden_states,
                        multimodal_outputs,
                        num_scheduled_tokens_np,
                        scheduler_output,
                        combined_hidden_states,
                        combined_multimodal_outputs,
                        req_ids_filter=downstream_req_id_set,
                        req_ids=req_ids_output_copy,
                        query_start_loc_cpu=query_start_loc_cpu,
                    )

            pooler_output = []
            with record_function_or_nullcontext("omni_output_builder:build_pooler_payloads"):
                for rid in req_ids_output_copy:
                    if rid not in downstream_req_id_set:
                        pooler_output.append({})
                        continue
                    idx = req_id_to_index_output_copy[rid]
                    start = int(query_start_loc_cpu[idx])
                    sched = int(num_scheduled_tokens_np[idx])
                    end = start + sched
                    payload = self._build_omni_pooler_payload(
                        rid=rid,
                        idx=idx,
                        start=start,
                        end=end,
                        hidden_states_cpu=hidden_states_cpu,
                        req_hidden_states_cpu=req_hidden_states_cpu,
                        combined_hidden_states=combined_hidden_states,
                        combined_multimodal_outputs=combined_multimodal_outputs,
                        mm_cpu=mm_cpu,
                        audio_sparse_output=audio_sparse_output,
                        sparse_mm_index=sparse_mm_index,
                        hidden_seq_len=hidden_seq_len,
                        scheduled_seq_len=scheduled_seq_len,
                    )
                    pooler_output.append(flatten_payload(payload))

        pooler_inter: Sequence[dict[str, object] | None] | None
        pooler_client: Sequence[dict[str, object] | None] | None
        pooler_output = pooler_output or []
        if self._async_chunk and stage_sends_async_output(self.model_config):
            pooler_inter, pooler_client = partition_payload_list(pooler_output)
        else:
            # Connector-less stages expose the same payload through the
            # orchestrator bridge; non-async stages preserve legacy behavior.
            pooler_inter, pooler_client = pooler_output, pooler_output
        client_output_keys = self._client_multimodal_output_keys()
        if client_output_keys:
            allowed = frozenset(client_output_keys)
            pooler_client = [
                {key: value for key, value in payload.items() if key in allowed} for payload in pooler_output
            ]

        if pooler_inter and self._should_accumulate_full_payload_output():
            with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"):
                for i, rid in enumerate(req_ids_output_copy):
                    req_state = self.requests.get(rid)
                    if req_state is not None and pooler_inter[i]:
                        self.accumulate_full_payload_output(rid, pooler_inter[i], req_state)

        with record_function_or_nullcontext("omni_output_builder:build_multimodal_outputs"):
            inter_stage_outputs, multimodal_outputs = self._build_omni_step_outputs(
                pooler_inter,
                pooler_client,
                defer_full_payload_d2h=defer_full_payload_d2h,
            )

        with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
            routed_experts_lists = None
            if self._should_return_omni_routed_experts():
                routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
            output = OmniModelRunnerOutput(
                req_ids=req_ids_output_copy,
                req_id_to_index=req_id_to_index_output_copy,
                sampled_token_ids=valid_sampled_token_ids,
                logprobs=logprobs_lists,
                prompt_logprobs_dict=prompt_logprobs_dict,
                pooler_output=None,
                multimodal_outputs=multimodal_outputs,
                inter_stage_outputs=inter_stage_outputs,
                kv_connector_output=kv_connector_output,
                ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
                num_nans_in_logits=num_nans_in_logits,
                cudagraph_stats=cudagraph_stats,
            )
            output.kv_extracted_req_ids = kv_extracted_req_ids
            with record_function_or_nullcontext("omni_output_builder:get_omni_connector_output"):
                output.omni_connector_output = self.get_omni_connector_output()
            output.routed_experts = routed_experts_lists
        return output

    @torch.inference_mode()
    def sample_tokens(
        self,
        grammar_output: GrammarOutput | None,
    ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
        kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None)
        self.kv_extracted_req_ids = None

        if self.execute_model_state is None:
            kv_connector_output = self.kv_connector_output
            self.kv_connector_output = None
            # receive sampled token ids from the last PP rank.
            if self.use_async_scheduling and not get_pp_group().is_last_rank:
                self._pp_receive_prev_sampled_token_ids_to_input_batch()
            # In case of PP with kv transfer, we need to pass through the
            # kv_connector_output
            return self.attach_omni_connector_output(
                OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
            )

        # Unpack ephemeral state.
        (
            scheduler_output,
            logits,
            spec_decode_metadata,
            spec_decode_common_attn_metadata,
            hidden_states,
            staged_hidden_states_cpu,
            sample_hidden_states,
            aux_hidden_states,
            ec_connector_output,
            cudagraph_stats,
            multimodal_outputs,
            slot_mappings,  # OMNI: unpack slot_mappings for drafter
        ) = self.execute_model_state
        self.execute_model_state = None

        # Apply structured output bitmasks if present.
        if grammar_output is not None:
            apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits)

        # Correct padding values of prompt_token_ids to match the logits vocabulary size.
        if logits is not None and not self.input_batch.sampling_metadata.no_penalties:
            smd = self.input_batch.sampling_metadata
            if smd.prompt_token_ids is not None:
                logits_vocab = logits.shape[-1]
                if self.input_batch.vocab_size > logits_vocab:
                    smd.prompt_token_ids = clamp_prompt_ids_to_penalty_padding(smd.prompt_token_ids, logits_vocab)

        # Drop min-tokens stop ids the head cannot emit (e.g. the text
        # tokenizer EOS folded into all_stop_token_ids on a narrow codec
        # talker head); they would index_put_ out of bounds (#4962).
        if logits is not None:
            sanitize_min_tokens_stop_ids(
                self.input_batch.sampling_metadata.logitsprocs,
                logits.shape[-1],
            )

        with record_function_or_nullcontext("gpu_model_runner: sample"):
            sampler_output = self._sample(logits, spec_decode_metadata)

        self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output)

        self._draft_token_ids = None
        self._draft_token_req_ids = None
        self.valid_sampled_token_count_gpu = None
        self.input_batch.prev_sampled_token_ids = None

        def propose_draft_token_ids(sampled_token_ids):
            assert spec_decode_common_attn_metadata is not None
            with record_function_or_nullcontext("gpu_model_runner: draft"):
                self._draft_token_ids = self.propose_draft_token_ids(
                    scheduler_output,
                    sampled_token_ids,
                    self.input_batch.sampling_metadata,
                    hidden_states,
                    sample_hidden_states,
                    aux_hidden_states,
                    spec_decode_metadata,
                    spec_decode_common_attn_metadata,
                    slot_mappings,  # OMNI: pass slot_mappings to drafter (upstream v1 API)
                )
                self._copy_draft_token_ids_to_cpu(scheduler_output)

        spec_config = self.speculative_config
        propose_drafts_after_bookkeeping = False
        if spec_config is not None:
            input_fits_in_drafter = self._input_fits_in_drafter(spec_decode_common_attn_metadata)
            use_gpu_toks = (
                spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states()
            ) and not spec_config.disable_padded_drafter_batch
            if use_gpu_toks:
                assert isinstance(
                    self.drafter,
                    EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer,
                )
                sampled_token_ids = sampler_output.sampled_token_ids
                if input_fits_in_drafter:
                    propose_draft_token_ids(sampled_token_ids)
                elif self.valid_sampled_token_count_event is not None:
                    assert spec_decode_common_attn_metadata is not None
                    next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded(
                        self.optimistic_seq_lens_cpu,
                        sampled_token_ids,
                        self.requests,
                        self.input_batch,
                        self.discard_request_mask.gpu,
                    )
                    self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
                    # Since we couldn't run the drafter,
                    # just use zeros for the draft tokens.
                    self._draft_token_ids = torch.zeros(1, device=self.device, dtype=torch.int32).expand(
                        len(self.input_batch.req_ids), self.num_spec_tokens
                    )
                    self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
            else:
                propose_drafts_after_bookkeeping = input_fits_in_drafter

        with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
            (
                num_nans_in_logits,
                num_nans,
                logprobs_lists,
                valid_sampled_token_ids,
                prompt_logprobs_dict,
                req_ids_output_copy,
                req_id_to_index_output_copy,
                invalid_req_indices,
            ) = self._bookkeeping_sync(
                scheduler_output,
                sampler_output,
                logits,
                hidden_states,
                scheduler_output.total_num_scheduled_tokens,
            )

        multimodal_outputs = self._run_post_sample_talker_mtp(
            req_ids=req_ids_output_copy,
            valid_sampled_token_ids=valid_sampled_token_ids,
            sampled_token_ids=sampler_output.sampled_token_ids,
            invalid_req_indices=invalid_req_indices,
            sample_hidden_states=sample_hidden_states,
            multimodal_outputs=multimodal_outputs,
        )

        if propose_drafts_after_bookkeeping:
            # ngram and other speculative decoding methods use the sampled
            # tokens on the CPU, so they are run after bookkeeping.
            propose_draft_token_ids(valid_sampled_token_ids)

        # Finalize KV connector (wait_for_save + clear metadata) after
        # draft model runs. Deferred from target model forward to allow
        # draft model to also save its KV cache.
        if self.speculative_config is not None:
            self.finalize_kv_connector()

        with record_function_or_nullcontext("gpu_model_runner: eplb"):
            self.eplb_step()

        # kv_connector_output may be modified during drafting
        kv_connector_output = self.kv_connector_output
        self.kv_connector_output = None

        num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
        if num_scheduled_tokens_np is None:
            num_scheduled_tokens_np = np.array(
                [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids_output_copy],
                dtype=np.int32,
            )
        else:
            num_scheduled_tokens_np = np.asarray(num_scheduled_tokens_np, dtype=np.int32).copy()

        query_start_loc_cpu = self._snapshot_query_start_loc_cpu()
        scheduler_output_snapshot = self._snapshot_scheduler_output_for_async_omni_output(scheduler_output)
        req_ids_output_snapshot = list(req_ids_output_copy)
        req_id_to_index_output_snapshot = dict(req_id_to_index_output_copy)
        valid_sampled_token_ids_snapshot = [list(token_ids) for token_ids in valid_sampled_token_ids]
        logprobs_lists_snapshot = copy(logprobs_lists) if logprobs_lists is not None else None
        prompt_logprobs_dict_snapshot = dict(prompt_logprobs_dict) if prompt_logprobs_dict is not None else {}
        num_nans_in_logits_snapshot = (
            dict(num_nans_in_logits) if isinstance(num_nans_in_logits, dict) else num_nans_in_logits
        )

        use_async_omni_output = self._should_use_async_omni_output()
        omni_postprocess_already_applied = False
        if use_async_omni_output:
            omni_postprocess_already_applied = self._maybe_run_eager_omni_postprocess_before_async_output(
                hidden_states=hidden_states,
                multimodal_outputs=multimodal_outputs,
                num_scheduled_tokens_np=num_scheduled_tokens_np,
                scheduler_output=scheduler_output,
                req_ids_output_copy=req_ids_output_copy,
                query_start_loc_cpu=query_start_loc_cpu,
            )
        output_tensor_snapshot = self._snapshot_omni_output_tensors_for_async_output(
            use_async_omni_output=use_async_omni_output,
            hidden_states=hidden_states,
            staged_hidden_states_cpu=staged_hidden_states_cpu,
            multimodal_outputs=multimodal_outputs,
        )

        def output_builder() -> OmniModelRunnerOutput:
            if output_tensor_snapshot.async_payload is not None:
                with record_function_or_nullcontext("omni_async_output:wait_cpu_payload"):
                    output_tensor_snapshot.async_payload.wait()
            with record_function_or_nullcontext("omni_output_builder:total"):
                return self._build_omni_model_runner_output_from_snapshot(
                    scheduler_output=scheduler_output_snapshot,
                    hidden_states=output_tensor_snapshot.hidden_states,
                    staged_hidden_states_cpu=output_tensor_snapshot.staged_hidden_states_cpu,
                    multimodal_outputs=output_tensor_snapshot.multimodal_outputs,
                    req_ids_output_copy=req_ids_output_snapshot,
                    req_id_to_index_output_copy=req_id_to_index_output_snapshot,
                    valid_sampled_token_ids=valid_sampled_token_ids_snapshot,
                    logprobs_lists=logprobs_lists_snapshot,
                    prompt_logprobs_dict=prompt_logprobs_dict_snapshot,
                    num_nans_in_logits=num_nans_in_logits_snapshot,
                    kv_connector_output=kv_connector_output,
                    ec_connector_output=ec_connector_output,
                    cudagraph_stats=cudagraph_stats,
                    kv_extracted_req_ids=kv_extracted_req_ids,
                    num_scheduled_tokens_np=num_scheduled_tokens_np,
                    query_start_loc_cpu=query_start_loc_cpu,
                    postprocess_already_applied=omni_postprocess_already_applied,
                )

        if not use_async_omni_output:
            output = output_builder()

            if not self.use_async_scheduling:
                return output
        with record_function_or_nullcontext("gpu_model_runner: AsyncGPUModelRunnerOutput"):
            async_output_cls = OmniAsyncGPUModelRunnerOutput if use_async_omni_output else AsyncGPUModelRunnerOutput
            async_output_kwargs = dict(
                sampled_token_ids=sampler_output.sampled_token_ids,
                logprobs_tensors=sampler_output.logprobs_tensors,
                invalid_req_indices=invalid_req_indices,
                async_output_copy_stream=self.async_output_copy_stream,
                vocab_size=self.input_batch.vocab_size,
                num_nans=num_nans,
            )
            if use_async_omni_output:
                async_output = async_output_cls(
                    model_runner_output_builder=output_builder,
                    cuda_device=self.device,
                    **async_output_kwargs,
                )
            else:
                async_output = async_output_cls(
                    model_runner_output=output,
                    **async_output_kwargs,
                )
        with record_function_or_nullcontext("gpu_model_runner: set_async_sampled_token_ids"):
            # Save ref of sampled_token_ids CPU tensor if the batch contains
            # any requests with sampling params that require output ids.
            self.input_batch.set_async_sampled_token_ids(
                async_output.sampled_token_ids_cpu,
                async_output.async_copy_ready_event,
            )

        return async_output

    def _resolve_global_request_id(self, req_id: str) -> str:
        """Resolve global request ID from request state."""
        req_state = self.requests.get(req_id)
        if not req_state:
            return req_id

        add_info = self.model_intermediate_buffer.get(req_id, {})
        global_id = add_info.get("global_request_id")
        if global_id:
            if isinstance(global_id, list) and global_id:
                global_id = global_id[0]
            if isinstance(global_id, bytes):
                return global_id.decode("utf-8")
            return str(global_id)
        return req_id