From 9869ea007ea08e431ba78bcb85b8f41b8d805097 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 20 Feb 2026 08:34:53 +0000 Subject: [PATCH] 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 --- docker/Dockerfile.ci | 4 +- .../distributed/omni_connectors/adapter.py | 16 ++- .../connectors/shm_connector.py | 2 +- vllm_omni/engine/input_processor.py | 93 +++++++------- vllm_omni/engine/output_processor.py | 9 +- vllm_omni/entrypoints/omni_llm.py | 27 +++- vllm_omni/entrypoints/openai/serving_chat.py | 120 ++++++++++++++---- vllm_omni/inputs/preprocess.py | 39 +++--- 8 files changed, 214 insertions(+), 96 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 9560a0dc8..e8c23f18f 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -14,11 +14,11 @@ RUN apt-get update && \ # Install vLLM wheel for the specific commit # URL format: https://wheels.vllm.ai// -# 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 diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 1d1dd6f0f..91a0146c9 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/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(), }, diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 341bbcdba..c654e2697 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/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: diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 29c74e2a9..37aadc7b0 100644 --- a/vllm_omni/engine/input_processor.py +++ b/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() diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index c2520f168..f87dc70ef 100644 --- a/vllm_omni/engine/output_processor.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 7b990e9d0..faf9bfce0 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/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. diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index b45b4069f..6874b4202 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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(" list[SamplingParams]: final_sampling_params_list = [] for sampling_params in sampling_params_list: diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 09b215bf9..53e09f08a 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/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]