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)