Merge branch 'main' into dev/vllm-align

Drop the old-base wheel reinstall and dependency repair block. Update source installation and configuration references to vLLM 0.29 while retaining explicitly matched 0.28 prebuilt Omni examples until its next release.

Validation: official image and CUDA/ROCm release artifacts checked; 6 documentation tests passed; diff whitespace checks passed. Markdown lint reports existing snippet/anchor issues. Container build not run: no builder available.
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/common/ci_mirror_hardwares.yml

@@ -28,6 +28,9 @@
 # packed onto g6.12xlarge L4 nodes (``vllm.ci/gpu-pool=l4x4``). GPU / CPU /
 # memory / shm requests scale with card count (1 GPU = 10 CPU / 40Gi / 24Gi shm).
 # Retry covers lost agents / killed pods, not pytest or OOM (exit 1 / 137).
+# The Kubernetes HF secret deliberately uses VLLM_CI_HF_TOKEN. Scheduled
+# build-level HF_TOKEN values can override pod env entries with the same name;
+# upload_pipeline.py restores the secret-backed value inside the job shell.
 
 x-npu-k8s-plugins: &npu_k8s_plugins
   - kubernetes:
@@ -86,7 +89,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /fsx/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -137,7 +140,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /fsx/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -188,7 +191,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /fsx/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -239,7 +242,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /fsx/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -273,7 +276,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -308,7 +311,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -343,7 +346,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -378,7 +381,7 @@ mirror_hardwares:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -428,7 +431,7 @@ mirror_hardwares:
                     value: "0"
                   - name: HF_HOME
                     value: /raid/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -479,7 +482,7 @@ mirror_hardwares:
                     value: "0"
                   - name: HF_HOME
                     value: /raid/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -530,7 +533,7 @@ mirror_hardwares:
                     value: "0"
                   - name: HF_HOME
                     value: /raid/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret
@@ -581,7 +584,7 @@ mirror_hardwares:
                     value: "0"
                   - name: HF_HOME
                     value: /raid/hf_cache
-                  - name: HF_TOKEN
+                  - name: VLLM_CI_HF_TOKEN
                     valueFrom:
                       secretKeyRef:
                         name: hf-token-secret

.buildkite/common/scripts/upload_pipeline.py

@@ -87,6 +87,8 @@ BOOTSTRAP_UPLOAD_IF_KEYS = {
 }
 E2E_GROUP_MARKER = "E2E Test"
 CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml"
+CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN"
+CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi'
 
 # Bootstrap Buildkite ``if`` expressions.
 # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
@@ -509,6 +511,17 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
 
     expanded = copy.deepcopy(preset)
     merged = {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded
+    # A scheduled build's HF_TOKEN may override a pod env entry with the same
+    # name. The preset injects the Kubernetes secret under a collision-proof
+    # alias; restore the conventional name after Buildkite has merged env.
+    if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()):
+        commands = merged.get("commands")
+        if isinstance(commands, list):
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands]
+        elif commands is None:
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT]
+        else:
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, commands]
     # Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
     if "retry" in step:
         merged["retry"] = step["retry"]

docker/Dockerfile.ci

@@ -1,27 +1,10 @@
-# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5).
-# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
-#
-# Release-boundary form: there is no precompiled-wheel pin here, because the
-# base image already ships exactly the target vLLM. Everything the runtime
-# needs — torch, flashinfer, nvJitLink, numpy — comes from that image as one
-# self-consistent set.
-#
-# This block reappears only when the rebase targets a vLLM commit that is NOT a
-# released tag. During the 0.27 rebase, targeting v0.27.0rc1 on top of the
-# v0.26.0 base image required reinstalling vLLM from wheels.vllm.ai and then
-# hand-repairing the fallout: pinning torch ahead of flashinfer (installing it
-# after bound flashinfer's TVM FFI runtime to the old torch and aborted engine
-# core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy
-# under numba's ceiling. All of it was compensation for a base image and a vLLM
-# that disagreed. At a real release boundary they agree, so none of it is
-# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels
-# rode on the v0.27.0 base image with the full reinstall block until the
-# official v0.28.0 image shipped, at which point this file dropped back to the
-# form you are reading. If you pin an untagged commit again, restore the
-# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather
-# than letting the Dockerfile silently ship the wrong vLLM.
+# Target vLLM version: v0.29.0 (tag @ 98dff2a81d74).
+# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer
+# runtime. Do not reinstall a commit wheel or repair dependencies on top of
+# an older base image at a release boundary. For a future unreleased target,
+# restore an explicit wheel pin and validate its complete runtime stack.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.28.0
+ARG VLLM_BASE_TAG=v0.29.0
 
 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine, please refer to the [vLLM 0.28 configuration guide](https://docs.vllm.ai/en/v0.28.0/configuration/index.html).
+For options within a vLLM Engine, please refer to the [vLLM 0.29 configuration guide](https://docs.vllm.ai/en/v0.29.0/configuration/index.html).
 
 Each model defines fixed topology in a registered `PipelineConfig` and runtime
 overrides in a deploy YAML.

docs/configuration/environment_variables.md

@@ -151,7 +151,7 @@ their keys only.
 ## Inherited vLLM variables
 
 vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
-the [vLLM 0.28 environment-variable reference](https://docs.vllm.ai/en/v0.28.0/configuration/env_vars.html)
+the [vLLM 0.29 environment-variable reference](https://docs.vllm.ai/en/v0.29.0/configuration/env_vars.html)
 for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
 XPU, ModelScope, and FlashInfer workspace settings.
 

docs/getting_started/installation/README.md

@@ -1,7 +1,7 @@
 # Installation
 
 !!! important
-    vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0.
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. The 0.29 development line requires vLLM 0.29.x; the source-install instructions pin vLLM 0.29.0. Published vLLM-Omni 0.28.0 wheels and images still require vLLM 0.28.x.
 
 vLLM-Omni supports the following hardware platforms:
 

docs/getting_started/installation/gpu.md

@@ -34,6 +34,8 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
+The source-install instructions target vLLM 0.29.0. Until vLLM-Omni 0.29.0 artifacts are published, the CUDA/ROCm pre-built examples below remain on the matching 0.28.0 release pair.
+
 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, 0.24.0, 0.25.0, 0.26.0, and 0.28.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

@@ -5,7 +5,7 @@
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
 
-vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.28.x release line uses vLLM 0.28.x.
+vLLM-Omni depends on the matching major/minor release of vLLM. The 0.29 development line uses vLLM 0.29.x. Published 0.28.0 wheels and images use vLLM 0.28.x.
 
 !!! note
     PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -18,6 +18,8 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 #### Installation of vLLM
 
+These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below.
+
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
 uv pip install vllm==0.28.0 --torch-backend=auto
@@ -26,12 +28,12 @@ uv pip install vllm==0.28.0 --torch-backend=auto
 #### Installation of vLLM-Omni
 
 ```bash
-uv pip install vllm-omni
+uv pip install vllm-omni==0.28.0
 ```
 
 To run Gradio demos, also install the optional extras:
 ```bash
-uv pip install 'vllm-omni[demo]'
+uv pip install 'vllm-omni[demo]==0.28.0'
 ```
 
 # --8<-- [end:pre-built-wheels]
@@ -39,13 +41,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.28.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library
 
 ```bash
-uv pip install vllm==0.28.0 --torch-backend=auto
+uv pip install vllm==0.29.0 --torch-backend=auto
 ```
 
-The 0.28.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.29.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 +68,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.28.0
+git checkout v0.29.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.28.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.29.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.29.0/vllm-0.29.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -131,7 +133,7 @@ If you want to specify the base vLLM version:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.29.0 \
   -t vllm-omni-cuda .
 ```
 

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

@@ -11,6 +11,8 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 #### Installation of vLLM
 
+These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below.
+
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
 uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
@@ -22,7 +24,7 @@ uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/roc
 # we need to add --no-build-isolation as the torch
 # is not obtained from pypi, we have to install using the
 # torch installed in our environment
-uv pip install vllm-omni
+uv pip install vllm-omni==0.28.0
 
 # Optional if want to run Qwen3 TTS
 uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
@@ -34,13 +36,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.28.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library
 
 ```bash
-uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
+uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.29.0/rocm723
 ```
 
-The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.29.0 vLLM wheel targets ROCm 7.2.3. 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 +60,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.28.0
+git checkout v0.29.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```

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.28.0 --torch-backend=auto
+uv pip install vllm==0.29.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
+uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.29.0/rocm723
 
 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.28.0` and vLLM Omni `0.28.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.29.0` and vLLM Omni `0.29.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

tests/attention/test_fish_kvcache_attn.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import pytest
 import torch
@@ -214,11 +214,20 @@ def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch):
     assert fish_kvcache_backend.get_fish_kvcache_attn_stats()["small_hit_count"] == 1
 
 
-def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size):
+def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size, dtype=torch.float16):
     # Source of truth for the paged KV cache layout the backend must unpack.
-    from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+    # vLLM 0.29 removed AttentionBackend.get_kv_cache_shape; the per-layer shape
+    # now comes from the spec, with the content dim expressed in bytes.
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, compute_layer_kv_cache_shape_bytes
 
-    return FlashAttentionBackend.get_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size)
+    spec = FullAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=num_kv_heads,
+        head_size=head_size,
+        dtype=dtype,
+    )
+    num_pages, num_heads, num_states, content_bytes = compute_layer_kv_cache_shape_bytes(spec, num_blocks)
+    return (num_pages, num_heads, num_states, content_bytes // dtype.itemsize)
 
 
 def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch):

tests/buildkite/test_upload_pipeline.py

@@ -12,6 +12,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "com
 
 from skip_ci import resolve_ci_decision  # noqa: E402
 from upload_pipeline import (  # noqa: E402
+    CUDA_HF_TOKEN_EXPORT,
     _expand_mirror_hardwares,
     _get_mirror_hw_selector,
     _load_bootstrap_steps,
@@ -143,6 +144,10 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest
     container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0]
     assert container["image"].endswith("$BUILDKITE_COMMIT")
     assert container["resources"]["limits"]["nvidia.com/gpu"] == 1
+    env_names = {item["name"] for item in container["env"]}
+    assert "VLLM_CI_HF_TOKEN" in env_names
+    assert "HF_TOKEN" not in env_names
+    assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"]
 
 
 def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
@@ -182,6 +187,33 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
     assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [
         {"name": "swr-secret"},
     ]
+    assert step["commands"] == ["pytest -sv tests/example"]
+
+
+def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None:
+    for name in (
+        "l4_1",
+        "l4_2",
+        "l4_3",
+        "l4_4",
+        "h100_1",
+        "h100_2",
+        "h100_3",
+        "h100_4",
+        "b200_1",
+        "b200_2",
+        "b200_3",
+        "b200_4",
+    ):
+        step = _expand_mirror_hardwares(
+            {"label": name, "mirror_hardwares": name, "commands": ["pytest -sv tests/example"]},
+        )
+        assert step is not None
+        container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0]
+        env_names = {item["name"] for item in container["env"]}
+        assert "VLLM_CI_HF_TOKEN" in env_names
+        assert "HF_TOKEN" not in env_names
+        assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT
 
 
 def _gpu_limit(step: dict) -> int:

tests/config/test_omni_config.py

@@ -215,6 +215,16 @@ def test_from_pipeline_config_normalizes_stage_engine_extras_without_expanding_s
     assert stage.diffusion_config.model_config["default_robot_embodiment"] == "roboarena"
 
 
+@pytest.mark.parametrize("disabled", [True, False])
+def test_frontend_log_stats_flag_is_not_an_unowned_stage_argument(disabled):
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+    config = _from_pipeline_key("dots_tts", cli_overrides={"disable_log_stats": disabled})
+    assert config.stage_configs
+    engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model")
+    assert "disable_log_stats" not in engine_args
+
+
 def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime_bridge():
     omni_config = _from_pipeline_key(
         "qwen3_tts",

tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The omni scheduler must stay on the v1 model runner, like every omni worker.
+
+vLLM 0.28 resolved this config to v1; 0.29 ends its fallback chain with v2. When
+only the workers forced v1, the scheduler took its v2 fast path -- which omits
+resumed-request token ids -- while the v1 runner still read them, so resuming a
+request raised ``KeyError: <req_id>`` in ``_update_states`` and killed the engine
+core (Multi-Replica Startup Test).
+
+The default is upstream's to change, so this pins Omni's side of the contract
+rather than the default itself.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Scheduler(OmniSchedulerMixin):
+    """Minimal carrier for the shared init hook both omni schedulers call."""
+
+    def __init__(self, *, use_v2_model_runner: bool) -> None:
+        self.use_v2_model_runner = use_v2_model_runner
+        self.vllm_config = SimpleNamespace(
+            model_config=SimpleNamespace(
+                async_chunk=False,
+                stage_id=0,
+                pooling_output_decoder=None,
+            )
+        )
+        self._init_omni_io_scheduling_state()
+
+
+@pytest.mark.parametrize("initial", [True, False])
+def test_scheduler_ends_on_v1_model_runner(initial: bool) -> None:
+    """Whatever upstream resolves, the omni scheduler runs v1."""
+    assert _Scheduler(use_v2_model_runner=initial).use_v2_model_runner is False
+
+
+def test_v2_default_does_not_reach_the_resumed_request_fast_path() -> None:
+    """The v2 fast path is what dropped resumed-request token ids.
+
+    ``OmniGenerationScheduler`` branches on ``use_v2_model_runner`` when it
+    assembles SchedulerOutput; forcing it false keeps the v1 assembly that
+    carries ``all_token_ids`` for resumed requests, which the v1 runner reads.
+    """
+    scheduler = _Scheduler(use_v2_model_runner=True)
+
+    assert not scheduler.use_v2_model_runner, (
+        "scheduler would take the v2 fast path while omni workers run v1; "
+        "resumed requests then KeyError in _update_states"
+    )

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -279,17 +279,17 @@ def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]
         group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None))
         return [[group], [group]], "cg-support", [4, 8]
 
+    # vLLM 0.29 dropped attn_groups and cache_dtype from init_kv_cache: the
+    # layout now rides on the resolved CacheConfig.kv_cache_layout instead.
     def init_cache(
         runner_kv_caches,
         forward_context,
         kv_cache_config,
-        attn_groups,
         device,
-        cache_dtype,
         kernel_block_sizes,
         vllm_config,
     ):
-        events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes))
+        events.append(("cache", kv_cache_config, kernel_block_sizes))
         runner_kv_caches.extend(["cache-0", "cache-1"])
         return {"layer-0": "tensor-0", "layer-1": "tensor-1"}
 

tests/diffusion/diffusion_kv/test_initialization.py

@@ -299,8 +299,10 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
     )
 
     assert vllm_config.model_config.original_max_model_len == -1
-    # Sixteen physical pages provide ``16 * block_size`` tokens.  The native
-    # backend chooses the block geometry, so keep this assertion backend
+    # Sixteen physical pages provide ``(16 - 1) * block_size`` tokens: vLLM 0.29
+    # subtracts the null block that BlockPool permanently holds back before
+    # sizing (check_enough_kv_cache_memory), so only 15 pages are usable.  The
+    # native backend chooses the block geometry, so keep this assertion backend
     # independent instead of assuming the generic 16-token block size.
-    expected_max_model_len = min(model_limit, spec.block_size * 16)
+    expected_max_model_len = min(model_limit, spec.block_size * (16 - 1))
     assert vllm_config.model_config.max_model_len == expected_max_model_len

tests/diffusion/diffusion_kv/test_layout.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Physical KV cache layout selection and propagation for the Diffusion cache."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.config import CacheConfig
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+from vllm_omni.diffusion.diffusion_kv.layout import (
+    DEFAULT_DIFFUSION_KV_CACHE_LAYOUT,
+    adopt_kv_cache_layout,
+    assert_backend_layout_supported,
+    build_kv_cache_tensor,
+    resolve_diffusion_kv_cache_layout,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _config():
+    return SimpleNamespace(cache_config=CacheConfig())
+
+
+class _DenseBackend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return False
+
+
+class _BlockStrideBackend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return True
+
+
+def test_default_layout_is_not_block_outermost() -> None:
+    """The default reproduces the pre-0.29 indexes_kv_by_block_stride=False."""
+    vllm_config = _config()
+    layout = resolve_diffusion_kv_cache_layout(vllm_config)
+
+    assert layout is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    assert layout.is_block_outermost is False
+    assert vllm_config.cache_config.get_resolved_kv_cache_layout() is layout
+
+
+def test_block_stride_backend_selects_block_outermost_layout() -> None:
+    layout = resolve_diffusion_kv_cache_layout(_config(), indexes_kv_by_block_stride=True)
+
+    assert layout.is_block_outermost is True
+
+
+def test_resolution_is_idempotent() -> None:
+    vllm_config = _config()
+
+    assert resolve_diffusion_kv_cache_layout(vllm_config) is resolve_diffusion_kv_cache_layout(vllm_config)
+
+
+def test_contradicting_requirement_raises_instead_of_misaddressing() -> None:
+    vllm_config = _config()
+    resolve_diffusion_kv_cache_layout(vllm_config)
+
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True)
+
+
+def test_preset_layout_is_honoured() -> None:
+    vllm_config = _config()
+    vllm_config.cache_config.kv_cache_layout = KVCacheLayout.BLNHC.name
+
+    layout = resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True)
+
+    assert layout is KVCacheLayout.BLNHC
+
+
+def test_backend_assertion_is_quiet_until_the_layout_is_resolved() -> None:
+    """Workers collect specs before the control plane hands them the layout."""
+    vllm_config = _config()
+
+    assert_backend_layout_supported(vllm_config, _BlockStrideBackend)  # unresolved: nothing to check
+
+    resolve_diffusion_kv_cache_layout(vllm_config)
+    assert_backend_layout_supported(vllm_config, _DenseBackend)
+    assert_backend_layout_supported(vllm_config, None)
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        assert_backend_layout_supported(vllm_config, _BlockStrideBackend)
+
+
+def test_worker_adopts_the_layout_carried_on_the_kv_cache_config() -> None:
+    """Diffusion workers receive only a KVCacheConfig, so the name rides on it."""
+    worker = _config()
+    shipped = SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name)
+
+    assert adopt_kv_cache_layout(worker, shipped) is KVCacheLayout.BLNHC
+    assert worker.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.BLNHC
+
+
+def test_worker_without_a_carried_layout_falls_back_to_the_same_default() -> None:
+    worker = _config()
+
+    assert adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=None)) is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+
+
+def test_worker_disagreeing_with_the_control_plane_raises() -> None:
+    worker = _config()
+    resolve_diffusion_kv_cache_layout(worker)
+
+    with pytest.raises(ValueError, match="disagrees with the layout"):
+        adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name))
+
+
+def test_block_stride_requirement_propagates_control_plane_to_worker() -> None:
+    """A block-stride backend must reach the worker as a block-outermost layout.
+
+    Covers the whole seam: the control plane resolves with the requirement, the
+    name rides on each rank-local KVCacheConfig, the worker adopts it, and the
+    backend contract check then passes instead of raising.
+    """
+    control_plane = _config()
+    layout = resolve_diffusion_kv_cache_layout(control_plane, indexes_kv_by_block_stride=True)
+    assert layout.is_block_outermost is True
+
+    shipped = SimpleNamespace(kv_cache_layout=layout.name)
+
+    worker = _config()
+    assert adopt_kv_cache_layout(worker, shipped) is layout
+    # The worker's own backend requirement now agrees with what it was handed.
+    assert_backend_layout_supported(worker, _BlockStrideBackend)
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        assert_backend_layout_supported(worker, _DenseBackend)
+
+
+def test_tensor_strides_come_from_the_upstream_derivation() -> None:
+    spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
+
+    tensor = build_kv_cache_tensor(spec, 4, ["layer0", "layer1"])
+
+    assert tensor.layers == ["layer0", "layer1"]
+    assert tensor.size == spec.page_size_bytes * 2 * 4
+    # Layer-outermost: each layer owns a contiguous run of its own blocks.
+    assert tensor.layer_stride == spec.page_size_bytes * 4
+    assert tensor.block_stride == spec.page_size_bytes

tests/diffusion/diffusion_kv/test_manager.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 
@@ -24,7 +25,7 @@ def _config(num_blocks: int) -> KVCacheConfig:
     )
     return KVCacheConfig(
         num_blocks=num_blocks,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
     )
 

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -15,6 +15,7 @@ import torch
 from torch import nn
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 
+from tests.helpers.kv_layout import layout_for_backend
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
@@ -1239,7 +1240,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
-    assert native_layer.spec.indexes_kv_by_block_stride is True
+    # 0.29: the flag no longer rides on the spec, it selects the layout.
+    assert layout_for_backend(native_layer.attn_backend).is_block_outermost is True
     assert native_layer._q_scale_float == 1.0
     assert native_layer._k_scale_float == 1.0
     assert native_layer._v_scale_float == 1.0

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -10,12 +10,14 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
 from vllm.v1.worker.gpu.block_table import BlockTables
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
     DiffusionPagedAttentionLayerAdapter,
@@ -88,12 +90,7 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
     canonical_spec = native_layer.spec
     kv_cache_config = KVCacheConfig(
         num_blocks=4,
-        kv_cache_tensors=[
-            KVCacheTensor(
-                size=canonical_spec.page_size_bytes * 4,
-                shared_by=[_LAYER_NAME],
-            )
-        ],
+        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])],
         kv_cache_groups=[
             KVCacheGroupSpec(
                 layer_names=[_LAYER_NAME],
@@ -101,6 +98,9 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
             )
         ],
     )
+    # 0.29 reads the resolved physical layout when allocating; the diffusion
+    # control plane pins it, so this standalone harness must too.
+    resolve_diffusion_kv_cache_layout(vllm_config)
     with set_current_vllm_config(vllm_config):
         attn_groups, _, kernel_block_sizes = init_attn_backend(
             kv_cache_config,
@@ -108,13 +108,13 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
             device,
         )
         runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+        # vLLM 0.29 dropped attn_groups and cache_dtype here: the layout now
+        # comes from the resolved CacheConfig.kv_cache_layout inside the call.
         init_kv_cache(
             runner_kv_caches,
             vllm_config.compilation_config.static_forward_context,
             kv_cache_config,
-            attn_groups,
             device,
-            vllm_config.cache_config.cache_dtype,
             kernel_block_sizes,
             vllm_config,
         )

tests/diffusion/diffusion_kv/test_request.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import pytest
 import torch
@@ -10,10 +10,10 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
-    KVCacheTensor,
 )
 from vllm.v1.request import RequestStatus
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -119,7 +119,7 @@ def _manager(*, enable_caching: bool) -> KVCacheManager:
     group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["layer0"])],
         kv_cache_groups=[group],
     )
     return KVCacheManager(

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -8,9 +8,10 @@ from unittest.mock import Mock
 import pytest
 import torch
 from torch import nn
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
 import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
+from tests.helpers.kv_layout import build_kv_cache_tensor, layout_for_backend
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
@@ -79,7 +80,9 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None:
     assert spec.num_kv_heads == 2
     assert spec.head_size == 8
     assert spec.dtype is torch.bfloat16
-    assert spec.indexes_kv_by_block_stride is True
+    # vLLM 0.29 moved the block-stride decision off the spec onto the resolved
+    # physical layout; the backend's declaration must still select one.
+    assert layout_for_backend(_Backend).is_block_outermost is True
     assert spec.non_causal is True
 
 
@@ -112,7 +115,7 @@ def test_runner_retains_matching_rank_local_config() -> None:
     spec = runner.get_kv_cache_spec()["image_attention"]
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["image_attention"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
     )
 
@@ -126,7 +129,7 @@ def test_runner_rejects_rank_local_config_for_different_layers() -> None:
     spec = runner.get_kv_cache_spec()["image_attention"]
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["other_attention"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
     )
 

tests/diffusion/layers/test_adalayernorm.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for shared AdaLayerNorm layers used by FLUX and other models."""
 
 import os
@@ -46,7 +46,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/layers/test_rope_broadcast.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 Shape broadcasting regression tests for RotaryEmbedding.forward_native.
 
@@ -16,7 +16,11 @@ import types
 import pytest
 import torch
 
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd
+from vllm_omni.diffusion.layers.rope import (
+    RotaryEmbedding,
+    RotaryEmbeddingWan,
+    apply_rotary_emb_mindiesd,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -239,6 +243,25 @@ class TestRotaryEmbeddingNativeShapeRegression:
         assert not torch.isnan(output).any()
 
 
+def test_cuda_paths_fall_back_without_vendored_rotary_kernel() -> None:
+    """Source/CPU installs need not contain the wheel-vendored CUDA kernel."""
+    x = torch.randn(1, 5, 2, 8)
+
+    rope = RotaryEmbedding(is_neox_style=True)
+    rope.apply_rotary_emb_vllm_flash_attn = None
+    output = rope.forward_cuda(x, torch.randn(5, 4), torch.randn(5, 4))
+    assert output.shape == x.shape
+
+    wan_rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
+    wan_rope.apply_rotary_emb_vllm_flash_attn = None
+    output = wan_rope.forward_cuda(
+        x,
+        torch.randn(1, 5, 1, 4),
+        torch.randn(1, 5, 1, 4),
+    )
+    assert output.shape == x.shape
+
+
 def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None:
     """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout."""
     captured: dict[str, object] = {}

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -59,7 +59,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/dreamzero/test_qkv_fusion.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Equivalence test for DreamZero self-attn QKV fusion.
 
@@ -52,7 +52,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/glm_image/test_glm_image_quantization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for GLM-Image quantization support (W4A16/AutoRound).
 
 These tests verify that the GLM-Image DiT transformer correctly accepts and uses
@@ -69,7 +69,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 from types import SimpleNamespace
 
 import pytest
@@ -64,6 +64,13 @@ def setup_tp_group(monkeypatch, mocker):
         lambda: 0,
     )
 
+    # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever
+    # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path.
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda",
+        lambda: False,
+    )
+
     mock_tp_group = mocker.MagicMock()
     mock_tp_group.world_size = 2
     mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group)
@@ -94,7 +101,7 @@ def setup_tp_group(monkeypatch, mocker):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
     with set_current_vllm_config(VllmConfig(device_config=device_config)):

tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py

@@ -1,12 +1,39 @@
 import pytest
 import torch
 
-from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
+    QwenEmbedRope,
+    _apply_qwen_image_rotary_emb,
+)
 from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_qwen_rotary_preserves_frequency_precision(dtype, noncontiguous):
+    from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
+
+    generator = torch.Generator().manual_seed(42)
+    x = torch.randn(2, 17, 3, 128, generator=generator).to(dtype)
+    if noncontiguous:
+        x = x.transpose(1, 2).contiguous().transpose(1, 2)
+    angles = torch.randn(17, 64, generator=generator)
+    freqs = torch.polar(torch.ones_like(angles), angles)
+    expected = apply_rotary_emb_qwen(x, freqs, use_real=False)
+    actual = _apply_qwen_image_rotary_emb(x, freqs)
+    assert actual.dtype == dtype
+    assert actual.shape == x.shape
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    if dtype == torch.bfloat16:
+        # A regression to early BF16 frequency rounding must be observable.
+        rounded_freqs = torch.complex(freqs.real.to(dtype).float(), freqs.imag.to(dtype).float())
+        rounded = _apply_qwen_image_rotary_emb(x, rounded_freqs)
+        assert not torch.equal(actual, rounded)
+
+
 def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count():
     prompt_embeds = torch.zeros(2, 64, 8)
     prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool)

tests/diffusion/models/sana_video/test_transformer_sana_video.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import os
 
@@ -22,7 +22,7 @@ def _init_distributed(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
     os.environ.setdefault("MASTER_ADDR", "localhost")
     os.environ.setdefault("MASTER_PORT", "29501")

tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py

@@ -41,8 +41,8 @@ def tp1_group():
 
     os.environ.setdefault("MASTER_ADDR", "localhost")
     os.environ.setdefault("MASTER_PORT", "29501")
-    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
-    initialize_model_parallel()
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://", backend="gloo")
+    initialize_model_parallel(backend="gloo")
     yield
     cleanup_dist_env_and_memory()
 
@@ -54,7 +54,7 @@ def force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 
@@ -482,8 +482,12 @@ class _LockstepP2PDist:
 def test_tiny_transformer_sp2_lockstep_matches_dense(tp1_group, force_default_gemm, mocker, task) -> None:
     """Two lockstep rank threads over uneven frame shards must reproduce the
     dense tiny-transformer output, and must actually have sharded the work."""
+    from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
     from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
 
+    # These CPU tensor tests must not probe a CUDA device during construction.
+    mocker.patch("vllm_omni.diffusion.attention.selector._cached_get_backend_cls", return_value=SDPABackend)
+
     torch.manual_seed(3)
     model = SanaVideoTransformer3DModel(**_TINY_SP_CONFIG).eval()
     for _, param in sorted(model.named_parameters()):

tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """TP unit tests for SANA-Video: distributed RMSNorm, sharded attention and
 sharded GLUMB temporal conv.
@@ -51,7 +51,7 @@ def force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import pytest
 import torch
 from transformers import T5Config
@@ -54,6 +54,13 @@ def setup_tp_group(monkeypatch, mocker):
         lambda: 0,
     )
 
+    # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever
+    # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path.
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda",
+        lambda: False,
+    )
+
     # TP group
     mock_tp_group = mocker.MagicMock()
     mock_tp_group.world_size = 2

tests/diffusion/quantization/test_component_routing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for component routing for quantization."""
 
 from unittest.mock import MagicMock
@@ -372,7 +372,8 @@ class TestComponentApplyVllmMapper:
                 "thinker.model": "thinker.language_model.model",
                 "talker.model": "talker.language_model.model",
             }
-        ).get_unstacked_mapper()
+            # vLLM 0.29 renamed get_unstacked_mapper to get_rename_mapper.
+        ).get_rename_mapper()
 
         config.apply_vllm_mapper(mapper)
 

tests/diffusion/test_diffusion_scheduler.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import asyncio
 import queue
@@ -10,8 +10,9 @@ import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from pytest_mock import MockerFixture
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -104,7 +105,7 @@ def _initialize_paged_scheduler(
     )
     config = KVCacheConfig(
         num_blocks=num_blocks,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
     )
     scheduler.initialize(

tests/e2e/accuracy/test_diffusers_backend_similarity.py

@@ -164,6 +164,7 @@ def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_imag
         )
         pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FLOW_SHIFT)
         pipe.to("cuda")
+        _set_matched_attention_backend(pipe)
 
         _diffusers_dummy_run(pipe)
 

tests/e2e/online_serving/test_dots_tts_expansion.py

@@ -85,5 +85,7 @@ def test_text_only_nonstreaming_pcm_001(omni_server, openai_client) -> None:
         "voice": "default",
         "min_audio_bytes": _MIN_AUDIO_BYTES,
         "min_hnr_db": -2.0,
+        # dots.tts-soar emits native 48 kHz PCM, which has no rate header.
+        "expected_sample_rate": 48_000,
     }
     openai_client.send_audio_speech_request(request_config)

tests/engine/test_async_omni_engine_stage_init.py

@@ -1593,73 +1593,58 @@ def test_model_path_resolver_is_generic_and_model_owned(tmp_path):
     assert "model_path_resolver" not in engine_args
 
 
-def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
+def test_build_stage0_input_processor_uses_omni_renderer_subclass(monkeypatch):
+    from vllm.renderers import BaseRenderer
+
     import vllm_omni.engine.stage_init_utils as init_mod
+    from vllm_omni.inputs.preprocess import OmniRenderer, omni_renderer_cls
 
-    class DummyInputProcessor:
-        def __init__(self, vllm_config, renderer=None):
-            self.vllm_config = vllm_config
-            self.renderer = renderer or object()
-            self.input_preprocessor = None
+    class _Base(BaseRenderer):
+        def __init__(self, config, tokenizer):
+            self.config, self.tokenizer = config, tokenizer
 
-    class DummyOmniInputPreprocessor:
-        def __init__(self, vllm_config, renderer=None):
-            self.vllm_config = vllm_config
-            self.renderer = renderer
+        def render_messages(self, messages, params):  # pragma: no cover - abstract stub
+            raise NotImplementedError
 
-    monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
-    monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)
-
-    input_processor = build_stage0_input_processor(
-        types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {}))
-    )
-
-    assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor)
-    assert input_processor.input_preprocessor.renderer is input_processor.renderer
-
-
-def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(
-    monkeypatch,
-):
-    import vllm_omni.engine.stage_init_utils as init_mod
-
-    token_only_renderer = object()
+    config = types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {}))
+    built = omni_renderer_cls(_Base)(config, "tok")
     seen = {}
 
     class DummyInputProcessor:
         def __init__(self, vllm_config, renderer=None):
-            if renderer is None:
-                raise AssertionError("skip_tokenizer_init must supply a renderer")
             seen["renderer"] = renderer
             self.renderer = renderer
-            self.input_preprocessor = None
-
-    class DummyOmniInputPreprocessor:
-        def __init__(self, vllm_config, renderer=None):
-            seen["preprocessor_renderer"] = renderer
 
     monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
-    monkeypatch.setattr(
-        init_mod,
-        "_build_token_only_renderer",
-        lambda _config: token_only_renderer,
-    )
-    monkeypatch.setattr(
-        init_mod,
-        "OmniInputPreprocessor",
-        DummyOmniInputPreprocessor,
-    )
+    monkeypatch.setattr(init_mod, "build_omni_renderer", lambda cfg: built if cfg is config else None)
+    processor = build_stage0_input_processor(config)
+    assert seen["renderer"] is built
+    assert isinstance(processor.renderer, OmniRenderer)
+    assert isinstance(processor.renderer, _Base)
+    assert not hasattr(processor, "input_preprocessor")
 
+
+def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(monkeypatch):
+    import vllm_omni.engine.stage_init_utils as init_mod
+
+    original = object()
+
+    class DummyInputProcessor:
+        def __init__(self, vllm_config, renderer=None):
+            assert renderer is original
+            self.renderer = renderer
+
+    def _must_not_resolve(_cfg):
+        raise AssertionError("tokenizer must not be resolved when skip_tokenizer_init=True")
+
+    monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
+    monkeypatch.setattr(init_mod, "_build_token_only_renderer", lambda _: original)
+    monkeypatch.setattr(init_mod, "build_omni_renderer", _must_not_resolve)
     config = types.SimpleNamespace(
-        model_config=types.SimpleNamespace(
-            skip_tokenizer_init=True,
-            try_get_generation_config=lambda: {},
-        )
+        model_config=types.SimpleNamespace(skip_tokenizer_init=True, try_get_generation_config=lambda: {})
     )
-    build_stage0_input_processor(config)
-
-    assert seen["renderer"] is token_only_renderer
-    assert seen["preprocessor_renderer"] is token_only_renderer
+    processor = build_stage0_input_processor(config)
+    assert processor.renderer is original
 
 
 def test_inject_kv_stage_info_infers_sender_tp_topology():

tests/entrypoints/openai/test_batch_chat_completions.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 from unittest.mock import AsyncMock
 
@@ -12,7 +15,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionResponseChoice,
     ChatMessage,
 )
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo
 
 from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
 

tests/entrypoints/openai_api/test_audio_format.py

@@ -248,7 +248,7 @@ class TestResolveAudioFormat:
         assert result == "pcm"
 
     def test_invalid_format_returns_error(self, serving_chat):
-        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+        from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
         request = self._make_request({"format": "aac", "voice": "alloy"})
         result = serving_chat._resolve_audio_format(request)
@@ -256,7 +256,7 @@ class TestResolveAudioFormat:
         assert "aac" in result.error.message
 
     def test_all_supported_formats_accepted(self, serving_chat):
-        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+        from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
         for fmt in SUPPORTED_CHAT_AUDIO_FORMATS:
             request = self._make_request({"format": fmt, "voice": "alloy"})

tests/entrypoints/openai_api/test_modality_validation.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for output modality validation in chat completions.
 
 Covers:
@@ -13,7 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 

tests/entrypoints/openai_api/test_serving_chat_metrics.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field."""
 
 from types import SimpleNamespace
@@ -19,7 +20,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 def test_omni_chat_completion_response_metrics():
     """Test OmniChatCompletionResponse metrics field works correctly."""
-    from vllm.entrypoints.openai.engine.protocol import UsageInfo
+    from vllm.entrypoints.serve.engine.protocol import UsageInfo
 
     from vllm_omni.entrypoints.openai.protocol.chat_completion import (
         OmniChatCompletionResponse,

tests/entrypoints/openai_api/test_serving_speech.py

@@ -25,7 +25,7 @@ from fastapi.responses import JSONResponse
 from fastapi.testclient import TestClient
 from pydantic import ValidationError
 from pytest_mock import MockerFixture
-from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse
 
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler

tests/entrypoints/openai_api/test_video_stream_handler.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Runtime behavior tests for VideoStreamHandler.
 
 Every test creates a mock WebSocket, drives handle_session through a
@@ -281,7 +281,7 @@ class TestHandleQueryErrors:
     @pytest.mark.asyncio
     async def test_query_returns_error_response(self):
         """create_chat_completion returns ErrorResponse → error forwarded."""
-        from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
+        from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse
 
         err = ErrorResponse(
             error=ErrorInfo(

tests/entrypoints/test_stream_finish_reason.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for multi-modal streaming finish_reason behavior (commit 44c799bc).
 
 Verifies that the /v1/chat/completions streaming endpoint emits exactly one
@@ -411,7 +412,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
 @pytest.mark.asyncio
 async def test_audio_choice_error_response_is_not_iterated_as_choices():
     """Any ErrorResponse from the audio path is skipped, not iterated."""
-    from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+    from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
     serving_chat = build_serving_chat()
     request = make_request(modalities=["text", "audio"])

tests/helpers/kv_layout.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shared access to the production ``KVCacheTensor`` layout derivation.
+
+Tests deliberately go through the same code production uses
+(``vllm_omni.diffusion.diffusion_kv.layout``) rather than recomputing strides,
+so a regression in that derivation fails the suite instead of being masked by a
+parallel test-only copy.
+"""
+
+from types import SimpleNamespace
+
+from vllm.config import CacheConfig
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+from vllm_omni.diffusion.diffusion_kv.layout import (
+    build_kv_cache_tensor,
+    resolve_diffusion_kv_cache_layout,
+)
+
+__all__ = ["build_kv_cache_tensor", "layout_for_backend"]
+
+
+def layout_for_backend(attn_backend) -> KVCacheLayout:
+    """The physical layout ``attn_backend``'s block-stride declaration selects.
+
+    Before 0.29 a backend's ``indexes_kv_by_block_stride()`` was copied onto
+    ``AttentionSpec``; asserting on the resolved layout is the surviving way to
+    check that declaration still reaches the cache.
+    """
+
+    vllm_config = SimpleNamespace(cache_config=CacheConfig())
+    return resolve_diffusion_kv_cache_layout(
+        vllm_config,
+        indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(),
+    )

tests/helpers/serving_chat.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Shared scaffolding for OmniOpenAIServingChat unit tests.
 
 Provides factories for a minimal serving-chat instance, chat requests, and
@@ -8,11 +9,11 @@ OmniRequestOutput objects, plus SSE stream collection/parsing utilities.
 import json
 from unittest.mock import MagicMock
 
+from vllm.entrypoints.generate.base.protocol import DeltaMessage
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponseStreamChoice,
 )
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.outputs import CompletionOutput, RequestOutput
 

tests/inputs/test_mm_processor.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+from transformers import BatchFeature
+
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("counts", [{}, {"audio": 1, "image": 0}])
+@pytest.mark.parametrize("original_text", [None, "", "voice cloning text"])
+def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts, original_text):
+    class Processor(OmniMultiModalProcessor):
+        def _get_mm_fields_config(self, *args):
+            return {}
+
+        def _get_prompt_updates(self, *args):
+            return []
+
+    processor = object.__new__(Processor)
+    processor.dummy_inputs = SimpleNamespace(get_dummy_text=Mock(return_value="dummy"))
+    processor._get_hf_mm_data = Mock(return_value=({"audios": ["wave"]}, {"embedding": "preserve"}))
+    processor._call_hf_processor = Mock(return_value=BatchFeature({"features": "processed"}))
+    items = Mock()
+    items.get_all_counts.return_value = counts
+    kwargs = {"sampling_rate": 16000}
+    if original_text is not None:
+        kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text
+    before = dict(kwargs)
+    result = processor._apply_hf_processor_main(items, kwargs)
+    items.select.assert_called_once_with({k for k, v in counts.items() if v > 0})
+    processor._call_hf_processor.assert_called_once_with(
+        "dummy" if original_text is None else original_text,
+        {"audios": ["wave"]},
+        {"sampling_rate": 16000},
+        {},
+    )
+    assert dict(result) == {"features": "processed", "embedding": "preserve"}
+    assert kwargs == before

tests/inputs/test_preprocess.py

@@ -1,35 +1,157 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Unit tests for OmniInputPreprocessor._process_text routing."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise Omni routing through the active upstream rendering pipeline."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
 
 import pytest
-from pytest_mock import MockerFixture
+import torch
+from vllm.inputs import tokens_input
+from vllm.pooling_params import PoolingParams
+from vllm.renderers import BaseRenderer
+from vllm.renderers.params import TokenizeParams
+from vllm.v1.engine.input_processor import InputProcessor
+
+from vllm_omni.inputs import preprocess as preprocess_mod
+from vllm_omni.inputs.preprocess import OmniRenderer, build_omni_renderer, omni_renderer_cls
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-class TestProcessTextMmProcessorKwargsRouting:
-    """Presence-based routing: an explicit empty ``mm_processor_kwargs`` dict
-    on the prompt still routes through ``_process_multimodal``. Required for
-    AR-based image generation (e.g. GLM-Image t2i) where the HF processor
-    supplies the image-generation scaffold from its own defaults and no extra
-    kwargs are needed from the caller.
-    """
+class _Renderer(BaseRenderer):
+    def __init__(self):
+        self.model_config = SimpleNamespace(is_encoder_decoder=False, enable_prompt_embeds=True)
+        self.tokenizer = None
+        self.default_cmpl_tok_params = TokenizeParams(max_total_tokens=128)
+        self._tokenize_prompt = lambda prompt, params: {**prompt, "prompt_token_ids": [1, 2, 3]}
+        self._tokenize_prompt_async = AsyncMock(side_effect=self._tokenize_prompt)
+        self._process_multimodal = Mock(return_value=tokens_input([1, 2, 3, 99]))
+        self._process_multimodal_async = AsyncMock(side_effect=self._process_multimodal)
 
-    @pytest.fixture
-    def preprocessor(self, mocker: MockerFixture):
-        from vllm_omni.inputs.preprocess import OmniInputPreprocessor
+    def render_messages(self, messages, params):
+        return messages, {"prompt": "hello"}
 
-        instance = object.__new__(OmniInputPreprocessor)
-        instance._process_multimodal = mocker.MagicMock(return_value={})
-        instance._tokenize_prompt = mocker.MagicMock(return_value=[1, 2, 3])
-        return instance
 
-    def test_empty_mm_processor_kwargs_routes_to_multimodal(self, preprocessor):
-        preprocessor._process_text({"prompt": "hello", "mm_processor_kwargs": {}})
-        assert preprocessor._process_multimodal.called
-        assert not preprocessor._tokenize_prompt.called
+@pytest.fixture
+def renderer():
+    return omni_renderer_cls(_Renderer)()
 
-    def test_missing_mm_processor_kwargs_routes_to_tokenize(self, preprocessor):
-        preprocessor._process_text({"prompt": "hello"})
-        assert preprocessor._tokenize_prompt.called
-        assert not preprocessor._process_multimodal.called
+
+@pytest.mark.parametrize("kwargs", [{}, {"target_h": 512, "target_w": 768}])
+@pytest.mark.parametrize("tokenized", [False, True])
+def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, tokenized):
+    # Keep real process_inputs/render_cmpl; stub unrelated config validation.
+    processor = object.__new__(InputProcessor)
+    processor.renderer = renderer
+    processor.model_config = renderer.model_config
+    processor.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            data_parallel_size=1,
+            data_parallel_size_local=1,
+            local_engines_only=False,
+        )
+    )
+    processor._validate_params = Mock()
+    processor._validate_lora = Mock()
+    processor._validate_model_inputs = Mock()
+    prompt: dict[str, object] = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"}
+    prompt.update(mm_processor_kwargs=kwargs, cache_salt="salt")
+    request = processor.process_inputs("image-request", prompt, PoolingParams(), ("embed",))
+    assert request.prompt_token_ids == [1, 2, 3, 99]
+    assert request.cache_salt == "salt"
+    renderer._process_multimodal.assert_called_once_with(
+        [1, 2, 3],
+        {},
+        mm_processor_kwargs=kwargs,
+        mm_uuids=None,
+        skip_mm_cache=False,
+    )
+
+
+@pytest.mark.parametrize("async_mode", [False, True])
+@pytest.mark.parametrize("kind", ["text", "tokens", "media", "kwargs", "embeds"])
+def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode, kind):
+    prompt: dict[str, object] = {
+        "prompt": "hello",
+        "additional_information": {"speaker": 1},
+        "model_intermediate_buffer": {"x": 2},
+    }
+    if kind == "tokens":
+        prompt["prompt_token_ids"] = [1, 2, 3]
+    elif kind == "media":
+        prompt["multi_modal_data"] = {"image": "image"}
+    elif kind == "kwargs":
+        prompt["mm_processor_kwargs"] = {}
+    elif kind == "embeds":
+        prompt["prompt_embeds"] = torch.ones(3, 4)
+    if async_mode:
+        (result,) = asyncio.run(renderer.render_cmpl_async([prompt], skip_mm_cache=True))
+    else:
+        (result,) = renderer.render_cmpl([prompt], skip_mm_cache=True)
+    assert result["additional_information"] == {"speaker": 1}
+    assert result["model_intermediate_buffer"] == {"x": 2}
+    assert result["prompt"] == "hello"
+    if kind in ("media", "kwargs"):
+        assert result["prompt_token_ids"] == [1, 2, 3, 99]
+        assert renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True
+    else:
+        renderer._process_multimodal.assert_not_called()
+        if kind == "embeds":
+            torch.testing.assert_close(result["prompt_embeds"], prompt["prompt_embeds"])
+        else:
+            assert result["prompt_token_ids"] == [1, 2, 3]
+
+
+def test_omni_renderer_is_a_real_subclass_created_once():
+    cls = omni_renderer_cls(_Renderer)
+    assert issubclass(cls, OmniRenderer) and issubclass(cls, _Renderer)
+    assert omni_renderer_cls(_Renderer) is cls
+    assert omni_renderer_cls(cls) is cls
+    assert cls.__name__ == "Omni_Renderer"
+    instance = cls()
+    assert isinstance(instance, _Renderer)
+    assert not hasattr(instance, "_renderer")
+
+
+def test_concrete_renderer_overrides_are_not_suppressed():
+    """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides."""
+
+    class _Custom(_Renderer):
+        def render_cmpl(self, prompts, *args, **kwargs):
+            return ["custom"]
+
+    assert omni_renderer_cls(_Custom)().render_cmpl([{"prompt": "x"}]) == ["custom"]
+
+
+class _CtorRenderer(BaseRenderer):
+    def __init__(self, config, tokenizer):
+        self.config = config
+        self.tokenizer = tokenizer
+
+    def render_messages(self, messages, params):
+        raise NotImplementedError
+
+
+def test_build_omni_renderer_with_explicit_class_skips_registry(monkeypatch):
+    monkeypatch.setattr(preprocess_mod, "cached_tokenizer_from_config", Mock(side_effect=AssertionError("resolved")))
+    config = SimpleNamespace(model_config=SimpleNamespace())
+    built = build_omni_renderer(config, renderer_cls=_CtorRenderer, tokenizer=None)
+    assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer)
+    assert built.config is config and built.tokenizer is None
+
+
+def test_build_omni_renderer_resolves_like_upstream(monkeypatch):
+    config = SimpleNamespace(model_config=SimpleNamespace())
+    tokenizer = object()
+    monkeypatch.setattr(
+        preprocess_mod, "cached_tokenizer_from_config", lambda mc: tokenizer if mc is config.model_config else None
+    )
+    monkeypatch.setattr(preprocess_mod, "tokenizer_args_from_config", lambda mc: ("fake-mode", None))
+    load_cls = Mock(return_value=_CtorRenderer)
+    monkeypatch.setattr(preprocess_mod.RENDERER_REGISTRY, "load_renderer_cls", load_cls)
+    built = build_omni_renderer(config)
+    load_cls.assert_called_once_with("fake-mode")
+    assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer)
+    assert built.tokenizer is tokenizer

tests/model_executor/models/glm_image/test_glm_image_ar.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for GLM-Image AR model: DataParser, processor, and M-RoPE."""
 
 import importlib.util
@@ -91,6 +91,41 @@ GlmImageRotaryEmbedding = _ar_mod.GlmImageRotaryEmbedding
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.parametrize("num_images", [0, 1, 2])
+@pytest.mark.parametrize("preformatted", [False, True])
+def test_generation_prompt_ids_preserve_hf_target_scaffold(num_images, preformatted):
+    from types import SimpleNamespace
+
+    from vllm.multimodal.parse import MultiModalDataItems
+    from vllm.multimodal.processing import BaseMultiModalProcessor, ProcessorInputs
+
+    processor = object.__new__(GlmImageMultiModalProcessor)
+    # Only the public token attributes of GlmImageProcessor are needed; the
+    # expected suffix below is what HF's own formatter produces for 512x768.
+    hf_processor = SimpleNamespace(grid_bos_token="<sop>", grid_eos_token="<eop>", bos_token="<image_start>")
+    tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode()))
+    processor.info = SimpleNamespace(
+        get_tokenizer=lambda: tokenizer,
+        get_hf_processor=lambda: hf_processor,
+        get_hf_config=lambda: _make_hf_config(),
+    )
+    items = MagicMock(spec=MultiModalDataItems)
+    items.get_all_counts.return_value = {"image": num_images}
+    suffix = "<sop>16 24<eop>" + ("<sop>13 19<eop>" if not num_images else "") + "<image_start>"
+    text_ids = list(b"a red bird")
+    expected = [167855] * num_images + text_ids + list(suffix.encode())
+    inputs = ProcessorInputs(
+        prompt=expected if preformatted else text_ids,
+        mm_data_items=items,
+        hf_processor_mm_kwargs={"target_h": 512, "target_w": 768},
+    )
+    with patch.object(BaseMultiModalProcessor, "apply", side_effect=lambda inputs, timing: inputs) as upstream:
+        result = processor.apply(inputs, MagicMock())
+    assert result.prompt == expected
+    assert inputs.prompt == (expected if preformatted else text_ids)
+    upstream.assert_called_once()
+
+
 def test_glm_image_ar_omits_hidden_pooler_payload():
     assert GlmImageForConditionalGeneration.omni_pooler_payload_include_hidden is False
 

tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for the MiniCPM-o 4.5 ``audio_embeds`` input path.
 
 Two defects made the OpenAI ``audio_embeds`` content part unusable with this model:
@@ -71,6 +71,12 @@ class _IdentityTokenizer:
     def decode(self, tokens: str) -> str:
         return tokens
 
+    def convert_tokens_to_ids(self, token: str) -> int:
+        # vLLM 0.29 prompt updates select placeholder positions by token id,
+        # so the processor resolves ``<unk>`` once; the value is irrelevant here.
+        assert token == "<unk>"
+        return 0
+
 
 def _processing_info() -> MiniCPMO45OmniLLMProcessingInfo:
     processor = _FakeProcessor()

tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for MiniCPM-o 4.5 >30s audio preprocessing.
 
 Covers:
@@ -43,6 +43,34 @@ _FEAT_DIM = 128
 _CHUNK_FRAMES = 3000
 
 
+@pytest.mark.parametrize("version", [(2, 5), (4, 5)])
+def test_prompt_batching_delegates_upstream_without_tts_kwargs(version):
+    from unittest.mock import MagicMock
+
+    info = SimpleNamespace(
+        get_model_version=lambda: version,
+        get_hf_processor=MagicMock(return_value=object()),
+        ctx=SimpleNamespace(call_hf_processor=MagicMock()),
+    )
+    info.ctx.call_hf_processor.side_effect = (
+        [{"features": [1]}, {"features": [2]}] if version == (2, 5) else [{"features": [1, 2]}]
+    )
+    kwargs = {"use_tts": True, "sampling_rate": 16000}
+    result = MiniCPMO45OmniLLMMultiModalProcessor._call_hf_processor_on_prompts(
+        SimpleNamespace(info=info), ["one", "two"], {"audios": [10, 20]}, kwargs, out_keys={"features"}
+    )
+    assert result == {"features": [1, 2]}
+    assert kwargs == {"use_tts": True, "sampling_rate": 16000}
+    for call in info.ctx.call_hf_processor.call_args_list:
+        assert call.args[2] == {"sampling_rate": 16000}
+    payloads = [call.args[1] for call in info.ctx.call_hf_processor.call_args_list]
+    assert payloads == (
+        [{"text": "one", "audios": 10}, {"text": "two", "audios": 20}]
+        if version == (2, 5)
+        else [{"text": ["one", "two"], "audios": [10, 20]}]
+    )
+
+
 class TestMiniCPMOFieldConfig:
     def test_short_audios_stay_batched(self) -> None:
         # two <=30s audios: one feature entry + one 1-element lens per audio
@@ -99,10 +127,10 @@ def _make_processor(
     processor.info = SimpleNamespace(audio_pattern="(<audio>./</audio>)")
     processor.data_parser = MiniCPMOMultiModalDataParser(target_sr=_SAMPLE_RATE)
 
-    def fake_base_call_hf_processor(prompts, mm_data, mm_kwargs, tok_kwargs, *, out_keys):
+    def fake_call_hf_processor_on_prompts(prompts, mm_data, mm_kwargs, *, out_keys):
         return {key: fake_hf_outputs[key] for key in out_keys}
 
-    processor._base_call_hf_processor = fake_base_call_hf_processor
+    processor._call_hf_processor_on_prompts = fake_call_hf_processor_on_prompts
     return processor
 
 
@@ -120,7 +148,7 @@ class TestProcessAudiosUnpadding:
         }
         processor = _make_processor(fake_outputs)
 
-        result = processor.process_audios({"audios": [_audio_seconds(45)]}, {}, {})
+        result = processor.process_audios({"audios": [_audio_seconds(45)]}, {})
 
         features = result["audio_features"]
         assert [feat.shape[-1] for feat in features] == chunk_lens
@@ -138,7 +166,7 @@ class TestProcessAudiosUnpadding:
         }
         processor = _make_processor(fake_outputs)
 
-        result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {}, {})
+        result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {})
 
         widths = [feat.shape[-1] for feat in result["audio_features"]]
         assert widths == [_CHUNK_FRAMES, 1200, _CHUNK_FRAMES, 800]
@@ -151,7 +179,7 @@ class TestProcessAudiosUnpadding:
         }
         processor = _make_processor(fake_outputs)
 
-        result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {}, {})
+        result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {})
 
         features = result["audio_features"]
         assert len(features) == 1

tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""OmniVoice processor: the prompt is tokenized exactly once per request.
+
+Upstream ``_cached_apply_hf_processor`` calls ``_apply_hf_processor_main`` for
+every request, including text-only requests and requests whose reference audio
+is already cached (empty ``mm_items``). These tests pin that only ``apply``
+touches the text tokenizer and that the media path only encodes the audio.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.processing import ProcessorInputs
+
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
+from vllm_omni.model_executor.models.omnivoice.omnivoice import OmniVoiceMultiModalProcessor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_PROMPT_IDS = [5, 6, 7]
+
+
+def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProcessor:
+    processor = object.__new__(OmniVoiceMultiModalProcessor)
+    tokenizer = SimpleNamespace(decode=Mock(return_value="hello world"))
+    processor.info = SimpleNamespace(
+        ctx=SimpleNamespace(
+            get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(sample_rate=24000),
+            model_config=SimpleNamespace(model="fake-model"),
+        ),
+        get_tokenizer=lambda: tokenizer,
+    )
+    processor.text_tokenizer = Mock(return_value=SimpleNamespace(input_ids=torch.tensor([_PROMPT_IDS])))
+    processor.audio_tokenizer = SimpleNamespace(encode=Mock(return_value=torch.zeros(8, 4, dtype=torch.long)))
+    processor.feature_extractor = None
+    monkeypatch.setattr(processor, "_ensure_cached_runtime_components", lambda *_args, **_kwargs: None)
+    return processor
+
+
+def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace:
+    return SimpleNamespace(
+        get_all_counts=lambda: counts,
+        select=lambda _keys: mm_data,
+    )
+
+
+def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    processor = _make_processor(monkeypatch)
+    captured: dict[str, ProcessorInputs] = {}
+
+    def fake_parent_apply(self, inputs, timing_ctx):
+        captured["inputs"] = inputs
+        return "sentinel"
+
+    monkeypatch.setattr(OmniMultiModalProcessor, "apply", fake_parent_apply)
+    inputs = ProcessorInputs(
+        prompt=[1, 2],
+        mm_data_items=MultiModalDataItems({}),
+        hf_processor_mm_kwargs={"ref_text": "reference"},
+    )
+
+    assert processor.apply(inputs, timing_ctx=None) == "sentinel"
+    assert captured["inputs"].prompt == _PROMPT_IDS
+    assert processor.text_tokenizer.call_count == 1
+    (text_prompt,), _ = processor.text_tokenizer.call_args
+    assert text_prompt.endswith("<|text_start|>reference hello world<|text_end|>")
+    assert captured["inputs"].hf_processor_mm_kwargs["ref_text"] == "reference"
+
+
+def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.MonkeyPatch) -> None:
+    processor = _make_processor(monkeypatch)
+    audio = np.zeros(2400, dtype=np.float32)
+    monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({"audios": [audio]}, {}))
+
+    out = processor._apply_hf_processor_main(_items({"audio": 1}, {}), {"ref_text": "reference"})
+
+    assert tuple(out["ref_audio_tokens"].shape) == (8, 4)
+    assert list(out["ref_audio_len"]) == [4]
+    assert "input_ids" not in out
+    assert processor.audio_tokenizer.encode.call_count == 1
+    assert processor.text_tokenizer.call_count == 0
+
+
+def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Text-only requests and fully cached audio both reach this path with empty items."""
+    processor = _make_processor(monkeypatch)
+    monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({}, {}))
+
+    out = processor._apply_hf_processor_main(_items({}, {}), {"ref_text": "reference"})
+
+    assert len(out.keys()) == 0
+    assert processor.text_tokenizer.call_count == 0
+    assert processor.audio_tokenizer.encode.call_count == 0
+
+
+def test_legacy_call_hf_processor_is_gone() -> None:
+    assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor")

tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py

@@ -266,8 +266,6 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     inputs = MagicMock()
     inputs.mm_data_items = mm_data_items
     inputs.hf_processor_mm_kwargs = {"use_audio_in_video": True}
-    inputs.tokenization_kwargs = {}
-    inputs.prompt = [1, 2, 3]
     inputs.get_mm_hashes = MagicMock(return_value={"video": ["v0"], "audio": ["a0"]})
 
     timing_ctx = MagicMock()
@@ -307,7 +305,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     )
     fake_self._paired_cache_keys = Qwen2_5OmniThinkerMultiModalProcessor._paired_cache_keys
     fake_self.info.parse_mm_data = MagicMock(return_value=SimpleNamespace(name="missing"))
-    fake_self._apply_hf_processor_main = MagicMock(return_value=([9], {"video_grid_thw": object()}, False))
+    fake_self._apply_hf_processor_main = MagicMock(return_value={"video_grid_thw": object()})
     fake_self._get_mm_fields_config = MagicMock(return_value={})
     fake_self._get_mm_prompt_updates = MagicMock(return_value={"video": [[object()]], "audio": [[object()]]})
     fake_self._merge_mm_kwargs = capture_merge
@@ -316,13 +314,12 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
         "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs",
         return_value={"video": [object()], "audio": [object()]},
     ):
-        prompt_ids, mm_info, _ = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor(
+        mm_info = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor(
             fake_self,
             inputs,
             timing_ctx,
         )
 
-    assert prompt_ids == [9]
     # Both modalities were reprocessed as one unit after pair-miss expansion.
     assert fake_self.info.parse_mm_data.call_args.args[0] == {
         "video": ["video-0"],
@@ -623,10 +620,17 @@ def test_derive_audio_from_video_placeholders_only_pairs_true_videos():
 
 
 def test_qwen3_processor_inherits_vllm_omni_per_video_helpers():
+    from vllm.model_executor.models.qwen3_omni_moe_thinker import (
+        Qwen3OmniMoeThinkerMultiModalProcessor as UpstreamQwen3Processor,
+    )
+
     assert issubclass(
         Qwen3OmniMoeThinkerMultiModalProcessor,
         Qwen2_5OmniThinkerMultiModalProcessor,
     )
+    assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data
+    mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__
+    assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor)
 
 
 def test_qwen3_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs():

tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py

@@ -113,7 +113,9 @@ def _save_peft_adapter(tmp_path: Path) -> Path:
     return adapter_dir
 
 
-def test_peft_expert_lora_loads_into_thinker(tmp_path: Path):
+def test_peft_expert_lora_loads_into_thinker(tmp_path: Path, monkeypatch):
+    # This CPU harness tests weight mapping, not CUDA pinned-memory allocation.
+    monkeypatch.setattr("vllm.lora.lora_model.PIN_MEMORY", False)
     adapter_dir = _save_peft_adapter(tmp_path)
     peft_helper = PEFTHelper.from_local_dir(str(adapter_dir), max_position_embeddings=128)
     lora_model = LoRAModel.from_local_checkpoint(

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import patch
@@ -157,7 +157,10 @@ def _load_weights_noop(model: Qwen3TTSCode2Wav) -> set[str]:
         def __init__(self, *_: object, **__: object):
             pass
 
-        def load_weights(self, _weights: object) -> set[str]:
+        def load_weights(self, _weights: object, *, mapper: object | None = None) -> set[str]:
+            # vLLM 0.29 dropped skip_prefixes from AutoWeightsLoader; the model
+            # now passes the stage exclusions as a WeightsMapper instead.
+            del mapper
             return {"decoder.fake_weight"}
 
     with (

tests/model_executor/models/test_weight_loader.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from torch import nn
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _TinyModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.kept = nn.Linear(1, 1, bias=False)
+        self.skipped = nn.Linear(1, 1, bias=False)
+        self.rotary_embed = nn.Module()
+        self.rotary_embed.register_parameter("inv_freq", nn.Parameter(torch.zeros(1)))
+
+
+def test_upstream_mapper_preserves_stage_exclusions_after_renaming() -> None:
+    model = _TinyModel()
+    model.skipped.weight.data.zero_()
+
+    loaded = AutoWeightsLoader(model).load_weights(
+        [
+            ("checkpoint.kept.weight", torch.ones(1, 1)),
+            ("checkpoint.skipped.weight", torch.full((1, 1), 2.0)),
+            ("rotary_embed.inv_freq", torch.ones(1)),
+        ],
+        mapper=WeightsMapper(orig_to_new_prefix={"checkpoint.": ""})
+        | WeightsMapper(orig_to_new_prefix={"skipped.": None}, orig_to_new_substr={"rotary_embed.inv_freq": None}),
+    )
+
+    assert loaded == {"kept.weight"}
+    torch.testing.assert_close(model.kept.weight, torch.ones(1, 1))
+    torch.testing.assert_close(model.skipped.weight, torch.zeros(1, 1))
+    torch.testing.assert_close(model.rotary_embed.inv_freq, torch.zeros(1))

tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for the NVFP4 W4A4 weight_scale NaN clamp installed by
 ``vllm_omni.patch``.
 
@@ -198,7 +198,7 @@ def test_reload_does_not_nest_wrapper():
         print("NESTED" if nested else "FLAT")
         """
     )
