vllm-omni/vllm_omni/worker/gpu_model_runner.py

import contextlib
from collections.abc import Callable
from dataclasses import replace
from typing import TYPE_CHECKING, Any, cast

import numpy as np
import torch
from vllm.compilation.cuda_graph import CUDAGraphWrapper
from vllm.config import CUDAGraphMode
from vllm.distributed.parallel_state import get_pp_group
from vllm.forward_context import set_forward_context
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import supports_mrope
from vllm.model_executor.models.interfaces_base import VllmModelForPooling
from vllm.sampling_params import SamplingType
from vllm.tracing import instrument
from vllm.utils.import_utils import LazyLoader
from vllm.utils.math_utils import cdiv
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.spec_decode.ngram_proposer_gpu import (
    update_ngram_gpu_tensors_incremental,
    update_scheduler_for_invalid_drafts,
)
from vllm.v1.worker.gpu_input_batch import CachedRequestState
from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices

from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
from vllm_omni.engine.serialization import deserialize_additional_information
from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding
from vllm_omni.model_executor.models.output_templates import OmniOutput

if TYPE_CHECKING:
    from vllm.v1.core.sched.output import SchedulerOutput
else:
    xgr = LazyLoader("xgr", globals(), "xgrammar")
    xgr_torch_compile = LazyLoader(
        "xgr_torch_compile",
        globals(),
        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile",
    )

logger = init_logger(__name__)


class OmniGPUModelRunner(GPUModelRunner):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_intermediate_buffer: dict[str, dict[str, Any]] = {}
        self._omni_num_scheduled_tokens_np: np.ndarray | None = None
        self._omni_last_model_output: object | None = None
        # The Omni tensor prefix cache will be allocated
        # when we initialize the metadata builders if enabled
        self.omni_prefix_cache = None

    def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes):
        """Override to fix scheduler_metadata buffer size for FA3 + CUDA graph.

        The upstream FlashAttentionMetadataBuilder pre-allocates
        scheduler_metadata with (max_num_seqs + 1) entries, but FA3's
        get_scheduler_metadata() can return up to
        (max_num_seqs * max_num_splits + 1) entries, causing a RuntimeError
        during CUDA graph capture.  After calling the parent implementation
        we resize any too-small buffers.
        """
        super().initialize_metadata_builders(kv_cache_config, kernel_block_sizes)

        for kv_cache_group in self.attn_groups:
            for attn_group in kv_cache_group:
                for builder in attn_group.metadata_builders:
                    sm = getattr(builder, "scheduler_metadata", None)
                    max_num_splits = getattr(builder, "max_num_splits", 0)
                    if sm is not None and max_num_splits > 1:
                        required = self.scheduler_config.max_num_seqs * max_num_splits + 1
                        if sm.shape[0] < required:
                            builder.scheduler_metadata = torch.zeros(
                                required,
                                dtype=sm.dtype,
                                device=sm.device,
                            )

        # Initialize the wrapper for both multimodal output tensors
        # and for hidden states to be passed between stages
        if self.cache_config.enable_prefix_caching:
            self.omni_prefix_cache = OmniTensorPrefixCache(
                num_blocks=kv_cache_config.num_blocks,
                block_size=self.cache_config.block_size,
                hidden_size=self.model_config.get_hidden_size(),
                hs_dtype=self.dtype,
            )

    @instrument(span_name="Loading (GPU)")
    def load_model(self, *args, **kwargs) -> None:
        super().load_model(*args, **kwargs)
        self._prewarm_attention_capture_workspaces()

        # TODO move this model specific logic to a separate class
        # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
        self.has_talker_mtp = False
        talker_mtp = getattr(self.model, "talker_mtp", None)
        if talker_mtp is not None:
            self.talker_mtp = talker_mtp  # type: ignore[assignment]
            self.has_talker_mtp = True
            cudagraph_mode = self.compilation_config.cudagraph_mode
            assert cudagraph_mode is not None
            has_separate_talker = getattr(self.model, "talker", None) is not None
            talker_mtp_graph_safe = getattr(self.model, "talker_mtp_graph_safe", False)
            if cudagraph_mode.has_full_cudagraphs() and (has_separate_talker or talker_mtp_graph_safe):
                self.talker_mtp = CUDAGraphWrapper(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
            # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
            hidden_size = int(
                getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
            )
            max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
            self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
            self.talker_mtp_inputs_embeds = self._make_buffer(
                max_batch_size, hidden_size, dtype=self.dtype, numpy=False
            )
            self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
            self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)

    def _prewarm_attention_capture_workspaces(self) -> None:
        capture_sizes = getattr(self.compilation_config, "cudagraph_capture_sizes", None)
        if not capture_sizes:
            return
        from vllm_omni.attention.fish_kvcache_backend import prewarm_fish_kvcache_attn_capture_workspaces

        prewarm_fish_kvcache_attn_capture_workspaces(
            model_config=self.model_config,
            device=self.device,
            dtype=self.dtype,
            capture_sizes=capture_sizes,
        )

    def _maybe_attach_attention_metadata_extensions(
        self,
        *,
        attn_metadata: Any,
        num_reqs: int,
        num_reqs_padded: int,
        max_query_len: int,
        pad_attn: bool,
        for_cudagraph_capture: bool = False,
        num_scheduled_tokens_np: np.ndarray | None = None,
    ) -> None:
        from vllm_omni.attention.fish_kvcache_backend import maybe_attach_fish_kvcache_seq_lens_upper_bound

        maybe_attach_fish_kvcache_seq_lens_upper_bound(
            model_config=self.model_config,
            attn_metadata=attn_metadata,
            input_batch=self.input_batch,
            optimistic_seq_lens_cpu=self.optimistic_seq_lens_cpu,
            num_reqs=num_reqs,
            num_reqs_padded=num_reqs_padded,
            max_query_len=max_query_len,
            pad_attn=pad_attn,
            for_cudagraph_capture=for_cudagraph_capture,
            num_scheduled_tokens_np=num_scheduled_tokens_np,
        )

    def _build_model_sampler_output_token_ids(self) -> list[list[int]]:
        """Build decoded-token history for ``prefer_model_sampler`` models.

        vLLM only populates ``sampling_metadata.output_token_ids`` when penalties
        or logits processors require it. Models that opt into a custom sampler
        (e.g. CosyVoice3 RAS sampler, HunyuanImage3 stage-transition sampler)
        also depend on this history, so we reconstruct it directly from the
        input batch. Shared by GPU and NPU AR runners.
        """
        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]

        return output_token_ids

    def _sampling_metadata_for_model_sampler(self, sampling_metadata):
        output_token_ids = self._build_model_sampler_output_token_ids()
        if output_token_ids == sampling_metadata.output_token_ids:
            return sampling_metadata
        return replace(sampling_metadata, output_token_ids=output_token_ids)

    def _init_mrope_positions(self, req_state: CachedRequestState):
        """Initialize M-RoPE positions for multimodal inputs.

        Extracts multimodal feature metadata (image grids, video grids,
        audio features) and computes M-RoPE positions for proper positional
        encoding of multimodal tokens.

        Args:
            req_state: Cached request state containing multimodal features

        Raises:
            AssertionError: If the model does not support M-RoPE
        """
        image_grid_thw = []
        video_grid_thw = []
        second_per_grid_ts = []
        audio_feature_lengths = []
        use_audio_in_video = False
        for mm_feature in req_state.mm_features:
            mm_item = mm_feature.data
            if mm_item is None:
                continue
            mm_input = mm_item.get_data()
            if (t := mm_input.get("image_grid_thw")) is not None:
                image_grid_thw.append(t.tolist())
            if (t := mm_input.get("video_grid_thw")) is not None:
                video_grid_thw.append(t.tolist())
            if (t := mm_input.get("second_per_grid_ts")) is not None:
                second_per_grid_ts.append(t)
            if (t := mm_input.get("audio_feature_lengths")) is not None:
                audio_feature_lengths.append(t)
            # Check for use_audio_in_video
            use_audio_in_video_value = mm_input.get("use_audio_in_video")
            if use_audio_in_video_value is not None:
                use_audio_in_video = bool(use_audio_in_video_value.item())

        if supports_mrope(self.get_model()):
            # Model implements SupportsMRoPE interface
            # Pass all extracted metadata; models use what they need via **kwargs
            sp_extra_args = getattr(req_state.sampling_params, "extra_args", {}) if req_state.sampling_params else {}
            target_h = sp_extra_args.get("target_h") if isinstance(sp_extra_args, dict) else None
            target_w = sp_extra_args.get("target_w") if isinstance(sp_extra_args, dict) else None
            kwargs = dict(
                mm_features=req_state.mm_features,
                hf_config=self.model_config.hf_config,
                image_grid_thw=image_grid_thw,
                video_grid_thw=video_grid_thw,
                second_per_grid_ts=second_per_grid_ts,
                audio_feature_lengths=audio_feature_lengths,
                use_audio_in_video=use_audio_in_video,
            )
            if target_h is not None:
                kwargs["target_h"] = target_h
            if target_w is not None:
                kwargs["target_w"] = target_w
            req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
                req_state.prompt_token_ids,
                **kwargs,
            )
        else:
            req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(
                req_state.prompt_token_ids,
                hf_config=self.model_config.hf_config,
                image_grid_thw=image_grid_thw,
                video_grid_thw=video_grid_thw,
                second_per_grid_ts=second_per_grid_ts,
                audio_feature_lengths=audio_feature_lengths,
                use_audio_in_video=use_audio_in_video,
            )

    def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"):
        """Calculate M-RoPE positions for scheduled tokens.

        Delegates to the upstream implementation first, then applies a fixup
        pass for models that pre-compute 2D spatial decode positions (e.g.
        GLM-Image).  This avoids duplicating the full upstream method while
        still supporting non-linear decode position patterns.

        Models opt-in by declaring ``precomputed_mrope_decode = True`` as a
        class attribute.  When set, ``get_mrope_input_positions`` is expected
        to return positions covering **both** prefill and decode tokens.
        """
        # Run upstream logic (handles prompt positions + linear decode fallback)
        super()._calc_mrope_positions(scheduler_output)

        # Only run the fixup if the model pre-computes decode M-RoPE positions
        if not getattr(self.get_model(), "precomputed_mrope_decode", False):
            return

        self._fixup_precomputed_mrope_decode_positions(scheduler_output)

    def _fixup_precomputed_mrope_decode_positions(self, scheduler_output: "SchedulerOutput") -> None:
        """Overwrite linear decode M-RoPE positions with pre-computed ones.

        For image-generation models (like GLM-Image) that output tokens in 2D
        grid order, ``get_mrope_input_positions`` returns positions for the
        full sequence (prefill + decode).  The upstream runner only uses the
        prefill portion and falls back to linear increments for decode.  This
        method patches the decode slice with the correct pre-computed values.
        """
        from vllm.utils import length_from_prompt_token_ids_or_embeds

        mrope_pos_ptr = 0
        for index, req_id in enumerate(self.input_batch.req_ids):
            req = self.requests[req_id]
            assert req.mrope_positions is not None

            num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index]
            num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
            num_prompt_tokens = length_from_prompt_token_ids_or_embeds(req.prompt_token_ids, req.prompt_embeds)

            if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens:
                prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens)
                completion_part_len = max(0, num_scheduled_tokens - prompt_part_len)
            else:
                prompt_part_len = num_scheduled_tokens
                completion_part_len = 0

            mrope_pos_ptr += prompt_part_len

            if completion_part_len > 0:
                dst_start = mrope_pos_ptr
                decode_start = num_computed_tokens + prompt_part_len
                decode_end = decode_start + completion_part_len
                total_precomputed = req.mrope_positions.shape[1]

                if decode_end <= total_precomputed:
                    # Overwrite the linear positions written by upstream with
                    # the correct pre-computed 2D spatial positions.
                    self.mrope_positions.cpu[:, dst_start : dst_start + completion_part_len] = req.mrope_positions[
                        :, decode_start:decode_end
                    ]

                mrope_pos_ptr += completion_part_len

    def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None:
        """Update the cached states and the persistent batch with the scheduler
        output.

        The updated states are used by the `_prepare_inputs` function to create
        the input GPU tensors for the model.

        The SamplingMetadata is updated and copied to the GPU if there is a
        new/resumed/paused/finished request in the batch.
        """
        # Used for prefix cache
        if self.omni_prefix_cache is not None:
            self.omni_prefix_cache.reset_prefix_cached_new_req_ids()

        # Remove finished requests from the cached states.
        # cleanup_finished_request lives on OmniConnectorModelRunnerMixin and
        # is only safe to call once init_omni_connectors() has finished
        # populating mixin state (it sets ``_omni_connector_initialized = True``
        # at the very end).  Archs that inherit the method via MRO without
        # running that init must be skipped, so gate on the explicit flag
        # rather than probing private attribute names.
        cleanup_finished_request = (
            getattr(self, "cleanup_finished_request", None)
            if getattr(self, "_omni_connector_initialized", False)
            else None
        )
        for req_id in scheduler_output.finished_req_ids:
            self.requests.pop(req_id, None)
            self.model_intermediate_buffer.pop(req_id, None)
            self.num_prompt_logprobs.pop(req_id, None)
            if self.omni_prefix_cache is not None:
                self.omni_prefix_cache.discard_deferred_mm_outputs(req_id)
            if hasattr(self, "_downstream_payload_cache"):
                self._downstream_payload_cache.pop(req_id, None)
            if hasattr(self, "_talker_mtp_generators"):
                self._talker_mtp_generators.pop(req_id, None)
            if cleanup_finished_request is not None:
                cleanup_finished_request(req_id)

        self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids)
        # Remove the finished requests from the persistent batch.
        # NOTE(woosuk): There could be an edge case where finished_req_ids and
        # scheduled_req_ids overlap. This happens when a request is aborted and
        # then resubmitted with the same ID. In this case, we treat them as two
        # distinct requests - clearing the cached states for the first request
        # and handling the second as a new request.
        for req_id in scheduler_output.finished_req_ids:
            self.input_batch.remove_request(req_id)

        # Zero GPU memory for freshly allocated cache blocks to prevent
        # stale NaN/data from corrupting attention or SSM computation.
        if scheduler_output.new_block_ids_to_zero:
            self._zero_block_ids(scheduler_output.new_block_ids_to_zero)

        # Free the cached encoder outputs.
        for mm_hash in scheduler_output.free_encoder_mm_hashes:
            self.encoder_cache.pop(mm_hash, None)

        # Remove the unscheduled requests from the persistent batch.
        # NOTE(woosuk): The unscheduled requests are either preempted requests
        # or running requests that are not scheduled in this step. We remove
        # them from the persistent batch but keep their cached states since
        # they will be scheduled again sometime in the future.
        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
        cached_req_ids = self.input_batch.req_id_to_index.keys()
        resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
        # NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint,
        # so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds
        # apart from the forced-preemption case in reset_prefix_cache. And in
        # that case we include the resumed_req_ids in the unscheduled set so
        # that they get cleared from the persistent batch before being re-scheduled
        # in the normal resumed request path.
        unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids)
        # NOTE(woosuk): The persistent batch optimization assumes that
        # consecutive batches contain mostly the same requests. If batches
        # have low request overlap (e.g., alternating between two distinct
        # sets of requests), this optimization becomes very inefficient.
        for req_id in unscheduled_req_ids:
            self.input_batch.remove_request(req_id)

        is_ngram_gpu = self.speculative_config is not None and self.speculative_config.use_ngram_gpu()
        if is_ngram_gpu:
            ngram_gpu_new_reqs: list[CachedRequestState] = []

        reqs_to_add: list[CachedRequestState] = []
        deferred_spec_decode_corrections = []
        # Add new requests to the cached states.
        for new_req_data in scheduler_output.scheduled_new_reqs:
            req_id = new_req_data.req_id
            if req_id in self.requests:
                self._update_streaming_input_additional_info(new_req_data, req_id)
                req_state = self._update_streaming_request(req_id, new_req_data)
                reqs_to_add.append(req_state)
                continue

            # Since this is the first time the request has been scheduled,
            # num_computed_tokens > 0 means that we have a hit in prefix
            # caching; mark it so that we can manage the hidden states
            # later on as needed.
            if self.omni_prefix_cache is not None and new_req_data.num_computed_tokens > 0:
                self.omni_prefix_cache.add_prefix_cached_new_req_id(req_id)

            sampling_params = new_req_data.sampling_params
            pooling_params = new_req_data.pooling_params

            if sampling_params and sampling_params.sampling_type == SamplingType.RANDOM_SEED:
                generator = torch.Generator(device=self.device)
                generator.manual_seed(sampling_params.seed)
            else:
                generator = None

            if self.is_pooling_model:
                assert pooling_params is not None
                task = pooling_params.task
                assert task is not None, "You did not set `task` in the API"

                model = cast(VllmModelForPooling, self.get_model())
                to_update = model.pooler.get_pooling_updates(task)
                to_update.apply(pooling_params)

            req_state = CachedRequestState(
                req_id=req_id,
                prompt_token_ids=new_req_data.prompt_token_ids,
                prompt_embeds=new_req_data.prompt_embeds,
                mm_features=new_req_data.mm_features,
                sampling_params=sampling_params,
                pooling_params=pooling_params,
                generator=generator,
                block_ids=new_req_data.block_ids,
                num_computed_tokens=new_req_data.num_computed_tokens,
                output_token_ids=[],
                lora_request=new_req_data.lora_request,
            )
            self.requests[req_id] = req_state
            self.late_interaction_runner.register_request(req_id, pooling_params)

            # If prompt embeddings are provided, decode and attach to inter_data
            try:
                if getattr(new_req_data, "prompt_embeds", None) is not None:
                    payload = new_req_data.prompt_embeds
                    dtype = getattr(np, payload.dtype)
                    arr = np.frombuffer(payload.data, dtype=dtype)
                    arr = arr.reshape(payload.shape)
                    pe_cpu = torch.from_numpy(arr)
                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
                    try:
                        new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
                    except Exception:
                        pass
            except Exception as e:
                logger.error(f"Error decoding prompt embeds: {e}")
            # Decode additional_information payloads (dictionary)
            try:
                if getattr(new_req_data, "additional_information", None) is not None:
                    logger.warning_once(
                        "additional_information on request data is deprecated, use model_intermediate_buffer"
                    )
                    info_dict = deserialize_additional_information(new_req_data.additional_information)
                    if info_dict:
                        self.model_intermediate_buffer[req_id] = info_dict
                        setattr(
                            self.requests[req_id],
                            "additional_information_cpu",
                            info_dict,
                        )
            except Exception as e:
                logger.error(f"Error decoding additional information: {e}")

            if sampling_params and sampling_params.prompt_logprobs is not None:
                self.num_prompt_logprobs[req_id] = (
                    self.input_batch.vocab_size
                    if sampling_params.prompt_logprobs == -1
                    else sampling_params.prompt_logprobs
                )
            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
            if self.uses_mrope:
                self._init_mrope_positions(req_state)

            # Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
            if self.uses_xdrope_dim > 0:
                self._init_xdrope_positions(req_state)

            reqs_to_add.append(self.requests[req_id])
            # Track new requests for ngram_gpu full tensor copy
            if is_ngram_gpu:
                ngram_gpu_new_reqs.append(req_state)

        # Update the states of the running/resumed requests.
        is_last_rank = get_pp_group().is_last_rank
        req_data = scheduler_output.scheduled_cached_reqs
        scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens

        # Save scheduler-allocated spec lengths before trimming so
        # prev_num_draft_len keeps the optimistic count for rejection correction.
        original_num_spec_per_req: dict[str, int] = {}
        if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
            for req_id, toks in scheduled_spec_tokens.items():
                original_num_spec_per_req[req_id] = len(toks)
            update_scheduler_for_invalid_drafts(
                self._num_valid_draft_tokens_event,
                self._num_valid_draft_tokens_cpu,
                scheduler_output,
                self.input_batch.req_id_to_index,
            )
        if self.use_async_spec_decode:
            self.prev_num_draft_tokens.np.fill(0)

        for i, req_id in enumerate(req_data.req_ids):
            req_state = self.requests[req_id]
            num_computed_tokens = req_data.num_computed_tokens[i]
            new_block_ids = req_data.new_block_ids[i]
            resumed_from_preemption = req_id in req_data.resumed_req_ids
            num_output_tokens = req_data.num_output_tokens[i]
            req_index = self.input_batch.req_id_to_index.get(req_id)

            if req_state.prev_num_draft_len and self.use_async_scheduling:
                # prev_num_draft_len is used in async scheduling mode with
                # spec decode. it indicates if need to update num_computed_tokens
                # of the request. for example:
                # first step: num_computed_tokens = 0, spec_tokens = [],
                # prev_num_draft_len = 0.
                # second step: num_computed_tokens = 100(prompt length),
                # spec_tokens = [a,b], prev_num_draft_len = 0.
                # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],
                # prev_num_draft_len = 2.
                # num_computed_tokens in first step and second step doesn't contain
                # the spec tokens length, but in third step it contains the
                # spec tokens length. we only need to update num_computed_tokens
                # when prev_num_draft_len > 0.
                if req_index is None:
                    req_state.prev_num_draft_len = 0
                else:
                    optimistic_num_accepted = req_state.prev_num_draft_len
                    req_state.output_token_ids.extend([-1] * optimistic_num_accepted)

                    deferred_spec_decode_corrections.append((req_id, optimistic_num_accepted, req_state))

                    prev_req_index = (
                        self.input_batch.prev_req_id_to_index.get(req_id)
                        if self.input_batch.prev_req_id_to_index
                        else None
                    )
                    if prev_req_index is not None:
                        self.prev_num_draft_tokens.np[prev_req_index] = optimistic_num_accepted

                    if is_ngram_gpu and optimistic_num_accepted > 0:
                        self.input_batch.num_tokens_no_spec[req_index] += optimistic_num_accepted

            # Update the cached states.
            req_state.num_computed_tokens = num_computed_tokens

            if not is_last_rank:
                if not req_data.new_token_ids:
                    new_token_ids: list[int] = []
                else:
                    new_token_ids = req_data.new_token_ids[i]
                    num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens
                    if num_new_tokens == 1:
                        req_state.output_token_ids.append(new_token_ids[-1])
                    elif num_new_tokens > 0:
                        req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
            elif num_output_tokens < len(req_state.output_token_ids):
                # Some output tokens were discarded due to a sync-KV-load
                # failure, or output_token_ids was inflated by the optimistic
                # extend above (async spec decode). Align the cached state.
                del req_state.output_token_ids[num_output_tokens:]
                if req_index is not None:
                    end_idx = self.input_batch.num_prompt_tokens[req_index] + num_output_tokens
                    self.input_batch.num_tokens_no_spec[req_index] = end_idx

            # Update the block IDs.
            if not resumed_from_preemption:
                if new_block_ids is not None:
                    # Append the new blocks to the existing block IDs.
                    for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
                        block_ids.extend(new_ids)
            else:
                assert req_index is None
                assert new_block_ids is not None
                # The request is resumed from preemption.
                # Replace the existing block IDs with the new ones.
                req_state.block_ids = new_block_ids

            req_index = self.input_batch.req_id_to_index.get(req_id)
            if req_index is None:
                # The request is not in the persistent batch.
                # The request was either preempted and resumed later, or was not
                # scheduled in the previous step and needs to be added again.

                if self.use_async_scheduling and num_output_tokens > 0:
                    # We must recover the output token ids for resumed requests in the
                    # async scheduling case, so that correct input_ids are obtained.
                    resumed_token_ids = req_data.all_token_ids[req_id]
                    req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]

                reqs_to_add.append(req_state)
                # Track resumed requests for ngram_gpu full tensor copy
                if is_ngram_gpu:
                    ngram_gpu_new_reqs.append(req_state)
                continue

            # Update the persistent batch.
            self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
            if new_block_ids is not None:
                self.input_batch.block_table.append_row(new_block_ids, req_index)

            # For the last rank, we don't need to update the token_ids_cpu
            # because the sampled tokens are already cached.
            if not is_last_rank:
                # Add new_token_ids to token_ids_cpu.
                start_token_index = num_computed_tokens
                end_token_index = num_computed_tokens + len(new_token_ids)
                self.input_batch.token_ids_cpu[req_index, start_token_index:end_token_index] = new_token_ids
                self.input_batch.num_tokens_no_spec[req_index] = end_token_index

            # Add spec_token_ids to token_ids_cpu.
            self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens)
            # Restore scheduler-side draft count after ngram trimming.
            if original_num_spec_per_req:
                orig = original_num_spec_per_req.get(req_id, 0)
                if orig != req_state.prev_num_draft_len:
                    req_state.prev_num_draft_len = orig

        # Add the new or resumed requests to the persistent batch.
        # The smaller empty indices are filled first.
        for request in reqs_to_add:
            self.input_batch.add_request(request)
            self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens)

        # Condense the batched states if there are gaps left by removed requests
        self.input_batch.condense()
        # Allow attention backend to reorder the batch, potentially
        self._may_reorder_batch(scheduler_output)
        # Refresh batch metadata with any pending updates.
        self.input_batch.refresh_metadata()

        # Incrementally update ngram_gpu tensors after batch is stable
        if is_ngram_gpu:
            update_ngram_gpu_tensors_incremental(
                self.input_batch,
                self.token_ids_gpu_tensor,
                self.num_tokens_no_spec_gpu,
                ngram_gpu_new_reqs,
                self.device,
                _pinned_idx_buf=self._ngram_pinned_idx_buf,
                _pinned_val_buf=self._ngram_pinned_val_buf,
            )

        if deferred_spec_decode_corrections:

            def correct_spec_decode_token_counts():
                valid_sampled_token_count = self._get_valid_sampled_token_count()
                if not valid_sampled_token_count:
                    return
                prev_req_id_to_index = self.input_batch.prev_req_id_to_index
                if not prev_req_id_to_index:
                    return
                for (
                    req_id,
                    optimistic_num_accepted,
                    req_state,
                ) in deferred_spec_decode_corrections:
                    prev_req_index = prev_req_id_to_index.get(req_id)
                    if prev_req_index is None:
                        continue
                    num_accepted = valid_sampled_token_count[prev_req_index] - 1
                    correction = optimistic_num_accepted - num_accepted
                    req_state.num_computed_tokens -= correction
                    cur_req_index = self.input_batch.req_id_to_index.get(req_id)
                    if cur_req_index is None:
                        continue
                    self.input_batch.num_computed_tokens_cpu[cur_req_index] -= correction
                    if is_ngram_gpu and correction > 0:
                        self.input_batch.num_tokens_no_spec[cur_req_index] -= correction
                        self.num_tokens_no_spec_gpu[cur_req_index] -= correction

            return correct_spec_decode_token_counts
        else:
            return None

    @torch.inference_mode()
    def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict:
        if (
            hasattr(self.model, "have_multimodal_outputs")
            and self.model.have_multimodal_outputs
            and isinstance(hidden_states, OmniOutput)
        ):
            text_hidden_states = hidden_states.text_hidden_states
            multimodal_outputs = hidden_states.multimodal_outputs

        elif isinstance(hidden_states, torch.Tensor):
            text_hidden_states = hidden_states
            multimodal_outputs = {}
        elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple):
            text_hidden_states = hidden_states[0]
            multimodal_outputs = {}
        else:
            raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
        return text_hidden_states, multimodal_outputs

    @torch.inference_mode()
    def _dummy_run(
        self,
        num_tokens: int,
        cudagraph_runtime_mode: CUDAGraphMode | None = None,
        force_attention: bool = False,
        uniform_decode: bool = False,
        allow_microbatching: bool = True,
        skip_eplb: bool = False,
        is_profile: bool = False,
        create_mixed_batch: bool = False,
        remove_lora: bool = True,
        is_graph_capturing: bool = False,
        num_active_loras: int = 0,
        profile_seq_lens: int | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Run a dummy forward pass to warm up/profile run or capture the
        CUDA graph for the model.

        Args:
            num_tokens: Number of tokens to run the dummy forward pass.
            cudagraph_runtime_mode: used to control the behavior.
                - if not set will determine the cudagraph mode based on using
                    the self.cudagraph_dispatcher.
                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
                    needed.
            force_attention: If True, always create attention metadata. Used to
                warm up attention backend when mode is NONE.
            uniform_decode: If True, the batch is a uniform decode batch.
            skip_eplb: If True, skip EPLB state update.
            is_profile: If True, this is a profile run.
            create_mixed_batch: If True, create a mixed batch with both decode
                (1 token) and prefill (multiple tokens) requests.
            remove_lora: If False, dummy LoRAs are not destroyed after the run
            num_active_loras: Number of distinct active LoRAs to capture for.
                LoRA is activated when num_active_loras > 0.
            profile_seq_lens: If provided, use this value for seq_lens instead
                of max_query_len. Used to profile attention workspace that
                scales with context length.
        """
        mm_config = self.vllm_config.model_config.multimodal_config
        if mm_config and mm_config.mm_encoder_only:
            # The current dummy run only covers LM execution, so we can skip it.
            # mm encoder dummy run may need to add in the future.
            return torch.tensor([]), torch.tensor([])

        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode()

        # If cudagraph_mode.decode_mode() == FULL and
        # cudagraph_mode.separate_routine(). This means that we are using
        # different graphs and/or modes for mixed prefill-decode batches vs.
        # uniform decode batches. A uniform decode batch means that all
        # requests have identical query length, except a potential virtual
        # request (shorter) in the batch account for padding.
        # Uniform decode batch could either be common pure decode, where
        # max_query_len == 1, or speculative decode, where
        # max_query_len == 1 + num_spec_decode_tokens.

        # When setting max_query_len = 1, we switch to and capture the optimized
        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
        # for GQA/MQA.
        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens

        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
        # for dummy run with LoRA so that the num_reqs collectively
        # has num_tokens in total.
        assert num_tokens <= self.max_num_tokens
        max_num_reqs = self.scheduler_config.max_num_seqs
        if create_mixed_batch:
            assert not uniform_decode
            # Create mixed batch:
            # first half decode tokens, second half one prefill
            num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2)
            num_prefill_tokens = num_tokens - num_decode_tokens
            num_reqs = num_decode_tokens + 1

            # Create decode requests (1 token each) followed by prefill request
            num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
            # Note: Overriding max_query_len to be the prefill tokens
            max_query_len = num_prefill_tokens
        elif uniform_decode:
            assert not create_mixed_batch
            num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))
            num_scheduled_tokens_list = [max_query_len] * num_reqs
            if num_tokens % max_query_len != 0:
                num_scheduled_tokens_list[-1] = num_tokens % max_query_len
        else:
            num_reqs = min(num_tokens, max_num_reqs)
            min_tokens_per_req = num_tokens // num_reqs
            num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
            num_scheduled_tokens_list[-1] += num_tokens % num_reqs

        assert sum(num_scheduled_tokens_list) == num_tokens
        assert len(num_scheduled_tokens_list) == num_reqs
        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
        num_tokens_unpadded = int(num_scheduled_tokens.sum())

        num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)

        _cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = (
            self._determine_batch_execution_and_padding(
                num_tokens=num_tokens_unpadded,
                num_reqs=num_reqs,
                num_scheduled_tokens_np=num_scheduled_tokens,
                max_num_scheduled_tokens=max_query_len,
                use_cascade_attn=False,
                allow_microbatching=allow_microbatching,
                force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
                # `force_uniform_decode` is used for cudagraph capture; because for
                # capturing mixed prefill-decode batches, we sometimes use
                # num_tokens == num_reqs which looks like a uniform decode batch to the
                # dispatcher; but we actually want to capture a piecewise cudagraph
                force_uniform_decode=uniform_decode,
                # `force_has_lora` is used for cudagraph capture; because LoRA is
                # activated later in the context manager, but we need to know the
                # LoRA state when determining the batch descriptor for capture
                force_has_lora=num_active_loras > 0,
                # `force_num_active_loras` is used for cudagraph capture; because we
                # need to capture graphs for specific num_active_loras counts
                force_num_active_loras=num_active_loras,
            )
        )

        if cudagraph_runtime_mode is None:
            cudagraph_runtime_mode = _cudagraph_mode
        else:
            assert cudagraph_runtime_mode == _cudagraph_mode, (
                f"Cudagraph runtime mode mismatch in dummy_run. "
                f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}."
            )

        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
        ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
            should_ubatch,
            num_scheduled_tokens,
            num_tokens_padded,
            num_reqs_padded,
            self.vllm_config.parallel_config.num_ubatches,
        )
        logger.debug(
            "ubatch_slices: %s, ubatch_slices_padded: %s",
            ubatch_slices,
            ubatch_slices_padded,
        )

        attn_metadata: PerLayerAttnMetadata | None = None

        slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
            num_tokens_padded=num_tokens_padded,
            num_reqs_padded=num_reqs_padded,
            num_tokens_unpadded=num_tokens_unpadded,
            ubatch_slices=ubatch_slices_padded,
        )

        if slot_mappings_by_group is not None:
            for sm in slot_mappings_by_group.values():
                sm.fill_(-1)

        with self.synchronize_input_prep():
            # If force_attention is True, we always capture attention.
            # Otherwise, it only happens for cudagraph_runtime_mode=FULL.
            if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
                if profile_seq_lens is not None:
                    seq_lens = profile_seq_lens  # type: ignore[assignment]
                elif create_mixed_batch:
                    seq_lens = torch.tensor(  # type: ignore[assignment]
                        [1] * num_decode_tokens + [num_prefill_tokens + 1],
                        dtype=torch.int,
                    )
                else:
                    seq_lens = max_query_len  # type: ignore[assignment]
                self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens
                self.optimistic_seq_lens_cpu[num_reqs:].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, self.query_pos.np)
                self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
                self.query_start_loc.copy_to_gpu()

                self.input_batch.block_table.commit_block_table(num_reqs_padded)

                pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
                attn_metadata, _ = self._build_attention_metadata(
                    num_tokens=num_tokens_unpadded,
                    num_tokens_padded=num_tokens_padded if pad_attn else None,
                    num_reqs=num_reqs_padded,
                    max_query_len=max_query_len,
                    ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices),
                    for_cudagraph_capture=is_graph_capturing,
                    slot_mappings=slot_mappings_by_group,
                    use_spec_decode=self.speculative_config is not None,
                )
                self._maybe_attach_attention_metadata_extensions(
                    attn_metadata=attn_metadata,
                    num_reqs=num_reqs_padded,
                    num_reqs_padded=num_reqs_padded,
                    max_query_len=max_query_len,
                    pad_attn=True,
                    for_cudagraph_capture=is_graph_capturing,
                )

        with self.maybe_dummy_run_with_lora(
            self.lora_config,
            num_scheduled_tokens,
            num_sampled_tokens,
            remove_lora,
            num_active_loras,
        ):
            # Make sure padding doesn't exceed max_num_tokens
            assert num_tokens_padded <= self.max_num_tokens
            model_kwargs = self._init_model_kwargs()
            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
                input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)

                model_kwargs = {
                    **model_kwargs,
                    **self._dummy_mm_kwargs(num_reqs),
                }
            elif self.enable_prompt_embeds:
                input_ids = None
                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
                model_kwargs = self._init_model_kwargs()
            elif getattr(getattr(self, "model", None), "has_preprocess", False):
                # Capture CUDA graph with inputs_embeds path so replay reads
                # from the same buffer that _preprocess writes into.
                input_ids = self.input_ids.gpu[:num_tokens_padded]
                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
            else:
                input_ids = self.input_ids.gpu[:num_tokens_padded]
                inputs_embeds = None

            if self.uses_mrope:
                positions = self.mrope_positions.gpu[:, :num_tokens_padded]
            elif self.uses_xdrope_dim > 0:
                positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
            else:
                positions = self.positions[:num_tokens_padded]

            if get_pp_group().is_first_rank:
                intermediate_tensors = None
            else:
                if self.intermediate_tensors is None:
                    self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
                        batch_size=self.max_num_tokens,
                        dtype=self.model_config.dtype,
                        device=self.device,
                    )

                intermediate_tensors = self.sync_and_gather_intermediate_tensors(num_tokens_padded, None, False)

            if ubatch_slices_padded is not None:
                # Adjust values to reflect a single ubatch.
                # TODO(sage,lucas): this is cruft that should be addressed in
                #  the padding refactor.
                num_tokens_padded = ubatch_slices_padded[0].num_tokens
                if num_tokens_across_dp is not None:
                    num_tokens_across_dp[:] = num_tokens_padded

            with (
                self.maybe_randomize_inputs(input_ids, inputs_embeds),
                set_forward_context(
                    attn_metadata,
                    self.vllm_config,
                    num_tokens=num_tokens_padded,
                    num_tokens_across_dp=num_tokens_across_dp,
                    cudagraph_runtime_mode=cudagraph_runtime_mode,
                    batch_descriptor=batch_desc,
                    ubatch_slices=ubatch_slices_padded,
                    slot_mapping=slot_mappings,
                ),
            ):
                if getattr(self.model, "talker", None) is not None and self.has_talker_mtp:
                    num_tokens_padded_talker_mtp = num_tokens_padded
                    if num_tokens_padded_talker_mtp == self.max_num_tokens:
                        num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
                    outputs = self.talker_mtp(
                        self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
                        self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
                        self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp],
                        self.text_step.gpu[:num_tokens_padded_talker_mtp],
                    )
                    self.compilation_config.cache_dir = None
                outputs = self.model(
                    input_ids=input_ids,
                    positions=positions,
                    intermediate_tensors=intermediate_tensors,
                    inputs_embeds=inputs_embeds,
                    **model_kwargs,
                )

            if self.use_aux_hidden_state_outputs:
                hidden_states, _ = outputs
            else:
                hidden_states = outputs
            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
            if self.speculative_config and (
                self.speculative_config.use_eagle()
                or self.speculative_config.uses_draft_model()
                or self.speculative_config.uses_extract_hidden_states()
            ):
                assert isinstance(
                    self.drafter,
                    EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer,
                )
                assert self.speculative_config is not None
                # Eagle currently only supports PIECEWISE cudagraphs.
                # Therefore only use cudagraphs if the main model uses PIECEWISE
                # NOTE(lucas): this is a hack, need to clean up.
                use_cudagraphs = (
                    (is_graph_capturing and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE)
                    or (not is_graph_capturing and cudagraph_runtime_mode != CUDAGraphMode.NONE)
                ) and not self.speculative_config.enforce_eager

                # Note(gnovack) - We need to disable cudagraphs for one of the two
                # lora cases when cudagraph_specialize_lora is enabled. This is a
                # short term mitigation for issue mentioned in
                # https://github.com/vllm-project/vllm/issues/28334
                if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0:
                    use_cudagraphs = False

                self.drafter.dummy_run(
                    num_tokens,
                    use_cudagraphs=use_cudagraphs,
                    is_graph_capturing=is_graph_capturing,
                    slot_mappings=slot_mappings,
                )

        # We register layerwise NVTX hooks here after the first dynamo tracing is
        # done to avoid nvtx operations in hook functions being traced by
        # torch dynamo and causing graph breaks.
        # Note that for DYNAMO_ONCE and VLLM_COMPILE mode,
        # compiled model's dynamo tracing is only done once and the compiled model's
        # __call__ function is replaced by calling the compiled function.
        # So it's safe to register hooks here. Hooks will be registered to
        # both compiled and uncompiled models but they will never
        # be called on the compiled model execution path.
        self._register_layerwise_nvtx_hooks()

        # This is necessary to avoid blocking DP.
        # For dummy runs, we typically skip EPLB since we don't have any real
        # requests to process.
        # However, in DP settings, there may be cases when some DP ranks do
        # not have any requests to process, so they're executing dummy batches.
        # In such cases, we still have to trigger EPLB to make sure
        # ranks execute the rearrangement in synchronization.
        if not skip_eplb:
            self.eplb_step(is_dummy=True, is_profile=is_profile)

        logit_indices = np.cumsum(num_scheduled_tokens) - 1
        logit_indices_device = torch.from_numpy(logit_indices).to(self.device, non_blocking=True)
        return hidden_states, hidden_states[logit_indices_device]

    # ------------------------------------------------------------------
    # Payload decoding helpers (torch.Tensor passthrough + legacy
    # PromptEmbedsPayload / AdditionalInformationPayload support)
    # ------------------------------------------------------------------

    @staticmethod
    def _resolve_prompt_embeds_cpu(
        pe: "torch.Tensor | object | None",
    ) -> torch.Tensor | None:
        """Convert *prompt_embeds* to a contiguous CPU tensor.

        Accepts:
        - ``torch.Tensor`` – moved to CPU as-is (the normal path after
          upstream added ``prompt_embeds`` to ``EngineCoreRequest``).
        - Legacy ``PromptEmbedsPayload`` (or any duck-typed object with
          ``.data``, ``.shape``, ``.dtype``) – decoded via numpy.
        - ``None`` – returns ``None``.
        """
        if pe is None:
            return None
        try:
            if isinstance(pe, torch.Tensor):
                return pe.detach().cpu().contiguous()
            data = getattr(pe, "data", None)
            shape = getattr(pe, "shape", None)
            if data is not None and shape is not None:
                dt = np.dtype(getattr(pe, "dtype", "float32"))
                arr = np.frombuffer(data, dtype=dt).reshape(shape)
                return torch.from_numpy(arr.copy())
        except Exception:
            logger.exception("Failed to decode prompt_embeds payload")
        return None

    def _decode_and_store_request_payloads(
        self,
        scheduler_output: "SchedulerOutput",
    ) -> None:
        """Decode per-request prompt_embeds and additional_information for
        newly scheduled requests and store them on CPU in the request state.
        """
        new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
        if not new_reqs:
            return
        for nr in new_reqs:
            req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
            if req_id is None or req_id not in self.requests:
                continue
            pe_cpu = self._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None))
            if pe_cpu is not None:
                setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
            info_payload = getattr(nr, "additional_information", None)
            if info_payload is not None:
                logger.warning_once(
                    "additional_information on request data is deprecated, use model_intermediate_buffer"
                )
            info_dict = deserialize_additional_information(info_payload)
            if info_dict:
                self.model_intermediate_buffer[req_id] = info_dict
                setattr(self.requests[req_id], "additional_information_cpu", info_dict)

    def _gather_runtime_additional_information(self) -> list[dict]:
        """Gather per-request model_intermediate_buffer in batch order."""
        per_req_runtime_info = []
        for req_id in self.input_batch.req_ids:
            req_state = self.requests.get(req_id)
            # MammothModa2 AR grid constraint: the model must emit a special
            # end-of-line (EOL) token at the end of each image row.  To determine
            # whether the current decoding step falls on a row boundary, the
            # constraint logic (see MammothModa2ARForConditionalGeneration.
            # _apply_t2i_token_constraints) computes:
            #   column_id = generated_len % (ar_width + 1)
            # and forces the EOL token when column_id == ar_width.
            generated_len = len(req_state.output_token_ids) if req_state is not None else 0
            info = self.model_intermediate_buffer.get(req_id, {})
            if info:
                info["generated_len"] = generated_len
                per_req_runtime_info.append(info)
                if "thinker_reply_part_per_request" in info:
                    q = info["thinker_reply_part_per_request"]
                    if hasattr(q, "shape"):
                        logger.debug(f"[OMNI] req={req_id} has thinker_reply_part_per_request queue shape: {q.shape}")
            else:
                per_req_runtime_info.append({})
        return per_req_runtime_info

    def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]:
        """Compute (start, end) token spans for each request within the flattened step sequence."""
        req_token_spans: list[tuple[int, int]] = []
        for req_index in range(len(self.input_batch.req_ids)):
            start_offset = int(self.query_start_loc.cpu[req_index])
            sched_tokens = int(num_scheduled_tokens_np[req_index])
            req_token_spans.append((start_offset, start_offset + sched_tokens))
        return req_token_spans

    def _sync_local_stage_payloads(self) -> None:
        """Move received full-payload stage inputs into model_intermediate_buffer."""
        cache = getattr(self, "_local_stage_payload_cache", None)
        if not cache:
            return
        lock = getattr(self, "_lock", None)
        ctx = lock if lock is not None else contextlib.nullcontext()
        with ctx:
            if not cache:
                return
            active_req_ids = set(getattr(self, "requests", {}))
            pending = set(getattr(self, "_full_payload_pending_broadcast_req_ids", set()))
            staged = {
                req_id: payload
                for req_id, payload in cache.items()
                if req_id not in pending and req_id in active_req_ids and isinstance(payload, dict)
            }
            for req_id in staged:
                cache.pop(req_id, None)
        for req_id, payload in staged.items():
            self._update_intermediate_buffer(req_id, payload)

    def _build_model_kwargs_extra(self) -> dict:
        """Build extra keyword arguments passed to the model for this step."""
        self._sync_local_stage_payloads()
        model_kwargs_extra: dict[str, object] = {}
        try:
            buffer_map = self._gather_runtime_additional_information()
            model_kwargs_extra["model_intermediate_buffer"] = buffer_map
            # Backward compatible: also emit old name
            model_kwargs_extra["runtime_additional_information"] = buffer_map
        except Exception as e:
            logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}")
            import traceback

            traceback.print_exc()

        if getattr(self.model_config, "has_sampling_extra_args", False):
            extra_args_list: list[dict] = []
            for req_id in self.input_batch.req_ids:
                req = self.requests[req_id]
                sp = req.sampling_params if req else None
                extra_args_list.append(sp.extra_args if sp and sp.extra_args else {})
            model_kwargs_extra["sampling_extra_args"] = extra_args_list

        return model_kwargs_extra

    def _process_additional_information_updates(
        self,
        hidden_states: torch.Tensor,
        multimodal_outputs: object,
        num_scheduled_tokens_np: np.ndarray,
        scheduler_output: "SchedulerOutput",
        combined_hidden_states: dict[str, torch.Tensor] | None = None,
        combined_multimodal_outputs: dict[str, object] | None = None,
        req_ids_filter: set[str] | None = None,
    ) -> None:
        """Process model-provided per-request updates and merge into model_intermediate_buffer."""
        try:
            # execute the custom postprocess function
            # TODO(Peiqi): do we have a more elegant way to do this?
            if hasattr(self.model, "has_postprocess") and self.model.has_postprocess:
                for req_index, req_id in enumerate(self.input_batch.req_ids):
                    if req_ids_filter is not None and req_id not in req_ids_filter:
                        continue
                    req_infos = self.model_intermediate_buffer.get(req_id, {})
                    if combined_hidden_states:
                        # Combined hidden states contains all hidden states for every request
                        hidden_states_slice = combined_hidden_states[req_id]
                    else:
                        start_offset = int(self.query_start_loc.cpu[req_index])
                        sched_tokens = int(num_scheduled_tokens_np[req_index])
                        s, e = start_offset, start_offset + sched_tokens
                        # only consider to store data into update dict.
                        hidden_states_slice = hidden_states[s:e]

                    if combined_multimodal_outputs:
                        # NOTE this is a bit ugly, but the mm data is structured as a list of
                        # keys mapping to request IDs, and if enabled, we will always have all
                        # request IDs in every subdict, including for cache misses.
                        mm_out = {k: v[req_id] for k, v in combined_multimodal_outputs.items()}
                    else:
                        mm_out = multimodal_outputs
                    # Exclude 'hidden_states' from kwargs to avoid clash with
                    # the positional arg. The buffer entry must be preserved
                    # because preprocess reads hidden_states['last'] from it.
                    # TODO: pass req_infos as a single payload arg instead of **unpacking
                    # to avoid key collisions with positional args.
                    postprocess_kwargs = {k: v for k, v in req_infos.items() if k != "hidden_states"}
                    update_dict = self.model.postprocess(
                        hidden_states_slice,
                        multimodal_outputs=mm_out,
                        **postprocess_kwargs,
                    )
                    self._update_intermediate_buffer(req_id, update_dict)
        except Exception as e:
            logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}")
            import traceback

            traceback.print_exc()

    def _collect_additional_information_for_prefill(
        self,
        num_scheduled_tokens_np: np.ndarray,
    ) -> dict[str, dict]:
        """Overlay per-request prompt_embeds for the prefill portion and collect
        additional_information slices for this step. Returns a map req_id -> dict."""
        for req_index, req_id in enumerate(self.input_batch.req_ids):
            req_state = self.requests[req_id]
            pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
            num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
            prompt_len = len(req_state.prompt_token_ids)
            prompt_remaining = max(0, prompt_len - num_computed_tokens)
            sched_tokens = int(num_scheduled_tokens_np[req_index])
            overlay_len = min(sched_tokens, prompt_remaining)
            if overlay_len <= 0:
                continue
            if overlay_len > 0 and pe_cpu is not None:
                src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
                    dtype=self.dtype, device=self.device, non_blocking=True
                )
                start_offset = int(self.query_start_loc.cpu[req_index])
                self.inputs_embeds.gpu[start_offset : start_offset + overlay_len].copy_(src)

    def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None:
        for new_req in scheduler_output.scheduled_new_reqs:
            payload_info = getattr(new_req, "additional_information", None)
            if isinstance(payload_info, dict):
                logger.warning_once(
                    "additional_information on request data is deprecated, use model_intermediate_buffer"
                )
                self._update_intermediate_buffer(new_req.req_id, payload_info)

        if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
            logger.warning_once(
                "additional_information on scheduled_cached_reqs is deprecated, use model_intermediate_buffer"
            )
            cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
            if isinstance(cached_infos, dict):
                for req_id, req_infos in cached_infos.items():
                    self._update_intermediate_buffer(req_id, req_infos)

    def _maybe_attach_mimo_audio_req_infos(
        self,
        req_state: CachedRequestState | None,
        req_infos: dict | None,
        req_id: str,
    ) -> dict | None:
        """Attach MiMoAudio-specific fields into req_infos if applicable.

        This helper is intentionally small and self-contained so that it can be
        unit-tested to prevent regressions when updating MiMoAudio handling.
        """
        if req_state is None or self.model.__class__.__name__ != "MiMoAudioForConditionalGeneration":
            return req_infos

        # Always operate on a dict copy to avoid mutating shared instances.
        req_infos = dict(req_infos) if isinstance(req_infos, dict) else {}
        mm_features = getattr(req_state, "mm_features", None)
        if mm_features and (not req_infos.get("mm_features")):
            req_infos["mm_features"] = mm_features
        req_infos["req_id"] = req_id

        return req_infos

    def _maybe_run_batch_preprocess(self, req_ids: list[str], device: torch.device) -> None:
        """Run an optional model-specific batch preprocess hook.

        The generic runner only supplies current request ids and the runner-owned
        intermediate buffer; model-specific code decides whether there is any
        batchable work.
        """
        preprocess_batch = getattr(self.model, "preprocess_batch", None)
        if not callable(preprocess_batch):
            return
        preprocess_batch(
            req_ids=req_ids,
            model_intermediate_buffer=self.model_intermediate_buffer,
            device=device,
        )

    def _preprocess(
        self,
        scheduler_output: "SchedulerOutput",
        num_input_tokens: int,
        intermediate_tensors: IntermediateTensors | None = None,
    ):
        """Align with v0.14.0 preprocess and omni's additional information handling."""
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        is_first_rank = get_pp_group().is_first_rank
        is_encoder_decoder = self.model_config.is_encoder_decoder

        # _prepare_inputs may reorder the batch, so we must gather multi
        # modal outputs after that to ensure the correct order
        ec_connector_output = None

        if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder:
            # Run the multimodal encoder if any.
            with self.maybe_get_ec_connector_output(
                scheduler_output,
                encoder_cache=self.encoder_cache,
            ) as ec_connector_output:
                self._execute_mm_encoder(scheduler_output)
                mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)

            # NOTE(woosuk): To unify token ids and soft tokens (vision
            # embeddings), we always use embeddings (rather than token ids)
            # as input to the multimodal model, even when the input is text.
            inputs_embeds_scheduled = self.model.embed_input_ids(
                self.input_ids.gpu[:num_scheduled_tokens],
                multimodal_embeddings=mm_embeds,
                is_multimodal=is_mm_embed,
            )

            # TODO(woosuk): Avoid the copy. Optimize.
            self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)

            input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens)
            model_kwargs = {
                **self._init_model_kwargs(),
                **self._extract_mm_kwargs(scheduler_output),
            }
        elif self.enable_prompt_embeds and is_first_rank:
            # Get the input embeddings for the tokens that are not input embeds,
            # then put them into the appropriate positions.
            # TODO(qthequartermasterman): Since even when prompt embeds are
            # enabled, (a) not all requests will use prompt embeds, and (b)
            # after the initial prompt is processed, the rest of the generated
            # tokens will be token ids, it is not desirable to have the
            # embedding layer outside of the CUDA graph all the time. The v0
            # engine avoids this by "double compiling" the CUDA graph, once
            # with input_ids and again with inputs_embeds, for all num_tokens.
            # If a batch only has token ids, then including the embedding layer
            # in the CUDA graph will be more performant (like in the else case
            # below).
            token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
            # Some tokens ids may need to become embeds
            if token_ids_idx.numel() > 0:
                token_ids = self.input_ids.gpu[token_ids_idx]
                tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids)
                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds

            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
            model_kwargs = self._init_model_kwargs()
            input_ids = None
        elif getattr(self.model, "has_preprocess", False):
            # Use pre-allocated buffer for CUDA graph compatibility.
            input_ids = self.input_ids.gpu[:num_input_tokens]
            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
            model_kwargs = self._init_model_kwargs()
        else:
            # For text-only models, we use token ids as input.
            # While it is possible to use embeddings as input just like the
            # multimodal models, it is not desirable for performance since
            # then the embedding layer is not included in the CUDA graph.
            input_ids = self.input_ids.gpu[:num_input_tokens]
            inputs_embeds = None
            model_kwargs = self._init_model_kwargs()

        if self.uses_mrope:
            positions = self.mrope_positions.gpu[:, :num_input_tokens]
        elif self.uses_xdrope_dim > 0:
            positions = self.xdrope_positions.gpu[:, :num_input_tokens]
        else:
            positions = self.positions[:num_input_tokens]
            if num_input_tokens > num_scheduled_tokens:
                self.positions[num_scheduled_tokens:num_input_tokens].zero_()

        if is_first_rank:
            intermediate_tensors = None
        else:
            assert intermediate_tensors is not None
            intermediate_tensors = self.sync_and_gather_intermediate_tensors(
                num_input_tokens, intermediate_tensors, True
            )

        if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
            # Run the encoder, just like we do with other multimodal inputs.
            # For an encoder-decoder model, our processing here is a bit
            # simpler, because the outputs are just passed to the decoder.
            # We are not doing any prompt replacement. We also will only
            # ever have a single encoder input.
            encoder_outputs = self._execute_mm_encoder(scheduler_output)
            model_kwargs.update({"encoder_outputs": encoder_outputs})

        req_ids = self.input_batch.req_ids
        num_scheduled_tokens_np = np.array(
            [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids],
            dtype=np.int32,
        )
        self._omni_num_scheduled_tokens_np = num_scheduled_tokens_np

        # Note: only prefill need collect additional_information for now.
        # Decode don't need per_req_additional_information anymore.
        if inputs_embeds is not None:
            # Prefill: overlay prompt_embeds and collect additional_information
            self._collect_additional_information_for_prefill(num_scheduled_tokens_np)

        # Keep per-request additional_information in sync for both new and
        # cached requests. This is required for stages without preprocess
        # (e.g., code2wav) so runtime_additional_information can be refreshed
        # from scheduler cached infos on every step.
        if hasattr(self.model, "has_preprocess") or hasattr(self.model, "enable_update_additional_information"):
            if self.vllm_config.model_config.async_chunk:
                self._update_additional_information(scheduler_output)
            else:
                # In full-payload (non-async-chunk) mode, connector-delivered
                # stage payloads must override any earlier engine-level
                # additional_information written by the legacy
                # custom_process_input_func codec, so talker_preprocess reads
                # the full thinker payload.
                self._sync_local_stage_payloads()

        if hasattr(self.model, "has_preprocess") and self.model.has_preprocess:
            preprocess_device = input_ids.device if input_ids is not None else inputs_embeds.device
            self._maybe_run_batch_preprocess(self.input_batch.req_ids, preprocess_device)

            # Overlay custom prompt_embeds per request for the prompt portion;
            # collect additional_information (tensor/list) for prefill portion only
            decode_req_ids = []
            decode_start_offsets = []
            decode_batch_items = []
            batch_decode_preprocess = getattr(self.model, "preprocess_decode_batch", None)

            def flush_decode_batch() -> None:
                nonlocal inputs_embeds
                if not decode_batch_items:
                    return

                req_ids_b = [item[0] for item in decode_batch_items]
                start_offsets_b = [item[1] for item in decode_batch_items]
                req_infos_b = [item[2] for item in decode_batch_items]
                ids_b = torch.stack([input_ids[offset : offset + 1].reshape(-1)[0] for offset in start_offsets_b])
                req_input_ids, req_embeds, last_talker_hidden, text_step, updates = batch_decode_preprocess(
                    input_ids=ids_b,
                    req_infos=req_infos_b,
                )
                if inputs_embeds is None:
                    inputs_embeds = torch.empty(
                        (input_ids.shape[0], req_embeds.shape[-1]),
                        device=req_embeds.device,
                        dtype=req_embeds.dtype,
                    )

                offsets_t = torch.tensor(start_offsets_b, device=req_embeds.device, dtype=torch.long)
                inputs_embeds.index_copy_(0, offsets_t, req_embeds)
                input_ids.index_copy_(0, offsets_t, req_input_ids.reshape(-1).to(dtype=input_ids.dtype))

                dst = slice(len(decode_req_ids), len(decode_req_ids) + len(req_ids_b))
                self.talker_mtp_input_ids.gpu[dst].copy_(req_input_ids.reshape(-1))
                self.talker_mtp_inputs_embeds.gpu[dst].copy_(req_embeds)
                self.last_talker_hidden.gpu[dst].copy_(last_talker_hidden)
                self.text_step.gpu[dst].copy_(text_step)

                for req_id_b, update_dict_b in zip(req_ids_b, updates, strict=True):
                    self._merge_additional_information_update(req_id_b, update_dict_b)

                decode_req_ids.extend(req_ids_b)
                decode_start_offsets.extend(start_offsets_b)
                decode_batch_items.clear()

            for req_index, req_id in enumerate(self.input_batch.req_ids):
                req_infos = self.model_intermediate_buffer.get(req_id, {})

                # mimo-audio check
                req_state = self.requests.get(req_id)
                req_infos = self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id)

                start_offset = int(self.query_start_loc.cpu[req_index])
                sched_tokens = int(num_scheduled_tokens_np[req_index])
                s, e = start_offset, start_offset + sched_tokens
                span_len = int(e) - int(s)

                # call the custom process function
                req_infos["request_id"] = req_id
                prompt_token_ids = getattr(req_state, "prompt_token_ids", ()) if req_state is not None else ()
                prompt_len = len(prompt_token_ids or ())
                num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
                is_prefill = num_computed_tokens < prompt_len
                req_infos["_omni_prompt_len"] = prompt_len
                req_infos["_omni_num_computed_tokens"] = num_computed_tokens
                req_infos["_omni_is_prefill"] = is_prefill
                if callable(batch_decode_preprocess) and self.has_talker_mtp and span_len == 1 and not is_prefill:
                    decode_batch_items.append((req_id, s, req_infos))
                    continue

                flush_decode_batch()

                embed_slice = inputs_embeds[s:e] if inputs_embeds is not None else None
                req_input_ids, req_embeds, update_dict = self.model.preprocess(
                    input_ids=input_ids[s:e], input_embeds=embed_slice, **req_infos
                )
                if inputs_embeds is None:
                    inputs_embeds = torch.empty(
                        (input_ids.shape[0], req_embeds.shape[-1]),
                        device=req_embeds.device,
                        dtype=req_embeds.dtype,
                    )

                if self.has_talker_mtp and span_len == 1 and not is_prefill:
                    last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
                    decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1)
                    self.talker_mtp_input_ids.gpu[decode_slice].copy_(req_input_ids)
                    self.talker_mtp_inputs_embeds.gpu[decode_slice].copy_(req_embeds)
                    self.last_talker_hidden.gpu[decode_slice].copy_(last_talker_hidden)
                    self.text_step.gpu[decode_slice].copy_(text_step)
                    decode_req_ids.append(req_id)
                    decode_start_offsets.append(s)

                # TODO(Peiqi): the merge stage could move out from the critical path
                self._merge_additional_information_update(req_id, update_dict)

                # update the inputs_embeds and input_ids
                seg_len = min(span_len, req_embeds.shape[0])
                inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
                if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
                    input_ids[s : s + seg_len] = req_input_ids

            flush_decode_batch()

            # run talker mtp decode
            if self.has_talker_mtp:
                self._talker_mtp_forward(decode_req_ids, inputs_embeds, decode_start_offsets)

        return (
            input_ids,
            inputs_embeds,
            positions,
            intermediate_tensors,
            model_kwargs,
            ec_connector_output,
        )

    def _talker_mtp_forward(
        self,
        decode_req_ids: list[str],
        inputs_embeds: torch.Tensor,
        start_offsets: list[int] | None = None,
    ) -> None:
        decode_batch_size = len(decode_req_ids)
        if decode_batch_size == 0:
            return
        _cudagraph_mode, batch_desc, _, _, _ = self._determine_batch_execution_and_padding(
            num_tokens=decode_batch_size,
            num_reqs=decode_batch_size,
            num_scheduled_tokens_np=np.ones(decode_batch_size, dtype=np.int32),
            max_num_scheduled_tokens=1,
            use_cascade_attn=False,
        )
        # Force eager for unwrapped code predictors (AR loops / multinomial).
        # When talker_mtp is not a CUDAGraphWrapper, it manages its own CUDA
        # graphs internally (code_predictor has its own bucket sizes).
        if not isinstance(self.talker_mtp, CUDAGraphWrapper):
            _cudagraph_mode = CUDAGraphMode.NONE
            num_tokens_padded = decode_batch_size
        else:
            num_tokens_padded = batch_desc.num_tokens
        req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded]
        req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded]
        last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded]
        text_step = self.text_step.gpu[:num_tokens_padded]
        subtalker_params = getattr(self.vllm_config.model_config, "subtalker_sampling_params", None)
        if not isinstance(subtalker_params, dict):
            subtalker_params = {}

        def _explicit_talker_seed(req_id: str) -> int | None:
            sampling_params = getattr(self.requests[req_id], "sampling_params", None)
            extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
            seed = extra_args.get("qwen3_tts_request_seed") if isinstance(extra_args, dict) else None
            return int(seed) if seed is not None else None

        if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids):
            # A torch.Generator is a single stream. Using one generator for a
            # multi-row batch would make explicitly-seeded requests depend on
            # other rows in the same scheduler step, so keep that path scalar.
            saved_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size].clone()
            saved_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].clone()
            saved_hidden = self.last_talker_hidden.gpu[:decode_batch_size].clone()
            saved_text = self.text_step.gpu[:decode_batch_size].clone()
            try:
                for row, req_id in enumerate(decode_req_ids):
                    self.talker_mtp_input_ids.gpu[:1].copy_(saved_input_ids[row : row + 1])
                    self.talker_mtp_inputs_embeds.gpu[:1].copy_(saved_embeds[row : row + 1])
                    self.last_talker_hidden.gpu[:1].copy_(saved_hidden[row : row + 1])
                    self.text_step.gpu[:1].copy_(saved_text[row : row + 1])
                    row_offsets = None if start_offsets is None else [start_offsets[row]]
                    self._talker_mtp_forward([req_id], inputs_embeds, row_offsets)
            finally:
                self.talker_mtp_input_ids.gpu[:decode_batch_size].copy_(saved_input_ids)
                self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].copy_(saved_embeds)
                self.last_talker_hidden.gpu[:decode_batch_size].copy_(saved_hidden)
                self.text_step.gpu[:decode_batch_size].copy_(saved_text)
            return

        generator = None
        if decode_req_ids:
            first_req_id = decode_req_ids[0]
            seed = _explicit_talker_seed(first_req_id)
            if seed is not None:
                generators = getattr(self, "_talker_mtp_generators", None)
                if generators is None:
                    generators = {}
                    self._talker_mtp_generators = generators
                generator = generators.get(first_req_id)
                if generator is None or generator.device != req_input_ids.device:
                    generator = torch.Generator(device=req_input_ids.device)
                    generator.manual_seed(int(seed))
                    generators[first_req_id] = generator
        talker_kwargs = {
            "do_sample": subtalker_params.get("do_sample"),
            "temperature": subtalker_params.get("temperature"),
            "top_k": subtalker_params.get("top_k"),
            "top_p": subtalker_params.get("top_p"),
        }
        if generator is not None:
            talker_kwargs["generator"] = generator
        with set_forward_context(
            None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
        ):
            req_embeds, code_predictor_codes = self.talker_mtp(
                req_input_ids,
                req_embeds,
                last_talker_hidden,
                text_step,
                **talker_kwargs,
            )
        # update the inputs_embeds and code_predictor_codes
        out_key = getattr(self.model, "talker_mtp_output_key", ("codes", "audio"))
        if not isinstance(out_key, tuple) or len(out_key) != 2:
            raise TypeError(f"talker_mtp_output_key must be a 2-tuple, got {type(out_key).__name__}: {out_key!r}")
        if start_offsets is None:
            start_offsets = []
            for req_id in decode_req_ids:
                req_index = self.input_batch.req_ids.index(req_id)
                start_offsets.append(int(self.query_start_loc.cpu[req_index]))
        for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)):
            inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
            update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
            self._merge_additional_information_update(req_id, update_dict)

    def _model_forward(
        self,
        input_ids: torch.Tensor | None = None,
        positions: torch.Tensor | None = None,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **model_kwargs: dict[str, Any],
    ):
        """Inject omni-specific kwargs into forward and cache model output"""
        model_kwargs_extra = self._build_model_kwargs_extra()

        model_output = super()._model_forward(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
            **model_kwargs,
            **model_kwargs_extra,
        )
        if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
            model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra)
        # Cache model output so later sample_tokens can consume multimodal results.
        self._omni_last_model_output = model_output
        return model_output

    def _store_value(self, dest: dict, key: str, value: Any, gpu_keys: set) -> None:
        if isinstance(value, torch.Tensor):
            if key in gpu_keys:
                dest[key] = value.detach().clone()
            else:
                dest[key] = value.detach().to("cpu").contiguous()
        elif isinstance(value, list):
            dest[key] = [
                (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in value
            ]
        else:
            dest[key] = value

    def _update_intermediate_buffer(self, req_id: str, upd: dict) -> None:
        if not isinstance(upd, dict) or not upd:
            return
        req_state = self.requests.get(req_id)
        if req_state is None:
            return
        # Check if the model declares keys that should stay on GPU (tuples of (type_key, qualifier))
        gpu_keys: set[tuple[str, str]] = set()
        if hasattr(self, "model") and hasattr(self.model, "gpu_resident_buffer_keys"):
            gpu_keys = self.model.gpu_resident_buffer_keys
        existing = self.model_intermediate_buffer.setdefault(req_id, {})
        for k, v in upd.items():
            if isinstance(v, dict):
                existing_sub = existing.setdefault(k, {})
                for qual, val in v.items():
                    self._store_value(existing_sub, qual, val, {q for tk, q in gpu_keys if tk == k})
            else:
                self._store_value(existing, k, v, set())
        # Backward compatible: mirror to old setattr location
        setattr(req_state, "additional_information_cpu", existing)

    def _merge_additional_information_update(self, req_id, upd):
        logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer")
        return self._update_intermediate_buffer(req_id, upd)

    def _update_streaming_input_additional_info(self, new_req_data, req_id):
        # For streaming input prefill case only. Update buffer from last segment input
        cached_additional_info = self.model_intermediate_buffer.get(req_id, {})
        if cached_additional_info:
            payload_info = getattr(new_req_data, "additional_information", None)
            inc_info = deserialize_additional_information(payload_info)
            if isinstance(inc_info, dict) and inc_info:
                accumulated_keys: set[tuple[str, str]] = set()
                if hasattr(self, "model") and hasattr(self.model, "streaming_accumulated_keys"):
                    accumulated_keys = self.model.streaming_accumulated_keys
                merged_info = dict(cached_additional_info)
                for key, value in inc_info.items():
                    if isinstance(value, dict):
                        existing_sub = merged_info.get(key)
                        merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
                        for sk, sv in value.items():
                            if (key, sk) in accumulated_keys and isinstance(sv, torch.Tensor):
                                inc_tensor = sv.detach().to("cpu").contiguous()
                                old_tensor = merged_sub.get(sk)
                                if old_tensor is None:
                                    merged_sub[sk] = inc_tensor
                                else:
                                    merged_sub[sk] = torch.cat((old_tensor, inc_tensor), dim=0)
                            else:
                                merged_sub[sk] = sv
                        merged_info[key] = merged_sub
                    else:
                        merged_info[key] = value
                merged_info.setdefault("meta", {})["num_processed_tokens"] = 0
                self.model_intermediate_buffer[req_id] = merged_info
                setattr(self.requests[req_id], "additional_information_cpu", merged_info)