From 75770c93097ebfa33f16c0722e2037a93c6b9f3b Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Fri, 13 Feb 2026 18:05:07 +0800 Subject: [PATCH 1/2] [Rebase] Rebase to vllm v0.16.0 (#1357) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Taichang Zhou Signed-off-by: tzhouam Signed-off-by: princepride Co-authored-by: 汪志鹏 Co-authored-by: Hongsheng Liu Co-authored-by: Isotr0py Co-authored-by: ZJY0516 Co-authored-by: youkaichao --- .buildkite/pipeline.yml | 2 +- .gitignore | 3 + docker/Dockerfile.ci | 29 +- docs/api/README.md | 5 +- .../installation/gpu/cuda.inc.md | 10 +- docs/getting_started/quickstart.md | 4 +- .../offline_inference/qwen3_omni/end2end.py | 31 - pyproject.toml | 1 + requirements/cuda.txt | 4 +- tests/e2e/online_serving/test_qwen3_tts.py | 1 + .../test_gpu_generation_model_runner.py | 1 + vllm_omni/benchmarks/patch/patch.py | 5 +- vllm_omni/config/model.py | 247 ++----- vllm_omni/core/sched/omni_ar_scheduler.py | 15 +- .../core/sched/omni_generation_scheduler.py | 29 +- vllm_omni/core/sched/output.py | 14 +- .../attention/backends/ring/ring_globals.py | 9 +- vllm_omni/diffusion/diffusion_engine.py | 4 + .../connectors/shm_connector.py | 62 +- vllm_omni/engine/arg_utils.py | 284 +++++--- vllm_omni/engine/input_processor.py | 17 +- vllm_omni/engine/output_processor.py | 19 +- vllm_omni/entrypoints/async_omni_llm.py | 13 +- vllm_omni/entrypoints/chat_utils.py | 337 ++------- vllm_omni/entrypoints/omni_stage.py | 3 + vllm_omni/entrypoints/openai/api_server.py | 16 +- vllm_omni/entrypoints/openai/serving_chat.py | 321 ++++----- vllm_omni/inputs/data.py | 8 +- vllm_omni/inputs/parse.py | 42 -- vllm_omni/inputs/preprocess.py | 52 +- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 85 ++- .../models/qwen2_5_omni/qwen2_old.py | 2 +- .../qwen3_omni/qwen3_omni_moe_thinker.py | 646 +++++++++++------- vllm_omni/outputs.py | 7 +- vllm_omni/request.py | 1 + vllm_omni/worker/gpu_ar_model_runner.py | 6 + vllm_omni/worker/gpu_ar_worker.py | 2 + .../worker/gpu_generation_model_runner.py | 96 ++- vllm_omni/worker/gpu_generation_worker.py | 2 + vllm_omni/worker/gpu_model_runner.py | 39 +- 40 files changed, 1210 insertions(+), 1264 deletions(-) delete mode 100644 vllm_omni/inputs/parse.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index c435f311e..88cfc22d5 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -232,7 +232,7 @@ steps: type: DirectoryOrCreate - label: "Qwen3-TTS E2E Test" - timeout_in_minutes: 10 + timeout_in_minutes: 20 depends_on: image-build commands: - export VLLM_LOGGING_LEVEL=DEBUG diff --git a/.gitignore b/.gitignore index 7b0a9c368..12486f4a7 100644 --- a/.gitignore +++ b/.gitignore @@ -242,3 +242,6 @@ configs/development.yaml Dockerfile.dev discussion tmp_test + +# output files +*.wav diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index c927e139a..560661761 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -12,9 +12,34 @@ RUN apt-get update && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* -# Install vllm-omni into the same uv-managed Python environment used by the base image. -RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" +# Install vLLM wheel for the specific commit +# URL format: https://wheels.vllm.ai// +# Current vLLM wheel commit: 2d5be1dd5ce2e44dfea53ea03ff61143da5137eb +# 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/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb + +RUN uv pip uninstall --system nvidia-cublas-cu12 +# Install vllm-omni into the same uv-managed Python environment used by the base image. +# Use bash -c so that $(python3 -c ...) is expanded inside the container. +RUN uv pip install --system --no-cache-dir ".[dev]" + +# pplx_kernels is optional; the base image may already ship a version. +# Reinstall attempt may fail if no compatible wheel exists for the current +# PyTorch/CUDA combination – allow the build to continue regardless. +RUN pip install --upgrade --force-reinstall "pplx_kernels" || true + +RUN /bin/bash -c 'FLASHINFER_CUDA_TAG="$(python3 -c "import torch; print((torch.version.cuda or \"12.4\").replace(\".\", \"\"))")" && \ + uv pip install --system --upgrade --force-reinstall \ + "flashinfer-python==0.6.3" \ + "flashinfer-cubin==0.6.3" \ + "flashinfer-jit-cache==0.6.3" \ + --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"' +# Pin nvidia-cublas-cu12 after vllm-omni install so dependency resolution does not overwrite it +RUN uv pip install --system --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4" +RUN uv pip install --system --upgrade --force-reinstall "numpy==2.2.6" RUN ln -sf /usr/bin/python3 /usr/bin/python ENTRYPOINT [] diff --git a/docs/api/README.md b/docs/api/README.md index b4ec398f2..efcb92041 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -7,9 +7,7 @@ Main entry points for vLLM-Omni inference and serving. - [vllm_omni.entrypoints.async_omni.AsyncOmni][] - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][] - [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][] -- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][] -- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][] -- [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][] +- [vllm_omni.entrypoints.chat_utils.extract_audio_from_video_async][] - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][] - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][] - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][] @@ -31,7 +29,6 @@ Input data structures for multi-modal inputs. - [vllm_omni.inputs.data.OmniTextPrompt][] - [vllm_omni.inputs.data.OmniTokenInputs][] - [vllm_omni.inputs.data.OmniTokensPrompt][] -- [vllm_omni.inputs.parse.parse_singleton_prompt_omni][] - [vllm_omni.inputs.preprocess.OmniInputPreprocessor][] ## Engine diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 09323cd25..18348c792 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.15.0 --torch-backend=auto +uv pip install vllm==0.16.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -34,10 +34,10 @@ uv pip install vllm-omni # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.15.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library ```bash -uv pip install vllm==0.15.0 --torch-backend=auto +uv pip install vllm==0.16.0 --torch-backend=auto ``` The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment. @@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.15.0 +git checkout v0.16.0 ``` Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file. ```bash -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.15.0/vllm-0.15.0-cp38-abi3-manylinux_2_31_x86_64.whl +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index c65dd32fb..a3eece127 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,10 +19,10 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.15.0 --torch-backend=auto +uv pip install vllm==0.16.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.15.0 --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700 +uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 263377762..509886a02 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -226,37 +226,6 @@ def get_multi_audios_query() -> QueryResult: ) -# def get_use_audio_in_video_query(video_path: str | None = None) -> QueryResult: -# question = ( -# "Describe the content of the video in details, then convert what the " -# "baby say into text." -# ) -# prompt = ( -# f"<|im_start|>system\n{default_system}<|im_end|>\n" -# "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" -# f"{question}<|im_end|>\n" -# f"<|im_start|>assistant\n" -# ) -# if video_path: -# if not os.path.exists(video_path): -# raise FileNotFoundError(f"Video file not found: {video_path}") -# video_frames = video_to_ndarrays(video_path, num_frames=16) -# else: -# video_frames = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays -# audio = extract_video_audio(video_path, sampling_rate=16000) -# return QueryResult( -# inputs={ -# "prompt": prompt, -# "multi_modal_data": { -# "video": video_frames, -# "audio": audio, -# }, -# "mm_processor_kwargs": { -# "use_audio_in_video": True, -# }, -# }, -# limit_mm_per_prompt={"audio": 1, "video": 1}, -# ) def get_use_audio_in_video_query() -> QueryResult: question = "Describe the content of the video in details, then convert what the baby say into text." prompt = ( diff --git a/pyproject.toml b/pyproject.toml index b2293da7c..965a3e4b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,6 +137,7 @@ warn_unreachable = true strict_equality = true [tool.pytest.ini_options] +asyncio_mode = "auto" testpaths = ["tests"] python_files = ["test_*.py", "*_test.py"] python_classes = ["Test*"] diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 823f47235..7ed74ca83 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -1,3 +1,5 @@ -r common.txt onnxruntime>=1.23.2 -fa3-fwd==0.0.1 +fa3-fwd==0.0.2 +# nvidia-cublas-cu12 is pinned to 12.9.1.4 via force-reinstall in Dockerfile.ci +# because torch pins its own cublas version, making it unsatisfiable at resolve time. diff --git a/tests/e2e/online_serving/test_qwen3_tts.py b/tests/e2e/online_serving/test_qwen3_tts.py index 9a1ff8b97..b4ca57835 100644 --- a/tests/e2e/online_serving/test_qwen3_tts.py +++ b/tests/e2e/online_serving/test_qwen3_tts.py @@ -46,6 +46,7 @@ def omni_server(): "120", "--trust-remote-code", "--enforce-eager", + "--disable-log-stats", ], ) as server: print("OmniServer started successfully") diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 5c44889a4..8ecb4b0f8 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -27,6 +27,7 @@ def _make_runner(multimodal_outputs): None, None, multimodal_outputs, + None, ) runner.kv_connector_output = None runner.input_batch = _DummyInputBatch() diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 4c030e25f..5a03a8069 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -5,6 +5,7 @@ import io import json import os import random +import ssl import sys import time import traceback @@ -240,6 +241,7 @@ async def benchmark( ramp_up_start_rps: int | None = None, ramp_up_end_rps: int | None = None, ready_check_timeout_sec: int = 600, + ssl_context: ssl.SSLContext | bool | None = None, ): try: request_func = ASYNC_REQUEST_FUNCS[endpoint_type] @@ -247,6 +249,7 @@ async def benchmark( raise ValueError(f"Unknown backend: {endpoint_type}") from None # Reuses connections across requests to reduce TLS handshake overhead. + ssl_setting = ssl_context if ssl_context is not None else ("https://" in api_url) connector = aiohttp.TCPConnector( limit=max_concurrency or 0, limit_per_host=max_concurrency or 0, @@ -255,7 +258,7 @@ async def benchmark( keepalive_timeout=60, enable_cleanup_closed=True, force_close=False, - ssl=("https://" in api_url), + ssl=ssl_setting, ) session = aiohttp.ClientSession( diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index b099c19a0..a9ffa015f 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -1,28 +1,12 @@ -import warnings from dataclasses import field from typing import Any -import torch from pydantic import ConfigDict -from pydantic.dataclasses import dataclass -from vllm.config import ModelConfig, config -from vllm.config.model import ( - _RUNNER_CONVERTS, - _get_and_verify_dtype, - get_served_model_name, -) -from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MultiModalConfig -from vllm.config.pooler import PoolerConfig +from vllm.config import ModelConfig +from vllm.config.multimodal import MMCacheType, MMEncoderTPMode +from vllm.config.utils import config from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.transformers_utils.config import ( - get_config, - get_hf_image_processor_config, - get_hf_text_config, - get_pooling_config, -) -from vllm.transformers_utils.gguf_utils import is_gguf, maybe_patch_hf_config_from_gguf -from vllm.transformers_utils.utils import maybe_model_redirect +from vllm.transformers_utils.config import get_hf_text_config from vllm.v1.attention.backends.registry import AttentionBackendEnum import vllm_omni.model_executor.models as me_models @@ -30,8 +14,7 @@ import vllm_omni.model_executor.models as me_models logger = init_logger(__name__) -@config -@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +@config(config=ConfigDict(arbitrary_types_allowed=True)) class OmniModelConfig(ModelConfig): """Configuration for Omni models, extending the base ModelConfig. @@ -84,6 +67,15 @@ class OmniModelConfig(ModelConfig): def architectures(self) -> list[str]: return [self.model_arch] + @property + def embedding_size(self): + if self.hf_config_name is not None: + stage_config = getattr(self.hf_config, self.hf_config_name, None) + override = getattr(stage_config, "embedding_size", None) + if override is not None: + return override + return super().embedding_size + def draw_hf_text_config(self): # transformers' get_text_config method is used to get the text config from thinker_config. # to handle the case that each model stage has their own text config, @@ -119,189 +111,32 @@ class OmniModelConfig(ModelConfig): skip_mm_profiling: bool | None, video_pruning_rate: float | None, ) -> None: - # Keep set served_model_name before maybe_model_redirect(self.model) - self.served_model_name = get_served_model_name(self.model, self.served_model_name) - self.model = maybe_model_redirect(self.model) - # The tokenizer is consistent with the model by default. - if self.tokenizer is None: - self.tokenizer = self.model - if self.tokenizer_revision is None: - self.tokenizer_revision = self.revision - self.tokenizer = maybe_model_redirect(self.tokenizer) - - if isinstance(self.hf_config_path, str): - self.hf_config_path = maybe_model_redirect(self.hf_config_path) - - if callable(self.hf_overrides): - hf_overrides_kw = {} - hf_overrides_fn = self.hf_overrides - dict_overrides: dict[str, Any] = {} - else: - # Separate dict overrides from flat ones - # We'll determine how to apply dict overrides after loading the config - hf_overrides_kw = {} - dict_overrides = {} - for key, value in self.hf_overrides.items(): - if isinstance(value, dict): - dict_overrides[key] = value - else: - hf_overrides_kw[key] = value - hf_overrides_fn = None - - self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer) - - if self.override_attention_dtype is not None and not current_platform.is_rocm(): - warnings.warn( - "override-attention-dtype is set but not using ROCm platform", - stacklevel=2, - ) - - if self.enable_sleep_mode and not current_platform.is_sleep_mode_available(): - raise ValueError("Sleep mode is not supported on current platform.") - - hf_config = get_config( - self.hf_config_path or self.model, - self.trust_remote_code, - self.revision, - self.code_revision, - self.config_format, - hf_overrides_kw=hf_overrides_kw, - hf_overrides_fn=hf_overrides_fn, - ) - hf_config = maybe_patch_hf_config_from_gguf( - self.model, - hf_config, + # Call parent's __post_init__ to handle all standard ModelConfig initialization + super().__post_init__( + limit_mm_per_prompt=limit_mm_per_prompt, + enable_mm_embeds=enable_mm_embeds, + media_io_kwargs=media_io_kwargs, + mm_processor_kwargs=mm_processor_kwargs, + mm_processor_cache_gb=mm_processor_cache_gb, + mm_processor_cache_type=mm_processor_cache_type, + mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, + mm_encoder_only=mm_encoder_only, + mm_encoder_tp_mode=mm_encoder_tp_mode, + mm_encoder_attn_backend=mm_encoder_attn_backend, + interleave_mm_strings=interleave_mm_strings, + skip_mm_profiling=skip_mm_profiling, + video_pruning_rate=video_pruning_rate, ) - self.hf_config = hf_config - if dict_overrides: - self._apply_dict_overrides(hf_config, dict_overrides) - self.hf_text_config = self.draw_hf_text_config() - self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None) - self.encoder_config = self._get_encoder_config() - self.hf_image_processor_config = get_hf_image_processor_config( - self.model, hf_token=self.hf_token, revision=self.revision - ) - self.model_arch_config = self.get_model_arch_config() - - if self.convert == "mm_encoder_only": - logger.warning_once( - "`--convert mm_encoder_only` is deprecated and " - "will be removed in v0.15. " - "Please use --mm-encoder-only` instead." - ) - mm_encoder_only = True - self.convert = "none" - - architectures = self.architectures - registry = self.registry - is_generative_model = registry.is_text_generation_model(architectures, self) - is_pooling_model = registry.is_pooling_model(architectures, self) - - self.runner_type = self._get_runner_type(architectures, self.runner) - self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert) - - if self.runner_type == "generate" and not is_generative_model: - generate_converts = _RUNNER_CONVERTS["generate"] - if self.convert_type not in generate_converts: - # Currently we don't have any converters for generative models - raise ValueError("This model does not support `--runner generate`.") - if self.runner_type == "pooling" and not is_pooling_model: - pooling_converts = _RUNNER_CONVERTS["pooling"] - if self.convert_type not in pooling_converts: - convert_option = "<" + "|".join(pooling_converts) + ">" - raise ValueError( - "This model does not support `--runner pooling`. " - f"You can pass `--convert {convert_option} to adapt " - "it into a pooling model." - ) - - # Note: Initialize these attributes early because transformers fallback - # may fail to load dynamic modules in child processes - model_info, arch = registry.inspect_model_cls(architectures, self) - self._model_info = model_info - self._architecture = arch - logger.info("Resolved architecture: %s", arch) - - # Init pooler config if needed - if self.runner_type == "pooling": - if self.pooler_config is None: - self.pooler_config = PoolerConfig() - - base_config = get_pooling_config(self.model, self.revision) - if base_config is not None: - # Only set values that are not overridden by the user - for k, v in base_config.items(): - if getattr(self.pooler_config, k) is None: - setattr(self.pooler_config, k, v) - - default_seq_pooling_type = self._model_info.default_seq_pooling_type - if self.pooler_config.seq_pooling_type is None: - self.pooler_config.seq_pooling_type = default_seq_pooling_type - default_tok_pooling_type = self._model_info.default_tok_pooling_type - if self.pooler_config.tok_pooling_type is None: - self.pooler_config.tok_pooling_type = default_tok_pooling_type - - self.dtype: torch.dtype = _get_and_verify_dtype( - self.model, - self.hf_config, - self.dtype, - is_pooling_model=self.runner_type == "pooling", - revision=self.revision, - ) - - self.original_max_model_len = self.max_model_len - self.max_model_len = self.get_and_verify_max_len(self.max_model_len) - - if self.is_encoder_decoder: - self.mm_processor_cache_gb = 0 - logger.info("Encoder-decoder model detected, disabling mm processor cache.") - - # Init multimodal config if needed - if self._model_info.supports_multimodal: - if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data: - logger.warning_once( - "This model does not support `--mm-encoder-tp-mode data`. " - "Falling back to `--mm-encoder-tp-mode weights`." - ) - mm_encoder_tp_mode = "weights" - - mm_config_kwargs = dict( - limit_per_prompt=limit_mm_per_prompt, - enable_mm_embeds=enable_mm_embeds, - media_io_kwargs=media_io_kwargs, - mm_processor_kwargs=mm_processor_kwargs, - mm_processor_cache_gb=mm_processor_cache_gb, - mm_processor_cache_type=mm_processor_cache_type, - mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, - mm_encoder_only=mm_encoder_only, - mm_encoder_tp_mode=mm_encoder_tp_mode, - mm_encoder_attn_backend=mm_encoder_attn_backend, - interleave_mm_strings=interleave_mm_strings, - skip_mm_profiling=skip_mm_profiling, - video_pruning_rate=video_pruning_rate, - ) - - mm_config_kwargs = {k: v for k, v in mm_config_kwargs.items() if v is not None} - - self.multimodal_config = MultiModalConfig(**mm_config_kwargs) - - # Multimodal GGUF models must use original repo for mm processing - if is_gguf(self.tokenizer) and self.is_multimodal_model: - raise ValueError( - "Loading a multimodal GGUF model needs to use original " - "tokenizer. Please specify the unquantized hf model's " - "repo name or path using the --tokenizer argument." - ) - - if self.disable_sliding_window: - # Set after get_and_verify_max_len to ensure that max_model_len - # can be correctly capped to sliding window size - self.hf_text_config.sliding_window = None - - # Avoid running try_verify_and_update_config multiple times - self.config_updated = False - self._try_verify_and_update_model_config() - self._verify_quantization() - self._verify_cuda_graph() - self._verify_bnb_config() + # Override hf_text_config with omni-specific logic for multi-stage models + # (e.g., thinker_config, talker_config) + new_hf_text_config = self.draw_hf_text_config() + if new_hf_text_config is not self.hf_text_config: + self.hf_text_config = new_hf_text_config + # Recalculate dependent attributes + self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None) + # Recalculate max_model_len since it depends on hf_text_config + self.max_model_len = self.get_and_verify_max_len(self.original_max_model_len) + # Reset sliding_window if needed + if self.disable_sliding_window: + self.hf_text_config.sliding_window = None diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index dc33f9cbc..7b0a1e54d 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -236,10 +236,10 @@ class OmniARScheduler(VLLMScheduler): # Skip requests that were recovered from KV load failure continue request = self.requests.get(req_id) - if request is None: + if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., - # in pipeline parallelism). + # in pipeline parallelism or async scheduling). continue req_index = model_runner_output.req_id_to_index[req_id] @@ -347,6 +347,7 @@ class OmniARScheduler(VLLMScheduler): kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, + num_external_computed_tokens=request.num_external_computed_tokens, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, ) @@ -413,7 +414,7 @@ class OmniARScheduler(VLLMScheduler): # publish collected KV cache events if events: - batch = KVEventBatch(ts=time.time(), events=events) + batch = KVEventBatch(ts=time(), events=events) self.kv_event_publisher.publish(batch) # Create EngineCoreOutputs for all clients that have requests with @@ -469,15 +470,16 @@ class OmniARScheduler(VLLMScheduler): return engine_core_outputs - def _free_request(self, request: Request) -> dict[str, Any] | None: + def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: # TODO(wzliu)! for offline mode, we should not end process until all data is transferred """Mark a request as finished and free its resources.""" + assert request.is_finished() # 1. Standard cleanup parts from base _free_request - delay_free_blocks = False + connector_delay_free_blocks = False kv_xfer_params = None if self.connector is not None: - delay_free_blocks, kv_xfer_params = self._connector_finished(request) + connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) self.encoder_cache_manager.free(request) request_id = request.request_id @@ -537,6 +539,7 @@ class OmniARScheduler(VLLMScheduler): return kv_xfer_params # 3. Standard Freeing + delay_free_blocks |= connector_delay_free_blocks if not delay_free_blocks: self._free_blocks(request) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index b86d49ed5..684aab9ce 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -317,17 +317,17 @@ class OmniGenerationScheduler(VLLMScheduler): # Skip requests that were recovered from KV load failure continue request = self.requests.get(req_id) - if request is None: + if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., - # in pipeline parallelism). + # in pipeline parallelism or async scheduling). continue req_index = model_runner_output.req_id_to_index[req_id] generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id) - if scheduled_spec_token_ids: + if scheduled_spec_token_ids and generated_token_ids: num_draft_tokens = len(scheduled_spec_token_ids) num_accepted = len(generated_token_ids) - 1 num_rejected = num_draft_tokens - num_accepted @@ -338,10 +338,16 @@ class OmniGenerationScheduler(VLLMScheduler): # tokens. if request.num_computed_tokens > 0: request.num_computed_tokens -= num_rejected + # If async scheduling, num_output_placeholders also includes + # the scheduled spec tokens count and so is similarly adjusted. + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_rejected spec_decoding_stats = self.make_spec_decoding_stats( spec_decoding_stats, num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted, + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, ) stopped = False @@ -412,6 +418,7 @@ class OmniGenerationScheduler(VLLMScheduler): kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, + num_external_computed_tokens=request.num_external_computed_tokens, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, ) @@ -427,6 +434,22 @@ class OmniGenerationScheduler(VLLMScheduler): # This is a rare case and unlikely to impact performance. self.waiting.remove_requests(stopped_preempted_reqs) + # Handle failed KV load requests + if failed_kv_load_req_ids and not self.recompute_kv_load_failures: + requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids] + self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) + for request in requests: + outputs[request.client_index].append( + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=request.get_finished_reason(), + events=request.take_events(), + trace_headers=request.trace_headers, + num_cached_tokens=request.num_cached_tokens, + ) + ) + # KV Connector: update state for finished KV Transfers. if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index 933ae7629..b86a97534 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -15,12 +15,15 @@ class OmniNewRequestData(NewRequestData): Args: prompt_embeds: Optional serialized prompt embeddings payload + (overrides parent's torch.Tensor type with PromptEmbedsPayload + for cross-process serialization) + external_req_id: Optional external request ID for tracking additional_information: Optional serialized additional information dictionary containing tensors or lists """ - # Optional serialized prompt embeddings - prompt_embeds: PromptEmbedsPayload | None = None + # Optional serialized prompt embeddings (override parent type for serialization) + prompt_embeds: PromptEmbedsPayload | None = None # type: ignore[assignment] # Optional external request ID for tracking external_req_id: str | None = None # Optional serialized additional information @@ -38,13 +41,14 @@ class OmniNewRequestData(NewRequestData): Args: request: Request object to convert block_ids: Tuple of block ID lists for KV cache allocation + prefill_token_ids: Optional prefill token IDs for v2 model runner Returns: OmniNewRequestData instance with data from the request """ return cls( req_id=request.request_id, - external_req_id=request.external_req_id, + external_req_id=getattr(request, "external_req_id", None), prompt_token_ids=request.prompt_token_ids, mm_features=request.mm_features, sampling_params=request.sampling_params, @@ -52,9 +56,9 @@ class OmniNewRequestData(NewRequestData): block_ids=block_ids, num_computed_tokens=request.num_computed_tokens, lora_request=request.lora_request, - prompt_embeds=request.prompt_embeds, + prompt_embeds=getattr(request, "prompt_embeds", None), prefill_token_ids=prefill_token_ids, - additional_information=request.additional_information, + additional_information=getattr(request, "additional_information", None), ) diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py index d80d81fb4..4ac7d1456 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py @@ -3,6 +3,8 @@ # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention +import logging + # test if flash_attn (FA2) is available try: import flash_attn # noqa: F401 @@ -45,12 +47,17 @@ HAS_FLASH_ATTN_HOPPER = HAS_FA3 flash_attn_forward_hopper = fa3_fwd_func flash3_attn_func = fa3_attn_func +logger = logging.getLogger(__name__) + try: from flashinfer.prefill import single_prefill_with_kv_cache # noqa: F401 HAS_FLASHINFER = True -except (ImportError, ModuleNotFoundError): +except Exception as e: + # flashinfer may raise RuntimeError at import-time for version/binary mismatches. + # Treat it as unavailable so the runtime can gracefully fall back to other backends. HAS_FLASHINFER = False + logger.warning("FlashInfer is unavailable; falling back to other attention backends. Reason: %s", e) try: import aiter # noqa: F401 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 2ab94d5b7..8e4a9f7a2 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -331,6 +331,10 @@ class DiffusionEngine: height=height, width=width, num_inference_steps=num_inference_steps, + # Keep warmup path minimal and robust across text encoders. + # Some models may fail when warmup implicitly triggers + # classifier-free guidance with an empty negative prompt. + guidance_scale=0.0, num_outputs_per_prompt=1, ), ) diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 719dbeae7..341bbcdba 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -33,7 +33,13 @@ class SharedMemoryConnector(OmniConnectorBase): "inline_writes": 0, } - def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + def put( + self, + from_stage: str, + to_stage: str, + put_key: str, + data: Any, + ) -> tuple[bool, int, dict[str, Any] | None]: try: # Always serialize first to check size (and for SHM writing) # Note: For extremely large objects in "inline" mode (e.g. Ray), @@ -42,9 +48,7 @@ class SharedMemoryConnector(OmniConnectorBase): payload = self.serialize_obj(data) size = len(payload) - metadata = {} - # if size > self.threshold: - if True: # TODO: correct put & get logic + if True: # Use Shared Memory lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock" with open(lock_file, "w") as lockf: @@ -53,13 +57,13 @@ class SharedMemoryConnector(OmniConnectorBase): fcntl.flock(lockf, fcntl.LOCK_UN) # meta contains {'name': ..., 'size': ...} - metadata[put_key] = {"shm": meta, "size": size} + metadata = {"shm": meta, "size": size} self._metrics["shm_writes"] += 1 else: # Inline - pass bytes directly to avoid double serialization of the object # We already serialized it to check size, so we pass the bytes. # The Queue will pickle these bytes (fast), avoiding re-serializing the complex object. - metadata[put_key] = {"inline_bytes": payload, "size": size} + metadata = {"inline_bytes": payload, "size": size} self._metrics["inline_writes"] += 1 self._metrics["puts"] += 1 @@ -71,7 +75,47 @@ class SharedMemoryConnector(OmniConnectorBase): logger.error(f"SharedMemoryConnector put failed for req {put_key}: {e}") return False, 0, None - def get(self, from_stage: str, to_stage: str, get_key: str, metadata=None) -> tuple[Any, int] | None: + def get( + self, + from_stage: str, + to_stage: str, + get_key: str, + metadata=None, + ) -> tuple[Any, int] | None: + if metadata is not None: + # Some callers may wrap metadata by request id. + if isinstance(metadata, dict) and get_key in metadata: + metadata = metadata.get(get_key) + + if not isinstance(metadata, dict): + return None + + if "inline_bytes" in metadata: + try: + obj = self.deserialize_obj(metadata["inline_bytes"]) + return obj, int(metadata.get("size", 0)) + except Exception as e: + logger.error(f"SharedMemoryConnector inline get failed for req {get_key}: {e}") + return None + + if "shm" in metadata: + try: + shm_handle = metadata["shm"] + lock_file = f"/dev/shm/shm_{shm_handle['name']}_lockfile.lock" + with open(lock_file, "w") as lockf: + fcntl.flock(lockf, fcntl.LOCK_SH) + data_bytes = shm_read_bytes(shm_handle) + fcntl.flock(lockf, fcntl.LOCK_UN) + if os.path.exists(lock_file): + os.remove(lock_file) + obj = self.deserialize_obj(data_bytes) + return obj, int(metadata.get("size", 0)) + except Exception as e: + logger.error(f"SharedMemoryConnector shm get failed for req {get_key}: {e}") + return None + + return None + from multiprocessing import shared_memory as shm_pkg # Wait for shared memory to be available (with retry logic) @@ -89,10 +133,10 @@ class SharedMemoryConnector(OmniConnectorBase): else: # Max retries reached, return None logger.warning(f"Shared memory '{get_key}' not found after {max_retries} retries") - return None, 0 + return None if shm is None: - return None, 0 + return None try: lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock" diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 6e6ee465f..6f14aa264 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,11 +1,10 @@ from dataclasses import dataclass, field from typing import Any -from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig -from vllm.engine.arg_utils import EngineArgs +import vllm.envs as envs +from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.logger import init_logger -from vllm.transformers_utils.config import get_hf_text_config -from vllm.v1.engine.async_llm import AsyncEngineArgs +from vllm.transformers_utils.gguf_utils import is_gguf from vllm_omni.config import OmniModelConfig from vllm_omni.plugins import load_omni_general_plugins @@ -77,24 +76,6 @@ class OmniEngineArgs(EngineArgs): omni_kv_config: dict | None = None worker_type: str | None = None - def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: - # transformers' get_text_config method is used to get the text config from thinker_config. - # to handle the case that each model stage has their own text config, - # we need to draw the text config from the corresponding model stage. - hf_config = config_dict["hf_config"] - hf_config_name = config_dict["hf_config_name"] - try: - # Try to get the stage-specific config (e.g., thinker_config, talker_config) - stage_config = getattr(hf_config, hf_config_name) - return stage_config.get_text_config() - except AttributeError: - # Fallback: if the attribute doesn't exist, use the default get_hf_text_config - logger.warning( - f"Config attribute '{hf_config_name}' not found in hf_config, " - "falling back to default get_hf_text_config" - ) - return get_hf_text_config(hf_config) - def __post_init__(self) -> None: load_omni_general_plugins() super().__post_init__() @@ -111,43 +92,112 @@ class OmniEngineArgs(EngineArgs): Returns: OmniModelConfig instance with all configuration fields set """ + # GGUF files need a specific model loader path in vLLM. + if is_gguf(self.model): + self.quantization = self.load_format = "gguf" + + if not envs.VLLM_ENABLE_V1_MULTIPROCESSING: + logger.warning( + "The global random seed is set to %d. Since " + "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may " + "affect the random state of the Python process that " + "launched vLLM.", + self.seed, + ) + # register omni models to avoid model not found error self._ensure_omni_models_registered() - # First, get the base ModelConfig from the parent class - base_config = super().create_model_config() + # Keep compatibility when async args are constructed from partial payloads. + limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {}) + enable_mm_embeds = getattr(self, "enable_mm_embeds", False) + interleave_mm_strings = getattr(self, "interleave_mm_strings", False) + media_io_kwargs = getattr(self, "media_io_kwargs", {}) + skip_mm_profiling = getattr(self, "skip_mm_profiling", False) + mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None) + mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4) + mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None) + mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128) + mm_encoder_only = getattr(self, "mm_encoder_only", False) + mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights") + mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None) + video_pruning_rate = getattr(self, "video_pruning_rate", 0.0) - # Create OmniModelConfig by copying all base config attributes - # and adding the new omni-specific fields - config_dict = base_config.__dict__.copy() - # FIXME(Isotr0py): This is a temporary workaround for multimodal_config - config_dict = { - **(getattr(mm := config_dict.pop("multimodal_config", None), "__dict__", mm or {})), - **config_dict, - } - - # Add the new omni-specific fields - config_dict["stage_id"] = self.stage_id - config_dict["async_chunk"] = self.async_chunk - config_dict["model_stage"] = self.model_stage - config_dict["model_arch"] = self.model_arch - config_dict["worker_type"] = self.worker_type - config_dict["engine_output_type"] = self.engine_output_type # Build stage_connector_config from stage_connector_spec stage_connector_config = { "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"), "extra": self.stage_connector_spec.get("extra", {}).copy(), } stage_connector_config["extra"]["stage_id"] = self.stage_id - config_dict["stage_connector_config"] = stage_connector_config - config_dict["hf_config_name"] = self.hf_config_name - config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func - config_dict["omni_kv_config"] = self.omni_kv_config - if self.hf_config_name is not None: - config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict) - # Create and return the OmniModelConfig instance - omni_config = OmniModelConfig(**config_dict) + # Create OmniModelConfig directly from engine args + # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config() + omni_config = OmniModelConfig( + # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config) + model=self.model, + model_weights=self.model_weights, + hf_config_path=self.hf_config_path, + runner=self.runner, + convert=self.convert, + tokenizer=self.tokenizer, + tokenizer_mode=self.tokenizer_mode, + trust_remote_code=self.trust_remote_code, + allowed_local_media_path=self.allowed_local_media_path, + allowed_media_domains=self.allowed_media_domains, + dtype=self.dtype, + seed=self.seed, + revision=self.revision, + code_revision=self.code_revision, + hf_token=self.hf_token, + hf_overrides=self.hf_overrides, + tokenizer_revision=self.tokenizer_revision, + max_model_len=self.max_model_len, + quantization=self.quantization, + allow_deprecated_quantization=self.allow_deprecated_quantization, + enforce_eager=self.enforce_eager, + enable_return_routed_experts=self.enable_return_routed_experts, + max_logprobs=self.max_logprobs, + logprobs_mode=self.logprobs_mode, + disable_sliding_window=self.disable_sliding_window, + disable_cascade_attn=self.disable_cascade_attn, + skip_tokenizer_init=self.skip_tokenizer_init, + enable_prompt_embeds=self.enable_prompt_embeds, + served_model_name=self.served_model_name, + limit_mm_per_prompt=limit_mm_per_prompt, + enable_mm_embeds=enable_mm_embeds, + interleave_mm_strings=interleave_mm_strings, + media_io_kwargs=media_io_kwargs, + skip_mm_profiling=skip_mm_profiling, + config_format=self.config_format, + mm_processor_kwargs=mm_processor_kwargs, + mm_processor_cache_gb=mm_processor_cache_gb, + mm_processor_cache_type=mm_processor_cache_type, + mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, + mm_encoder_only=mm_encoder_only, + mm_encoder_tp_mode=mm_encoder_tp_mode, + mm_encoder_attn_backend=mm_encoder_attn_backend, + pooler_config=self.pooler_config, + logits_processor_pattern=self.logits_processor_pattern, + generation_config=self.generation_config, + override_generation_config=self.override_generation_config, + enable_sleep_mode=self.enable_sleep_mode, + model_impl=self.model_impl, + override_attention_dtype=self.override_attention_dtype, + logits_processors=self.logits_processors, + video_pruning_rate=video_pruning_rate, + io_processor_plugin=self.io_processor_plugin, + # Omni-specific fields + stage_id=self.stage_id, + async_chunk=self.async_chunk, + model_stage=self.model_stage, + model_arch=self.model_arch, + worker_type=self.worker_type, + engine_output_type=self.engine_output_type, + hf_config_name=self.hf_config_name, + custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, + stage_connector_config=stage_connector_config, + omni_kv_config=self.omni_kv_config, + ) omni_config.hf_config.architectures = omni_config.architectures return omni_config @@ -182,24 +232,6 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): omni_kv_config: dict | None = None worker_type: str | None = None - def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: - # transformers' get_text_config method is used to get the text config from thinker_config. - # to handle the case that each model stage has their own text config, - # we need to draw the text config from the corresponding model stage. - hf_config = config_dict["hf_config"] - hf_config_name = config_dict["hf_config_name"] - try: - # Try to get the stage-specific config (e.g., thinker_config, talker_config) - stage_config = getattr(hf_config, hf_config_name) - return stage_config.get_text_config() - except AttributeError: - # Fallback: if the attribute doesn't exist, use the default get_hf_text_config - logger.warning( - f"Config attribute '{hf_config_name}' not found in hf_config, " - "falling back to default get_hf_text_config" - ) - return get_hf_text_config(hf_config) - def __post_init__(self) -> None: load_omni_general_plugins() super().__post_init__() @@ -212,36 +244,116 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): return True def create_model_config(self) -> OmniModelConfig: + """Create an OmniModelConfig from these engine arguments. + Returns: + OmniModelConfig instance with all configuration fields set + """ + # GGUF files need a specific model loader path in vLLM. + if is_gguf(self.model): + self.quantization = self.load_format = "gguf" + + if not envs.VLLM_ENABLE_V1_MULTIPROCESSING: + logger.warning( + "The global random seed is set to %d. Since " + "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may " + "affect the random state of the Python process that " + "launched vLLM.", + self.seed, + ) + # register omni models to avoid model not found error self._ensure_omni_models_registered() - # First, get the base ModelConfig from the parent class - base_config = super().create_model_config() - # Create OmniModelConfig by copying all base config attributes - # and adding the new omni-specific fields - config_dict = base_config.__dict__.copy() + # Keep compatibility when async args are constructed from partial payloads. + limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {}) + enable_mm_embeds = getattr(self, "enable_mm_embeds", False) + interleave_mm_strings = getattr(self, "interleave_mm_strings", False) + media_io_kwargs = getattr(self, "media_io_kwargs", {}) + skip_mm_profiling = getattr(self, "skip_mm_profiling", False) + mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None) + mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4) + mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None) + mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128) + mm_encoder_only = getattr(self, "mm_encoder_only", False) + mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights") + mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None) + video_pruning_rate = getattr(self, "video_pruning_rate", 0.0) - # Add the new omni-specific fields - config_dict["stage_id"] = self.stage_id - config_dict["async_chunk"] = self.async_chunk - config_dict["model_stage"] = self.model_stage - config_dict["model_arch"] = self.model_arch - config_dict["worker_type"] = self.worker_type - config_dict["engine_output_type"] = self.engine_output_type + # Build stage_connector_config from stage_connector_spec stage_connector_config = { "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"), "extra": self.stage_connector_spec.get("extra", {}).copy(), } stage_connector_config["extra"]["stage_id"] = self.stage_id - config_dict["stage_connector_config"] = stage_connector_config - config_dict["hf_config_name"] = self.hf_config_name - config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func - config_dict["omni_kv_config"] = self.omni_kv_config - if self.hf_config_name is not None: - config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict) - # Create and return the OmniModelConfig instance - omni_config = OmniModelConfig(**config_dict) + # Create OmniModelConfig directly from engine args + # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config() + omni_config = OmniModelConfig( + # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config) + model=self.model, + model_weights=self.model_weights, + hf_config_path=self.hf_config_path, + runner=self.runner, + convert=self.convert, + tokenizer=self.tokenizer, + tokenizer_mode=self.tokenizer_mode, + trust_remote_code=self.trust_remote_code, + allowed_local_media_path=self.allowed_local_media_path, + allowed_media_domains=self.allowed_media_domains, + dtype=self.dtype, + seed=self.seed, + revision=self.revision, + code_revision=self.code_revision, + hf_token=self.hf_token, + hf_overrides=self.hf_overrides, + tokenizer_revision=self.tokenizer_revision, + max_model_len=self.max_model_len, + quantization=self.quantization, + allow_deprecated_quantization=self.allow_deprecated_quantization, + enforce_eager=self.enforce_eager, + enable_return_routed_experts=self.enable_return_routed_experts, + max_logprobs=self.max_logprobs, + logprobs_mode=self.logprobs_mode, + disable_sliding_window=self.disable_sliding_window, + disable_cascade_attn=self.disable_cascade_attn, + skip_tokenizer_init=self.skip_tokenizer_init, + enable_prompt_embeds=self.enable_prompt_embeds, + served_model_name=self.served_model_name, + limit_mm_per_prompt=limit_mm_per_prompt, + enable_mm_embeds=enable_mm_embeds, + interleave_mm_strings=interleave_mm_strings, + media_io_kwargs=media_io_kwargs, + skip_mm_profiling=skip_mm_profiling, + config_format=self.config_format, + mm_processor_kwargs=mm_processor_kwargs, + mm_processor_cache_gb=mm_processor_cache_gb, + mm_processor_cache_type=mm_processor_cache_type, + mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, + mm_encoder_only=mm_encoder_only, + mm_encoder_tp_mode=mm_encoder_tp_mode, + mm_encoder_attn_backend=mm_encoder_attn_backend, + pooler_config=self.pooler_config, + logits_processor_pattern=self.logits_processor_pattern, + generation_config=self.generation_config, + override_generation_config=self.override_generation_config, + enable_sleep_mode=self.enable_sleep_mode, + model_impl=self.model_impl, + override_attention_dtype=self.override_attention_dtype, + logits_processors=self.logits_processors, + video_pruning_rate=video_pruning_rate, + io_processor_plugin=self.io_processor_plugin, + # Omni-specific fields + stage_id=self.stage_id, + async_chunk=self.async_chunk, + model_stage=self.model_stage, + model_arch=self.model_arch, + worker_type=self.worker_type, + engine_output_type=self.engine_output_type, + hf_config_name=self.hf_config_name, + custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, + stage_connector_config=stage_connector_config, + omni_kv_config=self.omni_kv_config, + ) omni_config.hf_config.architectures = omni_config.architectures return omni_config diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 64ffab3a1..99dc71cc4 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -1,4 +1,3 @@ -import os import time from collections.abc import Mapping from typing import Any, cast @@ -25,7 +24,9 @@ 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.inputs import DictPrompt, TokPrompt from vllm.sampling_params import SamplingParams +from vllm.tasks import SupportedTask from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.engine.input_processor import InputProcessor @@ -100,7 +101,7 @@ class OmniInputProcessor(InputProcessor): def process_inputs( self, request_id: str, - prompt: PromptType, + prompt: PromptType | DictPrompt | TokPrompt, params: SamplingParams | PoolingParams, arrival_time: float | None = None, lora_request: LoRARequest | None = None, @@ -108,6 +109,7 @@ class OmniInputProcessor(InputProcessor): trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, + supported_tasks: tuple[SupportedTask, ...] | None = None, resumable: bool = False, ) -> OmniEngineCoreRequest: """Process input prompt into an engine core request. @@ -140,7 +142,7 @@ class OmniInputProcessor(InputProcessor): has incorrect shape """ self._validate_lora(lora_request) - self._validate_params(params) + self._validate_params(params, supported_tasks) parallel_config = self.vllm_config.parallel_config dp_size = parallel_config.data_parallel_size @@ -179,14 +181,7 @@ class OmniInputProcessor(InputProcessor): # 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. - num_threads = int(os.environ.get("OMP_NUM_THREADS", "1")) - if "OMP_NUM_THREADS" not in os.environ: - logger.debug_once( - "OMP_NUM_THREADS is not set; defaulting Torch threads to %d for input preprocessing.", - num_threads, - ) - - with set_request_id(request_id), set_default_torch_num_threads(num_threads): + with set_request_id(request_id), set_default_torch_num_threads(): processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess( prompt, tokenization_kwargs=tokenization_kwargs, diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index e39dc9fec..c2520f168 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,4 +1,3 @@ -from ast import Dict from collections.abc import Callable from typing import Any @@ -37,8 +36,9 @@ class OmniRequestState(RequestState): **kwargs, ): super().__init__(*args, **kwargs) + # Omni-specific: multimodal output accumulation self.mm_type: str | None = None - self.mm_accumulated: Dict[str, Any] | None = None + self.mm_accumulated: dict[str, Any] | None = None def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: if payload is None: @@ -57,7 +57,7 @@ class OmniRequestState(RequestState): return x if isinstance(payload, dict): - incoming: Dict[str, Any] = {} + incoming: dict[str, Any] = {} target_key = self.mm_type or "hidden" # Iterate directly without unnecessary dict copy @@ -205,17 +205,18 @@ class OmniRequestState(RequestState): new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] self.sent_tokens_offset = len(self.detokenizer.output_token_ids) - request_id = self.request_id + external_req_id = self.external_req_id output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) if self.parent_req is None: outputs = [output] else: - request_id, outputs, finished = self.parent_req.get_outputs(request_id, output) + outputs, finished = self.parent_req.get_outputs(self.request_id, output) if not outputs: return None + external_req_id = self.parent_req.external_req_id - return self._new_request_output(request_id, outputs, finished, kv_transfer_params) + return self._new_request_output(external_req_id, outputs, finished, kv_transfer_params) def _new_completion_output( self, @@ -257,8 +258,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): def __init__( self, - tokenizer: TokenizerLike, + tokenizer: TokenizerLike | None, log_stats: bool, + stream_interval: int = 1, engine_core_output_type: str | None = None, ): """Initialize the multimodal output processor. @@ -266,11 +268,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): Args: tokenizer: Tokenizer for detokenizing text outputs log_stats: Whether to log statistics + stream_interval: Stream interval for output generation engine_core_output_type: Optional output type specification (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) + super().__init__(tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval) 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/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index b557c07dd..02d175ad0 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -6,7 +6,6 @@ import socket from typing import TYPE_CHECKING import torch -import vllm.envs as envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry @@ -164,21 +163,23 @@ class AsyncOmniLLM(AsyncLLM): except RuntimeError: pass - if envs.VLLM_TORCH_PROFILER_DIR and not envs.VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM: + # Use profiler_config from vllm_config (new way, aligned with vllm v1) + if vllm_config.profiler_config.profiler == "torch" and not vllm_config.profiler_config.ignore_frontend: + profiler_dir = vllm_config.profiler_config.torch_profiler_dir logger.info( "Torch profiler enabled. AsyncOmniLLM CPU traces will be collected under %s", - envs.VLLM_TORCH_PROFILER_DIR, + profiler_dir, ) worker_name = f"{socket.gethostname()}_{os.getpid()}.async_omni_llm" self.profiler = torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, ], - with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK, + with_stack=vllm_config.profiler_config.torch_profiler_with_stack, on_trace_ready=torch.profiler.tensorboard_trace_handler( - envs.VLLM_TORCH_PROFILER_DIR, + profiler_dir, worker_name=worker_name, - use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP, + use_gzip=vllm_config.profiler_config.torch_profiler_use_gzip, ), ) else: diff --git a/vllm_omni/entrypoints/chat_utils.py b/vllm_omni/entrypoints/chat_utils.py index e2cb64505..8970e5898 100644 --- a/vllm_omni/entrypoints/chat_utils.py +++ b/vllm_omni/entrypoints/chat_utils.py @@ -1,307 +1,62 @@ -from collections.abc import Awaitable, Iterable -from typing import Any, cast - import numpy as np -from openai.types.chat import ChatCompletionContentPartTextParam -from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - AsyncMultiModalContentParser, - AsyncMultiModalItemTracker, - BaseMultiModalContentParser, - BaseMultiModalItemTracker, - ChatCompletionContentPartParam, - ChatCompletionMessageParam, - ChatTemplateContentFormat, - ConversationMessage, - MultiModalDataDict, - MultiModalUUIDDict, - _AssistantParser, - _ContentPart, - _get_full_multimodal_text_prompt, - _parse_chat_message_content_part, - _postprocess_messages, - _ToolParser, -) -from vllm.logger import init_logger - -logger = init_logger(__name__) -class OmniAsyncMultiModalItemTracker(AsyncMultiModalItemTracker): - def create_parser(self) -> "BaseMultiModalContentParser": - return OmniAsyncMultiModalContentParser(self) +async def extract_audio_from_video_async(video_url: str) -> tuple[np.ndarray, int | float]: + """Extract audio from a video URL using librosa. + Returns a (audio_array, sample_rate) tuple compatible with audio format. + All blocking I/O operations are run in a thread pool. + """ + import asyncio + import os + import tempfile + from urllib.parse import urlparse -class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser): - def __init__(self, tracker: AsyncMultiModalItemTracker) -> None: - super().__init__(tracker=tracker) - self._mm_processor_kwargs: dict[str, Any] | None = None + parsed_url = urlparse(video_url) + temp_video_file_path = None - def set_mm_processor_kwargs(self, mm_processor_kwargs: dict[str, Any] | None) -> None: - """Set mm_processor_kwargs for use in parsing.""" - self._mm_processor_kwargs = mm_processor_kwargs + def _download_video_sync(url: str) -> bytes: + from urllib.request import urlopen - def parse_video(self, video_url: str | None, uuid: str | None = None) -> None: - # OMNI: Follow upstream async pattern - create coroutine that resolves to (data, uuid) - coro = self._video_with_uuid_async(video_url, uuid) - placeholder = self._tracker.add("video", coro) - self._add_placeholder("video", placeholder) + return urlopen(url).read() - # Extract audio from video if use_audio_in_video is True - if video_url and self._mm_processor_kwargs and self._mm_processor_kwargs.get("use_audio_in_video", False): - audio_coro = self._audio_from_video_with_uuid_async(video_url, uuid) - audio_placeholder = self._tracker.add("audio", audio_coro) - self._add_placeholder("audio", audio_placeholder) + def _write_temp_file_sync(data: bytes, suffix: str) -> str: + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: + temp_file.write(data) + return temp_file.name - async def _video_with_uuid_async(self, video_url: str | None, uuid: str | None): - """Fetch video and return (video, uuid) tuple.""" - video = await self._connector.fetch_video_async(video_url=video_url) if video_url else None - return video, uuid + def _load_audio_sync(file_path: str) -> tuple[np.ndarray, int | float]: + import librosa - async def _audio_from_video_with_uuid_async(self, video_url: str, uuid: str | None): - """Extract audio from video and return (audio, uuid) tuple.""" - audio = await self._extract_audio_from_video_async(video_url) - return audio, uuid - - async def _extract_audio_from_video_async(self, video_url: str) -> tuple[np.ndarray, int | float]: - """ - Extract audio from video URL using librosa. - Returns tuple of (audio_array, sample_rate) compatible with audio format. - - All blocking I/O operations are run in a thread pool to avoid blocking the event loop. - """ - import asyncio - import os - import tempfile - from urllib.parse import urlparse - - # Parse URL to determine type - parsed_url = urlparse(video_url) - temp_video_file_path = None - - def _download_video_sync(url: str) -> bytes: - """Synchronous video download - runs in thread pool.""" - from urllib.request import urlopen - - return urlopen(url).read() - - def _write_temp_file_sync(data: bytes, suffix: str) -> str: - """Synchronous temp file write - runs in thread pool.""" - with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: - temp_file.write(data) - return temp_file.name - - def _load_audio_sync(file_path: str) -> tuple[np.ndarray, int | float]: - """Synchronous audio loading with librosa - runs in thread pool.""" - import librosa - - return librosa.load(file_path, sr=16000) - - def _cleanup_file_sync(file_path: str) -> None: - """Synchronous file deletion - runs in thread pool.""" - try: - if os.path.exists(file_path): - os.unlink(file_path) - except OSError: - pass + return librosa.load(file_path, sr=16000) + def _cleanup_file_sync(file_path: str) -> None: try: - if parsed_url.scheme in ("http", "https"): - # Download video from HTTP/HTTPS URL asynchronously - video_data = await asyncio.to_thread(_download_video_sync, video_url) - # Write temp file asynchronously - temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4") - elif parsed_url.scheme == "file": - # Use file path directly (handle Windows paths) - from urllib.request import url2pathname + if os.path.exists(file_path): + os.unlink(file_path) + except OSError: + pass - temp_video_file_path = url2pathname(parsed_url.path) - elif parsed_url.scheme == "data": - # Handle data URL (base64 encoded video) - import base64 + try: + if parsed_url.scheme in ("http", "https"): + video_data = await asyncio.to_thread(_download_video_sync, video_url) + temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4") + elif parsed_url.scheme == "file": + from urllib.request import url2pathname - header, data = video_url.split(",", 1) - video_data = base64.b64decode(data) - # Write temp file asynchronously - temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4") - else: - # Assume it's a local file path - temp_video_file_path = video_url + temp_video_file_path = url2pathname(parsed_url.path) + elif parsed_url.scheme == "data": + import base64 - # Extract audio using librosa asynchronously (CPU-intensive, runs in thread pool) - audio_array, sample_rate = await asyncio.to_thread(_load_audio_sync, temp_video_file_path) + header, data = video_url.split(",", 1) + video_data = base64.b64decode(data) + temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4") + else: + # Assume local file path + temp_video_file_path = video_url - return audio_array, sample_rate - finally: - # Clean up temporary file if we created one (asynchronously) - if temp_video_file_path and parsed_url.scheme in ("http", "https", "data"): - await asyncio.to_thread(_cleanup_file_sync, temp_video_file_path) - - -def _ensure_system_prompt( - messages: list[ChatCompletionMessageParam], - model_config: ModelConfig, -) -> list[ChatCompletionMessageParam]: - """ - Ensure a system prompt exists for Qwen-Omni models to preserve precision of the model. - Args: - messages: List of chat messages - model_config: Model configuration - - Returns: - Messages list with system prompt - """ - model_name = getattr(model_config, "model", "").lower() - hf_config = getattr(model_config, "hf_config", None) - architectures = getattr(hf_config, "architectures", []) if hf_config else [] - - is_qwen_omni = ("qwen" in model_name and "omni" in model_name) or any( - "qwen" in arch.lower() and "omni" in arch.lower() for arch in architectures - ) - - if not is_qwen_omni: - return messages - - if messages and messages[0].get("role") == "system": - return messages - - default_qwen_omni_system_prompt = ( - "You are Qwen, a virtual human developed by the Qwen Team, Alibaba " - "Group, capable of perceiving auditory and visual inputs, as well as " - "generating text and speech." - ) - - system_message: ChatCompletionMessageParam = { - "role": "system", - "content": default_qwen_omni_system_prompt, - } - - logger.info(f"injecting system prompt {default_qwen_omni_system_prompt} for Qwen-Omni model") - return [system_message] + list(messages) - - -def parse_chat_messages_futures( - messages: list[ChatCompletionMessageParam], - model_config: ModelConfig, - content_format: ChatTemplateContentFormat, - mm_processor_kwargs: dict[str, Any] | None = None, -) -> tuple[ - list[ConversationMessage], - Awaitable[tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]], -]: - """Parse chat messages and return conversation with multimodal data future. - - OMNI: Updated to use upstream vLLM v0.15.0 API where resolve_items() - returns both mm_data and mm_uuids together as a tuple. - - Returns: - Tuple of (conversation, mm_future) where mm_future resolves to - (mm_data, mm_uuids) when awaited. - """ - # auto-inject system prompt for Qwen-Omni models if missing - messages = _ensure_system_prompt(messages, model_config) - - conversation: list[ConversationMessage] = [] - mm_tracker = OmniAsyncMultiModalItemTracker(model_config) - - for msg in messages: - sub_messages = _parse_chat_message_content( - msg, - mm_tracker, - content_format, - interleave_strings=( - content_format == "string" - and model_config.multimodal_config is not None - and model_config.multimodal_config.interleave_mm_strings - ), - mm_processor_kwargs=mm_processor_kwargs, - ) - - conversation.extend(sub_messages) - - _postprocess_messages(conversation) - - # OMNI: Use upstream resolve_items() which returns (mm_data, mm_uuids) tuple - return conversation, mm_tracker.resolve_items() - - -def _parse_chat_message_content( - message: ChatCompletionMessageParam, - mm_tracker: BaseMultiModalItemTracker, - content_format: ChatTemplateContentFormat, - interleave_strings: bool, - mm_processor_kwargs: dict[str, Any] | None = None, -) -> list[ConversationMessage]: - role = message["role"] - content = message.get("content") - - if content is None: - content = [] - elif isinstance(content, str): - content = [ChatCompletionContentPartTextParam(type="text", text=content)] - result = _parse_chat_message_content_parts( - role, - content, # type: ignore - mm_tracker, - wrap_dicts=(content_format == "openai"), - interleave_strings=interleave_strings, - mm_processor_kwargs=mm_processor_kwargs, - ) - - for result_msg in result: - if role == "assistant": - parsed_msg = _AssistantParser(message) - - # The 'tool_calls' is not None check ensures compatibility. - # It's needed only if downstream code doesn't strictly - # follow the OpenAI spec. - if "tool_calls" in parsed_msg and parsed_msg["tool_calls"] is not None: - result_msg["tool_calls"] = list(parsed_msg["tool_calls"]) - elif role == "tool": - parsed_msg = _ToolParser(message) - if "tool_call_id" in parsed_msg: - result_msg["tool_call_id"] = parsed_msg["tool_call_id"] - - if "name" in message and isinstance(message["name"], str): - result_msg["name"] = message["name"] - - return result - - -def _parse_chat_message_content_parts( - role: str, - parts: Iterable[ChatCompletionContentPartParam], - mm_tracker: BaseMultiModalItemTracker, - *, - wrap_dicts: bool, - interleave_strings: bool, - mm_processor_kwargs: dict[str, Any] | None = None, -) -> list[ConversationMessage]: - content = list[_ContentPart]() - - mm_parser = mm_tracker.create_parser() - # Set mm_processor_kwargs if parser supports it - if hasattr(mm_parser, "set_mm_processor_kwargs"): - mm_parser.set_mm_processor_kwargs(mm_processor_kwargs) - - for part in parts: - parse_res = _parse_chat_message_content_part( - part, - mm_parser, - wrap_dicts=wrap_dicts, - interleave_strings=interleave_strings, - ) - if parse_res: - content.append(parse_res) - - if wrap_dicts: - # Parsing wraps images and texts as interleaved dictionaries - return [ConversationMessage(role=role, content=content)] # type: ignore - texts = cast(list[str], content) - mm_placeholder_storage = mm_parser.mm_placeholder_storage() - if mm_placeholder_storage: - text_prompt = _get_full_multimodal_text_prompt(mm_placeholder_storage, texts, interleave_strings) - else: - text_prompt = "\n".join(texts) - - return [ConversationMessage(role=role, content=text_prompt)] + audio_array, sample_rate = await asyncio.to_thread(_load_audio_sync, temp_video_file_path) + return audio_array, sample_rate + finally: + if temp_video_file_path and parsed_url.scheme in ("http", "https", "data"): + await asyncio.to_thread(_cleanup_file_sync, temp_video_file_path) diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index f53177ce3..6c9723b6b 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -1126,6 +1126,9 @@ async def _stage_worker_async( vllm_config=vllm_config, usage_context=usage_context, engine_args=omni_engine_args, + disable_log_stats=bool( + engine_args.get("disable_log_stats", True) or getattr(omni_engine_args, "disable_log_stats", True) + ), ) omni_stage.set_async_engine(stage_engine) if hasattr(omni_stage.async_engine, "log_stats") and omni_stage.async_engine.log_stats: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 521bdf0f3..3898da308 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -31,9 +31,15 @@ from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launcher import serve_http from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer -from vllm.entrypoints.openai.api_server import base, load_log_config from vllm.entrypoints.openai.api_server import build_app as build_openai_app from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server + +# vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. +# Keep a fallback for older/newer upstream layouts during rebase windows. +try: + from vllm.entrypoints.serve.instrumentator.basic import base +except ModuleNotFoundError: + from vllm.entrypoints.openai.basic.api_router import base from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -53,7 +59,8 @@ from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.orca_metrics import metrics_header from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses -from vllm.entrypoints.openai.translations.serving import ( +from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config +from vllm.entrypoints.openai.speech_to_text.serving import ( OpenAIServingTranscription, OpenAIServingTranslation, ) @@ -195,7 +202,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) # Load logging config for uvicorn if specified - log_config = load_log_config(getattr(args, "log_config_file", None)) + log_config = get_uvicorn_log_config(args) if log_config is not None: uvicorn_kwargs["log_config"] = log_config @@ -211,7 +218,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, if not supported_tasks: supported_tasks = ("generate",) - app = build_openai_app(args) + # OMNI: Pass supported_tasks to build_app (required by upstream vLLM) + app = build_openai_app(args, supported_tasks) # OMNI: Remove upstream routes that we override with omni-specific handlers _remove_route_from_app(app, "/v1/chat/completions", {"POST"}) _remove_route_from_app(app, "/v1/models", {"GET"}) # Remove upstream /v1/models to use omni's handler diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 4cc181e09..4271d799e 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -3,7 +3,7 @@ import base64 import json import time import uuid -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable from datetime import datetime, timedelta, timezone from io import BytesIO from typing import TYPE_CHECKING, Any, Final, Optional, cast @@ -13,7 +13,7 @@ import torch from fastapi import Request from PIL import Image from pydantic import TypeAdapter -from vllm.renderers import RendererLike +from vllm.renderers.protocol import BaseRenderer from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse @@ -63,20 +63,15 @@ from vllm.entrypoints.openai.parser.harmony_utils import ( from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls from vllm.entrypoints.utils import should_include_usage -from vllm.inputs.data import PromptType, TokensPrompt +from vllm.inputs.data import PromptType from vllm.logger import init_logger from vllm.outputs import RequestOutput -from vllm.renderers.hf import ( - resolve_chat_template_content_format, -) -from vllm.renderers.hf import ( - safe_apply_chat_template as apply_hf_chat_template, -) -from vllm.renderers.mistral import ( - safe_apply_chat_template as apply_mistral_chat_template, -) +from vllm.reasoning import ReasoningParser +from vllm.renderers import merge_kwargs +from vllm.renderers.inputs import TokPrompt from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike +from vllm.tokenizers import TokenizerLike as AnyTokenizer from vllm.tokenizers.mistral import ( MistralTokenizer, maybe_serialize_tool_calls, @@ -85,10 +80,8 @@ from vllm.tokenizers.mistral import ( ) from vllm.tool_parsers import ToolParser from vllm.tool_parsers.mistral_tool_parser import MistralToolCall -from vllm.transformers_utils.tokenizer import AnyTokenizer -from vllm.utils.collection_utils import as_list, is_list_of +from vllm.utils.collection_utils import as_list -from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio @@ -182,6 +175,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if tokenizer is None: tokenizer = await self.engine_client.get_tokenizer() + reasoning_parser: ReasoningParser | None = None + if self.reasoning_parser_cls: + chat_template_kwargs = self._prepare_extra_chat_template_kwargs( + request.chat_template_kwargs, + self.default_chat_template_kwargs, + ) + reasoning_parser = self.reasoning_parser_cls( + tokenizer, + chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] + ) + tool_parser = self.tool_parser if isinstance(tokenizer, MistralTokenizer): @@ -231,29 +235,29 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): chat_template_kwargs = request.chat_template_kwargs or {} chat_template_kwargs.update(reasoning_effort=request.reasoning_effort) - ( - conversation, - request_prompts, - engine_prompts, - ) = await self._preprocess_chat( + # Merge chat_template_kwargs with defaults + merged_template_kwargs = self._prepare_extra_chat_template_kwargs( + chat_template_kwargs, + self.default_chat_template_kwargs, + ) + conversation, engine_prompts = await self._preprocess_chat( request, - renderer, request.messages, - chat_template=request.chat_template or self.chat_template, - chat_template_content_format=self.chat_template_content_format, + default_template=request.chat_template or self.chat_template, + default_template_content_format=self.chat_template_content_format, + default_template_kwargs=merged_template_kwargs, + tool_dicts=tool_dicts, + tool_parser=tool_parser, + # OMNI: Additional parameters + renderer=renderer, add_generation_prompt=request.add_generation_prompt, continue_final_message=request.continue_final_message, - tool_dicts=tool_dicts, documents=getattr(request, "documents", None), - chat_template_kwargs=chat_template_kwargs, - default_chat_template_kwargs=self.default_chat_template_kwargs, - tool_parser=tool_parser, add_special_tokens=request.add_special_tokens, ) else: should_include_tools = tool_dicts is not None conversation, engine_prompts = self._make_request_with_harmony(request, should_include_tools) - request_prompts = [engine_prompt.get("prompt_token_ids", []) for engine_prompt in engine_prompts] except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e: logger.exception("Error in preprocessing prompt inputs") @@ -282,7 +286,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): self._log_inputs( request_id, - request_prompts[i], + engine_prompt, params_list=sampling_params_list, lora_request=lora_request, ) @@ -311,6 +315,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation, tokenizer, request_metadata, + reasoning_parser, ) try: @@ -322,6 +327,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation, tokenizer, request_metadata, + reasoning_parser, ) except ValueError as e: return self.create_error_response(e) @@ -329,88 +335,74 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): async def _preprocess_chat( self, request: ChatLikeRequest | ResponsesRequest, - renderer: RendererLike, messages: list[ChatCompletionMessageParam], - chat_template: str | None, - chat_template_content_format: ChatTemplateContentFormatOption, + default_template: str | None, + default_template_content_format: ChatTemplateContentFormatOption, + default_template_kwargs: dict[str, Any] | None = None, + tool_dicts: list[dict[str, Any]] | None = None, + tool_parser: Callable[[TokenizerLike], ToolParser] | None = None, + # OMNI: Additional parameters for backward compatibility + renderer: BaseRenderer | None = None, add_generation_prompt: bool = True, continue_final_message: bool = False, - tool_dicts: list[dict[str, Any]] | None = None, documents: list[dict[str, str]] | None = None, - chat_template_kwargs: dict[str, Any] | None = None, - default_chat_template_kwargs: dict[str, Any] | None = None, - tool_parser: Callable[[TokenizerLike], ToolParser] | None = None, add_special_tokens: bool = False, - ) -> tuple[ - list[ConversationMessage], - Sequence[PromptType], - list[TokensPrompt], - ]: - model_config = self.model_config - tokenizer = renderer.get_tokenizer() if renderer is not None else None + ) -> tuple[list[ConversationMessage], list[TokPrompt]]: + if renderer is None: + renderer = self.renderer - if tokenizer is None or isinstance(tokenizer, MistralTokenizer): - resolved_content_format = ( - chat_template_content_format if chat_template_content_format != "auto" else "string" - ) - else: - resolved_content_format = resolve_chat_template_content_format( - chat_template, - tool_dicts, - chat_template_content_format, - tokenizer, - model_config=model_config, - ) - # OMNI: Updated for vLLM v0.15.0 API - resolve_items() returns (mm_data, mm_uuids) tuple - conversation, mm_future = parse_chat_messages_futures( - messages, - model_config, - content_format=resolved_content_format, - mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), - ) - - # Merge default_chat_template_kwargs with request-provided kwargs - # Request kwargs take precedence over defaults - merged_kwargs = self._prepare_extra_chat_template_kwargs( - chat_template_kwargs, - default_chat_template_kwargs, - ) - - _chat_template_kwargs: dict[str, Any] = dict( - chat_template=chat_template, - add_generation_prompt=add_generation_prompt, - continue_final_message=continue_final_message, - tools=tool_dicts, - documents=documents, - ) - _chat_template_kwargs.update(merged_kwargs) - - request_prompt: str | list[int] - - if tokenizer is None: - request_prompt = "placeholder" - elif isinstance(tokenizer, MistralTokenizer): - request_prompt = apply_mistral_chat_template( - tokenizer, - messages=messages, - **_chat_template_kwargs, - ) - else: - hf_chat_template_kwargs = dict(_chat_template_kwargs) - hf_chat_template_kwargs.pop("tools", None) - hf_chat_template_kwargs.pop("chat_template", None) - request_prompt = apply_hf_chat_template( - model_config=model_config, - tokenizer=tokenizer, - conversation=conversation, + # Keep OMNI compatibility args wired while delegating rendering + # to the upstream async renderer pipeline. + default_template_kwargs = merge_kwargs( + default_template_kwargs, + dict( tools=tool_dicts, - chat_template=chat_template, - tokenize=False, - **hf_chat_template_kwargs, - ) + documents=documents, + add_generation_prompt=add_generation_prompt, + continue_final_message=continue_final_message, + add_special_tokens=add_special_tokens, + tokenize=isinstance(renderer.tokenizer, MistralTokenizer), + ), + ) - # OMNI: Await the combined future to get both mm_data and mm_uuids - mm_data, mm_uuids = await mm_future + tok_params = request.build_tok_params(self.model_config) + chat_params = request.build_chat_params( + default_template, + default_template_content_format, + ).with_defaults(default_template_kwargs) + + (conversation,), (engine_prompt,) = await renderer.render_chat_async( + [messages], + chat_params, + tok_params, + prompt_extras={ + k: v for k in ("mm_processor_kwargs", "cache_salt") if (v := getattr(request, k, None)) is not None + }, + ) + + # 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 # tool_choice is not "none" (if tool_choice is "none" but a tool_parser @@ -428,42 +420,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): request=request ) - if tokenizer is None: - assert isinstance(request_prompt, str), ( - "Prompt has to be a string", - "when the tokenizer is not initialised", - ) - prompt_inputs = TokensPrompt(prompt=request_prompt, prompt_token_ids=[1]) - elif isinstance(request_prompt, str): - prompt_inputs = await self._tokenize_prompt_input_async( - request, - tokenizer, - request_prompt, - add_special_tokens=add_special_tokens, - ) - else: - # For MistralTokenizer - assert is_list_of(request_prompt, int), "Prompt has to be either a string or a list of token ids" - prompt_inputs = TokensPrompt( - prompt=tokenizer.decode(request_prompt), - prompt_token_ids=request_prompt, - ) - - engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"]) - if mm_data is not None: - engine_prompt["multi_modal_data"] = mm_data - - if mm_uuids is not None: - engine_prompt["multi_modal_uuids"] = mm_uuids - - mm_processor_kwargs = getattr(request, "mm_processor_kwargs", None) - if mm_processor_kwargs is not None: - engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs - - if hasattr(request, "cache_salt") and request.cache_salt is not None: - engine_prompt["cache_salt"] = request.cache_salt - - return conversation, [request_prompt], [engine_prompt] + return conversation, [engine_prompt] def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]: final_sampling_params_list = [] @@ -563,29 +520,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): def _log_inputs( self, request_id: str, - inputs: PromptType, + inputs: PromptType | TokPrompt, params_list: list[SamplingParams] | None, lora_request: LoRARequest | None, ) -> None: if self.request_logger is None: return - prompt, prompt_token_ids, prompt_embeds = None, None, None - if isinstance(inputs, str): - prompt = inputs - elif isinstance(inputs, list): - prompt_token_ids = inputs - else: - prompt = getattr(inputs, "prompt", None) - prompt_token_ids = getattr(inputs, "prompt_token_ids", None) - - logger.info( - "Received request %s: prompt: %r, params_list: %s, prompt_token_ids: %s, prompt_embeds shape: %s, lora_request: %s.", # noqa: E501 + components = self._extract_prompt_components(inputs) + self.request_logger.log_inputs( request_id, - prompt, - params_list, - prompt_token_ids, - prompt_embeds.shape if prompt_embeds is not None else None, - lora_request, + components.text, + components.token_ids, + components.embeds, + params=params_list, + lora_request=lora_request, ) async def chat_completion_stream_generator( @@ -597,6 +545,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation: list[ConversationMessage], tokenizer: AnyTokenizer, request_metadata: RequestResponseMetadata, + reasoning_parser: ReasoningParser | None = None, ): created_time = int(time.time()) chunk_object_type: Final = "chat.completion.chunk" @@ -638,7 +587,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Only one of these will be used, thus previous_texts and # all_previous_token_ids will not be used twice in the same iteration. - if tool_choice_auto or self.reasoning_parser: + if tool_choice_auto or reasoning_parser: # These are only required in "auto" tool choice case all_previous_token_ids = [[]] * num_choices # For reasoning parser and tool call all enabled @@ -646,23 +595,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): reasoning_end_arr = [False] * num_choices else: all_previous_token_ids = None - - try: - if self.reasoning_parser: - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser( - tokenizer, - chat_template_kwargs=chat_template_kwargs, # type: ignore - ) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - data = self.create_streaming_error_response(e) - yield f"data: {data}\n\n" - yield "data: [DONE]\n\n" - return # Prepare the tool parser if it's needed try: if tool_choice_auto and self.tool_parser: @@ -818,7 +750,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): delta_message: DeltaMessage | None # just update previous_texts and previous_token_ids - if tool_choice_auto or self.reasoning_parser: + if tool_choice_auto or reasoning_parser: assert previous_texts is not None assert all_previous_token_ids is not None previous_text = previous_texts[i] @@ -885,7 +817,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # handle streaming deltas for tools with named tool_choice elif tool_choice_function_name: if ( - self.reasoning_parser + reasoning_parser and not reasoning_end_arr[i] and not reasoning_parser.is_reasoning_end(previous_token_ids) ): @@ -915,7 +847,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): current_text = "" else: # Just to add remaining `content` - if self.reasoning_parser: + if reasoning_parser: delta_text = previous_text + delta_text current_text = "" @@ -951,14 +883,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): output_token_ids = as_list(output.token_ids) if ( - self.reasoning_parser is not None + reasoning_parser is not None and not reasoning_end_arr[i] and res.prompt_token_ids and reasoning_parser.is_reasoning_end(res.prompt_token_ids) ): reasoning_end_arr[i] = True - if self.reasoning_parser and not reasoning_end_arr[i]: + if reasoning_parser and not reasoning_end_arr[i]: delta_message = reasoning_parser.extract_reasoning_streaming( previous_text, current_text, @@ -997,7 +929,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # handle streaming deltas for tools with "auto" tool choice # and reasoning parser - elif tool_choice_auto and self.reasoning_parser: + elif tool_choice_auto and reasoning_parser: assert tool_parser is not None assert reasoning_parser is not None assert added_content_delta_arr is not None @@ -1078,7 +1010,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tools_streamed[i] = True # when only reasoning - elif self.reasoning_parser: + elif reasoning_parser: delta_message = reasoning_parser.extract_reasoning_streaming( previous_text, current_text, @@ -1092,7 +1024,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): delta_message = DeltaMessage(content=delta_text) # update the previous values for the next iteration - if (tool_choice_auto or self.reasoning_parser) and not self.use_harmony: + if (tool_choice_auto or reasoning_parser) and not self.use_harmony: assert previous_texts is not None assert all_previous_token_ids is not None previous_texts[i] = current_text @@ -1332,6 +1264,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation: list[ConversationMessage], tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, + reasoning_parser: ReasoningParser | None = None, ) -> ErrorResponse | OmniChatCompletionResponse: created_time = int(time.time()) final_res: RequestOutput | None = None @@ -1378,7 +1311,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): prompt_logprobs, prompt_token_ids, kv_transfer_params, - ) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role) + ) = self._create_text_choice( + request, + omni_outputs, + tokenizer, + conversation, + role, + reasoning_parser, + ) elif omni_outputs.final_output_type == "audio": choices_data = self._create_audio_choice(omni_outputs, role, request, stream=False) elif omni_outputs.final_output_type == "image": @@ -1441,6 +1381,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer: TokenizerLike, conversation: list[ConversationMessage], role: str, + reasoning_parser: ReasoningParser | None = None, ): final_res = omni_outputs.request_output if self.tool_call_id_type == "kimi_k2": @@ -1508,19 +1449,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): choices.append(choice_data) continue - if self.reasoning_parser: - try: - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser( - tokenizer, - chat_template_kwargs=chat_template_kwargs, # type: ignore - ) - except RuntimeError as e: - logger.exception("Error in reasoning parser creation.") - return self.create_error_response(e) + if reasoning_parser: # If the reasoning parser is enabled, # tool calls are extracted exclusively from the content. reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index bc9434817..bf8563bd2 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -3,7 +3,8 @@ import pprint from dataclasses import asdict, dataclass, field from typing import Any, TypeAlias -from vllm import PromptType, SamplingParams +from vllm.inputs import PromptType +from vllm.sampling_params import SamplingParams from vllm_omni.lora.request import LoRARequest @@ -13,6 +14,7 @@ except ImportError: # Python < 3.11: use typing_extensions from typing_extensions import NotRequired + import torch from vllm.inputs.data import EmbedsPrompt, TextPrompt, TokenInputs, TokensPrompt @@ -104,8 +106,8 @@ class OmniEmbedsPrompt(EmbedsPrompt): # Must ensure that all additional prompt types are inherited from vLLM prompt types # Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime. # Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate() -OmniSingletonPrompt: TypeAlias = str | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt -"""Omni singleton prompt type extending vLLM's SingletonPrompt with additional fields.""" +OmniSingletonPrompt: TypeAlias = str | list[int] | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt +"""Omni singleton prompt type extending vLLM's SingletonPrompt.""" OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py deleted file mode 100644 index 3fe46ff53..000000000 --- a/vllm_omni/inputs/parse.py +++ /dev/null @@ -1,42 +0,0 @@ -from vllm.inputs.parse import ( - ParsedEmbedsPrompt, - ParsedSingletonPrompt, - ParsedStrPrompt, - ParsedTextPrompt, - ParsedTokensPrompt, - SingletonPrompt, -) - - -def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt: - """Parse a singleton prompt into a typed parsed prompt. - - Handles omni-specific prompt types including tokens prompts with - embeddings and additional information. Supports string, text, - tokens, and embeddings prompts. - - Args: - prompt: Singleton prompt to parse. Can be a string, TextPrompt, - TokensPrompt (with optional prompt_embeds and additional_information), - or EmbedsPrompt. - - Returns: - ParsedSingletonPrompt containing the parsed prompt with type information - - Raises: - TypeError: If the prompt type is not supported - """ - if isinstance(prompt, str): - return ParsedStrPrompt(type="str", content=prompt) - elif isinstance(prompt, dict): - # Type ignores are because mypy does not correctly infer the TypedDicts - # Pyright does succeed. - # Priority tokens: When both tokens and embeds exist, keep both and - # follow the tokens path - if "prompt_token_ids" in prompt: - return ParsedTokensPrompt(type="tokens", content=prompt) # type: ignore[typeddict-item] - elif "prompt_embeds" in prompt: - return ParsedEmbedsPrompt(type="embeds", content=prompt) # type: ignore[typeddict-item] - elif "prompt" in prompt: - return ParsedTextPrompt(type="text", content=prompt) - raise TypeError("inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt") diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index d04c66331..09b215bf9 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -1,18 +1,19 @@ from typing import Any from typing_extensions import assert_never -from vllm.inputs.data import SingletonInputs, SingletonPrompt +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.renderers.inputs import SingletonDictPrompt from vllm_omni.inputs.data import ( + OmniEmbedsPrompt, OmniTextPrompt, OmniTokenInputs, OmniTokensPrompt, token_inputs_omni, ) -from vllm_omni.inputs.parse import parse_singleton_prompt_omni logger = init_logger(__name__) @@ -101,9 +102,28 @@ class OmniInputPreprocessor(InputPreprocessor): return inputs + def _process_embeds( + self, + parsed_content: OmniEmbedsPrompt, + ) -> EmbedsInputs: + """Process embeddings prompt with omni-specific extensions. + + Extends base _process_embeds to handle additional_information payload + for direct transfer between pipeline stages. + """ + # Call parent implementation for base embeds processing + inputs = super()._process_embeds(parsed_content) + + # Add omni-specific additional_information if present + additional_information = parsed_content.get("additional_information") + if additional_information is not None: + inputs["additional_information"] = additional_information # type: ignore[typeddict-unknown-key] + + return inputs + def _prompt_to_llm_inputs( self, - prompt: SingletonPrompt, + prompt: SingletonDictPrompt, tokenization_kwargs: dict[str, Any] | None = None, *, mm_uuids: MultiModalUUIDDict | None = None, @@ -114,31 +134,25 @@ class OmniInputPreprocessor(InputPreprocessor): Arguments: * prompt: single encoder or decoder input prompt - * lora_request: this is only valid for decoder prompts - * return_mm_hashes: whether to return multimodal hashes Returns: - * Input container compatible with vLLM's singleton prompt handling. + * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance """ - parsed = parse_singleton_prompt_omni(prompt) - - if parsed["type"] == "tokens": + if "prompt_token_ids" in prompt: return self._process_tokens( - parsed["content"], + prompt, # type: ignore[arg-type] mm_uuids=mm_uuids, ) - if parsed["type"] == "text": + + if "prompt_embeds" in prompt: + return self._process_embeds(prompt) # type: ignore[arg-type] + + if "prompt" in prompt: return self._process_text( - parsed["content"], - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) - if parsed["type"] == "str": - return self._process_text( - OmniTextPrompt(prompt=parsed["content"]), + prompt, # type: ignore[arg-type] tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, ) - assert_never(parsed) + assert_never(prompt) # type: ignore[arg-type] 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 bb2de4f14..8ee516376 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 @@ -27,6 +27,8 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerDummyInputsBuilder, Qwen2_5OmniThinkerMultiModalProcessor, Qwen2_5OmniThinkerProcessingInfo, + check_interleaved_audio_video, + merge_interleaved_embeddings, ) from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase, @@ -43,6 +45,7 @@ from vllm.model_executor.models.qwen2_5_vl import ( from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, + _merge_multimodal_embeddings, init_vllm_registered_model, maybe_prefix, split_list_into_ranges, @@ -275,28 +278,32 @@ class Qwen2_5OmniThinkerForConditionalGeneration( "in the audio tower part." ) - if multimodal_config.get_limit_per_prompt("audio"): - self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config) - else: - self.audio_tower = None - - if multimodal_config.get_limit_per_prompt("image") or multimodal_config.get_limit_per_prompt("video"): - self.visual = Qwen2_5_VisionTransformer( - vision_config=thinker_config.vision_config, - norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), - quant_config=quant_config, - prefix=maybe_prefix(prefix, "visual"), - ) - else: - self.visual = None - self.quant_config = quant_config - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "language_model"), - hf_config=thinker_config.text_config, - architectures=["Qwen2ForCausalLM"], - ) + + with self._mark_tower_model(vllm_config, "audio"): + if multimodal_config.get_limit_per_prompt("audio"): + self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config) + else: + self.audio_tower = None + + with self._mark_tower_model(vllm_config, {"image", "video"}): + if multimodal_config.get_limit_per_prompt("image") or multimodal_config.get_limit_per_prompt("video"): + self.visual = Qwen2_5_VisionTransformer( + vision_config=thinker_config.vision_config, + norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + ) + else: + self.visual = None + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + hf_config=thinker_config.text_config, + architectures=["Qwen2ForCausalLM"], + ) self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors @@ -310,7 +317,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration( mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs) if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality: mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs) - if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality: + if input_key in ("input_audio_features",) and "audio" not in mm_input_by_modality: mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs) return mm_input_by_modality @@ -503,8 +510,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( multimodal_embeddings += tuple(audio_embeddings) return multimodal_embeddings - # TODO (ywang96): support overlapping modality embeddings so that - # `use_audio_in_video` will work on V1. def embed_input_ids( self, input_ids: torch.Tensor, @@ -517,13 +522,41 @@ class Qwen2_5OmniThinkerForConditionalGeneration( if multimodal_embeddings is None or is_multimodal is None: return super().embed_input_ids(input_ids) - return super().embed_input_ids( + inputs_embeds = self._embed_text_input_ids( input_ids, - multimodal_embeddings=multimodal_embeddings, + self.get_language_model().embed_input_ids, is_multimodal=is_multimodal, handle_oov_mm_token=handle_oov_mm_token, ) + if len(multimodal_embeddings) == 0: + return inputs_embeds + + # Check for audio-in-video: interleaved video and audio tokens + # in the multimodal region. + video_token_id = self.config.video_token_index + audio_token_id = self.config.audio_token_index + + is_video = is_multimodal & (input_ids == video_token_id) + is_audio = is_multimodal & (input_ids == audio_token_id) + + num_video = is_video.sum().item() + num_audio = is_audio.sum().item() + + if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio): + return merge_interleaved_embeddings( + inputs_embeds, + multimodal_embeddings, + is_video, + is_audio, + is_multimodal, + num_video, + num_audio, + ) + + # Default: standard merge (no interleaving) + return _merge_multimodal_embeddings(inputs_embeds, multimodal_embeddings, is_multimodal) + def forward( self, input_ids: torch.Tensor, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py index e04010196..27eb25720 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py @@ -3,12 +3,12 @@ from collections.abc import Iterable import torch from torch import nn from transformers import Qwen2Config -from vllm.attention.layer import Attention from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor 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 f54182eed..a0920c069 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 @@ -22,9 +22,9 @@ # limitations under the License. """Inference-only Qwen3-Omni-Moe model (thinker part).""" -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from functools import partial -from typing import Any +from typing import Any, Literal, cast import numpy as np import torch @@ -42,8 +42,9 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import ( ) from transformers.models.whisper import WhisperFeatureExtractor from vllm.compilation.decorators import support_torch_compile -from vllm.config import VllmConfig +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig from vllm.distributed import get_pp_group +from vllm.inputs.data import PromptType from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead @@ -52,12 +53,15 @@ from vllm.model_executor.models.interfaces import ( SupportsMRoPE, SupportsMultiModal, SupportsPP, + SupportsTranscription, ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniAudioFeatureInputs, Qwen2_5OmniThinkerDummyInputsBuilder, Qwen2_5OmniThinkerMultiModalProcessor, + check_interleaved_audio_video, + merge_interleaved_embeddings, ) from vllm.model_executor.models.qwen2_5_vl import ( Qwen2_5_VLProcessingInfo, @@ -76,9 +80,6 @@ from vllm.model_executor.models.utils import ( _merge_multimodal_embeddings, maybe_prefix, ) -from vllm.model_executor.models.vision import ( - get_llm_pos_ids_for_vision, -) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems @@ -90,6 +91,7 @@ from vllm.multimodal.processing.processor import ( PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.processor import cached_processor_from_config from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, @@ -102,6 +104,29 @@ except (ImportError, ModuleNotFoundError): logger = init_logger(__name__) +# Speech input languages supported by Qwen3-Omni +# From: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct +ISO639_1_SUPPORTED_LANGS = { + "en": "English", + "zh": "Chinese", + "ko": "Korean", + "ja": "Japanese", + "de": "German", + "ru": "Russian", + "it": "Italian", + "fr": "French", + "es": "Spanish", + "pt": "Portuguese", + "ms": "Malay", + "nl": "Dutch", + "id": "Indonesian", + "tr": "Turkish", + "vi": "Vietnamese", + "yue": "Cantonese", + "ar": "Arabic", + "ur": "Urdu", +} + @support_torch_compile( dynamic_arg_dims={ @@ -119,7 +144,6 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel): positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, - *, capture_layer_indices: Sequence[int] | None = None, return_hidden_states: bool = False, deepstack_input_embeds: IntermediateTensors | None = None, @@ -639,6 +663,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( SupportsPP, SupportsMRoPE, Qwen3OmniMoeConditionalGenerationMixin, + SupportsTranscription, ): hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -660,6 +685,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ], } + supported_languages = ISO639_1_SUPPORTED_LANGS + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): @@ -679,46 +706,57 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( multimodal_config = vllm_config.model_config.multimodal_config self.config = thinker_config self.multimodal_config = multimodal_config - - self.audio_tower = Qwen3OmniMoeAudioEncoder( - thinker_config.audio_config, - ) - - self.visual = Qwen3Omni_VisionTransformer( - vision_config=thinker_config.vision_config, - norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), - quant_config=quant_config, - prefix=maybe_prefix(prefix, "visual"), - ) self.quant_config = quant_config - self.language_model = Qwen3MoeLLMForCausalLM( - vllm_config=vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"]), - prefix=maybe_prefix(prefix, "language_model"), - ) - - self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + with self._mark_tower_model(vllm_config, "audio"): + self.audio_tower = Qwen3OmniMoeAudioEncoder( + thinker_config.audio_config, + prefix=maybe_prefix(prefix, "audio_tower"), + ) self.use_deepstack = hasattr(thinker_config.vision_config, "deepstack_visual_indexes") self.deepstack_num_level = ( len(thinker_config.vision_config.deepstack_visual_indexes) if self.use_deepstack else 0 ) - # register buffer for deepstack - self.deepstack_input_embeds = ( - [ - torch.zeros( - vllm_config.scheduler_config.max_num_batched_tokens, - thinker_config.text_config.hidden_size, - ) - for _ in range(self.deepstack_num_level) - ] - if self.use_deepstack - else None - ) self.visual_dim = thinker_config.vision_config.out_hidden_size self.multiscale_dim = self.visual_dim * self.deepstack_num_level - def _get_deepstack_input_embeds(self, num_tokens: int) -> IntermediateTensors: + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Qwen3Omni_VisionTransformer( + vision_config=thinker_config.vision_config, + norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + ) + + # register buffer for deepstack + if self.use_deepstack: + self.deepstack_input_embeds = [ + torch.zeros( + vllm_config.scheduler_config.max_num_batched_tokens, + thinker_config.text_config.hidden_size, + ) + for _ in range(self.deepstack_num_level) + ] + + with self._mark_language_model(vllm_config): + self.language_model = Qwen3MoeLLMForCausalLM( + vllm_config=vllm_config.with_hf_config( + thinker_config.text_config, + architectures=["Qwen3MoeForCausalLM"], + ), + prefix=maybe_prefix(prefix, "language_model"), + ) + + self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + + def _get_deepstack_input_embeds( + self, + num_tokens: int, + ) -> IntermediateTensors | None: + if not getattr(self, "deepstack_input_embeds", None): + return None # If vision tower is skipped + # get deepstack_input_embeds from buffer, and clear the buffer return IntermediateTensors( { @@ -728,6 +766,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None: + if not getattr(self, "deepstack_input_embeds", None): + return + # set deepstack_input_embeds to buffer num_tokens = deepstack_input_embeds.size(1) if num_tokens > self.deepstack_input_embeds[0].size(0): @@ -744,6 +785,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx]) def _clear_deepstack_input_embeds(self, num_tokens: int) -> None: + if not getattr(self, "deepstack_input_embeds", None): + return + # clear deepstack_input_embeds in buffer if num_tokens > 0: for idx in range(self.deepstack_num_level): @@ -759,7 +803,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs) if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality: mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs) - if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality: + if input_key in ("input_audio_features",) and "audio" not in mm_input_by_modality: mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs) return mm_input_by_modality @@ -808,6 +852,17 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( if multimodal_embeddings is None or len(multimodal_embeddings) == 0: return inputs_embeds + # Detect interleaved audio-in-video early, since it affects + # both the deepstack path and the final embedding merge. + video_token_id = self.config.video_token_id + audio_token_id = self.config.audio_token_id + is_video = is_multimodal & (input_ids == video_token_id) + is_audio = is_multimodal & (input_ids == audio_token_id) + num_video = is_video.sum().item() + num_audio = is_audio.sum().item() + + is_interleaved = check_interleaved_audio_video(is_video, is_audio, num_video, num_audio) + deepstack_input_embeds = None # split the feat dim to obtain multi-scale visual feature has_vision_embeddings = [ @@ -816,12 +871,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( if self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings): multiscale_len = len(self.visual.deepstack_visual_indexes) multimodal_embeddings_multiscale = [] - is_vision = torch.zeros_like(is_multimodal) - mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0] - mm_position_idx = 0 + + if is_interleaved: + # Use input_ids-based mask for correct vision positions + # when audio and video tokens are interleaved. + is_vision = is_video.clone() + else: + is_vision = torch.zeros_like(is_multimodal) + mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0] + mm_position_idx = 0 + for index, embeddings in enumerate(multimodal_embeddings): num_tokens = embeddings.shape[0] - current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens] # Vision embeddings if embeddings.shape[-1] != self.config.text_config.hidden_size: @@ -830,13 +891,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( embeddings_main, embeddings_multiscale = torch.split(embeddings, [visual_dim, multi_dim], dim=-1) multimodal_embeddings[index] = embeddings_main multimodal_embeddings_multiscale.append(embeddings_multiscale) - is_vision[current_positions] = True + if not is_interleaved: + current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens] + is_vision[current_positions] = True # Audio embeddings else: - is_vision[current_positions] = False + if not is_interleaved: + current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens] + is_vision[current_positions] = False - mm_position_idx += num_tokens + if not is_interleaved: + mm_position_idx += num_tokens deepstack_input_embeds = inputs_embeds.new_zeros( inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1) @@ -853,6 +919,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) self._set_deepstack_input_embeds(deepstack_input_embeds) + if is_interleaved: + return merge_interleaved_embeddings( + inputs_embeds, + multimodal_embeddings, + is_video, + is_audio, + is_multimodal, + num_video, + num_audio, + ) + + # Default: standard merge (no interleaving) inputs_embeds = _merge_multimodal_embeddings( inputs_embeds=inputs_embeds, multimodal_embeddings=multimodal_embeddings, @@ -910,214 +988,304 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( return loaded_weights + def _compute_audio_token_count(self, audio_feature_length: int) -> int: + """Compute audio tokens from feature length using Qwen3-Omni formula.""" + return _get_feat_extract_output_lengths(torch.tensor([audio_feature_length])).item() + + def _get_audio_for_video_mapping(self, mm_features: list[MultiModalFeatureSpec]) -> tuple[dict[int, int], set[int]]: + """ + Map video offset -> paired audio_feature_length for use_audio_in_video. + + When use_audio_in_video=True, audio is interleaved within video. + The pairing is based on feature order in mm_features. + + Returns: + Tuple of (video_offset -> audio_feature_length mapping, + set of paired audio offsets to skip) + """ + videos_with_audio = [ + f + for f in mm_features + if f.modality == "video" and f.data.get("use_audio_in_video") and f.data["use_audio_in_video"].data.item() + ] + audios = [f for f in mm_features if f.modality == "audio"] + + mapping: dict[int, int] = {} + paired_audio_offsets: set[int] = set() + for i, video_f in enumerate(videos_with_audio): + if i < len(audios): + audio_len = audios[i].data["audio_feature_lengths"].data.item() + mapping[video_f.mm_position.offset] = audio_len + paired_audio_offsets.add(audios[i].mm_position.offset) + return mapping, paired_audio_offsets + + def iter_mm_features(self, mm_features: list[MultiModalFeatureSpec]) -> Iterator[tuple[int, str, dict[str, Any]]]: + """ + Iterate over multimodal features sorted by position offset. + + Yields: (offset, modality, feature_data) where feature_data contains: + - image: {"grid_t", "grid_h", "grid_w", "t_factor"} + - video: {"grid_t", "grid_h", "grid_w", "t_factor", + "use_audio_in_video", "audio_feature_length"} + - audio: {"audio_feature_length"} + """ + config = self.config + spatial_merge_size = config.vision_config.spatial_merge_size + position_id_per_seconds = config.position_id_per_seconds + + sorted_features = sorted(mm_features, key=lambda f: f.mm_position.offset) + audio_for_video, paired_audio_offsets = self._get_audio_for_video_mapping(sorted_features) + + for mm_feature in sorted_features: + offset = mm_feature.mm_position.offset + modality = mm_feature.modality + + if modality == "image": + t, h, w = mm_feature.data["image_grid_thw"].data.tolist() + yield ( + offset, + "image", + { + "grid_t": t, + "grid_h": h // spatial_merge_size, + "grid_w": w // spatial_merge_size, + "t_factor": position_id_per_seconds, + }, + ) + elif modality == "video": + t, h, w = mm_feature.data["video_grid_thw"].data.tolist() + second_per_grid_ts = 2.0 + if mm_feature.data.get("second_per_grid_ts"): + second_per_grid_ts = mm_feature.data["second_per_grid_ts"].data.item() + use_audio_in_video = bool( + mm_feature.data.get("use_audio_in_video") and mm_feature.data["use_audio_in_video"].data.item() + ) + + yield ( + offset, + "video", + { + "grid_t": t, + "grid_h": h // spatial_merge_size, + "grid_w": w // spatial_merge_size, + "t_factor": second_per_grid_ts * position_id_per_seconds, + "use_audio_in_video": use_audio_in_video, + "audio_feature_length": audio_for_video.get(offset), + }, + ) + elif modality == "audio": + if offset not in paired_audio_offsets: + audio_len = mm_feature.data["audio_feature_lengths"].data.item() + yield offset, "audio", {"audio_feature_length": audio_len} + + def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) -> tuple[np.ndarray, int]: + """ + Compute positions for interleaved video+audio using Qwen3 token-by-token + interleaving logic. + + Returns: (position_ids [3, N], total_token_count) + """ + grid_t = data["grid_t"] + grid_h = data["grid_h"] + grid_w = data["grid_w"] + t_factor = data["t_factor"] + audio_feature_length = data["audio_feature_length"] + + audio_len = self._compute_audio_token_count(audio_feature_length) + + h_index = np.tile(np.arange(grid_h).reshape(1, -1, 1), (grid_t, 1, grid_w)).flatten() + w_index = np.tile(np.arange(grid_w).reshape(1, 1, -1), (grid_t, grid_h, 1)).flatten() + t_index_raw = np.arange(grid_t) + t_index_scaled = (t_index_raw * t_factor).astype(np.int64) + t_index = np.repeat(t_index_scaled, grid_h * grid_w) + + video_pos = np.stack([t_index, h_index, w_index]) + start_idx + audio_pos = np.broadcast_to(np.arange(audio_len), (3, audio_len)) + start_idx + + video_t_values = video_pos[0] + audio_t_values = audio_pos[0] + + pos_ids_list: list[np.ndarray] = [] + video_idx, audio_idx = 0, 0 + num_video = grid_t * grid_h * grid_w + + while video_idx < num_video and audio_idx < audio_len: + if video_t_values[video_idx] <= audio_t_values[audio_idx]: + pos_ids_list.append(video_pos[:, video_idx : video_idx + 1]) + video_idx += 1 + else: + pos_ids_list.append(audio_pos[:, audio_idx : audio_idx + 1]) + audio_idx += 1 + + if video_idx < num_video: + pos_ids_list.append(video_pos[:, video_idx:]) + if audio_idx < audio_len: + pos_ids_list.append(audio_pos[:, audio_idx:]) + + total_tokens = num_video + audio_len + return np.concatenate(pos_ids_list, axis=1), total_tokens + + @classmethod + def get_speech_to_text_config(cls, model_config: ModelConfig, task_type: str) -> SpeechToTextConfig: + processor = cached_processor_from_config(model_config, processor_cls=Qwen3OmniMoeProcessor) + return SpeechToTextConfig( + max_audio_clip_s=processor.feature_extractor.chunk_length, + sample_rate=processor.feature_extractor.sampling_rate, + min_energy_split_window_size=None, + ) + + @classmethod + def get_generation_prompt( + cls, + audio: np.ndarray, + stt_config: SpeechToTextConfig, + model_config: ModelConfig, + language: str | None, + task_type: Literal["transcribe", "translate"], + request_prompt: str, + to_language: str | None, + ) -> PromptType: + """ + Construct a transcription/translation prompt for Qwen3-Omni. + """ + # Transcribe this audio [into ] | for transcription + # Translate this audio [from into ] | for translation + instruction = "Transcribe" if task_type == "transcribe" else "Translate" + instruction += " this audio" + + # Default to_language to English for translation + if task_type == "translate" and to_language is None: + to_language = "en" + + # Get full language names from supported_languages mapping + full_lang_name = cls.supported_languages.get(language, "") + full_lang_name_to = cls.supported_languages.get(to_language, "") + + if task_type == "transcribe" and full_lang_name: + instruction += f" into {full_lang_name}" + elif task_type == "translate": + if full_lang_name: + instruction += f" from {full_lang_name}" + if full_lang_name_to: + instruction += f" into {full_lang_name_to}" + + instruction += "." + + if request_prompt: + instruction += f" {request_prompt}" + + processor = cached_processor_from_config(model_config, processor_cls=Qwen3OmniMoeProcessor) + # Audio placeholder format: <|audio_start|><|audio_pad|><|audio_end|> + audio_placeholder = "<|audio_start|><|audio_pad|><|audio_end|>" + user_content = f"{audio_placeholder}{instruction}" + + messages = [{"role": "user", "content": user_content}] + prompt = processor.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + + audio_data = (audio, stt_config.sample_rate) + prompts_dict = {"multi_modal_data": {"audio": audio_data}, "prompt": prompt} + return cast(PromptType, prompts_dict) + def get_mrope_input_positions( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - kwargs = MultiModalFeatureSpec.gather_kwargs( - mm_features, - { - "image_grid_thw", - "video_grid_thw", - "second_per_grid_ts", - "audio_feature_lengths", - "use_audio_in_video", - }, - ) - image_grid_thw = kwargs.get("image_grid_thw", []) - video_grid_thw = kwargs.get("video_grid_thw", []) - second_per_grid_ts = kwargs.get("second_per_grid_ts", []) - audio_feature_lengths = kwargs.get("audio_feature_lengths", []) - use_audio_in_video = any(kwargs.get("use_audio_in_video", [])) + """Compute M-RoPE input positions using mm_features directly.""" + seq_len = len(input_tokens) - image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(image_grid_thw) - video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(video_grid_thw) - - input_ids = torch.tensor(input_tokens) - if input_ids is None or input_ids.ndim != 1: - raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids") - - seq_len = input_ids.shape[0] - - if isinstance(audio_feature_lengths, list): - audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long) - - if not len(second_per_grid_ts) and len(video_grid_thw): - second_per_grid_ts = 2.0 - second_per_grids = torch.ones(len(video_grid_thw), dtype=torch.float32) * second_per_grid_ts - else: - second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32) - - config = self.config - spatial_merge_size = config.vision_config.spatial_merge_size - image_token_id = config.image_token_id - video_token_id = config.video_token_id - audio_token_id = config.audio_token_id - vision_start_token_id = config.vision_start_token_id - audio_start_token_id = config.audio_start_token_id - position_id_per_seconds = config.position_id_per_seconds - - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - if vision_start_indices.numel() > 0: - vision_tokens = input_ids[vision_start_indices + 1] - else: - vision_tokens = input_ids.new_empty((0,), dtype=input_ids.dtype) - audio_nums = torch.sum(input_ids == audio_start_token_id) - image_nums = (vision_tokens == image_token_id).sum() - video_nums = ( - (vision_tokens == audio_start_token_id).sum() - if use_audio_in_video - else (vision_tokens == video_token_id).sum() - ) - - llm_pos_ids_list: list[torch.Tensor] = [] + llm_pos_ids_list: list[np.ndarray] = [] st = 0 - image_idx = 0 - video_idx = 0 - audio_idx = 0 - remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums # noqa: E501 - multimodal_nums = image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums # noqa: E501 - for _ in range(multimodal_nums): - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - if (image_token_id in input_tokens or video_token_id in input_tokens) and ( - remain_videos > 0 or remain_images > 0 - ): - ed_vision_start = input_tokens.index(vision_start_token_id, st) - else: - ed_vision_start = len(input_tokens) + 1 - if audio_token_id in input_tokens and remain_audios > 0: - ed_audio_start = input_tokens.index(audio_start_token_id, st) - else: - ed_audio_start = len(input_tokens) + 1 - min_ed = min(ed_vision_start, ed_audio_start) + for offset, modality, data in self.iter_mm_features(mm_features): + text_len = offset - st + st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0 - if min_ed == ed_audio_start: - text_len = min_ed - st - if text_len != 0: - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - bos_len = 1 - llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) - llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx - llm_pos_ids_list.append(llm_pos_ids) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - eos_len = 1 - llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st += text_len + bos_len + audio_len + eos_len - audio_idx += 1 - remain_audios -= 1 - elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == image_token_id: - text_len = min_ed - st - if text_len != 0: - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - bos_len = 1 - llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - grid_t = image_grid_thw[image_idx][0] - grid_hs = image_grid_thw[:, 1] - grid_ws = image_grid_thw[:, 2] - t_index = torch.arange(grid_t) * position_id_per_seconds - llm_pos_ids = get_llm_pos_ids_for_vision( - st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws - ) - image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2) - llm_pos_ids_list.append(llm_pos_ids) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - eos_len = 1 - llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st += text_len + bos_len + image_len + eos_len - image_idx += 1 - remain_images -= 1 - elif ( - min_ed == ed_vision_start - and input_ids[ed_vision_start + 1] == video_token_id - and not use_audio_in_video - ): - text_len = min_ed - st - if text_len != 0: - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - bos_len = 1 - llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - grid_t = video_grid_thw[video_idx][0] - grid_hs = video_grid_thw[:, 1] - grid_ws = video_grid_thw[:, 2] - t_index = torch.arange(grid_t) * float(second_per_grids[video_idx].item()) * position_id_per_seconds - llm_pos_ids = get_llm_pos_ids_for_vision( - st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws - ) - video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) - llm_pos_ids_list.append(llm_pos_ids) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - eos_len = 1 - llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st += text_len + bos_len + video_len + eos_len - video_idx += 1 - remain_videos -= 1 - elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start and use_audio_in_video: - text_len = min_ed - st - if text_len != 0: - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - bos_len = 1 - bos_block = torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx - llm_pos_ids_list.append(bos_block) - llm_pos_ids_list.append(bos_block) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) - audio_llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx - grid_t = video_grid_thw[video_idx][0] - grid_hs = video_grid_thw[:, 1] - grid_ws = video_grid_thw[:, 2] - t_index = torch.arange(grid_t) * float(second_per_grids[video_idx].item()) * position_id_per_seconds - video_llm_pos_ids = get_llm_pos_ids_for_vision( - st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws - ) - video_data_index, audio_data_index = 0, 0 - while video_data_index < video_llm_pos_ids.shape[-1] and audio_data_index < audio_llm_pos_ids.shape[-1]: - if video_llm_pos_ids[0][video_data_index] <= audio_llm_pos_ids[0][audio_data_index]: - llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_data_index + 1]) - video_data_index += 1 - else: - llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_data_index + 1]) - audio_data_index += 1 - if video_data_index < video_llm_pos_ids.shape[-1]: - llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_llm_pos_ids.shape[-1]]) - if audio_data_index < audio_llm_pos_ids.shape[-1]: - llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_llm_pos_ids.shape[-1]]) - video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - eos_len = 1 - eos_block = torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx - llm_pos_ids_list.append(eos_block) - llm_pos_ids_list.append(eos_block) - st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2 # noqa: E501 - audio_idx += 1 - video_idx += 1 - remain_videos -= 1 - remain_audios -= 1 + if text_len > 0: + llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx) + st_idx += text_len - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) + llm_pos_ids_list.append(bos_pos) + st_idx += 1 - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + if modality == "audio": + audio_tokens = self._compute_audio_token_count(data["audio_feature_length"]) + audio_pos = np.broadcast_to(np.arange(audio_tokens), (3, audio_tokens)) + st_idx + llm_pos_ids_list.append(audio_pos) + st_idx = int(audio_pos.max()) + 1 + + eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) + llm_pos_ids_list.append(eos_pos) + st = offset + 1 + audio_tokens + 1 + + elif modality == "image": + grid_t = data["grid_t"] + grid_h = data["grid_h"] + grid_w = data["grid_w"] + t_factor = data["t_factor"] + + grid_indices = np.indices((grid_t, grid_h, grid_w)) + if t_factor != 1.0: + grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64) + llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx) + + image_len = grid_t * grid_h * grid_w + st_idx = int(llm_pos_ids_list[-1].max()) + 1 + + eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) + llm_pos_ids_list.append(eos_pos) + st = offset + 1 + image_len + 1 + + elif modality == "video": + grid_t = data["grid_t"] + grid_h = data["grid_h"] + grid_w = data["grid_w"] + t_factor = data["t_factor"] + + if not data["use_audio_in_video"]: + grid_indices = np.indices((grid_t, grid_h, grid_w)) + if t_factor != 1.0: + grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64) + llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx) + + video_len = grid_t * grid_h * grid_w + st_idx = int(llm_pos_ids_list[-1].max()) + 1 + + eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) + llm_pos_ids_list.append(eos_pos) + st = offset + 1 + video_len + 1 + else: + audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1)) + llm_pos_ids_list.append(audio_bos_pos) + + pos_ids, _ = self._compute_interleaved_positions(st_idx, data) + llm_pos_ids_list.append(pos_ids) + st_idx = int(pos_ids.max()) + 1 + + eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) + llm_pos_ids_list.append(eos_pos) + llm_pos_ids_list.append(eos_pos) + + video_len = grid_t * grid_h * grid_w + audio_len = self._compute_audio_token_count(data["audio_feature_length"]) + st = offset + 2 + video_len + audio_len + 2 + + if st < seq_len: + st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0 + text_len = seq_len - st + llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx) + + llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) if llm_positions.shape[1] != seq_len: raise RuntimeError("Position ids length mismatch with input ids length") - mrope_position_delta = llm_positions.max() + 1 - seq_len - return llm_positions, mrope_position_delta + mrope_position_delta = int(llm_positions.max()) + 1 - seq_len + return torch.from_numpy(llm_positions), mrope_position_delta def get_mm_mapping(self) -> MultiModelKeys: """ diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 64569846a..3708bc8f9 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -128,9 +128,10 @@ class OmniRequestOutput: For diffusion outputs, this returns the local _multimodal_output field. """ if self.request_output is not None: - # Check completion outputs first (where multimodal_output is attached) - if self.request_output.outputs: - for output in self.request_output.outputs: + # Check completion outputs first (where multimodal_output is attached). + outputs = getattr(self.request_output, "outputs", None) + if isinstance(outputs, list) and outputs: + for output in outputs: mm = getattr(output, "multimodal_output", None) if mm: return mm diff --git a/vllm_omni/request.py b/vllm_omni/request.py index c992816de..a869fa23c 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -91,4 +91,5 @@ class OmniRequest(Request): block_hasher=block_hasher, additional_information=request.additional_information, resumable=request.resumable, + reasoning_ended=request.reasoning_ended, ) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 9332abeb1..7b23b3b91 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -50,7 +50,10 @@ class ExecuteModelState(NamedTuple): aux_hidden_states: list[torch.Tensor] | None ec_connector_output: Any cudagraph_stats: Any + # OMNI: multimodal_outputs field for omni-specific multimodal handling multimodal_outputs: Any + # slot_mappings for attention/drafter (aligned with upstream v1 API) + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None class GPUARModelRunner(OmniGPUModelRunner): @@ -386,6 +389,7 @@ class GPUARModelRunner(OmniGPUModelRunner): ec_connector_output, cudagraph_stats, multimodal_outputs, + slot_mappings, # OMNI: pass slot_mappings for drafter ) self.kv_connector_output = kv_connector_output @@ -428,6 +432,7 @@ class GPUARModelRunner(OmniGPUModelRunner): ec_connector_output, cudagraph_stats, multimodal_outputs, + slot_mappings, # OMNI: unpack slot_mappings for drafter ) = self.execute_model_state self.execute_model_state = None @@ -454,6 +459,7 @@ class GPUARModelRunner(OmniGPUModelRunner): aux_hidden_states, spec_decode_metadata, spec_decode_common_attn_metadata, + slot_mappings, # OMNI: pass slot_mappings to drafter (upstream v1 API) ) self._copy_draft_token_ids_to_cpu(scheduler_output) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 30edba643..198d42c8a 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -4,6 +4,7 @@ import os import torch from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.tracing import instrument from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats @@ -25,6 +26,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): model runners for text generation stages (e.g., thinker stages). """ + @instrument(span_name="Init device") def init_device(self): if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index c52ab1e8b..0747db3ea 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -299,6 +299,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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 return None @@ -340,6 +341,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ec_connector_output, cudagraph_stats, multimodal_outputs, + slot_mappings, # OMNI: unpack slot_mappings for upstream v1 API compatibility ) = self.execute_model_state self.execute_model_state = None @@ -459,8 +461,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): is_profile: bool = False, create_mixed_batch: bool = False, remove_lora: bool = True, - activate_lora: bool = False, is_graph_capturing: bool = False, + num_active_loras: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -483,7 +485,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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 - activate_lora: If False, dummy_run is performed without LoRAs. + num_active_loras: Number of active LoRAs to capture for. """ mm_config = self.vllm_config.model_config.multimodal_config if mm_config and mm_config.mm_encoder_only: @@ -561,7 +563,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # `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=activate_lora, + force_has_lora=num_active_loras > 0, + # Capture shape specialization for specific active LoRA counts. + force_num_active_loras=num_active_loras, ) ) @@ -630,8 +634,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): self.lora_config, num_scheduled_tokens, num_sampled_tokens, - activate_lora, remove_lora, + num_active_loras, ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens @@ -718,13 +722,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # 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 activate_lora: + 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 @@ -761,55 +766,46 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): assert mm_budget is not None if (encoder_budget := mm_budget.get_encoder_budget()) > 0: - # 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.max_items_per_batch_by_modality[dummy_modality] + 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, - ) + 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, - ) + # 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) + # 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, - ) - - # NOTE: This happens when encoder cache needs to store - # the embeddings that encoder outputs are scattered onto. - # In this case we create dummy embeddings of size - # (max_tokens_for_modality, hidden_size) and scatter - # encoder output into it. - encoder_output_shape = dummy_encoder_outputs[0].shape - max_mm_tokens_per_item = mm_budget.max_tokens_by_modality[dummy_modality] - if encoder_output_shape[0] < max_mm_tokens_per_item: - encoder_hidden_size = encoder_output_shape[-1] - expanded_outputs = [] - for output in dummy_encoder_outputs: - expanded = output.new_zeros((max_mm_tokens_per_item, encoder_hidden_size)) - num_tokens = output.shape[0] - expanded[:num_tokens].copy_(output) - expanded_outputs.append(expanded) - - dummy_encoder_outputs = expanded_outputs - - # Cache the dummy encoder outputs. - self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs)) + 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) diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 8ba4efa3d..7de9a7922 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -4,6 +4,7 @@ import os import torch from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.tracing import instrument from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats @@ -25,6 +26,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): worker_cls: "vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner" """ + @instrument(span_name="Init device") def init_device(self): if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index d52a2014f..09a792bc8 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType +from vllm.tracing import instrument from vllm.utils.import_utils import LazyLoader from vllm.utils.math_utils import cdiv from vllm.v1.spec_decode.eagle import EagleProposer @@ -40,6 +41,33 @@ class OmniGPUModelRunner(GPUModelRunner): self._omni_num_scheduled_tokens_np: np.ndarray | None = None self._omni_last_model_output: object | None = None + def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes): + """Override to fix scheduler_metadata buffer size for FA3 + CUDA graph. + + The upstream FlashAttentionMetadataBuilder pre-allocates + scheduler_metadata with (max_num_seqs + 1) entries, but FA3's + get_scheduler_metadata() can return up to + (max_num_seqs * max_num_splits + 1) entries, causing a RuntimeError + during CUDA graph capture. After calling the parent implementation + we resize any too-small buffers. + """ + super().initialize_metadata_builders(kv_cache_config, kernel_block_sizes) + + for kv_cache_group in self.attn_groups: + for attn_group in kv_cache_group: + for builder in attn_group.metadata_builders: + sm = getattr(builder, "scheduler_metadata", None) + max_num_splits = getattr(builder, "max_num_splits", 0) + if sm is not None and max_num_splits > 1: + required = self.scheduler_config.max_num_seqs * max_num_splits + 1 + if sm.shape[0] < required: + builder.scheduler_metadata = torch.zeros( + required, + dtype=sm.dtype, + device=sm.device, + ) + + @instrument(span_name="Loading (GPU)") def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) # TODO move this model specific logic to a separate class @@ -422,8 +450,9 @@ class OmniGPUModelRunner(GPUModelRunner): is_profile: bool = False, create_mixed_batch: bool = False, remove_lora: bool = True, - activate_lora: bool = False, is_graph_capturing: bool = False, + num_active_loras: int = 0, + activate_lora: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -446,8 +475,12 @@ class OmniGPUModelRunner(GPUModelRunner): 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 - activate_lora: If False, dummy_run is performed without LoRAs. + num_active_loras: Number of active LoRAs to capture for. + activate_lora: Backward-compatible override for LoRA activation. """ + if activate_lora is None: + activate_lora = num_active_loras > 0 + 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. @@ -525,6 +558,8 @@ class OmniGPUModelRunner(GPUModelRunner): # activated later in the context manager, but we need to know the # LoRA state when determining the batch descriptor for capture force_has_lora=activate_lora, + # Capture shape specialization for specific active LoRA counts. + force_num_active_loras=num_active_loras, ) ) From a42b7487b0681afab4a5a4b305d7f59c0b158de9 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Sat, 14 Feb 2026 19:43:38 +0800 Subject: [PATCH 2/2] [Test] Fix expansion and example test case for qwen3-omni (#1358) Signed-off-by: yenuo26 <410167048@qq.com> --- .buildkite/pipeline.yml | 210 ++++----- .buildkite/test-nightly.yaml | 6 +- tests/conftest.py | 13 +- .../test_qwen3_omni_expansion.py | 436 +++++++++--------- tests/e2e/stage_configs/qwen2_5_omni_ci.yaml | 20 +- .../online_serving/test_qwen2_5_omni.py | 45 +- .../online_serving/test_qwen3_omni.py | 45 +- tests/perf/scripts/run_benchmark.py | 33 +- 8 files changed, 432 insertions(+), 376 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 88cfc22d5..d0566e18b 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -139,41 +139,41 @@ steps: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Benchmark&Engine Test" - timeout_in_minutes: 15 - depends_on: image-build - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/benchmarks/test_serve_cli.py - - pytest -s -v tests/engine/test_async_omni_engine_abort.py - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + # - label: "Benchmark&Engine Test" + # timeout_in_minutes: 15 + # depends_on: image-build + # commands: + # - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # - pytest -s -v tests/benchmarks/test_serve_cli.py + # - pytest -s -v tests/engine/test_async_omni_engine_abort.py + # agents: + # queue: "mithril-h100-pool" + # plugins: + # - kubernetes: + # podSpec: + # containers: + # - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # resources: + # limits: + # nvidia.com/gpu: 2 + # volumeMounts: + # - name: devshm + # mountPath: /dev/shm + # - name: hf-cache + # mountPath: /root/.cache/huggingface + # env: + # - name: HF_HOME + # value: /root/.cache/huggingface + # nodeSelector: + # node.kubernetes.io/instance-type: gpu-h100-sxm + # volumes: + # - name: devshm + # emptyDir: + # medium: Memory + # - name: hf-cache + # hostPath: + # path: /mnt/hf-cache + # type: DirectoryOrCreate - label: "Omni Model Test" timeout_in_minutes: 15 @@ -194,42 +194,42 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Omni Model Test with H100" - timeout_in_minutes: 20 - depends_on: image-build - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export VLLM_TEST_CLEAN_GPU_MEMORY="1" - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py - - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + # - label: "Omni Model Test with H100" + # timeout_in_minutes: 20 + # depends_on: image-build + # commands: + # - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # - export VLLM_TEST_CLEAN_GPU_MEMORY="1" + # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py + # - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model" + # agents: + # queue: "mithril-h100-pool" + # plugins: + # - kubernetes: + # podSpec: + # containers: + # - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # resources: + # limits: + # nvidia.com/gpu: 2 + # volumeMounts: + # - name: devshm + # mountPath: /dev/shm + # - name: hf-cache + # mountPath: /root/.cache/huggingface + # env: + # - name: HF_HOME + # value: /root/.cache/huggingface + # nodeSelector: + # node.kubernetes.io/instance-type: gpu-h100-sxm + # volumes: + # - name: devshm + # emptyDir: + # medium: Memory + # - name: hf-cache + # hostPath: + # path: /mnt/hf-cache + # type: DirectoryOrCreate - label: "Qwen3-TTS E2E Test" timeout_in_minutes: 20 @@ -251,40 +251,40 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Diffusion Image Edit Test with H100 (1 GPU)" - timeout_in_minutes: 20 - depends_on: image-build - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + # - label: "Diffusion Image Edit Test with H100 (1 GPU)" + # timeout_in_minutes: 20 + # depends_on: image-build + # commands: + # - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py + # agents: + # queue: "mithril-h100-pool" + # plugins: + # - kubernetes: + # podSpec: + # containers: + # - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # resources: + # limits: + # nvidia.com/gpu: 1 + # volumeMounts: + # - name: devshm + # mountPath: /dev/shm + # - name: hf-cache + # mountPath: /root/.cache/huggingface + # env: + # - name: HF_HOME + # value: /root/.cache/huggingface + # nodeSelector: + # node.kubernetes.io/instance-type: gpu-h100-sxm + # volumes: + # - name: devshm + # emptyDir: + # medium: Memory + # - name: hf-cache + # hostPath: + # path: /mnt/hf-cache + # type: DirectoryOrCreate # - label: "Bagel Text2Img Model Test with H100" # timeout_in_minutes: 30 diff --git a/.buildkite/test-nightly.yaml b/.buildkite/test-nightly.yaml index f665de270..70c6b1fc4 100644 --- a/.buildkite/test-nightly.yaml +++ b/.buildkite/test-nightly.yaml @@ -5,8 +5,8 @@ steps: if: build.env("NIGHTLY") == "1" commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py - - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py + - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model" + - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model" agents: queue: "mithril-h100-pool" plugins: @@ -43,7 +43,7 @@ steps: commands: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py + - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model" agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: diff --git a/tests/conftest.py b/tests/conftest.py index 267386aa3..08d26a6d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1057,7 +1057,8 @@ def omni_server(request, run_level): Multi-stage initialization can take 10-20+ minutes. """ with _omni_server_lock: - model, stage_config_path = request.param + *port, model, stage_config_path = request.param + port = port[0] if port else None if run_level == "advanced_model": stage_config_path = modify_stage_config( stage_config_path, @@ -1070,7 +1071,8 @@ def omni_server(request, run_level): }, ) - with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server: + server_args = ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"] + with OmniServer(model, server_args, port=port) if port else OmniServer(model, server_args) as server: print("OmniServer started successfully") yield server print("OmniServer stopping...") @@ -1288,11 +1290,15 @@ class OpenAIClientHandler: responses = [] stream = request_config.get("stream", False) + modalities = request_config.get("modalities", ["text", "audio"]) if request_num == 1: # Send single request chat_completion = self.client.chat.completions.create( - model=request_config.get("model"), messages=request_config.get("messages"), stream=stream + model=request_config.get("model"), + messages=request_config.get("messages"), + stream=stream, + modalities=modalities, ) if stream: @@ -1314,6 +1320,7 @@ class OpenAIClientHandler: self.client.chat.completions.create, model=request_config.get("model"), messages=request_config.get("messages"), + modalities=modalities, stream=stream, ) futures.append(future) diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 9e84c7559..84cc85ba1 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -4,66 +4,55 @@ E2E Online tests for Qwen3-Omni model. """ -import concurrent.futures import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -import threading -import time from pathlib import Path -import openai import pytest from tests.conftest import ( - OmniServer, - convert_audio_to_text, - cosine_similarity_text, dummy_messages_from_mix_data, generate_synthetic_audio, generate_synthetic_image, + generate_synthetic_video, + modify_stage_config, ) +from tests.utils import hardware_test models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] -# CI stage config for 2*H100-80G GPUs -stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")] +AUDIO_KEY = ["water", "chirping", "crackling"] +IMAGE_KEY = ["square", "quadrate"] +VIDEO_KEY = ["sphere", "globe", "circle", "round", "ball"] + +def get_chunk_config(default_path): + path = modify_stage_config( + default_path, + updates={ + "async_chunk": True, + "stage_args": { + 0: { + "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" + }, + 1: { + "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk" + }, + }, + }, + deletes={"stage_args": {2: ["custom_process_input_func"]}}, + ) + return path + + +# CI stage config for 2*H100-80G GPUs +default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml") +stage_configs = [default_path, get_chunk_config(default_path)] # Create parameter combinations for model and stage config test_params = [(model, stage_config) for model in models for stage_config in stage_configs] -_omni_server_lock = threading.Lock() - - -@pytest.fixture(scope="module") -def omni_server(request): - """Start vLLM-Omni server as a subprocess with actual model weights. - Uses session scope so the server starts only once for the entire test session. - Multi-stage initialization can take 10-20+ minutes. - """ - with _omni_server_lock: - model, stage_config_path = request.param - - print(f"Starting OmniServer with model: {model}") - - with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") - - print("OmniServer stopped") - - -@pytest.fixture -def client(omni_server): - """OpenAI client for the running vLLM-Omni server.""" - return openai.OpenAI( - base_url=f"http://{omni_server.host}:{omni_server.port}/v1", - api_key="EMPTY", - ) - - def get_system_prompt(): return { "role": "system", @@ -83,7 +72,10 @@ def get_system_prompt(): def get_prompt(prompt_type="text_only"): prompts = { "text_only": "What is the capital of China?", - "mix": "What is recited in the audio? What is in this image? Describe the video briefly.", + "mix": "What is recited in the audio? What is in this image? What is in this video?", + "text_video": "What is in this video? ", + "text_image": "What is in this image? ", + "text_audio": "What is in this audio? ", } return prompts.get(prompt_type, prompts["text_only"]) @@ -93,196 +85,228 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -def get_deploy_config(deploy_type="TP1"): - result = { - "TP1": { - "stage_args": { - 0: { - "engine_args.gpu_memory_utilization": 0.95, - "engine_args.tensor_parallel_size": 1, - "runtime.devices": "0", - }, - 2: {"runtime.devices": "1"}, - } - } - } - return result.get(deploy_type, result["TP1"]) - - +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_001(client: openai.OpenAI, omni_server) -> None: - """Test processing text, generating text output via OpenAI API.""" +def test_text_to_audio_001(omni_server, openai_client) -> None: messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()) - # Test single completion - start_time = time.perf_counter() - chat_completion = client.chat.completions.create( - model=omni_server.model, messages=messages, max_tokens=20, modalities=["text"] - ) - # Verify E2E - print(f"the request e2e is: {time.perf_counter() - start_time}") - # TODO: Verify the E2E latency after confirmation baseline. + request_config = { + "model": omni_server.model, + "messages": messages, + "modalities": ["audio"], + "key_words": {"text": ["beijing"]}, + } - # Verify only output text - assert len(chat_completion.choices) == 1, "The generated content includes more than just text." - - # Verify text output success - text_choice = chat_completion.choices[0] - assert text_choice.message.content is not None, "No text output is generated" - assert chat_completion.usage.completion_tokens <= 20, "The output length more than the requested max_tokens." - assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords." + openai_client.send_request(request_config) +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_audio_to_text_001(client: openai.OpenAI, omni_server) -> None: - """Test processing text, generating text output via OpenAI API.""" - audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 1)['base64']}" - messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url) - # Test single completion - start_time = time.perf_counter() - chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, modalities=["text"]) - # Verify only output text - assert len(chat_completion.choices) == 1, "The generated content includes more than just text." +def test_text_to_text_audio_001(omni_server, openai_client) -> None: + messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()) - # Verify text output success - text_choice = chat_completion.choices[0] - assert text_choice.message.content is not None, "No text output is generated" - assert "water" in text_choice.message.content.lower(), "The output do not contain keywords." + request_config = { + "model": omni_server.model, + "messages": messages, + "key_words": {"text": ["beijing"]}, + } - # Verify E2E - print(f"the request e2e is: {time.perf_counter() - start_time}") - # TODO: Verify the E2E latency after confirmation baseline. + openai_client.send_request(request_config) +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_audio_to_text_audio_001(client: openai.OpenAI, omni_server) -> None: - """Test processing text, generating audio output via OpenAI API.""" - - audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 5)['base64']}" - - messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url) - num_concurrent_requests = get_max_batch_size() - # Test single completion - e2e_list = list() - with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor: - # Submit multiple completion requests concurrently - futures = [ - executor.submit(client.chat.completions.create, model=omni_server.model, messages=messages) - for _ in range(num_concurrent_requests) - ] - start_time = time.perf_counter() - # Wait for all requests to complete and collect results - chat_completions = list() - for future in concurrent.futures.as_completed(futures): - chat_completions.append(future.result()) - # Verify E2E - current_e2e = time.perf_counter() - start_time - print(f"the request e2e is: {current_e2e}") - # TODO: Verify the E2E latency after confirmation baseline. - e2e_list.append(current_e2e) - - print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}") - # Verify all completions succeeded - assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded." - for chat_completion in chat_completions: - # Verify audio output success - audio_message = chat_completion.choices[1].message - audio_data = audio_message.audio.data - assert audio_data is not None, "No audio output is generated" - assert audio_message.audio.expires_at > time.time(), "The generated audio has expired." - - # Verify text output success - text_choice = chat_completion.choices[0] - text_content = text_choice.message.content - assert text_choice.message.content is not None, "No text output is generated" - - # Verify text output same as audio output - audio_content = convert_audio_to_text(audio_data) - print(f"text content is: {text_content}") - print(f"audio content is: {audio_content}") - similarity = cosine_similarity_text(audio_content, text_content) - print(f"similarity between audio and text is: {similarity}") - assert similarity > 0.9, "The audio content is not same as the text" - - -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_image_to_text_001(client: openai.OpenAI, omni_server) -> None: - """Test processing text, generating text output via OpenAI API.""" +def test_image_to_text_001(omni_server, openai_client) -> None: image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" messages = dummy_messages_from_mix_data(image_data_url=image_data_url) - # Test single completion - start_time = time.perf_counter() - chat_completion = client.chat.completions.create( - model=omni_server.model, messages=messages, max_tokens=100, modalities=["text"] - ) - # Verify E2E - print(f"the request e2e is: {time.perf_counter() - start_time}") - # TODO: Verify the E2E latency after confirmation baseline. - # Verify only output text - assert len(chat_completion.choices) == 1, "The generated content includes more than just text." + request_config = { + "model": omni_server.model, + "messages": messages, + "modalities": ["text"], + "key_words": {"image": IMAGE_KEY}, + } - # Verify text output success - text_choice = chat_completion.choices[0] - text_content = text_choice.message.content - assert text_content is not None, "No text output is generated" - assert chat_completion.usage.completion_tokens <= 100, "The output length more than the requested max_tokens." - assert "square" in text_content.lower(), "The output do not contain keywords." + openai_client.send_request(request_config) +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_image_to_text_audio_001(client: openai.OpenAI, omni_server) -> None: - """Test processing text, generating audio output via OpenAI API.""" +def test_image_to_audio_001(omni_server, openai_client) -> None: + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + messages = dummy_messages_from_mix_data(image_data_url=image_data_url) - num_concurrent_requests = get_max_batch_size() + request_config = { + "model": omni_server.model, + "messages": messages, + "modalities": ["audio"], + "key_words": {"image": IMAGE_KEY}, + } + openai_client.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_image_to_text_audio_001(omni_server, openai_client) -> None: image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}" messages = dummy_messages_from_mix_data(image_data_url=image_data_url) - # Test single completion - e2e_list = list() - with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor: - # Submit multiple completion requests concurrently - futures = [ - executor.submit( - client.chat.completions.create, - model=omni_server.model, - messages=messages, - ) - for _ in range(num_concurrent_requests) - ] - start_time = time.perf_counter() - # Wait for all requests to complete and collect results - chat_completions = list() - for future in concurrent.futures.as_completed(futures): - chat_completions.append(future.result()) - # Verify E2E - current_e2e = time.perf_counter() - start_time - print(f"the request e2e is: {current_e2e}") - # TODO: Verify the E2E latency after confirmation baseline. - e2e_list.append(current_e2e) + request_config = { + "model": omni_server.model, + "messages": messages, + "key_words": {"image": IMAGE_KEY}, + } - print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}") - # Verify all completions succeeded - assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded." - for chat_completion in chat_completions: - # Verify audio output success - audio_message = chat_completion.choices[1].message - audio_data = audio_message.audio.data - assert audio_data is not None, "No audio output is generated" - assert audio_message.audio.expires_at > time.time(), "The generated audio has expired." + openai_client.send_request(request_config, request_num=get_max_batch_size()) - # Verify text output success - text_choice = chat_completion.choices[0] - text_content = text_choice.message.content - assert text_content is not None, "No text output is generated" - assert "square" in text_content.lower(), "The output do not contain keywords." - # Verify text output same as audio output - audio_content = convert_audio_to_text(audio_data) - print(f"text content is: {text_content}") - print(f"audio content is: {audio_content}") - similarity = cosine_similarity_text(audio_content, text_content) - print(f"similarity between audio and text is: {similarity}") - assert similarity > 0.9, "The audio content is not same as the text" +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_video_to_text_001(omni_server, openai_client) -> None: + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + messages = dummy_messages_from_mix_data(video_data_url=video_data_url) + + request_config = { + "model": omni_server.model, + "messages": messages, + "modalities": ["text"], + "key_words": {"video": VIDEO_KEY}, + } + + openai_client.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_video_to_audio_001(omni_server, openai_client) -> None: + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + messages = dummy_messages_from_mix_data(video_data_url=video_data_url) + + request_config = { + "model": omni_server.model, + "messages": messages, + "modalities": ["audio"], + "key_words": {"video": VIDEO_KEY}, + } + + openai_client.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_video_to_text_audio_001(omni_server, openai_client) -> None: + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + + messages = dummy_messages_from_mix_data(video_data_url=video_data_url) + + request_config = { + "model": omni_server.model, + "messages": messages, + "key_words": {"video": VIDEO_KEY}, + } + + openai_client.send_request(request_config, request_num=get_max_batch_size()) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None: + audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" + messages = dummy_messages_from_mix_data( + audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio") + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "modalities": ["text"], + "key_words": {"audio": AUDIO_KEY}, + } + + openai_client.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_image_to_text_audio_001(omni_server, openai_client) -> None: + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + + messages = dummy_messages_from_mix_data( + image_data_url=image_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_image") + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "modalities": ["audio"], + "key_words": {"image": IMAGE_KEY}, + } + + openai_client.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 30)['base64']}" + + messages = dummy_messages_from_mix_data( + video_data_url=video_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_video") + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "key_words": {"video": VIDEO_KEY}, + } + + openai_client.send_request(request_config, request_num=get_max_batch_size()) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_mix_to_text_audio_001(omni_server, openai_client) -> None: + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + video_data_url=video_data_url, + image_data_url=image_data_url, + audio_data_url=audio_data_url, + content_text=get_prompt("mix"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY}, + } + openai_client.send_request(request_config, request_num=get_max_batch_size()) diff --git a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml index f8655f77e..32b1ba15b 100644 --- a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml @@ -13,10 +13,10 @@ stage_args: model_arch: Qwen2_5OmniForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 2400 - max_num_batched_tokens: 2400 + max_model_len: 16384 + max_num_batched_tokens: 16384 max_num_seqs: 1 - gpu_memory_utilization: 0.8 + gpu_memory_utilization: 0.9 skip_mm_profiling: true enforce_eager: true # Now we only support eager mode trust_remote_code: true @@ -43,10 +43,10 @@ stage_args: model_arch: Qwen2_5OmniForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 2400 - max_num_batched_tokens: 2400 + max_model_len: 16384 + max_num_batched_tokens: 16384 max_num_seqs: 1 - gpu_memory_utilization: 0.8 + gpu_memory_utilization: 0.9 skip_mm_profiling: true enforce_eager: true trust_remote_code: true @@ -58,7 +58,7 @@ stage_args: temperature: 0.9 top_p: 0.8 top_k: 40 - max_tokens: 128 + max_tokens: 4096 seed: 42 detokenize: True repetition_penalty: 1.05 @@ -66,7 +66,7 @@ stage_args: - stage_id: 2 runtime: process: true - devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU + devices: "2" # Example: use a different GPU than the previous stage; use "0" if single GPU max_batch_size: 1 engine_args: model_stage: code2wav @@ -78,7 +78,7 @@ stage_args: trust_remote_code: true enable_prefix_caching: false engine_output_type: audio - max_num_batched_tokens: 4069 + max_num_batched_tokens: 100000 engine_input_source: [1] final_output: true final_output_type: audio @@ -86,7 +86,7 @@ stage_args: temperature: 0.0 top_p: 1.0 top_k: -1 - max_tokens: 128 + max_tokens: 4096 seed: 42 detokenize: True repetition_penalty: 1.1 diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index 5a529662b..a4199f18e 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -12,7 +12,8 @@ from pathlib import Path import pytest -from tests.conftest import OmniServer, convert_audio_file_to_text, cosine_similarity_text +from tests.conftest import convert_audio_file_to_text, cosine_similarity_text +from tests.utils import hardware_test models = ["Qwen/Qwen2.5-Omni-7B"] @@ -21,26 +22,7 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen2_5_omni") # Create parameter combinations for model and stage config -test_params = [(model, stage_config) for model in models for stage_config in stage_configs] - - -@pytest.fixture(scope="module") -def omni_server(request): - """Start vLLM-Omni server as a subprocess with actual model weights. - Uses module scope so the server starts only once per test module. - Multi-stage initialization can take 10-20+ minutes. - """ - model, stage_config_path = request.param - - print(f"Starting OmniServer with model: {model}") - print("This may take 10-20+ minutes for initialization...") - - with OmniServer( - model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"], port=8091 - ) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopped") +test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs] def run_cmd(command): @@ -67,6 +49,9 @@ def extract_content_after_keyword(keywords, text): return matches[0] +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: command = [ @@ -99,6 +84,9 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: command = [ @@ -131,6 +119,9 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: command = ["bash", os.path.join(example_dir, "run_curl_multimodal_generation.sh"), "mixed_modalities"] @@ -151,6 +142,9 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: command = [ @@ -178,6 +172,9 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: command = [ @@ -203,6 +200,9 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: command = [ @@ -237,6 +237,9 @@ def test_modality_control_003(omni_server) -> None: @pytest.mark.skip(reason="There is a known issue with stream error.") +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: command = [ diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py index 675146ff5..c7354905e 100644 --- a/tests/examples/online_serving/test_qwen3_omni.py +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -12,7 +12,8 @@ from pathlib import Path import pytest -from tests.conftest import OmniServer, convert_audio_file_to_text, cosine_similarity_text +from tests.conftest import convert_audio_file_to_text, cosine_similarity_text +from tests.utils import hardware_test models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -21,26 +22,7 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni") # Create parameter combinations for model and stage config -test_params = [(model, stage_config) for model in models for stage_config in stage_configs] - - -@pytest.fixture(scope="module") -def omni_server(request): - """Start vLLM-Omni server as a subprocess with actual model weights. - Uses module scope so the server starts only once per test module.. - Multi-stage initialization can take 10-20+ minutes. - """ - model, stage_config_path = request.param - - print(f"Starting OmniServer with model: {model}") - print("This may take 10-20+ minutes for initialization...") - - with OmniServer( - model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"], port=8091 - ) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopped") +test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs] def run_cmd(command): @@ -67,6 +49,9 @@ def extract_content_after_keyword(keywords, text): return matches[0] +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: command = [ @@ -92,6 +77,9 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: command = [ @@ -120,6 +108,9 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: command = ["bash", os.path.join(example_dir, "run_curl_multimodal_generation.sh"), "use_image"] @@ -134,6 +125,9 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: command = [ @@ -155,6 +149,9 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: command = [ @@ -175,6 +172,9 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: command = [ @@ -202,6 +202,9 @@ def test_modality_control_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: command = [ diff --git a/tests/perf/scripts/run_benchmark.py b/tests/perf/scripts/run_benchmark.py index 5895fe2a1..688666162 100644 --- a/tests/perf/scripts/run_benchmark.py +++ b/tests/perf/scripts/run_benchmark.py @@ -46,7 +46,8 @@ def modify_stage(default_path, updates, deletes): def create_unique_server_params(configs: list[dict[str, Any]]) -> list[tuple[str, str, str]]: - unique_params = set() + unique_params = [] + seen = set() for config in configs: test_name = config["test_name"] model = config["server_params"]["model"] @@ -55,9 +56,13 @@ def create_unique_server_params(configs: list[dict[str, Any]]) -> list[tuple[str delete = config["server_params"].get("delete", None) update = config["server_params"].get("update", None) stage_config_path = modify_stage(stage_config_path, update, delete) - unique_params.add((test_name, model, stage_config_path)) - return list(unique_params) + server_param = (test_name, model, stage_config_path) + if server_param not in seen: + seen.add(server_param) + unique_params.append(server_param) + + return unique_params def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]: @@ -95,6 +100,7 @@ def omni_server(request): print(f"Starting OmniServer with test: {test_name}, model: {model}") with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server: + server.test_name = test_name print("OmniServer started successfully") yield server print("OmniServer stopping...") @@ -150,10 +156,15 @@ def get_benchmark_params_for_server(test_name: str) -> list: def create_benchmark_indices(): indices = [] - for test_name, config_data in server_to_benchmark_mapping.items(): - params_list = config_data["benchmark_params"] - for idx in range(len(params_list)): - indices.append((test_name, idx)) + seen = set() + for config in BENCHMARK_CONFIGS: + test_name = config["test_name"] + if test_name not in seen: + seen.add(test_name) + params_list = get_benchmark_params_for_server(test_name) + for idx in range(len(params_list)): + indices.append((test_name, idx)) + return indices @@ -164,6 +175,10 @@ benchmark_indices = create_benchmark_indices() def benchmark_params(request, omni_server): """Benchmark parameters fixture with proper parametrization""" test_name, param_index = request.param + + if test_name != omni_server.test_name: + pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}") + all_params = get_benchmark_params_for_server(test_name) if not all_params: @@ -172,6 +187,10 @@ def benchmark_params(request, omni_server): if param_index >= len(all_params): raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}") + current = param_index + 1 + total = len(all_params) + print(f"\n Running benchmark {current}/{total} for {test_name}") + return {"test_name": test_name, "params": all_params[param_index]}