-    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120)
+    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600)
     if result.stdout.strip().startswith("SKIP:"):
         pytest.skip(result.stdout.strip())
     assert "FLAT" in result.stdout, (
@@ -265,9 +265,12 @@ def test_env_var_escape_hatch_disables_install(env_value, acceptable):
     )
     env = os.environ.copy()
     env["VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP"] = env_value
-    # 120s rather than 60s: cold CI containers can take 30-50s just to
-    # import vllm_omni → vllm → torch before the subprocess does any work.
-    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120, env=env)
+    # Generous timeout: cold CI containers can take 30-50s just to import
+    # vllm_omni → vllm → torch, and heavily loaded shared machines have
+    # been observed taking >2min for the same import under disk I/O
+    # contention. The subprocess does a single import + flag print, so 600s
+    # only bounds pathological hangs, not normal runtime.
+    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600, env=env)
     if result.stdout.strip().startswith("SKIP:"):
         pytest.skip("modelopt not available in subprocess")
     # Exact token match on the final stdout line — not a substring test, since

tests/model_executor/test_upstream_prompt_update_api.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Guard that Omni only calls prompt-update helpers upstream still provides.
+
+vLLM 0.29 removed ``PromptUpdateDetails.select_text`` in favour of the
+token-oriented ``select_token_id`` / ``select_token_ids``. Three MiniCPM-o call
+sites survived the rebase and only failed at engine-core startup, where the
+traceback surfaced as an opaque "Engine core initialization failed" across eight
+CI jobs.
+
+These helpers are referenced from replacement closures that no unit test drives,
+so a static check is what actually catches the next removal. It resolves the
+attribute against the installed class rather than a hard-coded list, so it
+follows upstream automatically.
+"""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import pytest
+from vllm.multimodal.processing.processor import PromptUpdateDetails
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_OMNI_ROOT = Path(__file__).resolve().parents[2] / "vllm_omni"
+_GUARDED = "PromptUpdateDetails"
+
+
+def _referenced_attributes() -> dict[str, set[str]]:
+    """Map ``PromptUpdateDetails.<attr>`` -> source files referencing it."""
+    found: dict[str, set[str]] = {}
+    for path in _OMNI_ROOT.rglob("*.py"):
+        try:
+            tree = ast.parse(path.read_text(encoding="utf-8"))
+        except (SyntaxError, UnicodeDecodeError):  # pragma: no cover - defensive
+            continue
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == _GUARDED:
+                found.setdefault(node.attr, set()).add(str(path.relative_to(_OMNI_ROOT.parent)))
+    return found
+
+
+def test_omni_only_uses_prompt_update_helpers_that_exist_upstream() -> None:
+    referenced = _referenced_attributes()
+    assert referenced, "expected Omni to reference PromptUpdateDetails somewhere"
+
+    missing = {attr: sorted(files) for attr, files in referenced.items() if not hasattr(PromptUpdateDetails, attr)}
+
+    assert not missing, (
+        "Omni calls PromptUpdateDetails helpers that the installed vLLM does not "
+        f"provide: {missing}. Upstream moved prompt updates to token ids; use "
+        "select_token_id/select_token_ids with the encoded text."
+    )
+
+
+def test_token_oriented_helpers_are_the_ones_upstream_exposes() -> None:
+    """Pin the replacements, so their removal fails here and not at startup."""
+    assert hasattr(PromptUpdateDetails, "select_token_id")
+    assert hasattr(PromptUpdateDetails, "select_token_ids")
+    assert not hasattr(PromptUpdateDetails, "select_text"), (
+        "select_text is back upstream; re-check the MiniCPM-o replacements."
+    )

tests/platforms/test_platform_detection.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import logging
+
+import pytest
+
+import vllm_omni.platforms as platforms
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_failing_oot_plugin_is_logged_and_falls_back(monkeypatch, caplog):
+    plugin_name = "raising_oot_plugin"
+
+    def raising_plugin():
+        raise RuntimeError("injected platform detection failure")
+
+    monkeypatch.setattr(platforms, "builtin_omni_platform_plugins", {})
+    monkeypatch.setattr(
+        platforms,
+        "load_omni_plugins_by_group",
+        lambda _group: {plugin_name: raising_plugin},
+    )
+
+    target_logger = logging.getLogger(platforms.__name__)
+    target_logger.addHandler(caplog.handler)
+    previous_level = target_logger.level
+    target_logger.setLevel(logging.DEBUG)
+    try:
+        platform_cls = platforms.resolve_current_omni_platform_cls_qualname()
+    finally:
+        target_logger.removeHandler(caplog.handler)
+        target_logger.setLevel(previous_level)
+
+    assert platform_cls == "vllm_omni.platforms.interface.UnspecifiedOmniPlatform"
+    assert plugin_name in caplog.text
+    assert "RuntimeError: injected platform detection failure" in caplog.text

vllm_omni/benchmarks/patch/patch.py

@@ -2975,7 +2975,9 @@ async def benchmark(
             "duration": benchmark_duration,
             "completed": metrics.completed,
             "total_input_tokens": metrics.total_input,
+            "total_input_sequences": metrics.total_input_sequences,
             "request_throughput": metrics.request_throughput,
+            "input_sequence_throughput": metrics.input_sequence_throughput,
             "total_token_throughput": metrics.total_token_throughput,
             "input_lens": [output.prompt_len for output in outputs],
             "errors": [output.error for output in outputs],

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -107,6 +107,17 @@ class OmniSchedulerMixin:
 
     def _init_omni_io_scheduling_state(self) -> None:
         """Initialize scheduler state shared by AR and generation stages."""
+        # Every omni worker forces the v1 model runner (the v2 runner carries no
+        # omni hooks), so the scheduler has to agree or the two disagree about
+        # what the SchedulerOutput carries. vLLM 0.29 defaults this to v2 where
+        # 0.28 did not, so the scheduler took its v2 fast path -- the one that
+        # deliberately omits resumed-request token ids -- while the v1 runner
+        # still read them, and resuming a request raised ``KeyError: <req_id>``
+        # in ``_update_states``.
+        if getattr(self, "use_v2_model_runner", False):
+            logger.warning("OMNI scheduler forces v1 model runner for omni hooks.")
+            self.use_v2_model_runner = False
+
         model_config = self.vllm_config.model_config
         self.chunk_transfer_adapter = (
             OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None

vllm_omni/diffusion/attention/backends/abstract.py

@@ -121,6 +121,13 @@ class AttentionBackend(ABC):
         when it unifies cache layouts across layers. Dense diffusion backends
         conservatively keep the default ``False``; a paged backend should
         override this only when its kernel actually follows that layout.
+
+        Since vLLM 0.29 this no longer rides on ``AttentionSpec``.  It selects
+        the block-outermost physical layout pinned on
+        ``CacheConfig.kv_cache_layout`` -- see
+        ``vllm_omni.diffusion.diffusion_kv.layout`` -- whose
+        ``is_block_outermost`` is what upstream reads when it sizes
+        ``KVCacheTensor`` regions.
         """
 
         return False

vllm_omni/diffusion/attention/layer.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.layout import assert_backend_layout_supported
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
 from vllm_omni.diffusion.forward_context import (
     get_forward_context,
@@ -261,15 +262,19 @@ class Attention(nn.Module):
             return None
         dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
         # Keep backend layout discovery under the same config context used by
-        # upstream vLLM's attention-spec collector.
+        # upstream vLLM's attention-spec collector.  vLLM 0.29 moved the
+        # block-stride decision off the spec and onto the single physical
+        # ``CacheConfig.kv_cache_layout`` resolved before the cache is built,
+        # so the backend's preference is enforced there (see
+        # ``vllm_omni.diffusion.diffusion_kv.initialization``) rather than
+        # carried per layer.
         with set_current_vllm_config(vllm_config):
-            indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride()
+            assert_backend_layout_supported(vllm_config, self.attn_backend)
         return FullAttentionSpec(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=self.num_kv_heads,
             head_size=self.head_size,
             dtype=dtype,
-            indexes_kv_by_block_stride=indexes_kv_by_block_stride,
             non_causal=not self.causal,
         )
 

vllm_omni/diffusion/diffusion_kv/initialization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -16,6 +16,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout
 from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
 from vllm_omni.platforms import current_omni_platform
 
@@ -28,14 +29,35 @@ def build_native_kv_cache_configs(
     vllm_config: VllmConfig,
     worker_specs: list[dict[str, KVCacheSpec]],
     available_memory: list[int],
+    *,
+    indexes_kv_by_block_stride: bool = False,
 ) -> tuple[list[KVCacheConfig], KVCacheConfig]:
-    """Build rank-local and Scheduler-native configs using vLLM utilities."""
+    """Build rank-local and Scheduler-native configs using vLLM utilities.
+
+    ``indexes_kv_by_block_stride`` carries the native backend requirement that
+    0.29 removed from ``AttentionSpec``.  Every diffusion backend in tree
+    declares ``False``, so the default reproduces current behaviour; a caller
+    that learns otherwise from its workers passes ``True`` here and the
+    block-outermost layout propagates to every rank-local config.
+    """
 
     from vllm.v1.core import single_type_kv_cache_manager
 
+    # 0.29 sizes every KVCacheTensor region from the resolved physical layout,
+    # so it must be pinned before get_kv_cache_configs.  Idempotent, so callers
+    # that already resolved it (the control plane) are unaffected.
+    layout = resolve_diffusion_kv_cache_layout(
+        vllm_config,
+        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
+    )
     single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
     worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)
-    return worker_configs, generate_scheduler_kv_cache_config(worker_configs)
+    scheduler_config = generate_scheduler_kv_cache_config(worker_configs)
+    # Workers only receive a KVCacheConfig, so the resolved layout travels on it
+    # and is adopted into each worker's own CacheConfig before allocation.
+    for config in (*worker_configs, scheduler_config):
+        config.kv_cache_layout = layout.name
+    return worker_configs, scheduler_config
 
 
 def initialize_diffusion_kv_control_plane(

vllm_omni/diffusion/diffusion_kv/layout.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Physical KV cache layout selection for the native Diffusion paged cache.
+
+vLLM 0.29 removed ``AttentionSpec.indexes_kv_by_block_stride`` and replaced it
+with a single physical layout per model, resolved once and recorded on
+``CacheConfig.kv_cache_layout``.  ``KVCacheLayout.is_block_outermost`` is the
+direct successor of the old per-spec flag: upstream derives
+``interleaved_block_stride`` from it when it lays out ``KVCacheTensor`` regions
+(``vllm/v1/core/kv_cache_utils.py``).
+
+Diffusion never reaches vLLM's engine-core resolution.  Its backends subclass
+Omni's own ``AttentionBackend`` ABC, which upstream's
+``get_supported_kv_cache_layouts`` never sees, so the layout is pinned here
+before ``get_kv_cache_configs`` runs.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from vllm.v1.kv_cache_interface import KVCacheSpec, KVCacheTensor, compute_layout_strides
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+# Layer-outermost, matching the head of upstream's ``_DEFAULT_LAYOUT_PREFERENCE``
+# and reproducing the ``indexes_kv_by_block_stride=False`` that every dense
+# diffusion backend declared before 0.29.
+DEFAULT_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.LBNHC
+# Block-outermost counterpart, for a paged backend whose kernel really does read
+# K/V pages by the runtime block stride.
+BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.BLNHC
+
+
+def resolve_diffusion_kv_cache_layout(
+    vllm_config: VllmConfig,
+    *,
+    indexes_kv_by_block_stride: bool = False,
+) -> KVCacheLayout:
+    """Pin the physical KV layout for the Diffusion cache and return it.
+
+    A layout already present on the config wins, mirroring upstream's
+    ``resolve_kv_cache_layout``; it is validated against the requirement rather
+    than silently overwritten, because an inconsistent layout reads K/V
+    from the wrong offsets instead of raising.
+    """
+
+    required_block_outermost = bool(indexes_kv_by_block_stride)
+    cache_config = getattr(vllm_config, "cache_config", None)
+    if cache_config is None:
+        # A config double with nowhere to record the choice; report what the
+        # requirement implies without pretending it was pinned.
+        return BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+
+    current = getattr(cache_config, "kv_cache_layout", None)
+    if current is not None:
+        layout = KVCacheLayout[current]
+        if layout.is_block_outermost != required_block_outermost:
+            raise ValueError(
+                "Diffusion KV cache layout "
+                f"{layout.name} (is_block_outermost="
+                f"{layout.is_block_outermost}) contradicts the attention "
+                "backend, which requires is_block_outermost="
+                f"{required_block_outermost}."
+            )
+        return layout
+
+    layout = BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    cache_config.kv_cache_layout = layout.name
+    return layout
+
+
+def adopt_kv_cache_layout(vllm_config: VllmConfig, kv_cache_config) -> KVCacheLayout:
+    """Adopt the control plane's resolved layout inside a worker.
+
+    The engine core resolves the layout once and every worker adopts it before
+    allocating (``CacheConfig.kv_cache_layout``).  Diffusion workers receive only
+    a ``KVCacheConfig``, so the name rides along on ``kv_cache_config`` and is
+    copied onto the worker's own config here.  Without this, anything calling
+    ``get_resolved_kv_cache_layout`` in a worker -- ``init_kv_cache`` does --
+    raises.
+    """
+
+    name = getattr(kv_cache_config, "kv_cache_layout", None)
+    if name is None:
+        # Older configs, or a standalone harness: fall back to the same default
+        # the control plane pins so both sides still agree.
+        return resolve_diffusion_kv_cache_layout(vllm_config)
+
+    cache_config = getattr(vllm_config, "cache_config", None)
+    if cache_config is None:
+        return KVCacheLayout[name]
+
+    current = getattr(cache_config, "kv_cache_layout", None)
+    if current is None:
+        cache_config.kv_cache_layout = name
+    elif current != name:
+        raise ValueError(
+            f"Worker KV cache layout {current!r} disagrees with the layout resolved by the control plane ({name!r})."
+        )
+    return KVCacheLayout[name]
+
+
+def build_kv_cache_tensor(
+    spec: KVCacheSpec,
+    num_blocks: int,
+    layers: Sequence[str],
+    *,
+    layout: KVCacheLayout | None = None,
+    offset: int = 0,
+) -> KVCacheTensor:
+    """A ``KVCacheTensor`` covering ``layers`` for ``num_blocks`` blocks.
+
+    Before 0.29 a tensor named the layers sharing it
+    (``KVCacheTensor(size=..., shared_by=[...])``); 0.29 replaced that with an
+    explicit layout, so layer ``l``'s block ``b`` starts at
+    ``offset + l * layer_stride + b * block_stride``.  Note this lays the layers
+    out at distinct offsets rather than aliasing them, which is what upstream
+    now does for a single cache group.
+
+    The strides come from upstream's own ``compute_layout_strides``, called the
+    way ``vllm.v1.core.kv_cache_utils`` calls it, because hand-computed strides
+    read K/V from the wrong offsets silently instead of raising.
+    """
+
+    layout = layout or DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    layer_names = list(layers)
+    # Mirrors kv_cache_utils: the block stride is pinned up front only for
+    # block-outermost layouts, where every block repeats the same packing.
+    bytes_per_block = spec.page_size_bytes * len(layer_names)
+    interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
+    layer_stride, block_stride, _, _, _ = compute_layout_strides(
+        spec,
+        num_blocks,
+        len(layer_names),
+        layout,
+        fixed_strides=(None, interleaved_block_stride, None, None, None),
+    )
+    return KVCacheTensor(
+        size=bytes_per_block * num_blocks,
+        layers=layer_names,
+        layer_stride=layer_stride,
+        block_stride=block_stride,
+        offset=offset,
+    )
+
+
+def assert_backend_layout_supported(vllm_config: VllmConfig, attn_backend: type | None) -> None:
+    """Fail loudly when a backend's block-stride need contradicts the layout.
+
+    Skipped while the layout is still unresolved -- worker processes collect
+    specs before the control plane hands them the resolved name -- so this
+    enforces the contract wherever it is knowable and never crashes a rank that
+    has simply not been told yet.
+    """
+
+    if attn_backend is None:
+        return
+    # Spec discovery runs against lightweight config doubles in tests and against
+    # worker configs that have not been handed the resolved name yet; in both
+    # cases there is simply nothing to check.
+    cache_config = getattr(vllm_config, "cache_config", None)
+    if cache_config is None or getattr(cache_config, "kv_cache_layout", None) is None:
+        return
+    indexes_by_block_stride = getattr(attn_backend, "indexes_kv_by_block_stride", None)
+    if indexes_by_block_stride is None:
+        return
+    resolve_diffusion_kv_cache_layout(
+        vllm_config,
+        indexes_kv_by_block_stride=indexes_by_block_stride(),
+    )

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -18,6 +18,10 @@ from vllm.v1.worker.gpu.block_table import BlockTables
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
+from vllm_omni.diffusion.diffusion_kv.layout import (
+    adopt_kv_cache_layout,
+    assert_backend_layout_supported,
+)
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
@@ -285,13 +289,16 @@ class DiffusionKVModelRunnerBackend:
                     layers=self._kv_cache_layer_adapters,
                     resolve_row=self._resolve_paged_attention_row,
                 )
+                # vLLM 0.29 reads the resolved physical layout when allocating,
+                # and dropped attn_groups/cache_dtype from the signature.
+                adopt_kv_cache_layout(self.vllm_config, kv_cache_config)
+                for layer_adapter in self._kv_cache_layer_adapters.values():
+                    assert_backend_layout_supported(self.vllm_config, getattr(layer_adapter, "attn_backend", None))
                 init_kv_cache(
                     kv_caches,
                     self.vllm_config.compilation_config.static_forward_context,
                     kv_cache_config,
-                    attn_groups,
                     self.device,
-                    self.vllm_config.cache_config.cache_dtype,
                     kernel_block_sizes,
                     self.vllm_config,
                 )

vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py

@@ -105,10 +105,12 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
             attn_backend,
             ulysses_degree=ulysses_degree,
         )
+        # vLLM 0.29 carries the block-stride decision on the resolved physical
+        # ``CacheConfig.kv_cache_layout`` instead of the spec; the backend's
+        # requirement is enforced against that layout rather than copied here.
         canonical_spec = replace(
             spec,
             num_kv_heads=num_kv_heads,
-            indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(),
         )
         self.layer_name = layer_name
         self.spec = canonical_spec

vllm_omni/diffusion/layers/rope.py

@@ -1,3 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from importlib import import_module
 from importlib.util import find_spec
 
 import torch
@@ -98,6 +102,7 @@ class RotaryEmbedding(CustomOp):
         self.interleaved = not is_neox_style
         self.half_head_dim = half_head_dim
         self.apply_rotary_emb_flash_attn = None
+        self.apply_rotary_emb_vllm_flash_attn = None
         self.has_mindie = False
         # ``find_spec("flash_attn")`` is True as long as *any* package publishes
         # the ``flash_attn`` namespace — including ``flash-attn-4``, which ships
@@ -111,6 +116,17 @@ class RotaryEmbedding(CustomOp):
                 self.apply_rotary_emb_flash_attn = apply_rotary
             except ImportError:
                 pass
+        # CUDA wheels vendor this kernel under vllm_flash_attn, while source
+        # and CPU-only installations may not contain the generated package.
+        # Resolve it lazily and retain the native implementation as a portable
+        # fallback for those environments.
+        if current_omni_platform.is_cuda():
+            try:
+                self.apply_rotary_emb_vllm_flash_attn = import_module(
+                    "vllm.vllm_flash_attn.layers.rotary"
+                ).apply_rotary_emb
+            except ImportError:
+                pass
         if find_spec("mindiesd") is not None:
             self.has_mindie = True
 
@@ -135,12 +151,18 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
-
         cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin)
 
+        if self.apply_rotary_emb_vllm_flash_attn is None:
+            return apply_rotary_emb_torch(
+                x,
+                cos,
+                sin,
+                interleaved=self.interleaved,
+            )
+
         x, squeezed = _ensure_batch_dim(x)
-        output = apply_rotary_emb(
+        output = self.apply_rotary_emb_vllm_flash_attn(
             x,
             cos,
             sin,
@@ -231,13 +253,14 @@ class RotaryEmbeddingWan(RotaryEmbedding):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+        if self.apply_rotary_emb_vllm_flash_attn is None:
+            return self.forward_native(x, cos, sin)
 
         if cos.dim() > 2:
             cos = cos.reshape(-1, cos.shape[-1])
             sin = sin.reshape(-1, sin.shape[-1])
 
-        return apply_rotary_emb(
+        return self.apply_rotary_emb_vllm_flash_attn(
             x,
             cos,
             sin,

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -454,11 +454,10 @@ class HunyuanImage3Pipeline(
 
         # Note: guidance_emb and timestep_r_emb are no longer skipped
         # to support HunyuanImage-3.0-Distil and MeanFlow distilled models
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=skip_prefixes,
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(
+            weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())})
         )
-        return loader.load_weights(weights)
 
     def prepare_seed(self, seed=None, batch_size=1):
         # random seed

vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py

@@ -24,7 +24,7 @@ from vllm.distributed import (
     tensor_model_parallel_all_reduce,
 )
 from vllm.logger import init_logger
-from vllm.model_executor.layers.batch_invariant import linear_batch_invariant
+from vllm.model_executor.determinism.batch_invariant import linear_batch_invariant
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -50,6 +50,17 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 logger = init_logger(__name__)
 
 
+def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """Rotate interleaved pairs before rounding back to the activation dtype.
+
+    Qwen-Image's reference uses complex FP32 multiplication. Rounding the
+    frequencies to BF16 before rotation loses positional precision; those
+    errors accumulate across the denoising steps.
+    """
+    paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype)
+
+
 def _normalize_qwen_image_weight_name(name: str) -> str:
     name = name.removeprefix("transformer.")
     if ".to_out.0." in name:
@@ -642,15 +653,23 @@ class QwenImageCrossAttention(nn.Module):
         txt_query = self.norm_added_q(txt_query)
         txt_key = self.norm_added_k(txt_key)
 
-        img_cos = torch.real(vid_freqs).to(img_query.dtype)
-        img_sin = torch.imag(vid_freqs).to(img_query.dtype)
-        txt_cos = torch.real(txt_freqs).to(txt_query.dtype)
-        txt_sin = torch.imag(txt_freqs).to(txt_query.dtype)
+        if img_query.device.type == "cuda":
+            img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs)
+            img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs)
+            txt_query = _apply_qwen_image_rotary_emb(txt_query, txt_freqs)
+            txt_key = _apply_qwen_image_rotary_emb(txt_key, txt_freqs)
+        else:
+            # Retain the platform-specific kernels on other accelerators,
+            # which may not support complex tensors.
+            img_cos = torch.real(vid_freqs).to(img_query.dtype)
+            img_sin = torch.imag(vid_freqs).to(img_query.dtype)
+            txt_cos = torch.real(txt_freqs).to(txt_query.dtype)
+            txt_sin = torch.imag(txt_freqs).to(txt_query.dtype)
 
-        img_query = self.rope(img_query, img_cos, img_sin)
-        img_key = self.rope(img_key, img_cos, img_sin)
-        txt_query = self.rope(txt_query, txt_cos, txt_sin)
-        txt_key = self.rope(txt_key, txt_cos, txt_sin)
+            img_query = self.rope(img_query, img_cos, img_sin)
+            img_key = self.rope(img_key, img_cos, img_sin)
+            txt_query = self.rope(txt_query, txt_cos, txt_sin)
+            txt_key = self.rope(txt_key, txt_cos, txt_sin)
 
         seq_len_txt = encoder_hidden_states.shape[1]
         joint_query = torch.cat([txt_query, img_query], dim=1)

vllm_omni/engine/arg_utils.py

@@ -443,8 +443,8 @@ class OmniEngineArgs(EngineArgs):
 @dataclass
 class OmniAsyncEngineArgs(AsyncEngineArgs, OmniEngineArgs):
     @classmethod
-    def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
-        parser = AsyncEngineArgs.add_cli_args(parser)
+    def add_cli_args(cls, parser: argparse.ArgumentParser, async_args_only: bool = False) -> argparse.ArgumentParser:
+        parser = AsyncEngineArgs.add_cli_args(parser, async_args_only=async_args_only)
         parser = OmniEngineArgs._add_omni_specific_args(parser)
         return parser
 

vllm_omni/engine/stage_engine_core_proc_manager.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Process manager for omni stage engine subprocesses.
 
@@ -83,6 +83,12 @@ class StageEngineCoreProcManager(CoreEngineProcManager):
         if local_engine_count <= 0:
             raise ValueError(f"local_engine_count must be > 0, got {local_engine_count}")
 
+        # Mirrors the vLLM 0.29 parent __init__: the inherited shutdown() reads
+        # this to bound how long in-flight requests may drain. Omitting it makes
+        # shutdown raise AttributeError, which leaves the engine core
+        # subprocesses alive and hangs interpreter exit until the job timeout.
+        self._request_shutdown_timeout = vllm_config.shutdown_timeout
+
         context = get_mp_context()
         common_kwargs: dict[str, object] = {
             "vllm_config": vllm_config,

vllm_omni/engine/stage_init_utils.py

@@ -54,7 +54,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
 from vllm_omni.entrypoints.utils import filter_dataclass_kwargs
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
-from vllm_omni.inputs.preprocess import OmniInputPreprocessor
+from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls
 from vllm_omni.outputs.output_processor import MultimodalOutputProcessor
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.inc_config import OmniINCConfig
@@ -1549,25 +1549,24 @@ class _TokenOnlyRenderer(BaseRenderer):
 
 
 def _build_token_only_renderer(stage_vllm_config: Any) -> BaseRenderer:
-    return _TokenOnlyRenderer(stage_vllm_config, tokenizer=None)
+    return omni_renderer_cls(_TokenOnlyRenderer)(stage_vllm_config, tokenizer=None)
 
 
 def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
-    """Build the shared stage-0 input processor."""
+    """Build the shared stage-0 input processor.
+
+    The renderer is the Omni subclass of the upstream renderer class that
+    ``renderer_from_config`` would have picked (or of the token-only renderer
+    when the stage skips tokenizer initialization), so upstream's
+    ``InputProcessor`` runs unmodified on top of it.
+    """
 
     patch_generation_config_if_needed(stage_vllm_config.model_config)
     if bool(getattr(stage_vllm_config.model_config, "skip_tokenizer_init", False)):
-        input_processor = InputProcessor(
-            vllm_config=stage_vllm_config,
-            renderer=_build_token_only_renderer(stage_vllm_config),
-        )
+        renderer = _build_token_only_renderer(stage_vllm_config)
     else:
-        input_processor = InputProcessor(vllm_config=stage_vllm_config)
-    input_processor.input_preprocessor = OmniInputPreprocessor(
-        vllm_config=stage_vllm_config,
-        renderer=input_processor.renderer,
-    )
-    return input_processor
+        renderer = build_omni_renderer(stage_vllm_config)
+    return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer)
 
 
 def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str:

vllm_omni/entrypoints/async_omni.py

@@ -51,9 +51,9 @@ from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetric
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
-    from vllm.inputs.preprocess import InputPreprocessor
     from vllm.tokenizers import TokenizerLike
     from vllm.v1.engine import PauseMode
+    from vllm.v1.engine.input_processor import InputProcessor
 
     from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType
 
@@ -1714,7 +1714,7 @@ class AsyncOmni(EngineClient, OmniBase):
 
     # ==================== EngineClient Interface ====================
 
-    async def get_input_preprocessor(self) -> InputPreprocessor:
+    async def get_input_preprocessor(self) -> InputProcessor:
         """Get input preprocessor."""
         return self.input_processor
 
@@ -1755,10 +1755,12 @@ class AsyncOmni(EngineClient, OmniBase):
         """
         logger.debug("Weight update start requested (no-op in omni)")
 
-    async def finish_weight_update(self) -> None:
+    async def finish_weight_update(self, weight_version: str | None = None) -> None:
         """Finish the current weight update.
 
         Omni does not currently support weight transfer, so this is a no-op.
+        ``weight_version`` is accepted for upstream ``EngineClient`` protocol
+        compatibility (RLHF weight-transfer routers pass it positionally).
         """
         logger.debug("Weight update finish requested (no-op in omni)")
 

vllm_omni/entrypoints/cli/serve.py

@@ -18,7 +18,7 @@ from typing import Any, cast
 
 import uvloop
 from vllm.entrypoints.cli.types import CLISubcommand
-from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
+from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
 

vllm_omni/entrypoints/duplex/chat_fallback.py

@@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator
 from typing import Any
 
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.duplex.protocol import DuplexSession

vllm_omni/entrypoints/openai/api_server.py

@@ -34,7 +34,10 @@ from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
 from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
-from vllm.entrypoints.launcher import serve_http, terminate_if_errored
+from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
+from vllm.entrypoints.launchers.cli_args import make_arg_parser
+from vllm.entrypoints.launchers.launcher import serve_http, terminate_if_errored
+from vllm.entrypoints.launchers.utils.server_utils import get_uvicorn_log_config
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
@@ -43,19 +46,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponse,
 )
-from vllm.entrypoints.openai.cli_args import make_arg_parser
 
 # yapf conflicts with isort for this block
 # yapf: disable
 # yapf: enable
 from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
-from vllm.entrypoints.openai.engine.protocol import (
-    ErrorResponse,
-    ModelCard,
-    ModelList,
-    ModelPermission,
-    RequestResponseMetadata,
-)
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
@@ -64,14 +59,13 @@ from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAISer
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
 from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
-
-# vLLM < 0.28 keeps create_error_response under serve.utils (it is not
-# re-exported from vllm.entrypoints.serve); 0.28+ moved it under
-# serve.exception_handling and re-exports it from the package root.
-try:
-    from vllm.entrypoints.serve import create_error_response
-except ImportError:
-    from vllm.entrypoints.serve.utils.error_response import create_error_response
+from vllm.entrypoints.serve import create_error_response
+from vllm.entrypoints.serve.engine.protocol import (
+    ErrorResponse,
+    ModelCard,
+    ModelList,
+    ModelPermission,
+)
 
 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
 # Keep a fallback for older/newer upstream layouts during rebase windows.
@@ -85,7 +79,6 @@ from vllm.entrypoints.serve.utils.api_utils import (
 )
 from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
-from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config
 from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime
 from vllm.entrypoints.speech_to_text.transcription.serving import (
     OpenAIServingTranscription,

vllm_omni/entrypoints/openai/batch_serving.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import asyncio
 import time
 
@@ -11,7 +11,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionResponse,
     ChatCompletionResponseChoice,
 )
-from vllm.entrypoints.openai.engine.protocol import (
+from vllm.entrypoints.serve.engine.protocol import (
     ErrorResponse,
     UsageInfo,
 )

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import asyncio
@@ -8,7 +11,7 @@ from typing import cast
 from uuid import uuid4
 
 import numpy as np
-from vllm.entrypoints.openai.engine.protocol import UsageInfo
+from vllm.entrypoints.serve.engine.protocol import UsageInfo
 from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection as VllmRealtimeConnection
 from vllm.entrypoints.speech_to_text.realtime.protocol import TranscriptionDelta, TranscriptionDone
 from vllm.logger import init_logger

vllm_omni/entrypoints/openai/serving_chat.py

@@ -54,8 +54,17 @@ except ImportError:
     soundfile = None
 
 
+from vllm.entrypoints.generate.base.protocol import (
+    DeltaFunctionCall,
+    DeltaMessage,
+    DeltaToolCall,
+    FunctionCall,
+    FunctionDefinition,
+    RequestResponseMetadata,
+    ToolCall,
+)
 from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs
-from vllm.entrypoints.launcher import terminate_if_errored
+from vllm.entrypoints.launchers.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionRequest,
@@ -69,22 +78,15 @@ from vllm.entrypoints.openai.chat_completion.serving import (
     _get_mm_token_counts,
     _make_prompt_tokens_details,
 )
-from vllm.entrypoints.openai.engine.protocol import (
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ErrorInfo,
-    ErrorResponse,
-    FunctionCall,
-    FunctionDefinition,
-    RequestResponseMetadata,
-    ToolCall,
-    UsageInfo,
-)
 from vllm.entrypoints.openai.parser.harmony_utils import (
     get_streamable_parser_for_assistant,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.entrypoints.serve.engine.protocol import (
+    ErrorInfo,
+    ErrorResponse,
+    UsageInfo,
+)
 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
@@ -700,7 +702,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 if negative_prompt is not None:
                     tprompt["negative_prompt"] = negative_prompt
                 # Always attach mm_processor_kwargs (possibly empty) so
-                # OmniInputPreprocessor._process_text routes through the
+                # OmniRenderer routes through the
                 # multimodal processor path. Without it, the preprocessor
                 # falls back to plain _tokenize_prompt and AR-based image-gen
                 # models like GLM-Image never see their image-generation

vllm_omni/entrypoints/openai/serving_speech.py

@@ -24,12 +24,10 @@ import soundfile as sf
 import torch
 from fastapi import HTTPException, Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
+from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
 from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
-from vllm.entrypoints.launcher import terminate_if_errored
-from vllm.entrypoints.openai.engine.protocol import (
-    ErrorResponse,
-    RequestResponseMetadata,
-)
+from vllm.entrypoints.launchers.launcher import terminate_if_errored
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
 from vllm.utils import random_uuid

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -156,7 +156,9 @@ def _decode_video_bytes(
         frames_array, metadata = loader.load_bytes(
             video_bytes,
             num_frames=num_frames,
-            backend="pyav",
+            # vLLM 0.29 removed the "pyav" decoder backend; "opencv" is
+            # upstream's default and the remaining CPU-only option.
+            backend="opencv",
             keep=keep,
         )
     except Exception as exc:

vllm_omni/entrypoints/openai/video_stream_session.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Video streaming session manager and WebSocket handler.
 
 Provides ``VideoStreamConfig``, ``VideoStreamSession`` (frame/audio buffer
@@ -20,7 +20,7 @@ from fastapi import WebSocket, WebSocketDisconnect
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
 )
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter

vllm_omni/experimental/ar_diffusion/kv_cache/manager.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Engine-level KV cache orchestration for AR-Diffusion models.
 
 This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a
@@ -22,10 +23,10 @@ from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheSpec,
-    KVCacheTensor,
 )
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor
 from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
@@ -134,7 +135,7 @@ def build_kv_manager(
     """
     layer_names = list(layer_names)
     group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
-    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
+    tensors = [build_kv_cache_tensor(spec, num_blocks, layer_names)]
     config = KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=tensors,

vllm_omni/inputs/mm_processor.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Multimodal data plumbing for Omni's custom, non-HF processors."""
+
+from collections.abc import Mapping
+from typing import TypeVar
+
+from transformers import BatchFeature
+from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.processing import BaseMultiModalProcessor, BaseProcessingInfo
+
+_I = TypeVar("_I", bound=BaseProcessingInfo)
+
+
+class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]):
+    """Share data/passthrough handling, leaving model processing in its owner.
+
+    These processors implement ``_call_hf_processor`` locally, sometimes
+    without any HF processor object. Prompt tokens are prepared separately
+    in ``apply``; the upstream base still owns caching and prompt updates.
+    """
+
+    _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text"
+
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        kwargs = dict(hf_processor_mm_kwargs)
+        prompt = kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, None)
+        if prompt is None:
+            prompt = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough = self._get_hf_mm_data(valid_items)
+        result = self._call_hf_processor(str(prompt), dict(mm_data), kwargs, {})
+        result.update(passthrough)
+        return result

vllm_omni/inputs/preprocess.py

@@ -1,200 +1,118 @@
-from typing import Any
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Omni renderer: the upstream renderer class with Omni engine-input routing.
 
-from typing_extensions import assert_never
-from vllm.inputs import EmbedsInput, MultiModalInput, SingletonInput
-from vllm.inputs.preprocess import InputPreprocessor
-from vllm.logger import init_logger
-from vllm.renderers.inputs import SingletonDictPrompt
+``OmniRenderer`` is a mixin. ``omni_renderer_cls`` puts it in front of a
+concrete ``BaseRenderer`` subclass, so the Omni instance *is* an upstream
+renderer (tokenizer, executors, multimodal cache and every public entry point
+come from upstream unchanged) and only ``_process_singleton`` /
+``_process_singleton_async`` are extended. ``build_omni_renderer`` resolves
+the concrete class the same way upstream ``renderer_from_config`` does.
+"""
 
-from vllm_omni.inputs.data import (
-    OmniEmbedsPrompt,
-    OmniTextPrompt,
-    OmniTokenInputs,
-    OmniTokensPrompt,
-    token_inputs_omni,
-)
+from __future__ import annotations
 
-logger = init_logger(__name__)
+from typing import TYPE_CHECKING, cast
+
+from vllm.renderers import BaseRenderer
+from vllm.renderers.registry import RENDERER_REGISTRY
+from vllm.tokenizers.registry import cached_tokenizer_from_config, tokenizer_args_from_config
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.tokenizers import TokenizerLike
 
 
-class OmniInputPreprocessor(InputPreprocessor):
-    """Input preprocessor for omni models.
+class OmniRenderer:
+    """Mixin that routes Omni prompts through the upstream renderer.
 
-    Extends the base InputPreprocessor to handle omni-specific input
-    types including prompt embeddings and additional information payloads.
-    Supports processing tokens, embeddings, text, and multimodal inputs.
+    It has no base class of its own; ``omni_renderer_cls`` prepends it to a
+    concrete ``BaseRenderer`` subclass, so ``self`` and ``super()`` are that
+    renderer at runtime (hence the casts below).
+
+    Two extensions over upstream ``_process_singleton``:
+
+    * a token prompt that carries ``mm_processor_kwargs`` but no
+      ``multi_modal_data`` still goes through the multimodal processor (AR
+      image-generation models such as GLM-Image read their target size from
+      those kwargs);
+    * Omni-only prompt keys (``prompt``, ``cache_salt``,
+      ``additional_information``, ``model_intermediate_buffer``) are copied
+      onto the engine input.
     """
 
-    def _process_text(
-        self,
-        parsed_content: OmniTextPrompt,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: Any | None = None,
-    ) -> OmniTokenInputs | MultiModalInput:
-        """Process text prompts with support for mm_processor_kwargs.
+    @staticmethod
+    def _with_omni_extras(inputs, prompt):
+        for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"):
+            if key in prompt:
+                inputs[key] = prompt[key]
+        return inputs
 
-        Extends base class to support mm_processor_kwargs without multi_modal_data.
-        This is needed for models like GLM-Image where text-to-image generation
-        requires processor kwargs (target_h, target_w) to format the prompt.
-        """
-        prompt_text = parsed_content["prompt"]
-        mm_processor_kwargs = parsed_content.get("mm_processor_kwargs") or {}
-        # When the deprecated raw-prompt path is used, process_inputs does
-        # not pass mm_uuids to preprocess().  Fall back to reading it from
-        # the prompt dict so the Renderer's _validate_mm_uuids can see it.
-        effective_mm_uuids = mm_uuids or parsed_content.get("multi_modal_uuids")
+    @staticmethod
+    def _routes_no_media_kwargs(prompt) -> bool:
+        return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data")
 
-        inputs: OmniTokenInputs | MultiModalInput
-        if multi_modal_data := parsed_content.get("multi_modal_data"):
-            inputs = self._process_multimodal(
-                prompt_text,
-                multi_modal_data,
-                mm_processor_kwargs,
-                tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=effective_mm_uuids,
-            )
-            prompt_embeds = parsed_content.get("prompt_embeds")
-            if prompt_embeds is not None:
-                inputs["prompt_embeds"] = prompt_embeds
-            additional_information = parsed_content.get("additional_information")
-            if additional_information is not None:
-                inputs["additional_information"] = additional_information
-            model_intermediate_buffer = parsed_content.get("model_intermediate_buffer")
-            if model_intermediate_buffer is not None:
-                inputs["model_intermediate_buffer"] = model_intermediate_buffer
-        elif "mm_processor_kwargs" in parsed_content:
-            # Presence — not truthiness. An explicitly-set empty dict still
-            # signals "route through the multimodal processor" (needed for
-            # AR-based image-gen where the HF processor supplies its own
-            # defaults and scaffold).
-            inputs = self._process_multimodal(
-                prompt_text,
+    def _process_singleton(self, prompt, *, skip_mm_cache: bool = False):
+        renderer = cast(BaseRenderer, self)
+        if self._routes_no_media_kwargs(prompt):
+            inputs = renderer._process_multimodal(
+                prompt["prompt_token_ids"],
                 {},
-                mm_processor_kwargs,
-                tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=effective_mm_uuids,
+                mm_processor_kwargs=prompt["mm_processor_kwargs"],
+                mm_uuids=prompt.get("multi_modal_uuids"),
+                skip_mm_cache=skip_mm_cache,
             )
         else:
-            prompt_token_ids = self._tokenize_prompt(
-                prompt_text,
-                tokenization_kwargs=tokenization_kwargs,
+            inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
+        return self._with_omni_extras(inputs, prompt)
+
+    async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False):
+        renderer = cast(BaseRenderer, self)
+        if self._routes_no_media_kwargs(prompt):
+            inputs = await renderer._process_multimodal_async(
+                prompt["prompt_token_ids"],
+                {},
+                mm_processor_kwargs=prompt["mm_processor_kwargs"],
+                mm_uuids=prompt.get("multi_modal_uuids"),
+                skip_mm_cache=skip_mm_cache,
             )
-            inputs = token_inputs_omni(
-                prompt_token_ids,
-                prompt_embeds=parsed_content.get("prompt_embeds"),
-                additional_information=parsed_content.get("additional_information"),
-                model_intermediate_buffer=parsed_content.get("model_intermediate_buffer"),
-            )
-        prompt_embeds = parsed_content.get("prompt_embeds")
-        if prompt_embeds is not None:
-            inputs["prompt_embeds"] = prompt_embeds
-        additional_information = parsed_content.get("additional_information")
-        if additional_information is not None:
-            inputs["additional_information"] = additional_information
-        model_intermediate_buffer = parsed_content.get("model_intermediate_buffer")
-        if model_intermediate_buffer is not None:
-            inputs["model_intermediate_buffer"] = model_intermediate_buffer
-        if cache_salt := parsed_content.get("cache_salt"):
-            inputs["cache_salt"] = cache_salt
-
-        return inputs
-
-    def _process_tokens(
-        self,
-        parsed_content: OmniTokensPrompt,
-        tokenization_kwargs: dict[str, Any] | None = None,
-    ) -> OmniTokenInputs | MultiModalInput:
-        prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
-        prompt_embeds = parsed_content.get("prompt_embeds")
-        additional_information = parsed_content.get("additional_information")
-        model_intermediate_buffer = parsed_content.get("model_intermediate_buffer")
-
-        multi_modal_data = parsed_content.get("multi_modal_data")
-
-        inputs: OmniTokenInputs | MultiModalInput
-        if multi_modal_data:
-            inputs = self._process_multimodal(
-                prompt_token_ids,
-                multi_modal_data,
-                parsed_content.get("mm_processor_kwargs"),
-                tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=parsed_content.get("multi_modal_uuids"),
-            )
-
         else:
-            inputs = token_inputs_omni(
-                prompt_token_ids=prompt_token_ids,
-                prompt_embeds=prompt_embeds,
-                additional_information=additional_information,
-                model_intermediate_buffer=model_intermediate_buffer,
-            )
-        if prompt_embeds is not None:
-            inputs["prompt_embeds"] = prompt_embeds
-        if additional_information is not None:
-            inputs["additional_information"] = additional_information
-        if model_intermediate_buffer is not None:
-            inputs["model_intermediate_buffer"] = model_intermediate_buffer
-        if prompt_text := parsed_content.get("prompt"):
-            inputs["prompt"] = prompt_text
-        if cache_salt := parsed_content.get("cache_salt"):
-            inputs["cache_salt"] = cache_salt
+            inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache)
+        return self._with_omni_extras(inputs, prompt)
 
-        return inputs
 
-    def _process_embeds(
-        self,
-        parsed_content: OmniEmbedsPrompt,
-    ) -> EmbedsInput:
-        """Process embeddings prompt with omni-specific extensions.
+_OMNI_RENDERER_CLASSES: dict[type[BaseRenderer], type[BaseRenderer]] = {}
 
-        Extends base _process_embeds to handle additional_information payload
-        for direct transfer between pipeline stages.
-        """
-        # Call parent implementation for base embeds processing
-        inputs = super()._process_embeds(parsed_content)
 
-        # Add omni-specific additional_information if present
-        additional_information = parsed_content.get("additional_information")
-        if additional_information is not None:
-            inputs["additional_information"] = additional_information  # type: ignore[typeddict-unknown-key]
-        model_intermediate_buffer = parsed_content.get("model_intermediate_buffer")
-        if model_intermediate_buffer is not None:
-            inputs["model_intermediate_buffer"] = model_intermediate_buffer  # type: ignore[typeddict-unknown-key]
+def omni_renderer_cls(base_cls: type[BaseRenderer]) -> type[BaseRenderer]:
+    """Return the Omni subclass of ``base_cls`` (created once per base class)."""
+    if issubclass(base_cls, OmniRenderer):
+        return base_cls
+    cls = _OMNI_RENDERER_CLASSES.get(base_cls)
+    if cls is None:
+        name = f"Omni{base_cls.__name__}"
+        cls = type(name, (OmniRenderer, base_cls), {"__module__": __name__, "__qualname__": name})
+        _OMNI_RENDERER_CLASSES[base_cls] = cls
+    return cls
 
-        return inputs
 
-    def _prompt_to_llm_inputs(
-        self,
-        prompt: SingletonDictPrompt,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: Any | None = None,
-    ) -> SingletonInput:
-        """
-        Extract the singleton inputs from a prompt.
+def build_omni_renderer(
+    vllm_config: VllmConfig,
+    *,
+    renderer_cls: type[BaseRenderer] | None = None,
+    tokenizer: TokenizerLike | None = None,
+) -> BaseRenderer:
+    """Build the Omni renderer for ``vllm_config``.
 
-        Arguments:
-
-        * prompt: single encoder or decoder input prompt
-
-        Returns:
-
-        * [`SingletonInput`][vllm.inputs.engine.SingletonInput] instance
-        """
-        if "prompt_embeds" in prompt:
-            return self._process_embeds(prompt)  # type: ignore[arg-type]
-
-        if "prompt_token_ids" in prompt:
-            return self._process_tokens(
-                prompt,  # type: ignore[arg-type]
-            )
-
-        if "prompt" in prompt:
-            return self._process_text(
-                prompt,  # type: ignore[arg-type]
-                tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
-            )
-
-        assert_never(prompt)  # type: ignore[arg-type]
+    With ``renderer_cls`` unset this mirrors upstream ``renderer_from_config``:
+    the tokenizer and renderer mode come from the model config and the class
+    from ``RENDERER_REGISTRY``. Pass ``renderer_cls`` (and optionally
+    ``tokenizer``) to skip that resolution, e.g. for tokenizer-less stages.
+    """
+    model_config = vllm_config.model_config
+    if renderer_cls is None:
+        if tokenizer is None:
+            tokenizer = cached_tokenizer_from_config(model_config)
+        renderer_mode, *_ = tokenizer_args_from_config(model_config)
+        renderer_cls = RENDERER_REGISTRY.load_renderer_cls(renderer_mode)
+    return omni_renderer_cls(renderer_cls)(vllm_config, tokenizer)

vllm_omni/model_executor/models/audex/audex_omni_processing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0):
 #   inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/audio_features.py
@@ -33,13 +33,14 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
 from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
-    BaseMultiModalProcessor,
     BaseProcessingInfo,
     PromptReplacement,
     PromptUpdate,
     PromptUpdateDetails,
 )
 
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
+
 SOUND_TOKEN = "<so_embedding>"
 SOUND_START_TOKEN = "<so_start>"
 SOUND_END_TOKEN = "<so_end>"
@@ -176,7 +177,7 @@ class AudexDummyInputsBuilder(BaseDummyInputsBuilder[AudexProcessingInfo]):
         return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios, overrides=overrides)}
 
 
-class AudexMultiModalProcessor(BaseMultiModalProcessor[AudexProcessingInfo]):
+class AudexMultiModalProcessor(OmniMultiModalProcessor[AudexProcessingInfo]):
     def _call_hf_processor(
         self,
         prompt: str,

vllm_omni/model_executor/models/bagel/bagel.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from math import isqrt
 from typing import Any
 
@@ -20,7 +21,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.bagel import BagelForConditionalGeneration
 from vllm.model_executor.models.interfaces import MultiModalEmbeddings
 from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP
-from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
@@ -221,6 +222,20 @@ class OmniBagelDataParser(MultiModalDataParser):
 class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]):
     IMG2IMG_PLACEHOLDER = "<|fim_middle|>"
 
+    def apply(self, inputs, timing_ctx):
+        num_img2img = inputs.mm_data_items.get_all_counts().get("img2img", 0)
+        if num_img2img:
+            token_id = self.info.get_tokenizer().get_vocab().get(self.IMG2IMG_PLACEHOLDER)
+            if token_id is not None:
+                prompt_ids = list(inputs.prompt)
+                missing = num_img2img - prompt_ids.count(token_id)
+                if missing > 0:
+                    inputs = replace(
+                        inputs,
+                        prompt=[token_id] * missing + prompt_ids,
+                    )
+        return super().apply(inputs, timing_ctx)
+
     @staticmethod
     def _mm_kwargs_for_bagel_img2img_hf(mm_kwargs: Mapping[str, object]) -> dict[str, object]:
         # OpenAI / GLM-style serving may pass target_h/target_w for output grid sizing.
@@ -241,69 +256,51 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn
             "pixel_values_img2img": MultiModalFieldConfig.batched("img2img"),
         }
 
-    def _call_hf_processor(
+    def _apply_hf_processor_main(
         self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
-    ) -> "BatchFeature":
-        has_image = "images" in mm_data
-        has_img2img = "pixel_values_img2img" in mm_data
-
-        if has_img2img and self.IMG2IMG_PLACEHOLDER not in prompt:
-            prompt = f"{self.IMG2IMG_PLACEHOLDER}{prompt}"
-
-        if has_image and has_img2img:
-            outputs = BatchFeature()
-
-            img_data = dict(mm_data)
-            if "pixel_values_img2img" in img_data:
-                del img_data["pixel_values_img2img"]
-            kwargs_img = dict(mm_kwargs)
-            kwargs_img["is_img2img"] = False
-            out_img = super()._call_hf_processor(prompt, img_data, kwargs_img, tok_kwargs)
-            if "pixel_values" in out_img:
-                outputs["pixel_values"] = out_img["pixel_values"]
-            for k, v in out_img.items():
-                if k != "pixel_values":
-                    outputs[k] = v
-
-            img2img_data = dict(mm_data)
-            if "images" in img2img_data:
-                del img2img_data["images"]
-            img2img_data["images"] = img2img_data.pop("pixel_values_img2img")
-            kwargs_img2img = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs)
-            kwargs_img2img["is_img2img"] = True
-            out_img2img = super()._call_hf_processor(prompt, img2img_data, kwargs_img2img, tok_kwargs)
-            if "pixel_values" in out_img2img:
-                outputs["pixel_values_img2img"] = out_img2img["pixel_values"]
-            for k, v in out_img2img.items():
-                if k not in outputs:
-                    outputs[k] = v
-
-            return outputs
-
-        elif has_img2img:
-            mm_data = dict(mm_data)
-            mm_data["images"] = mm_data.pop("pixel_values_img2img")
-            mm_kwargs = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs)
-            mm_kwargs["is_img2img"] = True
-            outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
-            if "pixel_values" in outputs:
-                outputs["pixel_values_img2img"] = outputs.pop("pixel_values")
-            return outputs
-
-        return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
-
-    def _hf_processor_applies_updates(
-        self,
-        prompt_text: str,
         mm_items: MultiModalDataItems,
         hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> bool:
-        return False
+    ) -> BatchFeature:
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+
+        has_image = "images" in mm_data
+        has_img2img = "pixel_values_img2img" in mm_data
+        processed_data = BatchFeature()
+
+        if has_image:
+            image_kwargs = {**hf_processor_mm_kwargs, "is_img2img": False}
+            image_outputs = self.info.ctx.call_hf_processor(
+                processor,
+                {"text": prompt_text, "images": mm_data["images"]},
+                image_kwargs,
+            )
+            processed_data.update(image_outputs)
+
+        if has_img2img:
+            img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_processor_mm_kwargs)
+            img2img_kwargs["is_img2img"] = True
+            img2img_outputs = self.info.ctx.call_hf_processor(
+                processor,
+                {
+                    "text": prompt_text,
+                    "images": mm_data["pixel_values_img2img"],
+                },
+                img2img_kwargs,
+            )
+            pixel_values = img2img_outputs.pop("pixel_values", None)
+            if pixel_values is not None:
+                processed_data["pixel_values_img2img"] = pixel_values
+            for key, value in img2img_outputs.items():
+                processed_data.setdefault(key, value)
+
+        if not has_image and not has_img2img:
+            processed_data = BatchFeature(dict(passthrough_data))
+        else:
+            processed_data.update(passthrough_data)
+        return processed_data
 
     def _get_prompt_updates(
         self,
@@ -372,7 +369,8 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn
                 embed_mask = [True] * num_vae_total + [False] + [True] * num_vit_total
                 return PromptUpdateDetails(
                     full=tokens,
-                    is_embed=lambda _tok, _seq, _m=embed_mask: torch.tensor(_m, dtype=torch.bool),
+                    # vLLM 0.29 calls is_embed with just the full token list.
+                    is_embed=lambda _full, _m=embed_mask: torch.tensor(_m, dtype=torch.bool),
                 )
 
             replacements.append(
@@ -1102,11 +1100,12 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
             filtered_weights.append((mapped_name, tensor))
 
         loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=["vit_pos_embed.pos_embed"],
-            ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."],
+            self, ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."]
+        )
+        loaded = loader.load_weights(
+            filtered_weights,
+            mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"vit_pos_embed.pos_embed": None}),
         )
-        loaded = loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper)
 
         loaded |= self._load_moe_gen_weights(moe_gen_weights)
 

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import os
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from functools import partial
 from math import gcd
 from threading import Lock
@@ -24,7 +25,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
 from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
-    BaseMultiModalProcessor,
     BaseProcessingInfo,
     ProcessorInputs,
     PromptIndexTargets,
@@ -38,6 +38,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import random_sample
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import (
     concat_text_with_prompt_ids,
@@ -108,7 +109,45 @@ class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo):
         )
 
 
-class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]):
+class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]):
+    def apply(self, inputs: ProcessorInputs, timing_ctx):
+        tokenizer = self.info.get_tokenizer()
+        prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False)
+        config = self.info.ctx.get_hf_config()
+        model_dir = self.info.ctx.model_config.model
+        self._ensure_cached_runtime_components(model_dir, config)
+
+        text_token, text_token_len = extract_text_token(
+            prompt_text,
+            self.tokenizer,
+            config.allowed_special,
+        )
+        if inputs.mm_data_items.get_all_counts().get("audio", 0):
+            reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text")
+            if not isinstance(reference_text, str):
+                raise ValueError(f"prompt text is None : {reference_text}")
+            prompt_text_token, prompt_text_token_len = extract_text_token(
+                reference_text,
+                self.tokenizer,
+                config.allowed_special,
+            )
+            text_token, _ = concat_text_with_prompt_ids(
+                text_token,
+                text_token_len,
+                prompt_text_token,
+                prompt_text_token_len,
+            )
+
+        inputs = replace(
+            inputs,
+            prompt=text_token.reshape(-1).tolist(),
+            hf_processor_mm_kwargs={
+                **inputs.hf_processor_mm_kwargs,
+                self._OMNI_PROMPT_TEXT_KEY: prompt_text,
+            },
+        )
+        return super().apply(inputs, timing_ctx)
+
     def _ensure_cached_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> None:
         cached_model_dir = getattr(self, "_cached_model_dir", None)
         if cached_model_dir == model_dir:
@@ -376,15 +415,6 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal
             "embedding": MultiModalFieldConfig.batched("audio"),
         }
 
-    def _hf_processor_applies_updates(
-        self,
-        prompt_text: str,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> bool:
-        return False
-
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,

vllm_omni/model_executor/models/covo_audio/covo_audio.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 Tencent.
 from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
@@ -17,7 +20,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig
 from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
-    BaseMultiModalProcessor,
     BaseProcessingInfo,
     ProcessorInputs,
     PromptReplacement,
@@ -26,6 +28,7 @@ from vllm.multimodal.processing import (
 from vllm.sequence import IntermediateTensors
 from vllm.v1.sample.metadata import SamplingMetadata
 
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
@@ -106,10 +109,7 @@ class CovoAudioDummyInputsBuilder(BaseDummyInputsBuilder[CovoAudioProcessingInfo
         )
 
 
-class CovoAudioMultiModalProcessor(BaseMultiModalProcessor[CovoAudioProcessingInfo]):
-    def _hf_processor_applies_updates(self, prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs) -> bool:
-        return False
-
+class CovoAudioMultiModalProcessor(OmniMultiModalProcessor[CovoAudioProcessingInfo]):
     def _call_hf_processor(
         self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], tok_kwargs: Mapping[str, object]
     ) -> BatchFeature:

vllm_omni/model_executor/models/gepard/gepard_talker.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Gepard-1.0 native-AR talker.
 
 Single-stage AR TTS on a vLLM-native Qwen3.5 backbone. Each step samples one
@@ -22,7 +22,7 @@ import torch.nn as nn
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLM
-from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
@@ -199,8 +199,10 @@ class GepardTalkerForConditionalGeneration(nn.Module):
 
         # NanoCodec loads lazily on first decode, so load_format=dummy and
         # NeMo-less environments never pay for it.
-        loader = AutoWeightsLoader(self, skip_prefixes=["mtp.", "ref_compressor."])
-        return loader.load_weights(iter(rest))
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(
+            iter(rest), mapper=WeightsMapper(orig_to_new_prefix={"mtp.": None, "ref_compressor.": None})
+        )
 
     def _get_or_create_state(self, request_id: str) -> _GepardState:
         st = self._active_states.get(request_id)

vllm_omni/model_executor/models/glm_image/glm_image_ar.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 # Adapted from
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm_image/modeling_glm_image.py
@@ -24,6 +24,7 @@
 import math
 import os
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from functools import lru_cache
 from typing import Annotated, Literal
 
@@ -81,7 +82,6 @@ from vllm.multimodal.inputs import (
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import (
-    ImageProcessorItems,
     MultiModalDataItems,
     MultiModalDataParser,
 )
@@ -341,6 +341,30 @@ class GlmImageDummyInputsBuilder(BaseDummyInputsBuilder[GlmImageProcessingInfo])
         }
 
 
+_GLM_IMAGE_GRID_FACTOR = 32
+
+
+def _build_target_shape_scaffold(processor, *, height: int, width: int, is_text_to_image: bool) -> str:
+    """Return the grid scaffold HF appends to a GLM-Image generation prompt.
+
+    Mirrors ``GlmImageProcessor._build_prompt_with_target_shape`` (identical in
+    transformers 5.13 and 5.14) using only the processor's public token
+    attributes, so this does not depend on a private HF method that a stub or
+    a future release may not provide.
+    """
+    factor = _GLM_IMAGE_GRID_FACTOR
+    token_h = (height // factor * factor) // factor
+    token_w = (width // factor * factor) // factor
+    grid_bos, grid_eos, bos = processor.grid_bos_token, processor.grid_eos_token, processor.bos_token
+    scaffold = f"{grid_bos}{token_h} {token_w}{grid_eos}"
+    if is_text_to_image:
+        ratio = token_h / token_w
+        prev_token_h = int(math.sqrt(ratio) * (factor // 2))
+        prev_token_w = int(math.sqrt(1 / ratio) * (factor // 2))
+        scaffold += f"{grid_bos}{prev_token_h} {prev_token_w}{grid_eos}"
+    return scaffold + bos
+
+
 class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo]):
     """
     Multimodal processor for GLM-Image.
@@ -358,265 +382,38 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
             return self._apply_hf_processor(inputs, timing_ctx)
         return super()._cached_apply_hf_processor(inputs, timing_ctx)
 
-    def _call_hf_processor(
-        self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        """
-        Call the HuggingFace processor.
-
-        For text-to-image mode (no images), we need to:
-        1. Build the prompt with target grid dimensions
-        2. Build the image_grid_thw tensor for M-RoPE position encoding
-
-        For image-to-image mode:
-        1. Process source images through the image processor
-        2. Build prompt with image placeholders expanded
-        3. Build image_grid_thw including source and target grids
-        """
-        processor = self.info.get_hf_processor()
-
-        # Get target dimensions from mm_kwargs or use defaults
-        target_h = mm_kwargs.get("target_h", 1024) if mm_kwargs else 1024
-        target_w = mm_kwargs.get("target_w", 1024) if mm_kwargs else 1024
-
-        logger.debug(
-            f"_call_hf_processor: target dimensions for generation: {target_h}x{target_w}, mm_kwargs={mm_kwargs}"
-        )
-
-        if not mm_data or not mm_data.get("images"):
-            # Text-to-image mode
-            if processor is not None:
-                # Build messages format expected by processor
-                messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
-
-                # Use apply_chat_template which handles target dimensions
-                hf_inputs = processor.apply_chat_template(
-                    messages,
-                    tokenize=True,
-                    target_h=target_h,
-                    target_w=target_w,
-                    return_dict=True,
-                    return_tensors="pt",
-                )
-
-                return hf_inputs
-            else:
-                # Fallback: just tokenize (this won't work properly for generation)
-                tokenizer = self.info.get_tokenizer()
-                prompt_ids = tokenizer.encode(prompt)
-                return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
-
-        # Image-to-image mode
-        # NOTE: Use "images" (plural) - this is what vLLM's ImageProcessorItems.get_processor_data() returns
-        images = mm_data.get("images")
-        if not isinstance(images, list):
-            images = [images]
-
-        logger.debug(
-            f"_call_hf_processor i2i: num_images={len(images)}, image_types={[type(img).__name__ for img in images]}"
-        )
-
-        if processor is not None:
-            # Build messages with image objects directly in content
-            # This is how GlmImageProcessor expects images - embedded in the content dict
-            # NOT as a separate images= parameter
-            #
-            # IMPORTANT: Remove <|image|> placeholders from prompt since apply_chat_template
-            # will automatically insert them for each image in content. Having both leads to
-            # index out of bounds when processing image_grid_thw.
-            clean_prompt = prompt.replace("<|image|>", "")
-            content = []
-            for img in images:
-                content.append({"type": "image", "image": img})
-            content.append({"type": "text", "text": clean_prompt})
-            messages = [{"role": "user", "content": content}]
-
-            logger.debug(f"_call_hf_processor: calling apply_chat_template with {len(images)} images in content")
-
-            # Use apply_chat_template - processor will process images when they're in content
-            hf_inputs = processor.apply_chat_template(
-                messages,
-                tokenize=True,
-                target_h=target_h,
-                target_w=target_w,
-                return_dict=True,
-                return_tensors="pt",
+    def apply(self, inputs, timing_ctx):
+        """Add the generation scaffold before upstream expands source images."""
+        num_images = inputs.mm_data_items.get_all_counts().get("image", 0)
+        prompt_ids = list(inputs.prompt)
+        if num_images:
+            image_token_id = getattr(
+                self.info.get_hf_config(),
+                "image_token_id",
+                167855,
             )
+            missing = num_images - prompt_ids.count(image_token_id)
+            if missing > 0:
+                prompt_ids = [image_token_id] * missing + prompt_ids
 
-            logger.debug(f"_call_hf_processor: apply_chat_template returned keys: {list(hf_inputs.keys())}")
-
-            # IMPORTANT (i2i): vLLM multimodal encoder must see source-only grids
-            # (matching pixel_values and number of images), but M-RoPE needs full
-            # grids (source + target) to compute correct decode positions.
-            image_grid_thw = hf_inputs.get("image_grid_thw")
-            if image_grid_thw is not None:
-                # Preserve full grids for M-RoPE.
-                hf_inputs["mrope_image_grid_thw"] = image_grid_thw
-
-                # Expose source-only grids for MM.
-                # In most i2i requests, we process one prompt at a time here,
-                # so `len(images)` is the number of source images.
-                num_source_images = len(images)
-                if image_grid_thw.shape[0] != num_source_images:
-                    source_grids = image_grid_thw[:num_source_images]
-                    hf_inputs["image_grid_thw"] = source_grids
-                    logger.debug(
-                        "_call_hf_processor: adjusted image_grid_thw for MM from %s to %s (num_source_images=%d)",
-                        tuple(image_grid_thw.shape),
-                        tuple(source_grids.shape),
-                        num_source_images,
-                    )
-
-            # Debug: Analyze input_ids for image tokens
-            input_ids = hf_inputs.get("input_ids")
-            if input_ids is not None:
-                if hasattr(input_ids, "tolist"):
-                    ids_list = input_ids.tolist()
-                    if isinstance(ids_list[0], list):
-                        ids_list = ids_list[0]  # Unbatch
-                else:
-                    ids_list = list(input_ids)
-
-                # Get image token ID from config
-                hf_config = self.info.get_hf_config()
-                image_token_id = getattr(hf_config, "image_token_id", 167855)
-
-                # Count image tokens
-                image_token_count = ids_list.count(image_token_id)
-                logger.debug(
-                    f"_call_hf_processor: input_ids length={len(ids_list)}, "
-                    f"image_token_id={image_token_id}, "
-                    f"image_token_count={image_token_count}"
-                )
-
-                # Log first/last few tokens to understand structure
-                logger.debug(f"_call_hf_processor: first 20 tokens: {ids_list[:20]}")
-                logger.debug(f"_call_hf_processor: last 20 tokens: {ids_list[-20:]}")
-
-                # Find positions of image tokens
-                image_positions = [i for i, t in enumerate(ids_list) if t == image_token_id]
-                if image_positions:
-                    logger.debug(f"_call_hf_processor: image token positions (first 10): {image_positions[:10]}")
-
-            return hf_inputs
-        else:
-            # Fallback without processor - this is not ideal but prevents crashes
-            logger.warning("GlmImageProcessor not available, using fallback for i2i")
-            tokenizer = self.info.get_tokenizer()
-            hf_config = self.info.get_hf_config()
-
-            # Get image token
-            image_token_id = getattr(hf_config, "image_token_id", 167855)
-            try:
-                image_token = tokenizer.convert_ids_to_tokens(image_token_id)
-            except Exception:
-                image_token = "<|image|>"
-
-            # Build prompt with image placeholders
-            image_placeholders = image_token * len(images)
-            full_prompt = f"{image_placeholders}{prompt}"
-            prompt_ids = tokenizer.encode(full_prompt)
-            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
-
-    def _apply_hf_processor_mm_only(
-        self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        """
-        Apply the HF processor on the multi-modal data only.
-
-        GLM-Image requires special handling because apply_chat_template always
-        adds a target <|image|> placeholder in addition to source image placeholders.
-        This causes an IndexError when the HF processor tries to find grid info
-        for the target placeholder (which doesn't exist for source-only processing).
-
-        Solution: Call the image processor directly to get pixel_values and
-        image_grid_thw, bypassing apply_chat_template's target handling.
-        """
-        mm_counts = mm_items.get_all_counts()
-        num_images = mm_counts.get("image", 0)
-
-        if num_images == 0:
-            # No images - call parent implementation
-            return super()._apply_hf_processor_mm_only(
-                mm_items=mm_items,
-                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-                tokenization_kwargs=tokenization_kwargs,
-            )
-
-        # For i2i mode, we need to process images directly with the image processor
-        # to avoid the apply_chat_template target placeholder issue
         processor = self.info.get_hf_processor()
-        image_processor = processor.image_processor
-
-        # Get images from mm_items
-        images = mm_items.get_items("image", ImageProcessorItems)
-        image_list = [images.get(i) for i in range(images.get_count())]
-
-        logger.debug(f"_apply_hf_processor_mm_only: processing {len(image_list)} images directly")
-
-        # Process images directly with image processor
-        image_inputs = image_processor(
-            images=image_list,
-            return_tensors="pt",
-        )
-
-        # Get grid info for source images only (no target)
-        pixel_values = image_inputs.get("pixel_values")
-        image_grid_thw = image_inputs.get("image_grid_thw")
-        if image_grid_thw is not None and image_grid_thw.shape[0] != num_images:
-            # Be defensive: some processors may include extra target grids.
-            image_grid_thw = image_grid_thw[:num_images]
-            image_inputs["image_grid_thw"] = image_grid_thw
-
-        logger.debug(
-            f"_apply_hf_processor_mm_only: pixel_values shape=\
-                {pixel_values.shape if pixel_values is not None else None}, "
-            f"image_grid_thw shape={image_grid_thw.shape if image_grid_thw is not None else None}"
-        )
-
-        # Build input_ids with image token placeholders
-        # The _get_prompt_updates returns PromptReplacement(target=[image_token_id], ...)
-        # which needs to find image tokens in input_ids to replace them.
-        # We need to include one image_token_id per image so the replacement can work.
         tokenizer = self.info.get_tokenizer()
-        image_token_id = tokenizer.convert_tokens_to_ids("<|image|>")
-
-        # Build input_ids: [image_token] * num_images + tokenized text
-        # This way _apply_prompt_updates can find the image tokens and replace them
-        dummy_text = self.dummy_inputs.get_dummy_text(mm_counts)
-        text_ids = tokenizer.encode(dummy_text, add_special_tokens=False)
-        input_ids = [image_token_id] * num_images + text_ids
-
-        logger.debug(
-            f"_apply_hf_processor_mm_only: built input_ids with {num_images} image tokens + {len(text_ids)} text tokens"
+        # Append HF's target-shape scaffold without decoding/re-tokenizing the
+        # user's tokens. Text-to-image needs both target and preview grids; i2i
+        # needs only the target. Grid metadata alone does not supply these AR
+        # tokens.
+        target_grid = self._build_generation_grids(inputs.hf_processor_mm_kwargs)[0]
+        suffix = _build_target_shape_scaffold(
+            processor,
+            height=int(target_grid[1]) * 32,
+            width=int(target_grid[2]) * 32,
+            is_text_to_image=num_images == 0,
         )
-
-        return BatchFeature(
-            dict(
-                input_ids=[input_ids],
-                pixel_values=pixel_values,
-                image_grid_thw=image_grid_thw,
-            ),
-            tensor_type="pt",
-        )
-
-    def _apply_hf_processor_text_only(
-        self, prompt_text: str, hf_processor_mm_kwargs: Mapping[str, object], tokenization_kwargs: Mapping[str, object]
-    ) -> list[int]:
-        prompt_ids, _, _ = super()._apply_hf_processor_text_mm(
-            prompt_text=prompt_text,
-            mm_items=MultiModalDataItems({}),
-            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-            tokenization_kwargs=tokenization_kwargs,
-        )
-        return prompt_ids
+        suffix_ids = tokenizer.encode(suffix, add_special_tokens=False)
+        if prompt_ids[-len(suffix_ids) :] != suffix_ids:
+            prompt_ids.extend(suffix_ids)
+        inputs = replace(inputs, prompt=prompt_ids)
+        return super().apply(inputs, timing_ctx)
 
     def _build_generation_grids(self, hf_processor_mm_kwargs: Mapping[str, object]) -> torch.Tensor:
         """Build generation grids for M-RoPE decode positions.
@@ -661,171 +458,39 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
 
     def _apply_hf_processor_main(
         self,
-        prompt: str | list[int],
         mm_items: MultiModalDataItems,
         hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-        *,
-        enable_hf_prompt_update: bool,
-    ) -> tuple[list[int], BatchFeature, bool]:
-        """
-        Override to handle GLM-Image i2i mode correctly.
-
-        Problem: When vLLM processes cached mm items (enable_hf_prompt_update=False),
-        the base implementation:
-        1. Gets prompt_ids from _apply_hf_processor_text_only (no image tokens)
-        2. Gets mm_data from _apply_hf_processor_mm_only
-        3. Returns is_update_applied=False
-
-        This causes _apply_prompt_updates to fail because prompt_ids has no image tokens.
-
-        Solution: For i2i mode, we build prompt_ids that include image placeholders,
-        and return is_update_applied=False so _apply_prompt_updates can expand them.
-        """
+    ) -> BatchFeature:
+        """Process source images under vLLM's token-only prompt contract."""
         mm_counts = mm_items.get_all_counts()
         num_images = mm_counts.get("image", 0)
+        if num_images == 0:
+            return BatchFeature({"mrope_image_grid_thw": self._build_generation_grids(hf_processor_mm_kwargs)})
 
-        logger.debug(f"_apply_hf_processor_main: mm_counts={mm_counts}, num_images={num_images}")
+        valid_mm_items = mm_items.select({key for key, count in mm_counts.items() if count > 0})
+        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        images = processor_data.get("images")
+        if not images:
+            return BatchFeature(dict(passthrough_data))
 
-        if num_images == 0 and isinstance(prompt, str):
-            # t2i mode or normal flow - use parent implementation
-            prompt_ids = self._apply_hf_processor_text_only(
-                prompt_text=prompt,
-                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-                tokenization_kwargs=tokenization_kwargs,
-            )
-            mm_processed_data = self._apply_hf_processor_mm_only(
-                mm_items=mm_items,
-                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-                tokenization_kwargs=tokenization_kwargs,
-            )
-
-            # t2i has no source images, so mm features cannot provide image_grid_thw.
-            # Provide explicit generation grids for M-RoPE to avoid fallback token parsing
-            # (which can degrade high-resolution spatial positions, e.g. 1920x1920).
-            try:
-                mrope_grid_thw = self._build_generation_grids(hf_processor_mm_kwargs)
-                mm_processed_data["mrope_image_grid_thw"] = mrope_grid_thw
-                logger.debug(
-                    "_apply_hf_processor_main t2i: mrope_image_grid_thw=%s",
-                    mrope_grid_thw.tolist(),
-                )
-            except Exception as e:
-                logger.warning("_apply_hf_processor_main t2i: failed to set mrope_image_grid_thw: %s", e)
-
-            return prompt_ids, mm_processed_data, False
-
-        # i2i mode: use unified HF processor path only.
-        # This avoids drift between duplicated manual/HF i2i implementations.
-        logger.debug(
-            "_apply_hf_processor_main: i2i mode (enable_hf_prompt_update=%s), num_images=%s",
-            enable_hf_prompt_update,
-            num_images,
+        processor = self.info.get_hf_processor()
+        image_inputs = processor.image_processor(
+            images=list(images),
+            return_tensors="pt",
         )
+        image_grid_thw = image_inputs.get("image_grid_thw")
+        if image_grid_thw is not None:
+            source_grid_thw = image_grid_thw[:num_images]
+            image_inputs["image_grid_thw"] = source_grid_thw
 
-        if not isinstance(prompt, str):
-            # Online OpenAI chat preprocessing can arrive here with tokenized
-            # prompts (list[int]) before serving_chat replaces engine prompt
-            # with the clean text prompt. Do not fail the whole request.
-            logger.warning(
-                "_apply_hf_processor_main i2i: got tokenized prompt type=%s; "
-                "using compatibility path for preprocessing",
-                type(prompt).__name__,
+            target_grid = self._build_generation_grids(hf_processor_mm_kwargs)[:1].to(dtype=source_grid_thw.dtype)
+            image_inputs["mrope_image_grid_thw"] = torch.cat(
+                [source_grid_thw, target_grid],
+                dim=0,
             )
 
-            prompt_ids = list(prompt)
-            mm_processed_data = self._apply_hf_processor_mm_only(
-                mm_items=mm_items,
-                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-                tokenization_kwargs=tokenization_kwargs,
-            )
-
-            # Preserve full grids for M-RoPE decode (source + target), while
-            # keeping image_grid_thw source-only for MM batching.
-            try:
-                source_grid_thw = mm_processed_data.get("image_grid_thw")
-                if source_grid_thw is not None and isinstance(source_grid_thw, torch.Tensor):
-                    target_h = (
-                        hf_processor_mm_kwargs.get("target_h")
-                        if isinstance(hf_processor_mm_kwargs.get("target_h"), int)
-                        else None
-                    )
-                    target_w = (
-                        hf_processor_mm_kwargs.get("target_w")
-                        if isinstance(hf_processor_mm_kwargs.get("target_w"), int)
-                        else None
-                    )
-                    if target_h is None or target_w is None:
-                        target_h = (
-                            hf_processor_mm_kwargs.get("height")
-                            if isinstance(hf_processor_mm_kwargs.get("height"), int)
-                            else 1024
-                        )
-                        target_w = (
-                            hf_processor_mm_kwargs.get("width")
-                            if isinstance(hf_processor_mm_kwargs.get("width"), int)
-                            else 1024
-                        )
-
-                    factor = 32
-                    token_h = max(1, target_h // factor)
-                    token_w = max(1, target_w // factor)
-                    target_grid = torch.tensor([[1, token_h, token_w]], dtype=source_grid_thw.dtype)
-                    mm_processed_data["mrope_image_grid_thw"] = torch.cat([source_grid_thw, target_grid], dim=0)
-            except Exception:
-                pass
-
-            # Prompt updates will expand image placeholders in this compatibility path.
-            return prompt_ids, mm_processed_data, False
-
-        images = mm_items.get_items("image", ImageProcessorItems)
-        image_list = [images.get(i) for i in range(images.get_count())]
-        if not image_list:
-            raise ValueError("GLM-Image i2i requires at least one source image in mm_items")
-
-        hf_inputs = self._call_hf_processor(
-            prompt=prompt,
-            mm_data={"images": image_list},
-            mm_kwargs=hf_processor_mm_kwargs,
-            tok_kwargs=tokenization_kwargs,
-        )
-
-        input_ids = hf_inputs.get("input_ids")
-        if input_ids is None:
-            raise ValueError("HF i2i processor returned no input_ids")
-
-        if isinstance(input_ids, torch.Tensor):
-            prompt_ids = input_ids[0].tolist() if input_ids.dim() > 1 else input_ids.tolist()
-        else:
-            prompt_ids = (
-                input_ids[0]
-                if isinstance(input_ids, list) and input_ids and isinstance(input_ids[0], list)
-                else list(input_ids)
-            )
-
-        mm_processed_data = BatchFeature(dict(), tensor_type="pt")
-        for key in ("pixel_values", "image_grid_thw", "mrope_image_grid_thw"):
-            value = hf_inputs.get(key)
-            if value is not None:
-                mm_processed_data[key] = value
-
-        image_grid_thw = mm_processed_data.get("image_grid_thw")
-        mrope_grid_thw = mm_processed_data.get("mrope_image_grid_thw")
-        hf_config = self.info.get_hf_config()
-        image_token_id = getattr(hf_config, "image_token_id", 167855)
-        image_token_count = prompt_ids.count(image_token_id)
-        logger.debug(
-            "_apply_hf_processor_main i2i(HF): num_images=%s, prompt_len=%s, image_token_count=%s, "
-            "source_grid_shape=%s, mrope_grid_shape=%s",
-            num_images,
-            len(prompt_ids),
-            image_token_count,
-            tuple(image_grid_thw.shape) if image_grid_thw is not None else None,
-            tuple(mrope_grid_thw.shape) if mrope_grid_thw is not None else None,
-        )
-
-        # HF processor already expanded image placeholders in input_ids.
-        return prompt_ids, mm_processed_data, True
+        image_inputs.update(passthrough_data)
+        return image_inputs
 
     def _get_mm_fields_config(
         self,
@@ -836,7 +501,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
         Get the multimodal field configuration.
 
         For GLM-Image i2i mode:
-        - image_grid_thw has been sliced in _call_hf_processor to only include source images
+        - image_grid_thw contains only source grids from _apply_hf_processor_main
         - pixel_values has shape [total_patches, C, H, W] - only for source images
 
         For t2i mode:
@@ -844,7 +509,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
         """
         result = {}
 
-        # Get image_grid_thw if present (already sliced in _call_hf_processor)
+        # Get the source-only image_grid_thw if present.
         image_grid_thw = hf_inputs.get("image_grid_thw")
 
         if "pixel_values" in hf_inputs and image_grid_thw is not None:
@@ -861,45 +526,21 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
 
                 result["pixel_values"] = MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes)
 
-                # Register image_grid_thw - it's been sliced in _call_hf_processor
+                # Register image_grid_thw - it's been sliced in _apply_hf_processor_main
                 # to only include source image grids, so batching will work correctly
                 result["image_grid_thw"] = MultiModalFieldConfig.batched("image")
 
+                if "mrope_image_grid_thw" in hf_inputs:
+                    result["mrope_image_grid_thw"] = MultiModalFieldConfig.shared(
+                        "image",
+                        num_source_images,
+                        keep_on_cpu=True,
+                    )
+
         logger.debug(f"_get_mm_fields_config: result keys: {list(result.keys())}")
 
         return result
 
-    def _hf_processor_applies_updates(
-        self,
-        prompt_text: str,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> bool:
-        """
-        Return whether the HF processor applies prompt updates.
-
-        For GLM-Image i2i mode, the HF processor's apply_chat_template already
-        expands <|image|> to N tokens (e.g., 4096 for 64x64 grid).
-
-        By returning True, we tell vLLM that HF processor DID apply prompt updates,
-        so vLLM will use _find_mm_placeholders to locate the expanded tokens
-        instead of trying to apply replacements.
-
-        For t2i mode (no images), there are no image placeholders to expand.
-        """
-        # Check if we have images (i2i mode)
-        num_images = mm_items.get_all_counts().get("image", 0)
-        if num_images > 0:
-            logger.debug(
-                f"_hf_processor_applies_updates: returning True for i2i mode "
-                f"(num_images={num_images}) - HF processor already expanded tokens"
-            )
-            return True
-
-        # For t2i mode (no images), use default behavior
-        return True
-
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
@@ -909,17 +550,8 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
         """
         Get prompt updates for image tokens.
 
-        For GLM-Image image-to-image mode, the HF processor's apply_chat_template
-        already expands each <|image|> placeholder to the correct number of
-        image tokens (grid_h * grid_w tokens per source image).
-
-        The HF processor does:
-        1. Replace each <|image|> with num_image_tokens copies of <|placeholder|>
-        2. Replace all <|placeholder|> back to <|image|>
-
-        So the tokenized input already has the expanded tokens. We use
-        target=[image_token_id] to match each occurrence of the image token,
-        similar to how Qwen2VL handles this pattern.
+        Expand each source-image placeholder into grid_h * grid_w image tokens
+        through vLLM's token-based prompt updates.
 
         We use image_grid_thw from out_mm_kwargs to get the actual processed grid
         size, following the Qwen2VL pattern. This is critical because the HF processor
@@ -930,8 +562,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
         """
         hf_config = self.info.get_hf_config()
 
-        # Get image token ID - this is the token that appears multiple times
-        # in the tokenized input after HF processor expansion
+        # Match the source-image placeholder in the tokenized input.
         image_token_id = getattr(hf_config, "image_token_id", 167855)
 
         # Debug: log mm_items info

vllm_omni/model_executor/models/glm_tts/glm_tts.py

@@ -39,7 +39,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
 from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
-    BaseMultiModalProcessor,
     BaseProcessingInfo,
     ProcessorInputs,
     PromptIndexTargets,
@@ -51,6 +50,7 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
 from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.platforms import current_omni_platform
@@ -364,7 +364,7 @@ class GLMTTSMultiModalProcessingInfo(BaseProcessingInfo):
         )
 
 
-class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessingInfo]):
+class GLMTTSMultiModalProcessor(OmniMultiModalProcessor[GLMTTSMultiModalProcessingInfo]):
     """GLM-TTS voice-clone processor.
 
     Unlike CosyVoice3, GLM-TTS prompt speech tokens are normal Llama vocab IDs
@@ -373,6 +373,53 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi
     handoff.
     """
 
+    def apply(self, inputs: ProcessorInputs, timing_ctx):
+        source_tokenizer = self.info.get_tokenizer()
+        prompt_text = source_tokenizer.decode(
+            inputs.prompt,
+            skip_special_tokens=False,
+        )
+        config = self.info.ctx.get_hf_config()
+        model_dir = self.info.ctx.model_config.model
+        self._ensure_cached_runtime_components(
+            model_dir,
+            config,
+            load_voice_clone=False,
+        )
+
+        normalized_text = _normalize_glm_tts_processor_text(
+            self.text_frontend,
+            prompt_text,
+        )
+        text_ids = self._encode_text(normalized_text)
+        prompt_parts = []
+        if inputs.mm_data_items.get_all_counts().get("audio", 0):
+            reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text")
+            if not isinstance(reference_text, str) or not reference_text.strip():
+                raise ValueError("GLM-TTS voice cloning requires mm_processor_kwargs['prompt_text'].")
+            normalized_reference = _normalize_glm_tts_processor_text(
+                self.text_frontend,
+                reference_text,
+                add_trailing_space=True,
+            )
+            prompt_parts.append(self._encode_text(normalized_reference))
+        prompt_parts.extend(
+            [
+                text_ids,
+                torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long),
+            ]
+        )
+        prompt_ids = torch.cat(prompt_parts, dim=1).reshape(-1).tolist()
+        inputs = replace(
+            inputs,
+            prompt=prompt_ids,
+            hf_processor_mm_kwargs={
+                **inputs.hf_processor_mm_kwargs,
+                self._OMNI_PROMPT_TEXT_KEY: prompt_text,
+            },
+        )
+        return super().apply(inputs, timing_ctx)
+
     def _ensure_cached_runtime_components(
         self,
         model_dir: str,
@@ -599,15 +646,6 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi
             if key in hf_inputs
         }
 
-    def _hf_processor_applies_updates(
-        self,
-        prompt_text: str,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> bool:
-        return False
-
     def _cached_apply_hf_processor(self, inputs: ProcessorInputs, timing_ctx: Any):
         # GLM-TTS builds the actual AR prompt from both the request text and
         # mm_processor_kwargs["prompt_text"]. The base cache path separates text

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py

@@ -1669,8 +1669,10 @@ class _BackboneWrapper(nn.Module):
         self.config = config
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        from vllm.model_executor.models.utils import AutoWeightsLoader
+        from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
         skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None
-        loader = AutoWeightsLoader(self, skip_prefixes=skip)
-        return loader.load_weights(weights)
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(
+            weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip or ())})
+        )

vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py

@@ -0,0 +1,625 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+#
+# Vendored from vllm/model_executor/models/hunyuan_v1.py at vllm@2cf0a6915c
+# (the upstream base vllm-omni origin/main was rebased against). Upstream
+# removed the native HunYuan V1 implementation in #53272 / #53615 (migrated
+# to the Transformers modeling backend). vllm-omni's HunyuanImage3 model
+# subclasses the removed HunYuanModel / HunYuanSparseMoeBlock / HunYuanMLP,
+# so those internals are preserved here to keep the omni diffusion model
+# working against current upstream.
+# Weight loading, expert mapping, and QKV splitting are owned by Omni's
+# HunyuanModel subclass; their unused upstream implementations are omitted.
+"""Inference-only HunYuan V1 internals vendored for vllm-omni HunyuanImage3."""
+
+from itertools import islice
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_ep_group,
+    get_pp_group,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.models.interfaces import EagleModelMixin
+from vllm.model_executor.models.utils import (
+    PPMissingLayer,
+    make_layers,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backend import AttentionType
+
+
+def _is_moe(config: PretrainedConfig) -> bool:
+    num_experts = getattr(config, "num_experts", None)
+    if isinstance(num_experts, int):
+        return num_experts > 1
+    if isinstance(num_experts, list) and num_experts:
+        # Ensure all elements are integers before calling max.
+        if all(isinstance(e, int) for e in num_experts):
+            return max(num_experts) > 1
+        else:
+            return False
+    return False
+
+
+def _get_cla_factor(config: PretrainedConfig) -> int:
+    if not getattr(config, "use_cla", False):
+        return 1
+    return getattr(config, "cla_share_factor", 1)
+
+
+class HunYuanMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+        reduce_results: bool = True,
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+            reduce_results=reduce_results,
+        )
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.")
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class HunYuanAttention(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position_embeddings: int = 8192,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        cache_config: CacheConfig | None = None,
+        prefix: str = "",
+        layer_id: int = -1,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        if hasattr(config, "head_dim") and config.head_dim:
+            self.head_dim = config.head_dim
+        elif hasattr(config, "attention_head_dim"):
+            self.head_dim = config.attention_head_dim
+        else:
+            self.head_dim = self.hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.use_qk_norm = getattr(config, "use_qk_norm", False)
+        self.layer_id = layer_id
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=True,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+        )
+
+        if self.use_qk_norm:
+            self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_states: tuple[torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        ori_k = k
+        if self.use_qk_norm:
+            q = self.query_layernorm(
+                q.view(-1, self.num_heads, self.head_dim),
+            )
+            k = self.key_layernorm(
+                k.view(-1, self.num_kv_heads, self.head_dim),
+            )
+
+        attn_output = self.attn(q, k, v)
+        # For o_proj
+        attn_output = attn_output.view(q.shape[0], -1)
+        output, _ = self.o_proj(attn_output)
+        return output, (ori_k, v)
+
+
+class HunYuanCrossAttention(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position_embeddings: int = 8192,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        cache_config: CacheConfig | None = None,
+        prefix: str = "",
+        layer_id: int = -1,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
+        if hasattr(config, "head_dim"):
+            self.head_dim = config.head_dim
+        elif hasattr(config, "attention_head_dim"):
+            self.head_dim = config.attention_head_dim
+        else:
+            self.head_dim = self.hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.use_qk_norm = getattr(config, "use_qk_norm", False)
+        self.layer_id = layer_id
+
+        self.q_proj = ColumnParallelLinear(
+            hidden_size,
+            hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.q_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=True,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            attn_type=AttentionType.ENCODER_DECODER,
+        )
+
+        if self.use_qk_norm:
+            self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_states: tuple[torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        assert kv_states is not None
+        ori_k, v = kv_states  # use last layer kv,
+        k = ori_k
+        q, _ = self.q_proj(hidden_states)
+        k_tmp = torch.empty_like(k)  # Todo: reduant rotary embedding
+        q, _ = self.rotary_emb(positions, q, k_tmp)
+        if self.use_qk_norm:
+            q = self.query_layernorm(
+                q.view(-1, self.num_heads, self.head_dim),
+            )
+            k = self.key_layernorm(
+                k.view(-1, self.num_kv_heads, self.head_dim),
+            )
+
+        attn_output = self.attn(q, k, v)
+        # For o_proj
+        attn_output = attn_output.view(q.shape[0], -1)
+        output, _ = self.o_proj(attn_output)
+        return output, (ori_k, v)
+
+
+class HunYuanSparseMoeBlock(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: QuantizationConfig | None = None,
+        layer_id: int = -1,
+        prefix: str = "",
+        enable_eplb: bool = False,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        self.ep_group = get_ep_group().device_group
+        self.ep_size = self.ep_group.size()
+        self.n_routed_experts = config.num_experts
+
+        if self.tp_size > config.num_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than the number of experts {config.num_experts}."
+            )
+
+        # Get layer_id topk if config.moe_topk is a list
+        if isinstance(config.moe_topk, list):
+            assert layer_id >= 0
+            assert len(config.moe_topk) > layer_id
+            top_k = config.moe_topk[layer_id]
+        else:
+            top_k = config.moe_topk
+
+        # If it is moe, moe_intermediate_size is preferred
+        intermediate_size = config.intermediate_size
+        if config.moe_intermediate_size is not None:
+            intermediate_size = (
+                config.moe_intermediate_size
+                if isinstance(config.moe_intermediate_size, int)
+                else config.moe_intermediate_size[layer_id]
+            )
+
+        # Load balancing settings.
+        vllm_config = get_current_vllm_config()
+        eplb_config = vllm_config.parallel_config.eplb_config
+        self.enable_eplb = enable_eplb
+
+        self.n_logical_experts = self.n_routed_experts
+        self.n_redundant_experts = eplb_config.num_redundant_experts
+        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
+        self.n_local_physical_experts = self.n_physical_experts // self.ep_size
+        self.gate = ReplicatedLinear(
+            config.hidden_size,
+            config.num_experts,
+            bias=False,
+            quant_config=None,
+            prefix=f"{prefix}.gate",
+        )
+        if config.use_mixed_mlp_moe > 0:
+            # Get layer_id num_shared_expert if config.num_shared_expert is
+            # a list.
+            if isinstance(config.num_shared_expert, list):
+                assert layer_id >= 0
+                assert len(config.num_shared_expert) > layer_id
+                num_shared_expert = config.num_shared_expert[layer_id]
+            else:
+                num_shared_expert = config.num_shared_expert
+
+            self.shared_mlp = HunYuanMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size * num_shared_expert,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                reduce_results=False,
+                prefix=f"{prefix}.shared_mlp",
+            )
+        else:
+            self.shared_mlp = None
+
+        self.experts = FusedMoEFactory(
+            shared_experts=self.shared_mlp,
+            num_experts=self.n_routed_experts,
+            top_k=top_k,
+            hidden_size=config.hidden_size,
+            intermediate_size=intermediate_size,
+            renormalize=top_k > 1,
+            quant_config=quant_config,
+            prefix=f"{prefix}.experts",
+            enable_eplb=self.enable_eplb,
+            num_redundant_experts=self.n_redundant_experts,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # NOTE: hidden_states can have either 1D or 2D shape.
+        orig_shape = hidden_states.shape
+        hidden_dim = hidden_states.shape[-1]
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits)
+
+        return final_hidden_states.view(orig_shape)
+
+
+class HunYuanDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        layer_id: int = -1,
+        enable_eplb: bool = False,
+    ) -> None:
+        super().__init__()
+        assert layer_id >= 0
+        self.layer_id = layer_id
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = (
+            config.intermediate_size
+            if isinstance(config.intermediate_size, int)
+            else config.intermediate_size[layer_id]
+        )
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False)
+        cla_factor = _get_cla_factor(config)
+        attention_type = (
+            AttentionType.ENCODER_DECODER if layer_id >= 0 and layer_id % cla_factor != 0 else AttentionType.DECODER
+        )
+        if attention_type == AttentionType.DECODER:
+            self.self_attn = HunYuanAttention(
+                config=config,
+                hidden_size=self.hidden_size,
+                num_heads=config.num_attention_heads,
+                num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
+                max_position_embeddings=max_position_embeddings,
+                quant_config=quant_config,
+                bias=attention_bias,
+                cache_config=cache_config,
+                prefix=f"{prefix}.self_attn",
+                layer_id=layer_id,
+            )
+        elif attention_type == AttentionType.ENCODER_DECODER:
+            self.self_attn = HunYuanCrossAttention(
+                config=config,
+                hidden_size=self.hidden_size,
+                num_heads=config.num_attention_heads,
+                num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
+                max_position_embeddings=max_position_embeddings,
+                quant_config=quant_config,
+                bias=attention_bias,
+                cache_config=cache_config,
+                prefix=f"{prefix}.self_attn",
+                layer_id=layer_id,
+            )
+        else:
+            raise RuntimeError(f"Unsupported attention type: {attention_type}")
+
+        if _is_moe(config):
+            self.mlp = HunYuanSparseMoeBlock(
+                config=config,
+                quant_config=quant_config,
+                layer_id=layer_id,
+                prefix=f"{prefix}.mlp",
+                enable_eplb=enable_eplb,
+            )
+        else:
+            self.mlp = HunYuanMLP(
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                bias=getattr(config, "mlp_bias", False),
+                prefix=f"{prefix}.mlp",
+            )
+
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        kv_states: tuple[torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        hidden_states, ori_kv_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_states=kv_states,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual, ori_kv_states
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (xd, seq_len) if xdrope is enabled for hunyuan-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    }
+)
+class HunYuanModel(nn.Module, EagleModelMixin):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        eplb_config = vllm_config.parallel_config.eplb_config
+        enable_eplb = vllm_config.parallel_config.enable_eplb
+        self.num_redundant_experts = eplb_config.num_redundant_experts
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.vocab_size = config.vocab_size
+
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: HunYuanDecoderLayer(
+                config=config,
+                layer_id=int(prefix.split(".")[-1]),
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+                enable_eplb=enable_eplb,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        cla_factor = _get_cla_factor(self.config)
+        prev_kv_states = None
+        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
+        for i, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)):
+            hidden_states, residual, kv_states = layer(
+                positions,
+                hidden_states,
+                residual,
+                prev_kv_states,
+            )
+
+            if getattr(self.config, "use_cla", False) and i % cla_factor == 0:
+                prev_kv_states = kv_states
+            else:
+                prev_kv_states = None
+
+            self._maybe_add_hidden_state(aux_hidden_states, i + 1, hidden_states, residual)
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        if len(aux_hidden_states) > 0:
+            return hidden_states, aux_hidden_states
+        return hidden_states

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import gc
 import math
 import typing
@@ -45,13 +45,6 @@ from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
-from vllm.model_executor.models.hunyuan_v1 import (
-    HunYuanMLP,
-    HunYuanModel,
-    HunYuanSparseMoeBlock,
-    _get_cla_factor,
-    _is_moe,
-)
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsLoRA,
@@ -63,6 +56,7 @@ from vllm.model_executor.models.interfaces import (
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    WeightsMapper,
     _merge_multimodal_embeddings,
     is_pp_missing_parameter,
     maybe_prefix,
@@ -78,7 +72,6 @@ from vllm.multimodal.parse import (
 )
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
-    BaseMultiModalProcessor,
     BaseProcessingInfo,
     PromptReplacement,
     PromptUpdate,
@@ -91,6 +84,14 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
+from vllm_omni.model_executor.models.hunyuan_image3._hunyuan_v1_vendored import (
+    HunYuanMLP,
+    HunYuanModel,
+    HunYuanSparseMoeBlock,
+    _get_cla_factor,
+    _is_moe,
+)
 from vllm_omni.model_executor.models.hunyuan_image3.autoencoder_kl_3d import AutoencoderKLConv3D
 from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjector, Siglip2VisionTransformer
 
@@ -1027,7 +1028,7 @@ class HunyuanImage3DummyInputsBuilder(BaseDummyInputsBuilder[HunyuanImage3Proces
         }
 
 
-class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3ProcessingInfo]):
+class HunyuanImage3MultiModalProcessor(OmniMultiModalProcessor[HunyuanImage3ProcessingInfo]):
     """Multimodal processor for HunyuanImage3 model."""
 
     def _call_hf_processor(
@@ -1046,15 +1047,6 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
             batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long)
         return batch_feature
 
-    def _hf_processor_applies_updates(
-        self,
-        prompt_text: str,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> bool:
-        return False
-
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -2180,12 +2172,11 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
             "timestep_r_emb",
         ]
         skip_prefixes.extend(unexpected_keywords)
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=skip_prefixes,
-        )
+        loader = AutoWeightsLoader(self)
 
-        loaded_params = loader.load_weights(weights)
+        loaded_params = loader.load_weights(
+            weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())})
+        )
         return loaded_params
 
     def get_language_model(self) -> torch.nn.Module:

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2025 Xiaomi Corporation.
 import os
 from collections.abc import Iterable, Mapping, Sequence
@@ -27,14 +30,13 @@ from vllm.multimodal.inputs import (
 from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, MultiModalDataParser
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
-    BaseMultiModalProcessor,
     BaseProcessingInfo,
     ProcessorInputs,
     PromptReplacement,
     PromptUpdate,
     PromptUpdateDetails,
 )
-from vllm.multimodal.utils import group_mm_kwargs_by_modality
+from vllm.multimodal.utils import group_and_batch_mm_kwargs
 from vllm.sequence import IntermediateTensors
 from vllm.utils.cache import LRUCache
 from vllm.utils.collection_utils import is_list_of
@@ -42,6 +44,7 @@ from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import (
     NO_INTERLEAVE_NEXT_TOKEN_ID,
@@ -348,7 +351,7 @@ class MiMoAudioDataParser(MultiModalDataParser):
         return AudioProcessorItems(new_audios)
 
 
-class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProcessingInfo]):
+class MiMoAudioLLMMultiModalProcessor(OmniMultiModalProcessor[MiMoAudioLLMProcessingInfo]):
     def _call_hf_processor(
         self,
         prompt: str,
@@ -484,7 +487,7 @@ class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProces
         return [
             PromptReplacement(
                 modality="audio",
-                target=audio_token,
+                target=[audio_token_id],
                 replacement=get_replacement_mimo_audio,
             )
         ]
@@ -626,7 +629,7 @@ class MiMoAudioForConditionalGeneration(
                 )
             mm_kwargs.append((mm_feature.modality, mm_item))
 
-        for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
+        for modality, num_items, mm_kwargs_group in group_and_batch_mm_kwargs(
             mm_kwargs,
             device=self.device,
             pin_memory=self.pin_memory,

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -390,35 +390,25 @@ class MimoAudioDummyInputsBuilder(BaseDummyInputsBuilder[MimoAudioProcessingInfo
 
 
 class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingInfo]):
-    def _call_hf_processor(
+    def _preprocess_hf_mm_data(
         self,
-        prompt: str,
         mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, Any],
-        tok_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+        mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data(
+            mm_data,
+            hf_processor_mm_kwargs,
+        )
+        mm_data = dict(mm_data)
         audios = mm_data.pop("audios", [])
         if audios:
             mm_data["audio"] = audios
-
-        # Text-only input not supported in composite processor
-        if not mm_data.get("audio", []):
-            prompt_ids = self.info.get_tokenizer().encode(prompt)
-            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
-            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
-
-        feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
-        mm_kwargs = dict(
-            **mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
-        )
-
-        return super()._call_hf_processor(
-            prompt=prompt,
-            mm_data=mm_data,
-            mm_kwargs=mm_kwargs,
-            tok_kwargs=tok_kwargs,
-        )
+            feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
+            hf_processor_mm_kwargs = {
+                **hf_processor_mm_kwargs,
+                "sampling_rate": feature_extractor.sampling_rate,
+            }
+        return mm_data, hf_processor_mm_kwargs
 
     def _get_mm_fields_config(
         self,
@@ -480,7 +470,7 @@ class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingIn
         return [
             PromptReplacement(
                 modality="audio",
-                target=audio_token,
+                target=[audio_token_id],
                 replacement=get_replacement_mimo_audio,
             )
         ]