diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index b6298c037..6404e56d9 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -67,7 +67,7 @@ from vllm.entrypoints.openai.speech_to_text.serving import ( ) from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.pooling.classify.serving import ServingClassification -from vllm.entrypoints.pooling.embed.serving import ServingEmbedding +from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling from vllm.entrypoints.pooling.score.serving import ServingScores from vllm.entrypoints.serve.disagg.serving import ServingTokens @@ -536,7 +536,13 @@ async def omni_init_app_state( else vllm_config.model_config ) io_processor_plugin = model_config.io_processor_plugin - engine_client.io_processor = get_io_processor(vllm_config, io_processor_plugin) + renderer = getattr(engine_client, "renderer", None) + if renderer is None: + from vllm.renderers import renderer_from_config + + renderer = renderer_from_config(vllm_config) + engine_client.renderer = renderer + engine_client.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) logger.info("Initialized io_processor for AsyncOmni") else: logger.warning("Cannot initialize processors: tokenizer is None. OpenAIServingModels may fail.") @@ -627,7 +633,7 @@ async def omni_init_app_state( else None ) state.openai_serving_embedding = ( - ServingEmbedding( + OpenAIServingEmbedding( engine_client, state.openai_serving_models, request_logger=request_logger, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 1424ca775..c1b71e076 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -773,6 +773,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, ) -> torch.Tensor: # This is to satisfy the type checker for each overload if multimodal_embeddings is None or is_multimodal is None: @@ -782,6 +783,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, ) def base_local_forward( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py index 8c348a9c6..210d262f1 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py @@ -109,6 +109,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration( multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, ) -> torch.Tensor: # This is to satisfy the type checker for each overload if multimodal_embeddings is None or is_multimodal is None: @@ -118,6 +119,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration( input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, ) def forward( @@ -230,7 +232,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration( return [] # The result multimodal_embeddings is tuple of tensors, with each - # tensor correspoending to a multimodal data item (image or video). + # tensor corresponding to a multimodal data item (image or video). multimodal_embeddings: tuple[torch.Tensor, ...] = () # NOTE: It is important to iterate over the keys in this dictionary diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index c32b121c2..4560ede39 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -58,6 +58,10 @@ from vllm.multimodal.inputs import ( MultiModalKwargsItems, ) from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import ( + ProcessorInputs, + TimingContext, +) from vllm.multimodal.processing.processor import ( MultiModalPromptUpdates, PlaceholderFeaturesInfo, @@ -76,6 +80,16 @@ class Qwen2_5OmniThinkerMultiModalProcessor( ): """Override to fix use_audio_in_video detection when mm cache returns None.""" + def _cached_apply_hf_processor( + self, + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ): + mm_processor_kwargs = inputs.hf_processor_mm_kwargs + if mm_processor_kwargs.get("use_audio_in_video", False): + return self._apply_hf_processor(inputs, timing_ctx) + return super()._cached_apply_hf_processor(inputs, timing_ctx) + def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, @@ -586,6 +600,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration( multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, ) -> torch.Tensor: if multimodal_embeddings is None or is_multimodal is None: return super().embed_input_ids(input_ids) @@ -594,6 +609,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration( input_ids, self.get_language_model().embed_input_ids, is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, ) if len(multimodal_embeddings) == 0: @@ -616,6 +632,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration( input_ids, self.get_language_model().embed_input_ids, is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, ) return merge_interleaved_embeddings( inputs_embeds, @@ -632,6 +649,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration( input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, ) def forward( diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 9f5c20cf7..7b3d0b017 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -341,7 +341,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( ) # The result multimodal_embeddings is tuple of tensors, with each - # tensor correspoending to a multimodal data item (image or video). + # tensor corresponding to a multimodal data item (image or video). dummy_multimodal_embeddings: tuple[torch.Tensor, ...] = () # NOTE: It is important to iterate over the keys in this dictionary diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index a0baefca7..ca7bb58c7 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -966,11 +966,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, ) -> torch.Tensor: inputs_embeds = self._embed_text_input_ids( input_ids, self.language_model.embed_input_ids, is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, ) if multimodal_embeddings is None or len(multimodal_embeddings) == 0: @@ -1062,6 +1064,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, ) def forward( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 9dce9778d..d7d45031a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -141,7 +141,7 @@ class GPUARModelRunner(OmniGPUModelRunner): # Update persistent batch states. self._update_states(scheduler_output) - if has_ec_transfer() and not get_ec_transfer().is_consumer: + if has_ec_transfer() and get_ec_transfer().is_producer: with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -276,9 +276,9 @@ class GPUARModelRunner(OmniGPUModelRunner): # 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 + # When spec decode is enabled, delay clearing connector metadata + # until after draft model runs in sample_tokens. + clear_kv_metadata = self.speculative_config is None with ( set_forward_context( attn_metadata, @@ -292,8 +292,7 @@ class GPUARModelRunner(OmniGPUModelRunner): ), record_function_or_nullcontext("gpu_model_runner: forward"), self.maybe_get_kv_connector_output( - scheduler_output, - defer_finalize=defer_kv_connector_finalize, + scheduler_output, clear_metadata=clear_kv_metadata ) as kv_connector_output, ): model_output = self._model_forward( @@ -534,11 +533,11 @@ class GPUARModelRunner(OmniGPUModelRunner): # tokens on the CPU, so they are run after bookkeeping. propose_draft_token_ids(valid_sampled_token_ids) - # Finalize KV connector (wait_for_save + clear metadata) after - # draft model runs. Deferred from target model forward to allow - # draft model to also save its KV cache. + # Clear KV connector metadata after draft model runs (if spec decode). + # This was deferred from target model forward to allow draft model + # to also save its KV cache. if self.speculative_config is not None: - self.finalize_kv_connector() + self.clear_kv_connector_metadata() with record_function_or_nullcontext("gpu_model_runner: eplb"): self.eplb_step() diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 5a3d2f9ca..1b1d3d9cf 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -238,7 +238,8 @@ class OmniGPUModelRunner(GPUModelRunner): self.requests.pop(req_id, None) self.model_intermediate_buffer.pop(req_id, None) self.num_prompt_logprobs.pop(req_id, None) - self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids) + if hasattr(self, "late_interaction_runner"): + 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 @@ -250,7 +251,7 @@ class OmniGPUModelRunner(GPUModelRunner): # 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: + if hasattr(scheduler_output, "new_block_ids_to_zero") and scheduler_output.new_block_ids_to_zero: self._zero_block_ids(scheduler_output.new_block_ids_to_zero) # Free the cached encoder outputs. @@ -320,7 +321,8 @@ class OmniGPUModelRunner(GPUModelRunner): lora_request=new_req_data.lora_request, ) self.requests[req_id] = req_state - self.late_interaction_runner.register_request(req_id, pooling_params) + if hasattr(self, "late_interaction_runner"): + self.late_interaction_runner.register_request(req_id, pooling_params) # If prompt embeddings are provided, decode and attach to inter_data try: