# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Code2Wav GPU Model Runner for vLLM-Omni. Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav. This is a non-autoregressive model that doesn't require sampling or logits computation. """ from __future__ import annotations import gc import logging from collections.abc import Mapping from dataclasses import replace import numpy as np import torch from vllm.config import CUDAGraphMode from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.dflash import DFlashProposer from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer from vllm.v1.spec_decode.gemma4 import Gemma4Proposer from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( AsyncGPUModelRunnerOutput, PerLayerAttnMetadata, ) from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import partition_payload_list from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.mixins import maybe_unpad_input_ids from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, needs_omni_connector, ) logger = logging.getLogger(__name__) class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): """Generation model runner for vLLM-Omni (non-autoregressive). - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue. - Does not compute logits or perform token sampling. - Executes generation process and returns tensors via `pooler_output`. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._async_chunk = getattr(self.model_config, "async_chunk", False) if needs_omni_connector(self.model_config): self.init_omni_connectors( model_config=self.model_config, ) def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests for req_id in scheduler_output.finished_req_ids: self.input_batch.remove_request(req_id) 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 unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) for req_id in unscheduled_req_ids: self.input_batch.remove_request(req_id) cached_reqs = scheduler_output.scheduled_cached_reqs req_states = [] for req_id in cached_reqs.req_ids: req_state = self.requests.get(req_id) assert req_state is not None req_state.prompt_token_ids = cached_reqs.prompt_token_ids.get(req_id) req_states.append(req_state) # Remove the request from the current input batch only if it is still present. if req_id in self.input_batch.req_id_to_index: self.input_batch.remove_request(req_id) for req_state in req_states: # update the request state in self.input_batch self.input_batch.add_request(req_state) # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: self._init_mrope_positions(req_state) @torch.inference_mode() def execute_model( self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors: if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") if self.routed_experts_initialized: self.routed_experts_capturer.clear_buffer() if hasattr(self, "_omni_connector"): for request in getattr(scheduler_output, "pending_input_registrations", []): self.register_chunk_recv(request) self.recv_full_payload_inputs(scheduler_output) if self._pending_full_payload_send: flush_ids = set(getattr(scheduler_output, "finished_req_ids", set())) flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests}) if flush_ids: self.flush_full_payload_outputs(flush_ids) if self.routed_experts_initialized: capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() # If ngram_gpu is used, we need to copy the scheduler_output to avoid # the modification has influence on the scheduler_output in engine core process. # The replace is much faster than deepcopy. if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy() scheduler_output = replace( scheduler_output, num_scheduled_tokens=num_scheduled_tokens_copy, scheduled_spec_decode_tokens=spec_decode_tokens_copy, ) if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata if kv_connector_metadata is not None: get_kv_transfer_group().handle_preemptions(kv_connector_metadata) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with ( record_function_or_nullcontext("gpu_model_runner: preprocess"), self.synchronize_input_prep(), ): if self.model_config.async_chunk and num_scheduled_tokens: self._update_request_states(scheduler_output) deferred_state_corrections_fn = self._update_states(scheduler_output) # Notify stateful models of finished requests before the # zero-token early return. Request cleanup is a model lifecycle # concern and does not depend on the inter-stage transfer mode. if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): self.model.on_requests_finished(scheduler_output.finished_req_ids) if has_ec_transfer() and not get_ec_transfer().is_consumer: with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) return self.attach_omni_connector_output(make_empty_encoder_model_runner_output(scheduler_output)) # OMNI: keep this block in lock-step with the same block in # GPUARModelRunner.execute_model. # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). if num_scheduled_tokens <= 0: if ( self.parallel_config.distributed_executor_backend == "external_launcher" and self.parallel_config.data_parallel_size > 1 ): # this is a corner case when both external launcher # and DP are enabled, num_scheduled_tokens could be # 0, and has_unfinished_requests in the outer loop # returns True. before returning early here we call # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. self._dummy_run(1) if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) result = self.kv_connector_no_forward(scheduler_output, self.vllm_config) return self.attach_omni_connector_output(result) if self.cache_config.kv_sharing_fast_prefill: assert not self.num_prompt_logprobs, ( "--kv-sharing-fast-prefill produces incorrect " "logprobs for prompt tokens, tokens, please disable " "it when the requests need prompt logprobs" ) num_reqs = self.input_batch.num_reqs req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens logits_indices, spec_decode_metadata, max_num_sampled_tokens = self._prepare_inputs( scheduler_output, num_scheduled_tokens_np ) cascade_attn_prefix_lens = None # Disable cascade attention when using microbatching (DBO) if self.cascade_attn_enabled and not self.parallel_config.use_ubatching: # Pre-compute cascade attention prefix lengths cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens( num_scheduled_tokens_np, self.input_batch.num_computed_tokens_cpu[:num_reqs], scheduler_output.num_common_prefix_blocks, ) ( cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, cudagraph_stats, ) = self._determine_batch_execution_and_padding( num_tokens=num_tokens_unpadded, num_reqs=num_reqs, num_scheduled_tokens_np=num_scheduled_tokens_np, max_num_scheduled_tokens=max_num_scheduled_tokens, use_cascade_attn=cascade_attn_prefix_lens is not None, num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), ) logger.debug( "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " "should_ubatch: %s, num_tokens_across_dp: %s", cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, ) 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_np, num_tokens_padded, num_reqs_padded, self.parallel_config.num_ubatches, ) logger.debug( "ubatch_slices: %s, ubatch_slices_padded: %s", ubatch_slices, ubatch_slices_padded, ) pad_attn = cudagraph_mode == CUDAGraphMode.FULL use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices # OMNI: True if any attention backend handles KV cache update separately # from forward() (i.e., forward_includes_kv_cache_update=False). When true, # slot_mappings must use padded dimensions to match the key/value tensors. from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec has_separate_kv_update = not all( all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id]) for id, spec in enumerate(self.kv_cache_config.kv_cache_groups) if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec) ) slot_mappings_by_group, slot_mappings = self._get_slot_mappings( num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded, num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs), num_tokens_unpadded=num_tokens_unpadded, ubatch_slices=ubatch_slices_padded, ) attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, num_reqs=num_reqs, num_reqs_padded=num_reqs_padded if pad_attn else None, max_query_len=max_num_scheduled_tokens, ubatch_slices=ubatch_slices_attn, logits_indices=logits_indices, max_num_sampled_tokens=max_num_sampled_tokens, use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, slot_mappings=slot_mappings_by_group, ) self._maybe_attach_attention_metadata_extensions( attn_metadata=attn_metadata, num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, max_query_len=max_num_scheduled_tokens, pad_attn=pad_attn, num_scheduled_tokens_np=num_scheduled_tokens_np, ) ( input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ec_connector_output, ) = self._preprocess( scheduler_output, num_tokens_padded, intermediate_tensors, ) # [Omni] Exact-shape models need input_ids at its real token count (#6712). input_ids = maybe_unpad_input_ids(self.model, input_ids, num_tokens_unpadded) # [Omni] Pass token counts per request for code2wav output slicing model_kwargs["seq_token_counts"] = tokens if getattr(self.model, "requires_request_ids", False): model_kwargs["request_ids"] = list(req_ids) # Run the model. # Use persistent buffers for CUDA graphs. # When spec decode is enabled, defer connector finalization # (wait_for_save + clear metadata) until after draft model runs. defer_kv_connector_finalize = self.speculative_config is not None with ( 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_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output( scheduler_output, defer_finalize=defer_kv_connector_finalize, ) as kv_connector_output, ): outputs = self._run_generation_model( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, model_kwargs=model_kwargs, logits_indices=logits_indices, ) _, multimodal_outputs = self.extract_multimodal_outputs(outputs) self.execute_model_state = ExecuteModelState( scheduler_output, None, spec_decode_metadata, spec_decode_common_attn_metadata, None, None, None, None, ec_connector_output, cudagraph_stats, multimodal_outputs, slot_mappings, # OMNI: pass slot_mappings for upstream v1 API compatibility ) self.kv_connector_output = kv_connector_output if deferred_state_corrections_fn: deferred_state_corrections_fn() if self.routed_experts_initialized and hasattr(self, "_positions_cpu"): self._omni_routed_experts_d2h(scheduler_output) return None @torch.inference_mode() def sample_tokens( self, grammar_output: GrammarOutput | None = None, ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: # NOTE: Even though the model is non-autoregressive, we still need # this function to match the interface of the engine core. # In this case, this function kv_connector_output = self.kv_connector_output self.kv_connector_output = None if self.execute_model_state is None: # receive sampled token ids from the last PP rank. if self.use_async_scheduling and not get_pp_group().is_last_rank: self._pp_receive_prev_sampled_token_ids_to_input_batch() # In case of PP with kv transfer, we need to pass through the # kv_connector_output return self.attach_omni_connector_output( OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) ) # Unpack ephemeral state. ( scheduler_output, logits, spec_decode_metadata, spec_decode_common_attn_metadata, hidden_states, _hidden_states_cpu, sample_hidden_states, aux_hidden_states, ec_connector_output, cudagraph_stats, multimodal_outputs_raw, slot_mappings, # OMNI: unpack slot_mappings for upstream v1 API compatibility ) = self.execute_model_state self.execute_model_state = None # Finalize KV connector (wait_for_save + clear metadata) after # draft model runs. Deferred from target model forward. if self.speculative_config is not None: self.finalize_kv_connector() # Build per-request multimodal_outputs list (dedicated channel). # pooler_output is no longer used for multimodal data. per_req_payloads: list[dict[str, object]] = [] if isinstance(multimodal_outputs_raw, torch.Tensor): # One row per request. The old asserts (`shape[0] == 1` AND # `shape[0] == num_reqs`) jointly forced num_reqs == 1, silently # rejecting batched steps the generation scheduler does admit. num_reqs = self.input_batch.num_reqs if multimodal_outputs_raw.shape[0] != num_reqs: raise ValueError( f"Multimodal output tensor has leading dim {multimodal_outputs_raw.shape[0]} " f"but the batch has {num_reqs} requests (one row per request; " "to return multiple tensors per request, use a dict)." ) for i in range(num_reqs): per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) elif isinstance(multimodal_outputs_raw, list): # One entry per request. The old `len == 1` assert did not check # num_reqs, so a batched step built a length-1 payload list that # misaligned with `req_ids` downstream. num_reqs = self.input_batch.num_reqs if len(multimodal_outputs_raw) != num_reqs: raise ValueError( f"Multimodal output list has length {len(multimodal_outputs_raw)} " f"but the batch has {num_reqs} requests (one entry per request; " "to return multiple lists per request, use a dict)." ) for out in multimodal_outputs_raw: per_req_payloads.append( {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} ) elif isinstance(multimodal_outputs_raw, Mapping): num_reqs = self.input_batch.num_reqs for i in range(num_reqs): mm_payload = {} for key, out in multimodal_outputs_raw.items(): if isinstance(out, list): if len(out) != num_reqs: raise ValueError( f"Multimodal output list for key '{key}' has length {len(out)} " f"but expected {num_reqs} (one entry per request)." ) mm_payload[key] = out[i].detach().to("cpu").contiguous() elif isinstance(out, torch.Tensor): mm_payload[key] = out.detach().to("cpu").contiguous() else: logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") per_req_payloads.append(_ensure_tensor_values(mm_payload)) else: raise RuntimeError("Unsupported diffusion output type") if self._async_chunk: inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads) else: # See gpu_ar_model_runner: non-async-chunk ships the full payload to the next # stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792) inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() routed_experts_lists = None if self.routed_experts_initialized: routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) if inter_stage_outputs and self._should_accumulate_full_payload_output(): for i, rid in enumerate(req_ids_output_copy): req_state = self.requests.get(rid) if req_state is not None and inter_stage_outputs[i]: self.accumulate_full_payload_output(rid, inter_stage_outputs[i], req_state) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, sampled_token_ids=[], logprobs=None, prompt_logprobs_dict={}, pooler_output=None, multimodal_outputs=multimodal_outputs, inter_stage_outputs=inter_stage_outputs, kv_connector_output=kv_connector_output, num_nans_in_logits={}, cudagraph_stats=cudagraph_stats, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) output.omni_connector_output = self.get_omni_connector_output() output.routed_experts = routed_experts_lists if not self.use_async_scheduling: return output return AsyncGPUModelRunnerOutput( model_runner_output=output, sampled_token_ids=torch.tensor([], device=self.device), invalid_req_indices=[], async_output_copy_stream=self.async_output_copy_stream, vocab_size=self.input_batch.vocab_size, logprobs_tensors=None, ) def _run_generation_model( self, *, input_ids: torch.Tensor | None, positions: torch.Tensor | None, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None, model_kwargs: dict, logits_indices: torch.Tensor, ) -> torch.Tensor | list[torch.Tensor]: """Run generation from codec codes to waveforms. Args: scheduler_output: Contains codec codes in input_ids or additional info intermediate_tensors: PP intermediate tensors if applicable Returns: Audio waveforms: [batch, 1, waveform_len] or list of tensors """ # Keep inputs identical to AR runner kwargs = dict( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, **model_kwargs, sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, ) if hasattr(self.model, "forward"): return self._model_forward(**kwargs) raise RuntimeError( "The loaded model does not expose diffusion interfaces 'sample', " "'forward', or 'diffuse'. Please implement one of them or adapt the runner." ) @torch.inference_mode() def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> torch.Tensor: logger.warning("Dummy sampler run is not implemented for generation model") return torch.tensor([]) @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, randomize_inputs: bool = False, ) -> 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 # OMNI: Get slot mappings before building attention metadata 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.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1]) self.query_start_loc.copy_to_gpu() self.input_batch.block_table.commit_block_table(num_reqs_padded) 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 # Some generation-stage models (e.g. MammothModa2DiTPipeline) require # model-specific runtime information (such as image size and conditioning # embeddings) even during the dummy profiling run that vLLM uses to # estimate KV-cache capacity. get_dummy_runtime_additional_information # provides placeholder values of the correct shape so that the profiling # run does not raise an error due to missing inputs. if hasattr(self.model, "get_dummy_runtime_additional_information"): runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs) model_kwargs["runtime_additional_information"] = runtime_addi # Generation-stage code2wav models split flat input_ids using this # shared contract. Dummy/profile runs use padded synthetic inputs, # so expose one synthetic segment covering the actual tensor length. model_kwargs["seq_token_counts"] = [int(num_tokens_padded)] if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] elif getattr(self, "uses_xdrope_dim", 0) > 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 ( # Never randomize generation-stage dummy ids: they are # structured payloads (e.g. packed RVQ codec codes that get # code_offset added before embedding), so vocab-uniform random # ids index past the codebook tables — CUDA device-side assert # in code2wav warmup. The expert-balance rationale behind # randomize_inputs targets text-MoE models, which generation # stages are not. The kwarg stays accepted for vLLM API compat. 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, # OMNI: required for KV cache operations ), ): 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, # OMNI: pass slot_mappings (upstream v1 API) ) # 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) return hidden_states, None def profile_run(self) -> None: # Profile with multimodal encoder & encoder cache. if self.supports_mm_inputs: mm_config = self.model_config.multimodal_config if mm_config is not None and mm_config.skip_mm_profiling: logger.info("Skipping memory profiling for multimodal encoder and encoder cache.") else: mm_budget = self.mm_budget assert mm_budget is not None if (encoder_budget := mm_budget.get_encoder_budget()) > 0: if not mm_budget.mm_max_toks_per_item: # All modality limits are 0 — embedding-only mode. # Budget is non-zero for embedding storage, but # there's no encoder to profile. logger.info( "Skipping encoder profiling for embedding-only " "mode (all modality limits=0 with " "enable_mm_embeds=True).", ) else: # NOTE: Currently model is profiled with a single non-text # modality with the max possible input tokens even when # it supports multiple. dummy_modality = mm_budget.get_modality_with_max_tokens() max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[dummy_modality] logger.info( "Encoder cache will be initialized with a budget of " "%s tokens, and profiled with %s %s items of the " "maximum feature size.", encoder_budget, max_mm_items_per_batch, dummy_modality, ) # Create dummy batch of multimodal inputs. batched_dummy_mm_inputs = self._get_mm_dummy_batch( dummy_modality, max_mm_items_per_batch, ) # Run multimodal encoder. dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs) sanity_check_mm_encoder_outputs( dummy_encoder_outputs, expected_num_items=max_mm_items_per_batch, ) for i, output in enumerate(dummy_encoder_outputs): self.encoder_cache[f"tmp_{i}"] = output # Add `is_profile` here to pre-allocate communication buffers hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) output = None self._sync_device() del hidden_states self.encoder_cache.clear() gc.collect()