Update vLLM wheel version in Dockerfile and enhance input processing for omni models. Added functionality to reinject omni-specific fields into processed prompts and improved handling of multimodal audio extraction in serving_chat.py. Adjusted SharedMemoryConnector logic for payload size checks and refined input preprocessing methods for better compatibility with multimodal data.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

docker/Dockerfile.ci

@@ -14,11 +14,11 @@ RUN apt-get update && \
 
 # Install vLLM wheel for the specific commit
 # URL format: https://wheels.vllm.ai/<commit_hash>/
-# Current vLLM wheel commit: d5fe3f702c2f4392515cdfde6fd0442271c74dda
+# Current vLLM wheel commit: 16f24e87975ef4cd2c12879425062913ef62f6fd
 # Note: --prerelease=allow is required because nightly/commit wheels use
 # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
 # causing it to fall back to the older stable release from PyPI.
-RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/d5fe3f702c2f4392515cdfde6fd0442271c74dda
+RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/16f24e87975ef4cd2c12879425062913ef62f6fd
 
 
 RUN uv pip uninstall --system nvidia-cublas-cu12

vllm_omni/distributed/omni_connectors/adapter.py

@@ -35,12 +35,26 @@ def try_send_via_connector(
     try:
         t0 = time.time()
 
+        # Strip non-serializable multimodal feature fields from original_prompt
+        # before including it in metadata.  After stage-0 runs, the TokPrompt
+        # returned by render_chat_async may carry processed multimodal features
+        # (mm_kwargs, mm_placeholders, mm_hashes) that contain MultiModalKwargsItems
+        # objects, which are not supported by OmniMsgpackEncoder.  The receiving
+        # side (try_recv_via_connector) only extracts "engine_inputs" from the
+        # payload and never uses "original_prompt", so stripping these fields
+        # only affects debug metadata and is safe.
+        _MM_FEATURE_KEYS = frozenset({"mm_kwargs", "mm_placeholders", "mm_hashes"})
+        if isinstance(original_prompt, dict) and any(k in original_prompt for k in _MM_FEATURE_KEYS):
+            safe_prompt = {k: v for k, v in original_prompt.items() if k not in _MM_FEATURE_KEYS}
+        else:
+            safe_prompt = original_prompt
+
         # Prepare data for connector
         payload_data = {
             "engine_inputs": next_inputs,
             "sampling_params": sampling_params,
             "metadata": {
-                "original_prompt": original_prompt,
+                "original_prompt": safe_prompt,
                 "stage_transition": f"{stage_id}->{next_stage_id}",
                 "timestamp": time.time(),
             },

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -48,7 +48,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             payload = self.serialize_obj(data)
             size = len(payload)
 
-            if True:
+            if size > self.threshold:
                 # Use Shared Memory
                 lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"
                 with open(lock_file, "w") as lockf:

vllm_omni/engine/input_processor.py

@@ -1,6 +1,6 @@
 import time
 from collections.abc import Mapping
-from typing import Any, cast
+from typing import Any
 
 import numpy as np
 import torch
@@ -9,7 +9,7 @@ from vllm.inputs import ProcessorInputs, PromptType
 from vllm.inputs.parse import split_enc_dec_inputs
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict
+from vllm.multimodal.inputs import MultiModalFeatureSpec
 
 try:
     from vllm.multimodal.processing import set_request_id
@@ -22,8 +22,8 @@ except ImportError:  # vllm without set_request_id (older releases)
 
 
 from vllm.multimodal.utils import argsort_mm_positions
-from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
+from vllm.renderers import BaseRenderer
 from vllm.renderers.inputs import DictPrompt, TokPrompt
 from vllm.sampling_params import SamplingParams
 from vllm.tasks import SupportedTask
@@ -42,6 +42,35 @@ from vllm_omni.lora.request import LoRARequest
 
 logger = init_logger(__name__)
 
+_OMNI_EXTRA_KEYS = (
+    "additional_information",
+    "prompt_embeds",
+    "negative_prompt",
+    "negative_prompt_embeds",
+)
+
+
+def reinject_omni_fields(
+    results: list[ProcessorInputs],
+    original_prompts: list[dict],
+) -> None:
+    """Re-inject omni-specific fields that the upstream renderer discards.
+
+    The upstream renderer's ``process_for_engine`` creates new dicts that only
+    copy standard vLLM fields (prompt_token_ids, multi_modal_data, …).
+    Omni-specific fields such as ``additional_information`` and
+    ``prompt_embeds`` are silently dropped.  This helper copies them back from
+    the *original* parsed prompts into the renderer outputs so they survive
+    into ``OmniInputProcessor.process_inputs()``.
+    """
+    for result, orig in zip(results, original_prompts):
+        if not isinstance(orig, dict):
+            continue
+        for key in _OMNI_EXTRA_KEYS:
+            val = orig.get(key)
+            if val is not None and key not in result:
+                result[key] = val
+
 
 class OmniInputProcessor(InputProcessor):
     """Processor for omni models, handling multimodal inputs and embeddings.
@@ -88,9 +117,11 @@ class OmniInputProcessor(InputProcessor):
     def __init__(
         self,
         vllm_config: VllmConfig,
+        renderer: BaseRenderer | None = None,
+        *,
         mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
     ):
-        super().__init__(vllm_config, mm_registry=mm_registry)
+        super().__init__(vllm_config, renderer=renderer, mm_registry=mm_registry)
         self.input_preprocessor = OmniInputPreprocessor(
             vllm_config=vllm_config,
             renderer=self.renderer,
@@ -150,48 +181,24 @@ class OmniInputProcessor(InputProcessor):
         if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks):
             raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).")
 
-        if arrival_time is None:
-            arrival_time = time.time()
-
-        # Optionally generate multimodal hash overrides to avoid hashing
-        # multimodal data items by their content as their identifiers.
-
-        # NOTE: when users explicitly turn off BOTH prefix caching and input
-        # processing caching, no multimodal features or embeddings will be
-        # reused across requests, therefore identifying multimodal data items
-        # by their content is no longer necessary, and we create uuids with
-        # request id-modality-index as multimodal hash overrides.
-        if (
-            self.model_config.multimodal_config
-            and self.model_config.multimodal_config.mm_processor_cache_gb == 0
-            and not self.cache_config.enable_prefix_caching
-        ):
-            mm_uuids = self._maybe_build_mm_uuids(request_id, prompt)
+        # Short-circuit for prompts already processed by the renderer
+        # (they carry a "type" key).  Raw prompts must still go through the
+        # omni preprocessor which preserves additional_information, etc.
+        if isinstance(prompt, dict) and "type" in prompt:
+            if arrival_time is None:
+                arrival_time = prompt.get("arrival_time", time.time())
+            processed_inputs: ProcessorInputs = prompt  # type: ignore[assignment]
         else:
-            # Otherwise, use user-provided uuids as multimodal hash overrides
-            # if provided.
-            self._validate_mm_uuids(prompt)
-            if isinstance(prompt, dict):
-                mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids"))
-            else:
-                mm_uuids = None
+            if arrival_time is None:
+                arrival_time = time.time()
 
-        # Process inputs, which includes:
-        # 1. Tokenize text prompt, with LoRA request if one exists.
-        # 2. For multimodal models with a merged preprocessor, preprocess
-        #   multimodal data and expand prompt token ids accordingly.
-        with set_request_id(request_id), set_default_torch_num_threads():
-            processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
-                prompt,
-                tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
-            )
+            with set_request_id(request_id), set_default_torch_num_threads():
+                processed_inputs = self.input_preprocessor.preprocess(
+                    prompt,
+                    tokenization_kwargs=tokenization_kwargs,
+                )
 
-        current_platform.validate_request(
-            prompt=prompt,
-            params=params,
-            processed_inputs=processed_inputs,
-        )
+        self._platform_validate_request(processed_inputs, params)
 
         eos_token_id = self.renderer.get_eos_token_id()
 

vllm_omni/engine/output_processor.py

@@ -259,8 +259,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def __init__(
         self,
         tokenizer: TokenizerLike | None,
+        *,
         log_stats: bool,
         stream_interval: int = 1,
+        tracing_enabled: bool = False,
         engine_core_output_type: str | None = None,
     ):
         """Initialize the multimodal output processor.
@@ -273,7 +275,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 (e.g., "image", "audio", "latents"). Used to route outputs
                 to appropriate processors. If None, output type is inferred.
         """
-        super().__init__(tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval)
+        super().__init__(
+            tokenizer=tokenizer,
+            log_stats=log_stats,
+            stream_interval=stream_interval,
+            tracing_enabled=tracing_enabled,
+        )
         self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
         self._reqid_to_mm_type: dict[str, str] = {}
         self.engine_core_output_type = engine_core_output_type

vllm_omni/entrypoints/omni_llm.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import Any
 
 import cloudpickle
@@ -8,9 +8,11 @@ from tqdm import tqdm
 # External library imports (vLLM)
 from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
 from vllm.entrypoints.llm import LLM
+from vllm.inputs import ProcessorInputs, PromptType
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput, RequestOutput
 from vllm.plugins.io_processors import get_io_processor
+from vllm.renderers.inputs.preprocess import parse_model_prompt
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.counter import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
@@ -19,7 +21,7 @@ from vllm_omni.distributed.omni_connectors import initialize_orchestrator_connec
 
 # Internal imports (our code)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
-from vllm_omni.engine.input_processor import OmniInputProcessor
+from vllm_omni.engine.input_processor import OmniInputProcessor, reinject_omni_fields
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.entrypoints.utils import (
     load_stage_configs_from_model,
@@ -182,6 +184,27 @@ class OmniLLM(LLM):
         self.model_config = self.llm_engine.model_config
         self.input_processor = self.llm_engine.input_processor
 
+    # ------------------------------------------------------------------
+    # Override upstream _preprocess_cmpl so that omni-specific fields
+    # (additional_information, prompt_embeds, …) survive the renderer's
+    # process_for_engine step which only copies standard vLLM keys.
+    # ------------------------------------------------------------------
+
+    def _preprocess_cmpl(
+        self,
+        prompts: Sequence[PromptType],
+        tokenization_kwargs: dict[str, Any] | None = None,
+    ) -> Sequence[ProcessorInputs]:
+        renderer = self.renderer
+        model_config = self.model_config
+
+        parsed_prompts = [parse_model_prompt(model_config, prompt) for prompt in prompts]
+        tok_params = renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {}))
+        results = renderer.render_cmpl(parsed_prompts, tok_params)
+
+        reinject_omni_fields(results, parsed_prompts)
+        return results
+
     def close(self) -> None:
         """Close resources.
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -371,6 +371,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             default_template_content_format,
         ).with_defaults(default_template_kwargs)
 
+        # OMNI: When use_audio_in_video=True, the qwen2_5_omni_thinker mm
+        # processor asserts that audio items are present alongside video items
+        # (inside render_chat_async).  We must inject audio_url items into the
+        # messages BEFORE calling render_chat_async so the mm processor can
+        # count them correctly during tokenisation.
+        mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {}
+        if mm_proc_kw.get("use_audio_in_video", False):
+            messages = await self._inject_audio_from_video_urls(messages)
+
         (conversation,), (engine_prompt,) = await renderer.render_chat_async(
             [messages],
             chat_params,
@@ -380,28 +389,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             },
         )
 
-        # OMNI: When use_audio_in_video=True, the upstream renderer does not
-        # extract audio from video.  We do it here after rendering so that the
-        # audio data is present in multi_modal_data before the engine processes
-        # the request.
-        mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {}
-        if mm_proc_kw.get("use_audio_in_video", False) and isinstance(engine_prompt, dict):
-            mm_data = engine_prompt.get("multi_modal_data")
-            if mm_data is not None and "video" in mm_data and "audio" not in mm_data:
-                from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async
-
-                video_urls: list[str] = []
-                for msg in messages:
-                    for part in msg.get("content") or []:
-                        if isinstance(part, dict) and part.get("type") == "video_url":
-                            url = part.get("video_url", {}).get("url")
-                            if url:
-                                video_urls.append(url)
-
-                if video_urls:
-                    audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))
-                    engine_prompt.setdefault("multi_modal_data", {})["audio"] = list(audios)
-
         tokenizer = renderer.get_tokenizer()
 
         # tool parsing is done only if a tool_parser has been set and if
@@ -422,6 +409,95 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return conversation, [engine_prompt]
 
+    async def _inject_audio_from_video_urls(
+        self,
+        messages: list[ChatCompletionMessageParam],
+    ) -> list[ChatCompletionMessageParam]:
+        """Pre-extract audio from video URLs and inject as audio_url content items.
+
+        When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal
+        processor requires that the number of audio items equals the number of
+        video items (it subtracts mm_counts["video"] from mm_counts["audio"]).
+        The client only sends video_url items; this method adds the matching
+        audio_url items on the server side before the renderer processes them.
+        """
+        import io
+
+        from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async
+
+        new_messages: list[ChatCompletionMessageParam] = []
+        for msg in messages:
+            content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
+            if not isinstance(content, list):
+                new_messages.append(msg)
+                continue
+
+            video_urls = [
+                part.get("video_url", {}).get("url")
+                for part in content
+                if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url")
+            ]
+
+            if not video_urls:
+                new_messages.append(msg)
+                continue
+
+            audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))
+
+            audio_items: list[dict] = []
+            for audio_array, sample_rate in audios:
+                buf = io.BytesIO()
+                if soundfile is not None:
+                    soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV")
+                else:
+                    import struct
+
+                    import numpy as np
+
+                    audio_np = np.asarray(audio_array, dtype=np.float32)
+                    sr = int(sample_rate)
+                    num_channels = 1
+                    bits_per_sample = 32
+                    num_frames = len(audio_np)
+                    data_size = num_frames * num_channels * (bits_per_sample // 8)
+                    # Write minimal RIFF/WAV header
+                    buf.write(b"RIFF")
+                    buf.write(struct.pack("<I", 36 + data_size))
+                    buf.write(b"WAVE")
+                    buf.write(b"fmt ")
+                    buf.write(
+                        struct.pack(
+                            "<IHHIIHH",
+                            16,
+                            3,
+                            num_channels,
+                            sr,
+                            sr * num_channels * (bits_per_sample // 8),
+                            num_channels * (bits_per_sample // 8),
+                            bits_per_sample,
+                        )
+                    )
+                    buf.write(b"data")
+                    buf.write(struct.pack("<I", data_size))
+                    buf.write(audio_np.tobytes())
+
+                audio_b64 = base64.b64encode(buf.getvalue()).decode()
+                audio_items.append(
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"},
+                    }
+                )
+
+            new_content = list(content) + audio_items
+            if isinstance(msg, dict):
+                new_msg = {**msg, "content": new_content}
+            else:
+                new_msg = msg.model_copy(update={"content": new_content})
+            new_messages.append(new_msg)
+
+        return new_messages
+
     def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:
         final_sampling_params_list = []
         for sampling_params in sampling_params_list:

vllm_omni/inputs/preprocess.py

@@ -4,7 +4,7 @@ from typing_extensions import assert_never
 from vllm.inputs.data import EmbedsInputs, SingletonInputs
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
-from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
+from vllm.multimodal.inputs import MultiModalInputs
 from vllm.renderers.inputs import SingletonDictPrompt
 
 from vllm_omni.inputs.data import (
@@ -30,8 +30,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTextPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> OmniTokenInputs | MultiModalInputs:
         prompt_text = parsed_content["prompt"]
 
@@ -42,14 +40,8 @@ class OmniInputPreprocessor(InputPreprocessor):
                 multi_modal_data,
                 parsed_content.get("mm_processor_kwargs") or {},
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
-            prompt_embeds = parsed_content.get("prompt_embeds")
-            if prompt_embeds is not None:
-                inputs["prompt_embeds"] = prompt_embeds
-            additional_information = parsed_content.get("additional_information")
-            if additional_information is not None:
-                inputs["additional_information"] = additional_information
+
         else:
             prompt_token_ids = self._tokenize_prompt(
                 prompt_text,
@@ -60,7 +52,12 @@ class OmniInputPreprocessor(InputPreprocessor):
                 prompt_embeds=parsed_content.get("prompt_embeds"),
                 additional_information=parsed_content.get("additional_information"),
             )
-
+        prompt_embeds = parsed_content.get("prompt_embeds")
+        if prompt_embeds is not None:
+            inputs["prompt_embeds"] = prompt_embeds
+        additional_information = parsed_content.get("additional_information")
+        if additional_information is not None:
+            inputs["additional_information"] = additional_information
         if cache_salt := parsed_content.get("cache_salt"):
             inputs["cache_salt"] = cache_salt
 
@@ -70,8 +67,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTokensPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> OmniTokenInputs | MultiModalInputs:
         prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
         prompt_embeds = parsed_content.get("prompt_embeds")
@@ -82,21 +77,21 @@ class OmniInputPreprocessor(InputPreprocessor):
             inputs = self._process_multimodal(
                 prompt_token_ids,
                 multi_modal_data,
-                parsed_content.get("mm_processor_kwargs") or {},
+                parsed_content.get("mm_processor_kwargs"),
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
+                mm_uuids=parsed_content.get("multi_modal_uuids"),
             )
-            if prompt_embeds is not None:
-                inputs["prompt_embeds"] = prompt_embeds
-            if additional_information is not None:
-                inputs["additional_information"] = additional_information
+
         else:
             inputs = token_inputs_omni(
                 prompt_token_ids=prompt_token_ids,
                 prompt_embeds=prompt_embeds,
                 additional_information=additional_information,
             )
-
+        if prompt_embeds is not None:
+            inputs["prompt_embeds"] = prompt_embeds
+        if additional_information is not None:
+            inputs["additional_information"] = additional_information
         if cache_salt := parsed_content.get("cache_salt"):
             inputs["cache_salt"] = cache_salt
 
@@ -125,8 +120,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         prompt: SingletonDictPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> SingletonInputs:
         """
         Extract the singleton inputs from a prompt.
@@ -142,7 +135,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         if "prompt_token_ids" in prompt:
             return self._process_tokens(
                 prompt,  # type: ignore[arg-type]
-                mm_uuids=mm_uuids,
             )
 
         if "prompt_embeds" in prompt:
@@ -152,7 +144,6 @@ class OmniInputPreprocessor(InputPreprocessor):
             return self._process_text(
                 prompt,  # type: ignore[arg-type]
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
 
         assert_never(prompt)  # type: ignore[arg-type]