[Rebase] Rebase to vllm v0.16.0 (#1357)

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: ZJY0516 <zhu.jiangyun@foxmail.com>
Co-authored-by: youkaichao <youkaichao@gmail.com>

.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

.gitignore

@@ -242,3 +242,6 @@ configs/development.yaml
 Dockerfile.dev
 discussion
 tmp_test
+
+# output files
+*.wav

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/<commit_hash>/
+# 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 []

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

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

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

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 = (

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*"]

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.

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")

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()

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(

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,155 +111,9 @@ 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,
-        )
-
-        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,
+        # 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,
@@ -282,26 +128,15 @@ class OmniModelConfig(ModelConfig):
             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."
-            )
-
+        # 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:
-            # 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()

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)
 

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)

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),
         )
 
 

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

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,
             ),
         )

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"

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

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,

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

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:

vllm_omni/entrypoints/chat_utils.py

@@ -1,104 +1,36 @@
-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.
 
-
-class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
-    def __init__(self, tracker: AsyncMultiModalItemTracker) -> None:
-        super().__init__(tracker=tracker)
-        self._mm_processor_kwargs: dict[str, Any] | None = 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 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)
-
-        # 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)
-
-    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
-
-    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.
+    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
 
-        # 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)
@@ -107,201 +39,24 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
 
     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
 
             temp_video_file_path = url2pathname(parsed_url.path)
         elif parsed_url.scheme == "data":
-                # Handle data URL (base64 encoded video)
             import base64
 
             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
+            # Assume local file path
             temp_video_file_path = video_url
 
-            # 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)
-
         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)]

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:

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

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,
+        # 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,
                 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,
-                tools=tool_dicts,
-                chat_template=chat_template,
-                tokenize=False,
-                **hf_chat_template_kwargs,
+                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)

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
 

vllm_omni/inputs/parse.py

@@ -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")

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]

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,11 +278,15 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                 "in the audio tower part."
             )
 
+        self.quant_config = quant_config
+
+        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,
@@ -290,7 +297,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             else:
                 self.visual = None
 
-        self.quant_config = quant_config
+        with self._mark_language_model(vllm_config):
             self.language_model = init_vllm_registered_model(
                 vllm_config=vllm_config,
                 prefix=maybe_prefix(prefix, "language_model"),
@@ -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,

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

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.quant_config = quant_config
 
+        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
+        )
+        self.visual_dim = thinker_config.vision_config.out_hidden_size
+        self.multiscale_dim = self.visual_dim * self.deepstack_num_level
+
+        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"),
             )
-        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
-
-        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 = (
-            [
+            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)
                 ]
-            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_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 = []
+
+            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,12 +891,17 @@ 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)
+                    if not is_interleaved:
+                        current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens]
                         is_vision[current_positions] = True
 
                 # Audio embeddings
                 else:
+                    if not is_interleaved:
+                        current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens]
                         is_vision[current_positions] = False
 
+                if not is_interleaved:
                     mm_position_idx += num_tokens
 
             deepstack_input_embeds = inputs_embeds.new_zeros(
@@ -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 <language>] | for transcription
+        # Translate this audio [from <language> into <to_language>] | 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)
+        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 text_len > 0:
+                llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx)
+                st_idx += text_len
+
+            bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
+            llm_pos_ids_list.append(bos_pos)
+            st_idx += 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:
-                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)
+                    audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1))
+                    llm_pos_ids_list.append(audio_bos_pos)
 
-            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
+                    pos_ids, _ = self._compute_interleaved_positions(st_idx, data)
+                    llm_pos_ids_list.append(pos_ids)
+                    st_idx = int(pos_ids.max()) + 1
 
-        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)
+                    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)
 
-        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+                    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:
         """

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

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,
         )

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)
 

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.

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,11 +766,21 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 assert mm_budget is not None
 
                 if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
+                    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.max_items_per_batch_by_modality[dummy_modality]
+                        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 "
@@ -789,27 +804,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                             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))
+                        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)

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.

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,
             )
         )