anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Refactor OpenAI serving components and enhance multimodal handling

Browse Source 
- Updated import for OpenAIServingEmbedding in api_server.py.
- Enhanced omni_init_app_state to initialize renderer for engine_client.
- Added handle_oov_mm_token parameter to multiple model classes for better multimodal token handling.
- Improved comments for clarity in various model files.
- Adjusted GPUARModelRunner to ensure proper handling of late interaction runner attributes.

Signed-off-by: [Your Name] <[Your Email]>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-03-13 01:55:17 +00:00
1 parent 72cb4f661b
commit e17fb0e97e
8 files changed
+50 -18
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/entrypoints/openai/api_server.py
+9 -3
View File
@@ -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,
vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
+2
View File
@@ -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(
vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py
+3 -1
View File
@@ -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
vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+18
View File
@@ -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(
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
+1 -1
View File
@@ -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
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
+3
View File
@@ -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(
vllm_omni/worker/gpu_ar_model_runner.py
+9 -10
View File
@@ -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()
vllm_omni/worker/gpu_model_runner.py
+5 -3
View File
@@ -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: