[Rebase] Rebase to vllm v0.24.0 (#4709)

.buildkite/test-merge.yml

@@ -174,7 +174,7 @@ steps:
           - vllm_omni/deploy/qwen3_tts.yaml
         commands:
           - |
-            timeout 20m bash -c "
+            timeout 40m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
               pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -201,7 +201,7 @@ steps:
           - vllm_omni/deploy/qwen3_tts.yaml
         commands:
           - |
-            timeout 20m bash -c "
+            timeout 40m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
               pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -262,10 +262,10 @@ steps:
           - vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
           - vllm_omni/config/pipeline_registry.py
           - vllm_omni/model_executor/models/registry.py
-        timeout_in_minutes: 20
+        timeout_in_minutes: 60
         commands:
           - |
-            timeout 20m bash -c "
+            timeout 45m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
             "
@@ -290,7 +290,7 @@ steps:
           - vllm_omni/model_executor/models/qwen3_omni/
           - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
           - vllm_omni/deploy/qwen3_omni_moe.yaml
-        timeout_in_minutes: 30
+        timeout_in_minutes: 50
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         agents:

.buildkite/test-nightly.yml

@@ -12,7 +12,7 @@ steps:
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
     steps:
       - label: ":full_moon: Omni · Function Test with H100"
-        timeout_in_minutes: 90
+        timeout_in_minutes: 120
         commands:
           - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
         agents:
@@ -153,7 +153,7 @@ steps:
 
       - label: ":full_moon: Omni · Perf Test · No Async Chunk"
         key: nightly-omni-performance-no-async-chunk
-        timeout_in_minutes: 180
+        timeout_in_minutes: 300
         commands:
           - export BENCHMARK_DIR=tests/dfx/perf/results
           - |
@@ -376,7 +376,7 @@ steps:
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
     steps:
       - label: ":full_moon: TTS · Function Test with L4"
-        timeout_in_minutes: 90
+        timeout_in_minutes: 120
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"
@@ -826,7 +826,7 @@ steps:
 
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
         commands:
           - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:

docker/Dockerfile.ci

@@ -1,8 +1,8 @@
-# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
+# Target vllm version: v0.24.0 (releases/v0.24.0 branch).
 # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
 # be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.23.0
+ARG VLLM_BASE_TAG=v0.24.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.23.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.24.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.23.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.24.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docker/Dockerfile.xpu

@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.23.0
+ARG VLLM_VERSION=v0.24.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, and 0.23.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, and 0.24.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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.23.0 --torch-backend=auto
+uv pip install vllm==0.24.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --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.23.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library
 
 ```bash
-uv pip install vllm==0.23.0 --torch-backend=auto
+uv pip install vllm==0.24.0 --torch-backend=auto
 ```
 
-The 0.23.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.24.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ 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.23.0
+git checkout v0.24.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
-#For CUDA 13.0 (the default for v0.23.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.24.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.24.0/vllm-0.24.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.23.0 \
+    vllm/vllm-omni:v0.24.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
+uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --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.23.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library
 
 ```bash
-uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
+uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
 ```
 
-The pre-built 0.23.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.24.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ 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.23.0
+git checkout v0.24.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.23.0 \
+  vllm/vllm-omni-rocm:v0.24.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.23.0
+  vllm/vllm-omni-rocm:v0.24.0
 ```
 
 # --8<-- [end:pre-built-images]

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.23.0 --torch-backend=auto
+uv pip install vllm==0.24.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
+uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.23.0` and vLLM Omni `0.23.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.24.0` and vLLM Omni `0.24.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

requirements/common.txt

@@ -9,7 +9,7 @@
 # That breakage is now handled in code instead of by pinning:
 # MossTTSNanoForGeneration wraps weight loading in
 # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
-transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
+transformers >= 5.5.3
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
@@ -30,3 +30,7 @@ janus>=1.0.0
 msgpack>=1.0.0
 pydantic>=2.1.0
 pydantic-settings>=2.0.0
+# Required for GGUF-quantized diffusion model loading (gguf_adapters). The CI
+# base image does not ship it, so declare it explicitly. The loader also guards
+# this import so a missing gguf no longer breaks all diffusion model loading.
+gguf>=0.10.0

tests/diffusion/quantization/test_gguf_config.py

@@ -8,11 +8,11 @@ from types import SimpleNamespace
 
 import pytest
 import torch
-from vllm.model_executor.layers.linear import LinearBase
-from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES, UnquantizedLinearMethod
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
 
 from vllm_omni.quantization import build_quant_config
 from vllm_omni.quantization.gguf_config import (
+    UNQUANTIZED_TYPES,
     DiffusionGGUFConfig,
     DiffusionGGUFLinearMethod,
     dequant_gemm_gguf,

tests/diffusion/quantization/test_hsdp_fp8.py

@@ -38,13 +38,13 @@ class _ToyKernel:
 
 def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
     """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
-    ``quant_method`` is an ``Fp8OnlineLinearMethod`` with a ``_ToyKernel``.
+    ``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``.
     """
-    from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
+    from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
 
     # Create the quant-method instance without calling __init__ (avoids
     # heavyweight upstream config dependencies).  isinstance() still works.
-    qm = object.__new__(Fp8OnlineLinearMethod)
+    qm = object.__new__(Fp8LinearMethod)
     qm.fp8_linear = _ToyKernel()
 
     # randn does not support float8 on CPU; zeros does and the test only

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -72,6 +72,10 @@ class TestHiggsAudioV3OnlineHappyPath:
                 "stream": False,
                 "response_format": "wav",
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                # whisper-small mishears this short clip ~0.5% of the time; on a
+                # failed match re-verify with this stronger ASR before failing so
+                # the gate is not flaky. Local-only key (not forwarded to server).
+                "transcript_escalation_model": "large-v3",
             }
         )
 

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -40,11 +40,21 @@ def get_batch_token_config(default_path):
     Uses the new flat-stage schema (``stages.<id>.<field>``); the legacy
     ``stage_args.<id>.engine_args.<field>`` path no longer applies because
     the deploy YAML doesn't nest engine fields under ``engine_args:``.
+
+    Note:
+        64 was the original value but caused a GPU hang (NVIDIA watchdog killed
+        the process after ~2 s) when 5 concurrent requests arrived, due to an
+        upstream scheduler change (``throttle_prefills`` API in vLLM commit
+        0b131b16c). 512 was first tried but caused output quality degradation for
+        long audio (~120 s) inputs. 2048 gives the scheduler enough token budget
+        to handle 5 concurrent requests even with multi-second audio prefill,
+        while still being well below the default 32768.
+        If the upstream scheduler changes again, this value may need adjustment.
     """
     return modify_stage_config(
         default_path,
         updates={
-            "stages": {0: {"max_num_batched_tokens": 64}, 1: {"max_num_batched_tokens": 64}},
+            "stages": {0: {"max_num_batched_tokens": 2048}, 1: {"max_num_batched_tokens": 2048}},
         },
     )
 

tests/entrypoints/openai_api/test_serving_chat_speaker.py

@@ -79,6 +79,7 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M
     serving_chat.renderer.get_tokenizer.return_value = mocker.MagicMock()
     serving_chat.reasoning_parser_cls = None
     serving_chat.tool_parser = None
+    serving_chat.parser_cls = None
     serving_chat.use_harmony = False
     serving_chat.enable_auto_tools = False
     serving_chat.exclude_tools_when_tool_choice_none = False
@@ -86,8 +87,9 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M
     serving_chat.chat_template = None
     serving_chat.chat_template_content_format = "string"
     serving_chat.default_chat_template_kwargs = {}
-    serving_chat._validate_chat_template = mocker.MagicMock(return_value=None)
-    serving_chat._prepare_extra_chat_template_kwargs = mocker.MagicMock(return_value={})
+    serving_chat.openai_serving_render = mocker.MagicMock()
+    serving_chat.openai_serving_render.validate_chat_template.return_value = None
+    serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={})
     serving_chat._preprocess_chat = mocker.AsyncMock(
         side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian")
     )

tests/helpers/assertions.py

@@ -23,6 +23,15 @@ from tests.helpers.media import (
 
 _GENDER_PIPELINE = None
 _GENDER_PIPELINE_LOCK = threading.Lock()
+# Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
+# short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
+# leading SFX token), which flakes the deterministic similarity gate. A test can
+# opt in to ASR escalation by setting ``transcript_escalation_model`` to a
+# whisper model name (e.g. ``"large-v3"``) in its request_config: on a failed
+# fast pass the clip is re-transcribed with that stronger ASR before the test
+# fails, so a weak-ASR mishear is rescued while a genuine model artifact still
+# fails (the strong ASR mismatches too). Tests that do not opt in keep the
+# strict behaviour.
 _PCM_SPEECH_SAMPLE_RATE_HZ = 24_000
 _MIN_PCM_SPEECH_HNR_DB = 1.0
 _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
@@ -456,7 +465,15 @@ def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_l
     if "audio" not in modalities:
         return False
     keywords_dict = request_config.get("key_words", {}) or {}
-    if keywords_dict.get("audio") and "text" not in modalities:
+    # When text is not an output modality, the keyword loop validates keywords
+    # against the audio transcript -- for keywords under ANY word_type
+    # (text/image/audio/video), not just "audio". Mirror that here so the
+    # transcript is actually computed; otherwise the loop hits
+    # `assert transcript is not None` with transcript=None (e.g. an audio-only
+    # request carrying key_words={"text": [...]}).
+    if "text" not in modalities and any(
+        keywords_dict.get(word_type) for word_type in ("text", "image", "audio", "video")
+    ):
         return True
     if request_config.get("audio_ref_text"):
         return True
@@ -590,6 +607,56 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                 )
 
 
+def _assert_transcript_matches(
+    transcript: str,
+    audio_bytes: bytes | None,
+    expected_text: Any,
+    *,
+    threshold: float,
+    escalation_model: str | None = None,
+) -> None:
+    """Assert spoken audio matches ``expected_text``.
+
+    ``transcript`` is the fast whisper-``small`` result. If it clears
+    ``threshold`` the check passes immediately.
+
+    When ``escalation_model`` is set (opt-in via the ``transcript_escalation_model``
+    request_config key) and the fast check fails, the clip is re-transcribed with
+    that stronger ASR and the assertion is decided on its verdict -- so a weak
+    whisper-``small`` mishear on a short clip does not flake the gate, while a
+    genuine model artifact still fails (the strong ASR mismatches too). When
+    ``escalation_model`` is ``None`` the original strict behaviour is preserved,
+    so other tests are unaffected.
+    """
+    expected = str(expected_text).strip().lower()
+    similarity = cosine_similarity_text(transcript.strip().lower(), expected)
+    print(f"Cosine similarity: {similarity:.3f}")
+    if similarity > threshold:
+        return
+
+    if escalation_model and audio_bytes:
+        print(
+            f"whisper-small below threshold ({similarity:.2f} <= {threshold}); "
+            f"escalating to whisper-{escalation_model} to rule out an ASR mishear"
+        )
+        strong_transcript = convert_audio_bytes_to_text(audio_bytes, model_size=escalation_model)
+        strong_similarity = cosine_similarity_text(strong_transcript.strip().lower(), expected)
+        print(
+            f"audio content (whisper-{escalation_model}): {strong_transcript}\n"
+            f"Cosine similarity (whisper-{escalation_model}): {strong_similarity:.3f}"
+        )
+        assert strong_similarity > threshold, (
+            f"Transcript doesn't match input after ASR escalation: "
+            f"input={expected_text!r}; whisper-small='{transcript}' (sim={similarity:.2f}); "
+            f"whisper-{escalation_model}='{strong_transcript}' (sim={strong_similarity:.2f})"
+        )
+        return
+
+    assert similarity > threshold, (
+        f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
+    )
+
+
 def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None:
     """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`.
 
@@ -643,10 +710,12 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
             if expected_text:
                 print(f"audio content is: {transcript}")
                 print(f"input text is: {expected_text}")
-                similarity = cosine_similarity_text(transcript.strip().lower(), str(expected_text).lower())
-                print(f"Cosine similarity: {similarity:.3f}")
-                assert similarity > 0.9, (
-                    f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
+                _assert_transcript_matches(
+                    transcript,
+                    getattr(response, "audio_bytes", None),
+                    expected_text,
+                    threshold=0.9,
+                    escalation_model=request_config.get("transcript_escalation_model"),
                 )
         _assert_preset_voice_gender_from_audio(
             response.audio_bytes,

tests/helpers/media.py

@@ -617,11 +617,11 @@ def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer:
 
 
 @contextmanager
-def _serialize_whisper_small_model_download():
-    """Serialize Whisper ``small`` cache writes across processes (Linux/Unix)."""
+def _serialize_whisper_model_download(model_size: str = "small"):
+    """Serialize Whisper cache writes across processes (Linux/Unix), per model."""
     import fcntl
 
-    lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock"
+    lock_path = Path.home() / ".cache" / "whisper" / f".{model_size}_model_download.lock"
     lock_path.parent.mkdir(parents=True, exist_ok=True)
     f = open(lock_path, "a+b")
     try:
@@ -632,7 +632,7 @@ def _serialize_whisper_small_model_download():
         f.close()
 
 
-def _whisper_transcribe_in_current_process(output_path: str) -> str:
+def _whisper_transcribe_in_current_process(output_path: str, model_size: str = "small") -> str:
     import whisper
 
     device_index = None
@@ -656,8 +656,8 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
         use_accelerator = False
         device = "cpu"
 
-    with _serialize_whisper_small_model_download():
-        model = whisper.load_model("small", device=device)
+    with _serialize_whisper_model_download(model_size):
+        model = whisper.load_model(model_size, device=device)
     try:
         text = model.transcribe(
             output_path,
@@ -674,15 +674,15 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
     return text or ""
 
 
-def convert_audio_file_to_text(output_path: str) -> str:
+def convert_audio_file_to_text(output_path: str, model_size: str = "small") -> str:
     """Convert an audio file to text in an isolated subprocess."""
     ctx = multiprocessing.get_context("spawn")
     with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
-        future = executor.submit(_whisper_transcribe_in_current_process, output_path)
+        future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size)
         return future.result()
 
 
-def convert_audio_bytes_to_text(raw_bytes: bytes) -> str:
+def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") -> str:
     output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav")
     os.close(output_fd)
     if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"):
@@ -690,7 +690,7 @@ def convert_audio_bytes_to_text(raw_bytes: bytes) -> str:
     data, samplerate = sf.read(io.BytesIO(raw_bytes))
     sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16")
     print(f"audio data is saved: {output_path}")
-    return convert_audio_file_to_text(output_path)
+    return convert_audio_file_to_text(output_path, model_size)
 
 
 __all__ = [

tests/helpers/stage_config.py

@@ -482,6 +482,7 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             {
                 "stage_id": 1,
                 "max_num_seqs": 1,
+                "gpu_memory_utilization": 0.5,
             },
         ],
     },
@@ -577,10 +578,15 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             {
                 "stage_id": 0,
                 "max_num_seqs": 1,
-                "gpu_memory_utilization": 0.9,
+                # Tuned for the heavier rebased stack (vLLM v0.23.1rc1 + torch 2.11 + CUDA 13
+                # + flashinfer) on a 24 GiB L4: the 16.78 GiB model weights + non-KV overhead
+                # + a 16384-token activation peak left 0 KV cache at util 0.85 (build 2354
+                # OOM). The abort test only needs the engine to init, so a smaller batched-
+                # token budget is fine; outputs are unaffected (chunked prefill is on).
+                "gpu_memory_utilization": 0.90,
                 "enforce_eager": True,
                 "enable_prefix_caching": False,
-                "max_num_batched_tokens": 16384,
+                "max_num_batched_tokens": 2048,
                 "max_model_len": 16384,
                 "skip_mm_profiling": True,
                 "mm_processor_cache_gb": 0,

tests/model_executor/models/moss_tts/test_per_request_codes.py

@@ -156,14 +156,19 @@ class TestPerRequestRouting:
         assert torch.all(audio[0][1] == 0)  # this step's row 0 appended
         assert torch.all(audio[1] == 1)
 
-    def test_runs_without_spans_backward_compat(self):
-        """Missing request_token_spans must not crash (equal-split fallback)
-        and still yields a batch-aligned list."""
+    def test_runs_with_spans(self):
+        """request_token_spans are required for correct per-request routing;
+        verify that passing them produces a batch-aligned list."""
         talker = _make_bare_delay_talker()
         hidden = _row_indexed_hidden(2)
         info_dicts = [{}, {}]
+        spans = [(0, 1), (1, 2)]
 
-        result = talker.make_omni_output(hidden, runtime_additional_information=info_dicts)
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
         audio = _audio_payload(result)
         assert isinstance(audio, list)
         assert len(audio) == 2

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -41,6 +41,7 @@ def _make_minimal_talker(
     def _default_raise(**_kwargs):
         raise AssertionError("build_prompt_embeds was not stubbed in this test")
 
+    model._embedding_dtype = torch.bfloat16
     model._prompt_builder = SimpleNamespace(
         build_prompt_embeds=build_prompt_embeds if build_prompt_embeds is not None else _default_raise,
     )
@@ -89,6 +90,7 @@ def _make_minimal_builder(
     )
     builder._speaker_cache = None
     builder._text_tokenizer = None
+    builder._embedding_dtype = torch.bfloat16
     builder._ref_audio_artifact_cache_max_entries = 256
     builder._ref_audio_artifact_cache = OrderedDict()
     builder._resampler_cache = OrderedDict()

tests/worker/test_gpu_ar_model_runner.py

@@ -434,8 +434,12 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
 
     output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
 
-    assert output.inter_stage_outputs is None
+    # Non-async-chunk now ships the full payload to the next stage, so
+    # inter_stage_outputs mirrors multimodal_outputs (PR #4792).
+    assert output.inter_stage_outputs is not None
     assert output.multimodal_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["hidden"], output.multimodal_outputs[0]["hidden"])
+    assert torch.equal(output.inter_stage_outputs[1]["hidden"], output.multimodal_outputs[1]["hidden"])
     assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
     assert torch.equal(
         output.multimodal_outputs[1]["hidden"],

vllm_omni/benchmarks/patch/patch.py

@@ -1261,7 +1261,7 @@ if "daily-omni" not in OPENAI_COMPATIBLE_BACKENDS:
 # Prevent import order from causing patch failures
 from vllm.benchmarks import serve
 from vllm.benchmarks.lib.ready_checker import wait_for_endpoint
-from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request
+from vllm.benchmarks.serve import TaskType, _merge_overrides, calculate_metrics_for_embeddings, get_request
 
 from vllm_omni.benchmarks.metrics.metrics import (
     MultiModalsBenchmarkMetrics,
@@ -1334,6 +1334,8 @@ async def benchmark(
         input_requests[0].expected_output_len,
         input_requests[0].multi_modal_data,
     )
+    test_extra_body = _merge_overrides(extra_body, input_requests[0].request_overrides)
+    test_chat_messages = input_requests[0].chat_messages
 
     assert (
         test_mm_content is None
@@ -1351,7 +1353,8 @@ async def benchmark(
         multi_modal_content=test_mm_content,
         ignore_eos=ignore_eos,
         extra_headers=extra_headers,
-        extra_body=extra_body,
+        extra_body=test_extra_body,
+        chat_messages=test_chat_messages,
     )
     _attach_daily_omni_to_request_func_input(input_requests[0], test_input)
     _attach_seed_tts_to_request_func_input(input_requests[0], test_input)
@@ -1415,7 +1418,8 @@ async def benchmark(
             multi_modal_content=test_mm_content,
             ignore_eos=ignore_eos,
             extra_headers=extra_headers,
-            extra_body=extra_body,
+            extra_body=test_extra_body,
+            chat_messages=test_chat_messages,
         )
         _attach_daily_omni_to_request_func_input(input_requests[0], profile_input)
         _attach_seed_tts_to_request_func_input(input_requests[0], profile_input)
@@ -1481,6 +1485,7 @@ async def benchmark(
             request.multi_modal_data,
             request.request_id,
         )
+        per_request_extra_body = _merge_overrides(extra_body, request.request_overrides)
         req_model_id, req_model_name = model_id, model_name
         if lora_modules:
             req_lora_module = next(lora_modules)
@@ -1497,8 +1502,9 @@ async def benchmark(
             multi_modal_content=mm_content,
             ignore_eos=ignore_eos,
             extra_headers=extra_headers,
-            extra_body=extra_body,
+            extra_body=per_request_extra_body,
             request_id=request_id,
+            chat_messages=request.chat_messages,
         )
         _attach_daily_omni_to_request_func_input(request, request_func_input)
         _attach_seed_tts_to_request_func_input(request, request_func_input)
@@ -1670,6 +1676,8 @@ async def benchmark(
             prompt_len=test_prompt_len,
             output_len=test_output_len,
             logprobs=logprobs,
+            extra_body=test_extra_body,
+            chat_messages=test_chat_messages,
         )
         profile_output = await request_func(request_func_input=profile_input, session=session)
         if profile_output.success:

vllm_omni/config/config_factory.py

@@ -67,12 +67,42 @@ class StageConfigFactory:
         if trust_remote_code is None:
             trust_remote_code = False
 
-        # --- New path: check pipeline registry by model_type first ---
         model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
         if model_type == "vla":
             if _looks_like_dreamzero(model):
                 model_type = "dreamzero"
 
+        # --- 1. Explicit deploy-config pipeline override (highest precedence) ---
+        # A deploy YAML may set ``pipeline: <model_type>`` to force routing for
+        # models whose HF config has a generic/colliding ``model_type`` or no
+        # matching architectures. Ming-omni-tts, for example, reports
+        # model_type="dense" with arch BailingMMNative…, which matches no
+        # pipeline — without this it falls through to the diffusion default and
+        # dies with "Model class BailingMMNativeForConditionalGeneration not
+        # found in diffusion model registry". Honor the key before auto-detection.
+        explicit_pipeline = None
+        if deploy_config_path:
+            try:
+                explicit_pipeline = load_deploy_config(deploy_config_path).pipeline
+            except Exception:
+                logger.exception("Failed to read 'pipeline' key from deploy config %s", deploy_config_path)
+        if explicit_pipeline:
+            pipeline_cfg = cls.resolve_pipeline_config(explicit_pipeline, hf_config)
+            if pipeline_cfg is not None:
+                return cls._create_from_registry(
+                    explicit_pipeline,
+                    pipeline_cfg,
+                    cli_overrides,
+                    deploy_config_path,
+                )
+            logger.warning(
+                "Deploy config %s requested pipeline %r which is not in OMNI_PIPELINES; "
+                "falling back to auto-detection.",
+                deploy_config_path,
+                explicit_pipeline,
+            )
+
+        # --- 2. Auto-detected model_type registered in OMNI_PIPELINES ---
         if model_type and model_type in OMNI_PIPELINES:
             pipeline_cfg = cls.resolve_pipeline_config(model_type, hf_config)
             if pipeline_cfg is not None:

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -205,7 +205,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         return False
 
-    def schedule(self) -> SchedulerOutput:  # type: ignore[override]
+    def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         # Remove FINISHED_ABORTED requests before the upstream scheduler sees
         # them. Upstream vllm raises RuntimeError on this status; omni allows
         # async abort (e.g. client disconnect during TTS streaming) to leave
@@ -223,7 +223,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             )
 
         try:
-            scheduler_output = super().schedule()
+            scheduler_output = super().schedule(throttle_prefills)
         finally:
             if self.chunk_transfer_adapter:
                 # Add request waiting for chunk to the waiting and running queue

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -56,7 +56,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             )
         self._latest_omni_connector_output: OmniConnectorOutput | None = None
 
-    def schedule(self) -> SchedulerOutput:
+    def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         """Diffusion fast path:
         - Feed all input tokens of the request at once
           (if 0, allocate 1 placeholder token).
@@ -217,7 +217,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     scheduler_requests=self.requests,
                 )
             else:
-                res = super().schedule()
+                res = super().schedule(throttle_prefills)
                 if self.input_coordinator:
                     self.input_coordinator.restore_queues(self.waiting, self.running)
                 return self._wrap_omni_scheduler_output(res)

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -26,15 +26,15 @@ _STATS_INTERVAL_S = 1.0
 # Scope: this constant only covers the full-payload coordinator path
 # (``input_coordinator``).  The async-chunk path uses
 # ``chunk_transfer_adapter`` and is not affected by this constant.
-_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "300")
+_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600")
 try:
     DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW)
 except ValueError:
     logger.warning(
-        "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 300 seconds.",
+        "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.",
         _INPUT_WAIT_TIMEOUT_RAW,
     )
-    DEFAULT_INPUT_WAIT_TIMEOUT_S = 300.0
+    DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0
 
 
 class OmniSchedulerMixin:

vllm_omni/deploy/moss_sound_effect.yaml

@@ -4,6 +4,7 @@
 # no reference audio required.
 # Rate: ~12.5 tokens/second; pass `tokens` to control duration.
 async_chunk: true
+trust_remote_code: true  # OpenMOSS-Team/MOSS-SoundEffect ships custom modeling code
 
 connectors:
   shm:

vllm_omni/deploy/moss_tts.yaml

@@ -2,6 +2,7 @@
 # HF: OpenMOSS-Team/MOSS-TTS
 # General-purpose TTS with voice cloning. Production-recommended variant.
 async_chunk: true
+trust_remote_code: true  # OpenMOSS-Team/MOSS-TTS ships custom modeling code
 
 connectors:
   shm:

vllm_omni/deploy/moss_tts_realtime.yaml

@@ -3,6 +3,9 @@
 # Real-time streaming TTS with TTFB ~180 ms and LLM+TTS first-sentence ~377 ms.
 # Uses local depth transformer (no delay warm-up) — lowest latency variant.
 async_chunk: true
+# OpenMOSS-Team/MOSS-TTS-Realtime ships custom modeling code; the newer vLLM
+# ModelConfig hard-rejects it at validation without this (was lenient before).
+trust_remote_code: true
 
 connectors:
   shm:

vllm_omni/deploy/moss_ttsd.yaml

@@ -2,6 +2,7 @@
 # HF: OpenMOSS-Team/MOSS-TTSD-v1.0
 # Dialogue TTS — optimised for multi-turn conversational speech generation.
 async_chunk: true
+trust_remote_code: true  # OpenMOSS-Team/MOSS-TTSD-v1.0 ships custom modeling code
 
 connectors:
   shm:

vllm_omni/deploy/moss_voice_generator.yaml

@@ -2,6 +2,9 @@
 # HF: OpenMOSS-Team/MOSS-VoiceGenerator
 # Zero-shot voice design — create a new voice from a text description.
 async_chunk: true
+# OpenMOSS-Team/MOSS-VoiceGenerator ships custom modeling code; the newer vLLM
+# ModelConfig hard-rejects it at validation without this (was lenient before).
+trust_remote_code: true
 
 connectors:
   shm:
@@ -19,6 +22,11 @@ stages:
   - stage_id: 0
     max_num_seqs: 8
     gpu_memory_utilization: 0.60
+    # The MOSS talker (custom audio heads + per-request GPU buffers) is not
+    # CUDA-graph-safe; FULL-decode graph capture crashes the worker. Run the
+    # LLM stage eager, matching the other TTS talker configs (ming_tts,
+    # moss_tts_nano).
+    enforce_eager: true
     enable_prefix_caching: false
     async_scheduling: true
     max_num_batched_tokens: 4096

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -18,7 +18,6 @@ from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
-    download_gguf,
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
     filter_duplicate_safetensors_files,
@@ -42,6 +41,39 @@ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter imp
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import initialize_model
 
+
+# download_gguf was removed from upstream vLLM (commit 6635279d8).
+# Inlined from the last upstream version before the GGUF plugin migration.
+def download_gguf(
+    repo_id: str,
+    quant_type: str,
+    cache_dir: str | None = None,
+    revision: str | None = None,
+    ignore_patterns: str | list[str] | None = None,
+) -> str:
+    allow_patterns = [
+        f"*-{quant_type}.gguf",
+        f"*-{quant_type}-*.gguf",
+        f"*/*-{quant_type}.gguf",
+        f"*/*-{quant_type}-*.gguf",
+    ]
+    folder = download_weights_from_hf(
+        model_name_or_path=repo_id,
+        cache_dir=cache_dir,
+        allow_patterns=allow_patterns,
+        revision=revision,
+        ignore_patterns=ignore_patterns,
+    )
+    local_files: list[str] = []
+    for pattern in allow_patterns:
+        glob_pattern = os.path.join(folder, pattern)
+        local_files.extend(glob.glob(glob_pattern))
+    if not local_files:
+        raise ValueError(f"Downloaded GGUF files not found in {folder} for quant_type {quant_type}")
+    local_files.sort(key=lambda x: (x.count("-"), x))
+    return local_files[0]
+
+
 logger = init_logger(__name__)
 
 

vllm_omni/diffusion/model_loader/gguf_adapters/base.py

@@ -6,10 +6,18 @@ from collections.abc import Generator
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
-import gguf
 import numpy as np
 import torch
 
+try:
+    import gguf
+except ImportError:
+    # Optional dependency: only required for GGUF-quantized diffusion models.
+    # Keep this import lazy so a missing ``gguf`` does not break importing the
+    # diffusion model loader (and thus inspection/loading of every other
+    # diffusion model that imports through this package).
+    gguf = None
+
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
     from vllm_omni.diffusion.model_loader.diffusers_loader import (
@@ -65,6 +73,11 @@ def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Te
     layer with different quant types.
     """
 
+    if gguf is None:
+        raise ImportError(
+            "The 'gguf' package is required to load GGUF-quantized diffusion models. "
+            "Install it with `pip install gguf`."
+        )
     reader = gguf.GGUFReader(gguf_file)
 
     for tensor in reader.tensors:

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -53,7 +53,6 @@ from vllm_omni.diffusion.models.dreamzero.utils import (
 )
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 MAX_DREAMZERO_SESSIONS = 64
@@ -487,10 +486,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         latents: tuple[torch.Tensor, torch.Tensor],
         do_true_cfg: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Post-step sync: .contiguous() + platform synchronize."""
+        """Post-step sync: .contiguous() + cuda.synchronize()"""
         latents = tuple(t.contiguous() for t in latents)
         if do_true_cfg and get_classifier_free_guidance_world_size() > 1:
-            current_omni_platform.synchronize()
+            device = next((t.device for t in latents if t.is_cuda), None)
+            if device is not None:
+                torch.cuda.current_stream(device).synchronize()
         return latents
 
     # -----------------------------------------------------------------------

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -284,10 +284,17 @@ class HeliosPipeline(
     ) -> DiffusionRequestState:
         """Initialize Helios request state for chunk-wise step execution."""
         del kwargs
-        req = OmniDiffusionRequest(
-            prompt=state.prompt,
-            sampling_params=state.sampling,
-            request_id=state.request_id,
+        # Wrap the single request in a DiffusionRequestBatch so the batch
+        # compatibility properties (`prompts`, etc.) used below are available;
+        # OmniDiffusionRequest itself only exposes a singular `prompt`.
+        req = DiffusionRequestBatch(
+            requests=[
+                OmniDiffusionRequest(
+                    prompt=state.prompt,
+                    sampling_params=state.sampling,
+                    request_id=state.request_id,
+                )
+            ]
         )
         extra = getattr(state.sampling, "extra_args", {}) or {}
 

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -1282,10 +1282,14 @@ class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery):
             self.pipe.scheduler = new_scheduler
 
         # We only want to change num_inference_steps here, so no need
-        # to deep copy the whole request
+        # to deep copy the whole request. `req` is a DiffusionRequestBatch
+        # whose `sampling_params` is a read-only property, so clone the
+        # underlying request(s) and override their sampling params instead.
         stage_2_req = copy.copy(req)
-        stage_2_req.sampling_params = req.sampling_params.clone()
-        stage_2_req.sampling_params.num_inference_steps = 3
+        stage_2_req.requests = [copy.copy(r) for r in req.requests]
+        for stage_2_request in stage_2_req.requests:
+            stage_2_request.sampling_params = stage_2_request.sampling_params.clone()
+            stage_2_request.sampling_params.num_inference_steps = 3
 
         stage2_output = self.pipe(
             req=stage_2_req,

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py

@@ -870,9 +870,14 @@ class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsComponentDiscovery):
             )
             self.pipe.scheduler = new_scheduler
 
+        # `req` is a DiffusionRequestBatch whose `sampling_params` is a
+        # read-only property; clone the underlying request(s) and override
+        # their sampling params instead of assigning to the batch.
         stage_2_req = copy.copy(req)
-        stage_2_req.sampling_params = req.sampling_params.clone()
-        stage_2_req.sampling_params.num_inference_steps = 3
+        stage_2_req.requests = [copy.copy(r) for r in req.requests]
+        for stage_2_request in stage_2_req.requests:
+            stage_2_request.sampling_params = stage_2_request.sampling_params.clone()
+            stage_2_request.sampling_params.num_inference_steps = 3
 
         stage2_output = self.pipe(
             req=stage_2_req,

vllm_omni/diffusion/quantization/hsdp_fp8.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """HSDP/FSDP2 compatibility for online FP8 quantization.
 
-vllm's ``Fp8OnlineLinearMethod.process_weights_after_loading`` ends with
+vllm's ``Fp8LinearMethod.process_weights_after_loading`` ends with
 ``layer.weight = qweight.t()`` so that the Cutlass FP8 GEMM kernel sees its
 B operand as column-major ``[K, N]`` (the TN layout required by Hopper FP8
 ``wgmma`` instructions). The resulting tensor is a non-contiguous transpose
@@ -32,7 +32,7 @@ import types
 
 from torch import nn
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
+from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
 
 logger = init_logger(__name__)
 
@@ -55,7 +55,7 @@ def _build_transposed_get_layer_params(original_bound_method):
 def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int:
     """Make online-FP8 linear layers in ``model`` FSDP2-compatible.
 
-    For every layer whose ``quant_method`` is an :class:`Fp8OnlineLinearMethod`
+    For every layer whose ``quant_method`` is an :class:`Fp8LinearMethod`
     and whose weight is currently a non-contiguous transpose view, this
     function:
 
@@ -76,7 +76,7 @@ def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int:
     patched_kernel_ids: set[int] = set()
     for module in model.modules():
         qm = getattr(module, "quant_method", None)
-        if not isinstance(qm, Fp8OnlineLinearMethod):
+        if not isinstance(qm, Fp8LinearMethod):
             continue
 
         weight = getattr(module, "weight", None)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -81,6 +81,10 @@ class _DiffusionVllmModelConfig:
     def is_nvfp4_quantized(self) -> bool:
         return self.quantization == "modelopt_fp4"
 
+    @property
+    def is_diffusion(self) -> bool:
+        return False
+
 
 def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig:
     quant_config = getattr(od_config, "quantization_config", None)

vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py

@@ -692,12 +692,20 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
             return None
 
         # 1. Allocate Destination Buffer from Pool
+        # For non-fast-path (serialized objects), allocate a few extra bytes
+        # as padding.  Mooncake's TCP/memcpy transport may prepend a leading
+        # protocol byte (0x01) before the payload.  The extra padding allows
+        # the receiver to detect and skip this byte without truncating the
+        # actual payload.  The padding is harmless for fast-path (ManagedBuffer)
+        # because the caller reads exact-size slices.
+        _MOONCAKE_TCP_PADDING = 4 if not is_fast_path else 0
+        alloc_size = data_size + _MOONCAKE_TCP_PADDING
         try:
-            offset = self.allocator.alloc(data_size)
-            recv_buffer = ManagedBuffer(self.allocator, offset, data_size, self.pool)
+            offset = self.allocator.alloc(alloc_size)
+            recv_buffer = ManagedBuffer(self.allocator, offset, alloc_size, self.pool)
             dst_ptr = self.base_ptr + offset
         except MemoryError:
-            logger.error(f"Failed to allocate {data_size} bytes in receive pool")
+            logger.error(f"Failed to allocate {alloc_size} bytes in receive pool")
             return None
 
         _t2 = _time_mod.perf_counter()
@@ -775,8 +783,19 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
                         _t_copy_end = _time_mod.perf_counter()
                         _copy_ms = (_t_copy_end - _t_copy_start) * 1000
 
+                        # Try deserializing the first data_size bytes.  Mooncake's
+                        # TCP/memcpy transport may prepend a leading protocol byte
+                        # (0x01) before the payload — handle this by retrying from
+                        # offset 1 if the first attempt raises DecodeError.
                         _t_deser_start = _time_mod.perf_counter()
-                        val = OmniSerializer.deserialize(raw_bytes)
+                        payload = raw_bytes[:data_size]
+                        try:
+                            val = OmniSerializer.deserialize(payload)
+                        except msgspec.DecodeError:
+                            # Likely a leading Mooncake TCP protocol byte — retry
+                            # from offset 1, using the extra padding bytes if needed.
+                            payload = raw_bytes[1 : data_size + 1]
+                            val = OmniSerializer.deserialize(payload)
                         _t_deser_end = _time_mod.perf_counter()
                         _deser_ms = (_t_deser_end - _t_deser_start) * 1000
 

vllm_omni/engine/stage_engine_core_proc_manager.py

@@ -38,13 +38,11 @@ from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc
 
 logger = init_logger(__name__)
 
-try:
-    # ``set_device_control_env_var`` lives next to CoreEngineProcManager and
-    # is only required for non-CUDA DP, so we tolerate its absence on
-    # older / future vLLM revisions.
-    from vllm.v1.engine.utils import set_device_control_env_var  # type: ignore
-except ImportError:  # pragma: no cover - depends on vLLM build
-    set_device_control_env_var = None  # type: ignore[assignment]
+# ``set_device_control_env_var`` was removed from upstream vllm.
+# It was only required for non-CUDA DP; set to None so the existing
+# guard at the call-site (vllm_config is not None and ... is not None)
+# skips the call transparently.
+set_device_control_env_var = None
 
 
 class StageEngineCoreProcManager(CoreEngineProcManager):

vllm_omni/engine/stage_runtime.py

@@ -140,6 +140,13 @@ class StageRuntime:
         self.stage_pools: list[StagePool] = []
         self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None
         self._spawn_device_lock = threading.Lock()
+        # Serialize all LLM replica spawning + handshake across device groups
+        # to prevent ZMQ port-allocation races (get_engine_zmq_addresses) and
+        # CUDA-context conflicts when multiple engine core subprocesses
+        # initialize simultaneously on different GPUs.  Matches the old
+        # AsyncOmniEngine._initialize_llm_replica pattern which used a single
+        # ``llm_stage_launch_lock`` for all replicas.
+        self._replica_launch_lock = threading.Lock()
         self._init_visible_devices_baseline: str | None = None
 
     @staticmethod
@@ -559,19 +566,22 @@ class StageRuntime:
                     plan.engine_args_dict,
                     stage_init_timeout,
                 )
-            with launch_stage_replica(
-                vllm_config=vllm_config,
-                executor_class=executor_class,
-                log_stats=False,
-                stage_id=plan.metadata.stage_id,
-                replica_id=plan.replica_id,
-                stage_config=plan.stage_cfg,
-                omni_master_server=self._get_omni_master_server(),
-                omni_coordinator_address=self._get_coordinator_address(),
-                stage_visible_devices=physical_devices,
-                spawn_device_lock=self._spawn_device_lock,
-            ) as resources:
-                pass
+            # Serialize engine-core spawning across all LLM replicas to avoid
+            # ZMQ port-allocation races and simultaneous CUDA context init.
+            with self._replica_launch_lock:
+                with launch_stage_replica(
+                    vllm_config=vllm_config,
+                    executor_class=executor_class,
+                    log_stats=False,
+                    stage_id=plan.metadata.stage_id,
+                    replica_id=plan.replica_id,
+                    stage_config=plan.stage_cfg,
+                    omni_master_server=self._get_omni_master_server(),
+                    omni_coordinator_address=self._get_coordinator_address(),
+                    stage_visible_devices=physical_devices,
+                    spawn_device_lock=self._spawn_device_lock,
+                ) as resources:
+                    pass
 
             logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id)
             if resources is None:

vllm_omni/entrypoints/openai/api_server.py

@@ -66,7 +66,7 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens
 # Keep a fallback for older/newer upstream layouts during rebase windows.
 from vllm.entrypoints.serve.instrumentator.basic import base
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
-from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
 from vllm.entrypoints.serve.utils.api_utils import (
     load_aware_call,
     process_lora_modules,
@@ -807,19 +807,6 @@ async def omni_init_app_state(
 
     state.vllm_config = vllm_config
 
-    # Propagate enable_in_reasoning to the API-server process. The engine core
-    # runs in a separate process, so the contextvar that backs
-    # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers
-    # call `get_enable_structured_outputs_in_reasoning()` during request
-    # handling and need to see the real flag, otherwise they silently fall
-    # back to False and mismatch the engine-side bitmask gating.
-    if vllm_config is not None:
-        from vllm.tool_parsers.structural_tag_registry import (
-            set_enable_structured_outputs_in_reasoning,
-        )
-
-        set_enable_structured_outputs_in_reasoning(vllm_config.structured_outputs_config.enable_in_reasoning)
-
     # Get supported tasks
     supported_tasks: set[str] = {"generate"}
     if hasattr(engine_client, "get_supported_tasks"):
@@ -900,12 +887,8 @@ async def omni_init_app_state(
     )
     await state.openai_serving_models.init_static_loras()
 
-    # NOTE: kept aligned with vllm 0.20 `init_app_state`:
-    # - dropped the `io_processor` kwarg (no longer accepted by 0.20);
-    #   io_processor stays on `engine_client` and downstream serving classes
-    #   read it from there.
-    # - pass `reasoning_parser` so render-time `adjust_request` runs for
-    #   reasoning models (matches `vllm.entrypoints.openai.api_server`).
+    # NOTE: kept aligned with upstream `init_app_state`:
+    # Use OpenAIServingRender (replaces the old OnlineRenderer + OnlineDerenderer + ServingRender split).
     state.openai_serving_render = OpenAIServingRender(
         model_config=engine_client.model_config,
         renderer=engine_client.renderer,
@@ -1031,8 +1014,7 @@ async def omni_init_app_state(
         if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
         else None
     )
-    state.openai_serving_tokenization = OpenAIServingTokenization(
-        engine_client,
+    state.openai_serving_tokenization = ServingTokenization(
         state.openai_serving_models,
         state.openai_serving_render,
         request_logger=request_logger,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -64,12 +64,12 @@ from vllm.entrypoints.openai.engine.protocol import (
     ToolCall,
     UsageInfo,
 )
-from vllm.entrypoints.openai.engine.serving import ChatLikeRequest, clamp_prompt_logprobs
+from vllm.entrypoints.openai.engine.serving import clamp_prompt_logprobs
 from vllm.entrypoints.openai.parser.harmony_utils import (
     get_streamable_parser_for_assistant,
-    parse_chat_output,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.entrypoints.serve.engine.typing import ChatLikeRequest
 from vllm.entrypoints.serve.utils.api_utils import should_include_usage
 from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls
 from vllm.inputs import PromptType
@@ -90,6 +90,7 @@ from vllm.tokenizers.mistral import (
 )
 from vllm.tool_parsers import ToolParser
 from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
+from vllm.tool_parsers.streaming import extract_required_tool_call_streaming
 from vllm.utils.collection_utils import as_list
 from vllm.v1.engine.exceptions import EngineDeadError
 
@@ -146,6 +147,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     _diffusion_extra_body_params: frozenset[str] | None = None
     _diffusion_extra_output_params: frozenset[str] | None = None
 
+    # Harmony flag (always False for vllm-omni models)
+    use_harmony: bool = False
+
+    def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool:
+        """Check if streamed tokens should go through the tool-call parser.
+
+        We only want to do this IF user-provided tools are set, a tool parser
+        is configured, "auto" tool choice is enabled, and the request's tool
+        choice field indicates that "auto" tool choice should be used.
+
+        This method existed in upstream vLLM commit 91df0fad4 (OpenAIServingChat)
+        but was removed in the Harmony refactoring (PR #45171, #45104).  Omni's
+        independently-maintained ``chat_completion_stream_generator`` still calls
+        it, so we keep a local copy.
+        """
+        # parser_cls may be None (no tool parser configured); guard accordingly
+        tp_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None
+        return request.tools and tp_cls is not None and self.enable_auto_tools and request.tool_choice in ["auto", None]
+
     @classmethod
     def for_diffusion(
         cls,
@@ -303,17 +323,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 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(
+            if self.parser_cls is not None and self.parser_cls.reasoning_parser_cls is not None:
+                chat_template_kwargs = self._effective_chat_template_kwargs(request)
+                reasoning_parser = self.parser_cls.reasoning_parser_cls(
                     tokenizer,
                     chat_template_kwargs=chat_template_kwargs,  # type: ignore[call-arg]
                 )
 
-            tool_parser = self.tool_parser
+            tool_parser = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None
 
             if isinstance(tokenizer, MistralTokenizer):
                 # because of issues with pydantic we need to potentially
@@ -351,7 +368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tool_dicts = [tool.model_dump() for tool in request.tools]
 
             if not self.use_harmony:
-                error_check_ret = self._validate_chat_template(
+                error_check_ret = self.openai_serving_render.validate_chat_template(
                     request_chat_template=request.chat_template,
                     chat_template_kwargs=request.chat_template_kwargs,
                     trust_request_chat_template=self.trust_request_chat_template,
@@ -359,14 +376,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 if error_check_ret is not None:
                     return error_check_ret
 
-                chat_template_kwargs = request.chat_template_kwargs or {}
-                chat_template_kwargs.update(reasoning_effort=request.reasoning_effort)
-
-                # Merge chat_template_kwargs with defaults
-                merged_template_kwargs = self._prepare_extra_chat_template_kwargs(
-                    chat_template_kwargs,
-                    self.default_chat_template_kwargs,
-                )
+                # Effective kwargs fold request.chat_template_kwargs, reasoning_effort,
+                # and server defaults — mirrors OpenAIServingChat._effective_chat_template_kwargs.
+                merged_template_kwargs = self._effective_chat_template_kwargs(request)
                 conversation, engine_prompts = await self._preprocess_chat(
                     request,
                     request.messages,
@@ -1103,7 +1115,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         return (
             output.finish_reason is not None
             and self.enable_auto_tools
-            and self.tool_parser is not None
+            and self.parser_cls is not None
+            and self.parser_cls.tool_parser_cls is not None
             and delta_message is not None
             and delta_message.tool_calls
         )
@@ -1226,8 +1239,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             all_previous_token_ids = None
         # Prepare the tool parser if it's needed
         try:
-            if tool_choice_auto and self.tool_parser:
-                tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer, request.tools)] * num_choices
+            tool_parser_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None
+            if tool_choice_auto and tool_parser_cls is not None:
+                tool_parsers: list[ToolParser | None] = [
+                    tool_parser_cls(tokenizer, request.tools) for _ in range(num_choices)
+                ]
             else:
                 tool_parsers = [None] * num_choices
         except Exception as e:
@@ -1555,12 +1571,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                                 # either finished reasoning or no reasoning at all
                                 content = current_text
 
-                                delta_message, function_name_returned[i] = self.extract_tool_call_required_streaming(
+                                delta_message, function_name_returned[i] = extract_required_tool_call_streaming(
                                     previous_text=previous_text,
                                     current_text=content,
                                     delta_text=delta_text,
                                     function_name_returned=fn_name_returned,
                                     tool_call_idx=history_tool_call_cnt,
+                                    tool_call_id_type=self.tool_call_id_type,
                                 )
                                 if (
                                     delta_message
@@ -2256,12 +2273,24 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 logprobs = None
 
             if self.use_harmony:
-                reasoning, content, _ = parse_chat_output(token_ids)
+                # Use upstream's HarmonyParser (inherited as self.parser_cls for
+                # gpt_oss/harmony models) instead of vendoring a copy of the
+                # pre-refactor parse_chat_output. We take only reasoning+content;
+                # tool calls are extracted by omni's own path below.
+                if self.parser_cls is not None:
+                    parser = self.parser_cls(tokenizer, request.tools)
+                    reasoning, content, _ = parser.parse(
+                        "",
+                        request,
+                        model_output_token_ids=token_ids,
+                    )
+                else:
+                    reasoning, content = None, None
                 if not request.include_reasoning:
                     reasoning = None
 
-                if self.tool_parser is not None:
-                    tool_parser = self.tool_parser(tokenizer, request.tools)
+                if self.parser_cls is not None and self.parser_cls.tool_parser_cls is not None:
+                    tool_parser = self.parser_cls.tool_parser_cls(tokenizer, request.tools)
                     # NOTE: We use token_ids for openai tool parser
                     tool_call_info = tool_parser.extract_tool_calls(
                         "",
@@ -2309,7 +2338,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             auto_tools_called = False
             # if auto tools are not enabled, and a named tool choice using
             #   outlines is not being used
-            if (not self.enable_auto_tools or not self.tool_parser) and (
+            if (not self.enable_auto_tools or self.parser_cls is None or self.parser_cls.tool_parser_cls is None) and (
                 not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
                 and request.tool_choice != "required"
             ):
@@ -2375,10 +2404,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 request.tools
                 and (request.tool_choice == "auto" or request.tool_choice is None)
                 and self.enable_auto_tools
-                and self.tool_parser
+                and self.parser_cls is not None
+                and self.parser_cls.tool_parser_cls is not None
             ):
                 try:
-                    tool_parser = self.tool_parser(tokenizer, request.tools)
+                    tool_parser = self.parser_cls.tool_parser_cls(tokenizer, request.tools)
                 except RuntimeError as e:
                     logger.exception("Error in tool parser creation.")
                     return self.create_error_response(e)

vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py

@@ -32,7 +32,10 @@ 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.fused_moe import FusedMoE
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoE,
+    fused_moe_make_expert_params_mapping,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -341,7 +344,7 @@ class BailingMoeV2SparseMoeBlock(nn.Module):
             prefix=f"{prefix}.experts",
         )
 
-        self.experts.expert_mapping = FusedMoE.make_expert_params_mapping(
+        self.experts.expert_mapping = fused_moe_make_expert_params_mapping(
             self.experts,
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -140,9 +140,17 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         self.pad_token_id: int = getattr(self.config, "pad_token_id", 151643)
         self.im_end_token_id: int = getattr(self.config, "im_end_token_id", 151645)
 
-        # Qwen3 backbone — weights live under ``language_model.*``
+        # Qwen3 backbone — weights live under ``language_model.*``. vLLM's
+        # Qwen3Model reads num_hidden_layers/hidden_size/etc. from
+        # vllm_config.model_config.hf_config, but MOSS nests the LLM config under
+        # ``language_config`` (a Qwen3Config). Rebind the backbone's hf_config to
+        # it — same pattern as the Qwen3-Omni talker — so the inner model sees a
+        # real Qwen3Config instead of MossTTSDelayConfig (which lacks
+        # num_hidden_layers).
+        lm_vllm_config = vllm_config.with_hf_config(lm_cfg, architectures=["Qwen3ForCausalLM"])
+        lm_vllm_config.model_config.hf_text_config = lm_vllm_config.model_config.hf_config
         self.model = Qwen3Model(
-            vllm_config=vllm_config,
+            vllm_config=lm_vllm_config,
             prefix=_maybe_prefix(prefix, "model"),
         )
 

vllm_omni/patch.py

@@ -324,11 +324,13 @@ def _patch_chat_template_registry():
     try:
         from vllm.transformers_utils.chat_templates.registry import (
             _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK,
-            _get_qwen_chat_template_fallback,
+            CHAT_TEMPLATES_DIR,
         )
 
         if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK:
-            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = _get_qwen_chat_template_fallback
+            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = (
+                lambda _: CHAT_TEMPLATES_DIR / "template_chatml.jinja"
+            )
     except ImportError:
         pass
 

vllm_omni/platforms/__init__.py

@@ -7,7 +7,6 @@ from itertools import chain
 from typing import TYPE_CHECKING
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
-from vllm.utils.torch_utils import supports_xccl
 
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
 from vllm_omni.plugins import (
@@ -86,7 +85,7 @@ def xpu_omni_platform_plugin() -> str | None:
     try:
         import torch
 
-        if supports_xccl():
+        if torch.distributed.is_xccl_available():
             dist_backend = "xccl"
         else:
             dist_backend = "ccl"

vllm_omni/platforms/cuda/platform.py

@@ -224,8 +224,16 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
-        """Copied from vllm/platforms/cuda/platform.py v0.20.0 with force using vllm_c kernels"""
-        default = ["vllm_c", "native"]  # Originally using "native" here when compiling
+        """Copied from upstream CudaPlatformBase with inductor-aware logic.
+
+        When inductor is active (compiling) use native as the default;
+        otherwise prefer vllm_c kernels where available.
+        """
+        from vllm.config.compilation import CompilationMode
+
+        cc = vllm_config.compilation_config
+        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
+        default = ["native"] if using_inductor else ["vllm_c", "native"]
 
         # Use oink if enabled for rms_norm
         # TODO(Laurawly/luka): remove this env var,

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -1145,7 +1145,10 @@ class NPUARModelRunner(OmniNPUModelRunner):
         if self._async_chunk:
             pooler_inter, pooler_client = partition_payload_list(pooler_output)
         else:
-            pooler_inter, pooler_client = None, pooler_output
+            # Non-async-chunk ships the full payload to the next stage via
+            # inter_stage_outputs (the NPU runner has no separate full-payload
+            # accumulate). #4527's (None, pooler_output) starved it. (PR #4792)
+            pooler_inter, pooler_client = pooler_output, pooler_output
         inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
         multimodal_outputs = self._build_multimodal_outputs(pooler_client)
         model_runner_output = OmniModelRunnerOutput(

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -479,7 +479,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         if self._async_chunk:
             inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
         else:
-            inter_stage_outputs, multimodal_outputs = None, per_req_payloads
+            # See npu_ar_model_runner: non-async-chunk ships the full payload to the next
+            # stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792)
+            inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
         req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()

vllm_omni/platforms/rocm/platform.py

@@ -163,10 +163,18 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
-        """Copied from vllm/platforms/rocm/platform.py v0.20.0 with force using vllm_c kernels"""
+        """Copied from upstream RocmPlatform with inductor-aware logic.
+
+        When inductor is active (compiling) use native as the default;
+        otherwise prefer vllm_c kernels where available.
+        Preserves omni-specific is_custom_op_enabled('rms_norm') check.
+        """
+        from vllm.config.compilation import CompilationMode
+
         # TODO(luka/TJ) use aiter, vllm_c, native by default on ROCm
         cc = vllm_config.compilation_config
-        default = ["vllm_c", "native"]  # Originally using "native" here when compiling
+        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
+        default = ["native"] if using_inductor else ["vllm_c", "native"]
 
         # This (mostly) preserves previous CustomOp behavior
         # Necessary on ROCm because it's common that users

vllm_omni/platforms/xpu/platform.py

@@ -97,7 +97,15 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
-        """Copied from vllm/platforms/xpu/platform.py v0.20.0 with force using xpu_kernels kernels"""
-        default = ["xpu_kernels", "native"]  # Originally using "native" here when compiling
+        """Copied from upstream XPUPlatform with inductor-aware logic.
+
+        When inductor is active (compiling) use native as the default;
+        otherwise prefer xpu_kernels where available.
+        """
+        from vllm.config.compilation import CompilationMode
+
+        cc = vllm_config.compilation_config
+        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
+        default = ["native"] if using_inductor else ["xpu_kernels", "native"]
 
         return IrOpPriorityConfig.with_default(default)

vllm_omni/quantization/gguf_config.py

@@ -3,22 +3,260 @@
 """GGUF quantization config for diffusion transformers.
 
 Uses dequant+GEMM instead of the fused kernel path (which expects 2D inputs).
+
+The GGUF quantization module was migrated to an external plugin in upstream
+vLLM (commit 6635279d8).  This file inlines the base classes that
+DiffusionGGUFConfig / DiffusionGGUFLinearMethod depend on, reproduced from
+the last upstream version before the migration.
 """
 
 from __future__ import annotations
 
+from collections.abc import Mapping
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any
+
 import gguf
 import torch
+from gguf import GGMLQuantizationType as WeightType
+from torch.nn.parameter import Parameter, UninitializedParameter
 from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.gguf import (
-    UNQUANTIZED_TYPES,
-    GGUFConfig,
-    GGUFLinearMethod,
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import (
     LinearBase,
-    QuantizeMethodBase,
+    LinearMethodBase,
     UnquantizedLinearMethod,
-    is_layer_skipped_gguf,
 )
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from vllm.model_executor.utils import set_weight_attrs
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization import QuantizationMethods
+
+logger = init_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# GGUF base types (inlined from upstream vllm before migration to plugin)
+# ---------------------------------------------------------------------------
+
+UNQUANTIZED_TYPES: set[WeightType] = {WeightType.F32, WeightType.F16, WeightType.BF16}
+
+STANDARD_QUANT_TYPES: set[WeightType] = {
+    WeightType.Q4_0,
+    WeightType.Q4_1,
+    WeightType.Q5_0,
+    WeightType.Q5_1,
+    WeightType.Q8_0,
+    WeightType.Q8_1,
+}
+KQUANT_TYPES: set[WeightType] = {
+    WeightType.Q2_K,
+    WeightType.Q3_K,
+    WeightType.Q4_K,
+    WeightType.Q5_K,
+    WeightType.Q6_K,
+}
+IMATRIX_QUANT_TYPES: set[WeightType] = {
+    WeightType.IQ1_M,
+    WeightType.IQ1_S,
+    WeightType.IQ2_XXS,
+    WeightType.IQ2_XS,
+    WeightType.IQ2_S,
+    WeightType.IQ3_XXS,
+    WeightType.IQ3_S,
+    WeightType.IQ4_XS,
+    WeightType.IQ4_NL,
+}
+DEQUANT_TYPES: set[WeightType] = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES
+
+
+def is_layer_skipped_gguf(
+    prefix: str,
+    unquantized_modules: list[str],
+    fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
+) -> bool:
+    # Fused layers like gate_up_proj or qkv_proj will not be fused
+    # in the safetensors checkpoint. So, we convert the name
+    # from the fused version to unfused + check to make sure that
+    # each shard of the fused layer has the same scheme.
+    proj_name = prefix.split(".")[-1]
+    if proj_name in fused_mapping:
+        shard_prefixes = [prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name]]
+
+        is_skipped = None
+        for shard_prefix in shard_prefixes:
+            is_shard_skipped = any(shard_prefix in module_name for module_name in unquantized_modules)
+
+            if is_skipped is None:
+                is_skipped = is_shard_skipped
+            elif is_shard_skipped != is_skipped:
+                raise ValueError(
+                    f"Detected some but not all shards of {prefix} "
+                    "are quantized. All shards of fused layers "
+                    "to have the same precision."
+                )
+    else:
+        is_skipped = any(module_name in prefix for module_name in unquantized_modules)
+
+    assert is_skipped is not None
+    return is_skipped
+
+
+class _GGUFUninitializedParameter(UninitializedParameter):
+    """UninitializedParameter subclass that stores GGUF weight data in a list.
+
+    The actual tensor is materialized later via _create_padded_weight_param.
+    """
+
+    cls_to_become = Parameter
+    data_container: list[torch.Tensor]
+
+
+class GGUFConfig(QuantizationConfig):
+    """Config class for GGUF (inlined from upstream vllm)."""
+
+    def __init__(self, unquantized_modules: list[str] | None = None) -> None:
+        super().__init__()
+        self.unquantized_modules = unquantized_modules or []
+
+    def __repr__(self) -> str:
+        return "GGUFConfig()"
+
+    def get_name(self) -> QuantizationMethods:
+        return "gguf"
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        return [torch.half, torch.bfloat16, torch.float32]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 60
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []  # no extra configs.
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> GGUFConfig:
+        return cls()
+
+    @classmethod
+    def override_quantization_method(
+        cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config=None
+    ) -> QuantizationMethods | None:
+        if user_quant == "gguf":
+            return "gguf"
+        return None
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
+                return UnquantizedLinearMethod()
+            return GGUFLinearMethod(self)
+        return None
+
+
+class GGUFLinearMethod(LinearMethodBase):
+    """Linear method for GGUF (inlined from upstream vllm)."""
+
+    def __init__(self, quant_config: GGUFConfig):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        self.params_dtype = params_dtype
+        output_size_per_partition = sum(output_partition_sizes)
+
+        tensor_shape = (output_size_per_partition, input_size_per_partition)
+        qweight = _GGUFUninitializedParameter(requires_grad=False)
+        set_weight_attrs(
+            qweight,
+            {
+                "input_dim": 1,
+                "output_dim": 0,
+                "tensor_shape": tensor_shape,
+                "is_gguf_weight": True,
+                "data_container": [],
+                "shard_id": [],
+                "shard_id_map": {},
+            },
+        )
+        set_weight_attrs(qweight, extra_weight_attrs)
+        layer.register_parameter("qweight", qweight)
+
+        qweight_type = Parameter(
+            torch.empty(len(output_partition_sizes), dtype=torch.uint8),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            qweight_type,
+            {
+                "is_gguf_weight_type": True,
+                "weight_type": 0,
+                "shard_weight_type": {},
+                "ignore_warning": True,
+            },
+        )
+        set_weight_attrs(qweight_type, extra_weight_attrs)
+        layer.register_parameter("qweight_type", qweight_type)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        qweight_type = layer.qweight_type.weight_type
+        if not (qweight_type in UNQUANTIZED_TYPES or qweight_type in DEQUANT_TYPES):
+            qweight_type = WeightType(qweight_type)
+            raise ValueError(f"Unsupported GGUF quantization type {qweight_type} in layer {layer}.")
+        self._create_padded_weight_param(layer)
+
+    def _create_padded_weight_param(self, layer: torch.nn.Module):
+        """Create padded weight parameter for GGUF MergedLinear layer."""
+        qweight = layer.qweight
+        shard_id_map = qweight.shard_id_map
+        shard_id = qweight.shard_id
+        if len(data_container := qweight.data_container) > 1:
+            dtype = {data.dtype for data in data_container}
+            assert len(dtype) == 1, ValueError(f"Data container has mixed dtypes: {dtype}")
+            dtype = next(iter(dtype))
+            padded_side = max(x.size(1) for x in data_container)
+            concat_side = sum(x.size(0) for x in data_container)
+            padded_data = torch.zeros((concat_side, padded_side), dtype=dtype, device=qweight.device)
+            shard_offset_map = dict[str, tuple[int, int, int]]()
+            for idx in shard_id:
+                id_in_container = shard_id_map[idx]
+                start = sum(x.size(0) for x in data_container[:id_in_container])
+                end = start + data_container[id_in_container].size(0)
+                size = data_container[id_in_container].size(1)
+                padded_data[start:end, :size] = data_container[id_in_container]
+                shard_offset_map[idx] = (start, end, size)
+            qweight.data_container.clear()
+            padded_param = Parameter(padded_data, requires_grad=False)
+            set_weight_attrs(padded_param, vars(qweight))
+            set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map})
+            layer.register_parameter("qweight", padded_param)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Base apply: overridden by DiffusionGGUFLinearMethod.
+        raise NotImplementedError("Use DiffusionGGUFLinearMethod for diffusion GGUF")
+
+
+# ---------------------------------------------------------------------------
+# Diffusion-specific GGUF (dequant+GEMM, replaces fused kernel path)
+# ---------------------------------------------------------------------------
 
 
 def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor:
@@ -48,7 +286,13 @@ class DiffusionGGUFLinearMethod(GGUFLinearMethod):
             for idx in shard_id:
                 start, end, offset = layer.qweight.shard_offset_map[idx]
                 qweight_type = layer.qweight_type.shard_weight_type[idx]
-                result.append(dequant_gemm_gguf(x, qweight[start:end, :offset].contiguous(), qweight_type))
+                result.append(
+                    dequant_gemm_gguf(
+                        x,
+                        qweight[start:end, :offset].contiguous(),
+                        qweight_type,
+                    )
+                )
             out = torch.cat(result, axis=-1)
         else:
             qweight = layer.qweight

vllm_omni/worker/gpu_ar_model_runner.py

@@ -1654,7 +1654,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         if self._async_chunk:
             pooler_inter, pooler_client = partition_payload_list(pooler_output)
         else:
-            pooler_inter, pooler_client = None, pooler_output
+            # Non-async-chunk still ships the full payload to the next stage (via
+            # accumulate_full_payload_output and the inter_stage_outputs field); only
+            # client mm keys are split out when async_chunk is enabled. #4527 set this
+            # to (None, pooler_output), which skipped accumulation and starved the
+            # downstream stage (300s connector-input timeout / empty audio). (PR #4792)
+            pooler_inter, pooler_client = pooler_output, pooler_output
 
         if pooler_inter and self._should_accumulate_full_payload_output():
             with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"):

vllm_omni/worker/gpu_ar_worker.py

@@ -50,16 +50,32 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
                 # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
                 self.local_rank += dp_local_rank * tp_pp_world_size
+
+            # Publish the logical-to-physical mapping for topology queries
+            # such as NIC affinity and P2P checks (upstream PR #45026).
+            assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids
+            if assigned_physical_gpu_ids is not None:
+                from vllm.platforms.interface import set_assigned_physical_gpu_ids
+
+                set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
+                assert self.local_rank < len(assigned_physical_gpu_ids), (
+                    f"local_rank {self.local_rank} is out of bounds for "
+                    f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
+                )
+                if parallel_config.distributed_executor_backend not in ("ray", "external_launcher"):
+                    assert self.parallel_config.local_world_size <= len(assigned_physical_gpu_ids), (
+                        f"local_world_size ({self.parallel_config.local_world_size}) "
+                        "exceeds assigned_physical_gpu_ids count "
+                        f"({len(assigned_physical_gpu_ids)})"
+                    )
+            else:
                 assert self.local_rank < torch.accelerator.device_count(), (
-                    f"DP adjusted local rank {self.local_rank} is out of bounds. "
+                    f"DP adjusted local rank {self.local_rank} is out of "
+                    f"bounds for {torch.accelerator.device_count()} devices."
                 )
-                visible_device_count = torch.accelerator.device_count()
-                assert self.parallel_config.local_world_size <= visible_device_count, (
-                    f"local_world_size ({self.parallel_config.local_world_size}) must "
-                    f"be less than or equal to the number of visible devices "
-                    f"({visible_device_count})."
-                )
-            self.device = current_omni_platform.get_torch_device(self.local_rank)
+
+            visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank)
+            self.device = current_omni_platform.get_torch_device(visible_device_index)
             torch.accelerator.set_device_index(self.device)
 
             current_platform.check_if_supports_dtype(self.model_config.dtype)

vllm_omni/worker/gpu_generation_model_runner.py

@@ -459,7 +459,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         if self._async_chunk:
             inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
         else:
-            inter_stage_outputs, multimodal_outputs = None, per_req_payloads
+            # See gpu_ar_model_runner: non-async-chunk ships the full payload to the next
+            # stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792)
+            inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads
 
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()

vllm_omni/worker/gpu_model_runner.py

@@ -318,6 +318,12 @@ class OmniGPUModelRunner(GPUModelRunner):
 
         Raises:
             AssertionError: If the model does not support M-RoPE
+
+        Note:
+            Upstream vLLM (commit 470229c37) added a fallback for the case
+            where ``req_state.prompt_token_ids`` is None but
+            ``req_state.prompt_embeds`` is available.  Omni models always set
+            ``prompt_token_ids``, so this fallback is deliberately omitted.
         """
         image_grid_thw = []
         video_grid_thw = []
@@ -1524,7 +1530,16 @@ class OmniGPUModelRunner(GPUModelRunner):
         num_input_tokens: int,
         intermediate_tensors: IntermediateTensors | None = None,
     ):
-        """Align with v0.14.0 preprocess and omni's additional information handling."""
+        """Align with v0.14.0 preprocess and omni's additional information handling.
+
+        Note:
+            Upstream vLLM (commit c621af169) added a conditional in the
+            ``supports_mm_inputs`` path to handle precomputed ``prompt_embeds``
+            alongside multimodal inputs.  Omni models that reach this override
+            always go through the omni-specific ``model.preprocess`` /
+            ``has_preprocess`` code path below, so the upstream change is
+            deliberately not ported.
+        """
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         is_first_rank = get_pp_group().is_first_rank
         is_encoder_decoder = self.model_config.is_encoder_decoder

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -919,7 +919,7 @@ class OmniConnectorModelRunnerMixin:
         if not (finished_req_ids & pending_req_ids):
             return
 
-        logger.debug(
+        logger.info(
             "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s",
             self._stage_id,
             finished_req_ids,