Merge branch 'main' into codex-diffusion-shutdown-escalation

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

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

.buildkite/cuda/test-merge.yml

@@ -155,6 +155,7 @@ steps:
         source_file_dependencies:
           - tests/e2e/offline_inference/test_minicpmo_4_5.py
           - tests/e2e/online_serving/test_minicpmo_4_5.py
+          - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja
           - vllm_omni/model_executor/models/minicpmo_4_5/
           - vllm_omni/model_executor/models/step_audio2/
           - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/

.buildkite/cuda/test-ready.yml

@@ -188,6 +188,7 @@ steps:
       - label: "Omni · MiniCPM-o 4.5 Online Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_minicpmo_4_5.py
+          - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja
           - vllm_omni/model_executor/models/minicpmo_4_5/
           - vllm_omni/model_executor/models/step_audio2/
           - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/

docker/Dockerfile.ci

@@ -1,27 +1,10 @@
-# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5).
+# Target vLLM version: v0.29.0 (tag @ 98dff2a81d74).
-# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
+# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer
-#
+# runtime. Do not reinstall a commit wheel or repair dependencies on top of
-# Release-boundary form: there is no precompiled-wheel pin here, because the
+# an older base image at a release boundary. For a future unreleased target,
-# base image already ships exactly the target vLLM. Everything the runtime
+# restore an explicit wheel pin and validate its complete runtime stack.
-# 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.
 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)
+#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.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
+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
 

pyproject.toml

@@ -161,7 +161,7 @@ where = ["."]
 include = ["vllm_omni*"]
 
 [tool.setuptools.package-data]
-"vllm_omni" = ["_version.py", "py.typed"]
+"vllm_omni" = ["_version.py", "py.typed", "transformers_utils/chat_templates/*.jinja"]
 
 [tool.setuptools_scm]
 # no extra settings needed, presence enables setuptools-scm

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -28,13 +28,13 @@ also provided.
 
 - Default deploy configs (auto-loaded by HF `model_type=minicpmo` +
   `hf_config.version="4.5"`):
-  - Default single-GPU compatibility layout (auto-loaded):
+    - Default single-GPU compatibility layout (auto-loaded):
     [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
-  - Recommended 2-GPU continuous-batching layout:
+    - Recommended 2-GPU continuous-batching layout:
     [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml),
-  - 3-GPU layout:
+    - 3-GPU layout:
     [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
-  - 8x RTX 4090 layout:
+    - 8x RTX 4090 layout:
     [`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
 - Online example + Gradio demo:
   [`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/)
@@ -126,9 +126,20 @@ video. Use the two-GPU profile for production throughput.
 ```bash
 vllm serve openbmb/MiniCPM-o-4_5 --omni \
     --trust-remote-code \
+    --chat-template vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja \
+    --chat-template-content-format openai \
     --host 0.0.0.0 --port 8099
 ```
 
+Run these commands from the repository root. The bundled chat template
+matches the Hugging Face `chat(omni_mode=True)` content layout: media and text
+parts are concatenated without adding separators, while whitespace explicitly
+included in a text part is preserved. This matters for speech generation because
+an inserted newline changes the Thinker condition passed to Talker. For image
+questions, place the image before the question, as in the native image-chat
+layout. Continue to request `use_tts_template: true` and
+`enable_thinking: false` for spoken answers.
+
 The deploy config is auto-loaded by the model registry — no
 `--deploy-config` flag needed for this default single-GPU layout.
 
@@ -271,6 +282,8 @@ uses GPU 5. GPUs 6–7 are left free.
 vllm serve openbmb/MiniCPM-o-4_5 --omni \
     --deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
     --trust-remote-code \
+    --chat-template vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja \
+    --chat-template-content-format openai \
     --host 0.0.0.0 --port 8099
 ```
 

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
+    # Sixteen physical pages provide ``(16 - 1) * block_size`` tokens: vLLM 0.29
-    # backend chooses the block geometry, so keep this assertion backend
+    # 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=[
+        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])],
-            KVCacheTensor(
-                size=canonical_spec.page_size_bytes * 4,
-                shared_by=[_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://")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://", backend="gloo")
-    initialize_model_parallel()
+    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/diffusion/test_multiproc_engine_concurrency.py

@@ -4,8 +4,10 @@
 import asyncio
 import multiprocessing as mp
 import queue
+import signal
 import threading
 import time
+import weakref
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock
 
@@ -1330,6 +1332,7 @@ class TestExecutorShutdownCleaner:
         class FakeProcess:
             def __init__(self, name):
                 self.name = name
+                self.pid = 123
                 self.alive = True
                 self.terminated = False
                 self.join_timeouts = []
@@ -1358,6 +1361,132 @@ class TestExecutorShutdownCleaner:
         assert first.terminated and second.terminated
         assert not first.is_alive() and not second.is_alive()
 
+    @pytest.mark.parametrize("cooperative", [False, True])
+    def test_real_workers_exit_and_are_reaped(self, monkeypatch, cooperative):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        ctx = mp.get_context("fork")
+        ready, stop = ctx.Event(), ctx.Event()
+
+        def run_worker():
+            signal.signal(signal.SIGTERM, signal.SIG_IGN)
+            ready.set()
+            stop.wait()
+
+        dead = ctx.Process(target=lambda: None)
+        worker = ctx.Process(target=run_worker)
+        monkeypatch.setattr(executor_module, "_WORKER_SHUTDOWN_GRACE_S", 0.2)
+        monkeypatch.setattr(executor_module, "_WORKER_TERMINATE_GRACE_S", 0.1)
+        monkeypatch.setattr(executor_module, "_WORKER_KILL_GRACE_S", 5.0, raising=False)
+        dead.start()
+        worker.start()
+        try:
+            dead.join(5)
+            assert not dead.is_alive()
+            assert ready.wait(5), "worker did not install its signal handler"
+            mq = Mock()
+            if cooperative:
+                mq.enqueue.side_effect = lambda *args, **kwargs: stop.set()
+            else:
+                mq.enqueue.side_effect = OSError("shutdown queue unavailable")
+            cleaner = executor_module._ExecutorShutdownCleaner(mq, 2, [dead, worker])
+
+            cleaner()
+
+            assert not worker.is_alive()
+            assert worker.exitcode == (0 if cooperative else -signal.SIGKILL)
+            assert cleaner.processes == []
+            assert cleaner.broadcast_mq is None
+            cleaner()
+        finally:
+            for proc in (dead, worker):
+                if proc.is_alive():
+                    proc.kill()
+                proc.join(5)
+                proc.close()
+
+    @pytest.mark.parametrize("failed_action", ["terminate", "kill"])
+    def test_kill_phase_shares_deadline_and_continues_after_os_errors(self, monkeypatch, failed_action):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        first, second = Mock(pid=101), Mock(pid=102)
+        first.name, second.name = "first", "second"
+        for proc in (first, second):
+            proc.is_alive.side_effect = [True, True, True, False]
+        getattr(first, failed_action).side_effect = OSError("signal failed")
+        first.join.side_effect = [OSError("join failed"), None, None]
+        monotonic = Mock(side_effect=[100, 100, 110, 120, 120, 124, 130, 130, 134])
+        monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
+
+        cleaner()
+
+        first.kill.assert_called_once()
+        second.kill.assert_called_once()
+        assert [call.args[0] for call in first.join.call_args_list] == [15, 5, 5]
+        assert [call.args[0] for call in second.join.call_args_list] == [5, 1, 1]
+        assert cleaner.processes == []
+
+    def test_executor_retains_survivor_for_explicit_shutdown_retry(self, monkeypatch):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        survivor = Mock(pid=123)
+        survivor.name = "surviving-worker"
+        survivor.is_alive.return_value = True
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[survivor])
+        executor, _, _ = _make_executor()
+        executor._shutdown_cleaner = cleaner
+        executor._processes = [survivor]
+        executor._finalizer = weakref.finalize(executor, cleaner)
+        executor._pump_stop = threading.Event()
+        executor._futures_lock = threading.RLock()
+        executor._rpc_futures = {}
+        executor._output_futures = {}
+        executor._batch_split_map = {}
+        log_error = Mock()
+        monkeypatch.setattr(executor_module.logger, "error", log_error)
+
+        executor.shutdown()
+
+        assert not executor._finalizer.alive
+        assert executor._shutdown_cleaner is cleaner
+        assert executor._processes == [survivor]
+        log_error.assert_called_once()
+        assert log_error.call_args.args[1] == [("surviving-worker", 123)]
+        survivor.is_alive.return_value = False
+
+        executor.shutdown()
+        executor.shutdown()
+
+        assert executor._shutdown_cleaner is None
+        assert executor._processes == []
+
+    def test_concurrent_cleanup_does_not_repeat_process_operations(self):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        joining, release = threading.Event(), threading.Event()
+        proc = Mock(pid=123)
+        proc.is_alive.return_value = True
+
+        def join(timeout):
+            joining.set()
+            assert release.wait(5)
+            proc.is_alive.return_value = False
+
+        proc.join.side_effect = join
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[proc])
+        thread = threading.Thread(target=cleaner)
+        thread.start()
+        try:
+            assert joining.wait(5)
+            cleaner()
+            proc.join.assert_called_once()
+        finally:
+            release.set()
+            thread.join(5)
+        assert not thread.is_alive()
+        assert cleaner.processes == []
+
 
 # ───────── monitor thread & death sentinel integration tests ─────────
 
@@ -1406,7 +1535,7 @@ class TestMultiprocExecutorWorkerMonitor:
         executor._result_mq = None
         executor._shutdown_cleaner = None
         # Use a no-op so shutdown() doesn't crash on None resources.
-        executor._finalizer = lambda: None
+        executor._finalizer = weakref.finalize(executor, lambda: None)
         # ------------------------------------------------------------------
         # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates
         # over them, so they need to exist even when constructed via __new__.
@@ -1441,7 +1570,7 @@ class TestMultiprocExecutorWorkerMonitor:
         executor._broadcast_mq = None
         executor._result_mq = None
         executor._shutdown_cleaner = None
-        executor._finalizer = lambda: None
+        executor._finalizer = weakref.finalize(executor, lambda: None)
 
         proc = _make_short_lived_process()
         executor._processes = [proc]

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/e2e/online_serving/test_minicpmo_4_5.py

@@ -24,6 +24,9 @@ from tests.helpers.stage_config import get_deploy_config_path, modify_stage_conf
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 _MODEL = "openbmb/MiniCPM-o-4_5"
+_NATIVE_CHAT_TEMPLATE = (
+    Path(__file__).resolve().parents[3] / "vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja"
+)
 _CI_DEPLOY = modify_stage_config(
     get_deploy_config_path("minicpmo_4_5.yaml"),
     updates={
@@ -63,7 +66,14 @@ test_params = [
             model=_MODEL,
             stage_config_path=_CI_DEPLOY,
             use_stage_cli=False,
-            server_args=["--trust-remote-code", "--async-chunk"],
+            server_args=[
+                "--trust-remote-code",
+                "--async-chunk",
+                "--chat-template",
+                str(_NATIVE_CHAT_TEMPLATE),
+                "--chat-template-content-format",
+                "openai",
+            ],
             env_dict={
                 "VLLM_CONFIGURE_LOGGING": "1",
                 "VLLM_LOGGING_CONFIG_PATH": str(_PROMPT_LOGGING_CONFIG_PATH),
@@ -265,11 +275,18 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None:
     Input Setting: stream=True
     """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(24, 24)['base64']}"
-    messages = dummy_messages_from_mix_data(
+    # Match the native image-chat layout explicitly; do not rely on the
-        system_prompt=get_system_prompt(),
+    # renderer moving media placeholders ahead of text.
-        image_data_url=image_data_url,
+    messages = [
-        content_text=get_prompt("text_image"),
+        get_system_prompt(),
-    )
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": image_data_url}},
+                {"type": "text", "text": get_prompt("text_image")},
+            ],
+        },
+    ]
 
     request_config = {
         "model": omni_server.model,

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:
+    class _Base(BaseRenderer):
-        def __init__(self, vllm_config, renderer=None):
+        def __init__(self, config, tokenizer):
-            self.vllm_config = vllm_config
+            self.config, self.tokenizer = config, tokenizer
-            self.renderer = renderer or object()
-            self.input_preprocessor = None
 
-    class DummyOmniInputPreprocessor:
+        def render_messages(self, messages, params):  # pragma: no cover - abstract stub
-        def __init__(self, vllm_config, renderer=None):
+            raise NotImplementedError
-            self.vllm_config = vllm_config
-            self.renderer = renderer
 
-    monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
+    config = types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {}))
-    monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)
+    built = omni_renderer_cls(_Base)(config, "tok")
-
-    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()
     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(
+    monkeypatch.setattr(init_mod, "build_omni_renderer", lambda cfg: built if cfg is config else None)
-        init_mod,
+    processor = build_stage0_input_processor(config)
-        "_build_token_only_renderer",
+    assert seen["renderer"] is built
-        lambda _config: token_only_renderer,
+    assert isinstance(processor.renderer, OmniRenderer)
-    )
+    assert isinstance(processor.renderer, _Base)
-    monkeypatch.setattr(
+    assert not hasattr(processor, "input_preprocessor")
-        init_mod,
-        "OmniInputPreprocessor",
-        DummyOmniInputPreprocessor,
-    )
 
+
+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(
+        model_config=types.SimpleNamespace(skip_tokenizer_init=True, try_get_generation_config=lambda: {})
-            skip_tokenizer_init=True,
-            try_get_generation_config=lambda: {},
-        )
     )
-    build_stage0_input_processor(config)
+    processor = build_stage0_input_processor(config)
-
+    assert processor.renderer is original
-    assert seen["renderer"] is token_only_renderer
-    assert seen["preprocessor_renderer"] is token_only_renderer
 
 
 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:
+class _Renderer(BaseRenderer):
-    """Presence-based routing: an explicit empty ``mm_processor_kwargs`` dict
+    def __init__(self):
-    on the prompt still routes through ``_process_multimodal``. Required for
+        self.model_config = SimpleNamespace(is_encoder_decoder=False, enable_prompt_embeds=True)
-    AR-based image generation (e.g. GLM-Image t2i) where the HF processor
+        self.tokenizer = None
-    supplies the image-generation scaffold from its own defaults and no extra
+        self.default_cmpl_tok_params = TokenizeParams(max_total_tokens=128)
-    kwargs are needed from the caller.
+        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 render_messages(self, messages, params):
-    def preprocessor(self, mocker: MockerFixture):
+        return messages, {"prompt": "hello"}
-        from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 
-        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):
+@pytest.fixture
-        preprocessor._process_text({"prompt": "hello", "mm_processor_kwargs": {}})
+def renderer():
-        assert preprocessor._process_multimodal.called
+    return omni_renderer_cls(_Renderer)()
-        assert not preprocessor._tokenize_prompt.called
 
-    def test_missing_mm_processor_kwargs_routes_to_tokenize(self, preprocessor):
+
-        preprocessor._process_text({"prompt": "hello"})
+@pytest.mark.parametrize("kwargs", [{}, {"target_h": 512, "target_w": 768}])
-        assert preprocessor._tokenize_prompt.called
+@pytest.mark.parametrize("tokenized", [False, True])
-        assert not preprocessor._process_multimodal.called
+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_chat_template.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MiniCPM structured-content rendering must preserve native part boundaries."""
+
+from pathlib import Path
+
+import pytest
+from jinja2 import TemplateError
+from transformers.utils.chat_template_utils import _compile_jinja_template
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TEMPLATE = Path(__file__).resolve().parents[4] / "vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja"
+
+
+@pytest.mark.parametrize(
+    "parts,expected",
+    [
+        ([{"type": "image"}, {"type": "text", "text": "Question?"}], "(<image>./</image>)Question?"),
+        ([{"type": "text", "text": "Question?"}, {"type": "image"}], "Question?(<image>./</image>)"),
+        ([{"type": "image"}, {"type": "text", "text": "\n\nQuestion?"}], "(<image>./</image>)\n\nQuestion?"),
+        ([{"type": "text", "text": " A "}, {"type": "text", "text": " B "}], " A  B "),
+        ([{"type": "audio"}, {"type": "text", "text": "Transcribe."}], "(<audio>./</audio>)Transcribe."),
+        ([{"type": "video"}, {"type": "image"}], "(<video>./</video>)(<image>./</image>)"),
+        ("\nKeep this whitespace.\n", "\nKeep this whitespace.\n"),
+    ],
+)
+def test_native_content_concatenation(parts, expected):
+    rendered = _compile_jinja_template(TEMPLATE.read_text()).render(
+        messages=[{"role": "user", "content": parts}],
+        add_generation_prompt=True,
+        enable_thinking=False,
+        use_tts_template=True,
+    )
+    assert rendered == (
+        f"<|im_start|>user\n{expected}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n<|tts_bos|>"
+    )
+
+
+def test_unsupported_part_is_not_silently_discarded():
+    with pytest.raises(TemplateError, match="Unsupported MiniCPM-o content part"):
+        _compile_jinja_template(TEMPLATE.read_text()).render(
+            messages=[{"role": "user", "content": [{"type": "unknown"}]}],
+            add_generation_prompt=True,
+        )

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
+    # Generous timeout: cold CI containers can take 30-50s just to import
-    # import vllm_omni → vllm → torch before the subprocess does any work.
+    # vllm_omni → vllm → torch, and heavily loaded shared machines have
-    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120, env=env)
+    # 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/executor/multiproc_executor.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 concurrent.futures
@@ -10,7 +13,7 @@ import threading
 import time
 import weakref
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from multiprocessing.synchronize import Event
 from typing import TYPE_CHECKING, Any, cast
 
@@ -42,6 +45,7 @@ _DEQUEUE_TIMEOUT_S = 5.0
 _DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0))
 _WORKER_SHUTDOWN_GRACE_S = 15.0
 _WORKER_TERMINATE_GRACE_S = 5.0
+_WORKER_KILL_GRACE_S = 5.0
 _RESULT_PUMP_JOIN_TIMEOUT_S = 2.0
 
 
@@ -80,33 +84,61 @@ class _ExecutorShutdownCleaner:
     broadcast_mq: MessageQueue | None = None
     num_workers: int = 0
     processes: list[mp.Process] | None = None
+    _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
 
     def __call__(self) -> None:
         """Clean up background resources."""
+        # The worker monitor and an explicit shutdown may race. A retry must
+        # not signal/join the same Process objects while cleanup is in flight.
+        if not self._lock.acquire(blocking=False):
+            return
+        try:
+            self._cleanup()
+        finally:
+            self._lock.release()
+
+    def _cleanup(self) -> None:
         if self.broadcast_mq is not None:
             try:
                 for _ in range(self.num_workers):
                     self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0)
 
-                self.broadcast_mq = None
             except Exception as exc:
                 logger.warning("Failed to send shutdown signal: %s", exc)
+            finally:
+                self.broadcast_mq = None
 
         if self.processes:
-            join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S
-            for proc in self.processes:
-                if not proc.is_alive():
-                    continue
-                proc.join(max(0.0, join_deadline - time.monotonic()))
-
             alive = [proc for proc in self.processes if proc.is_alive()]
-            for proc in alive:
+            for action, grace in (
-                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
+                (None, _WORKER_SHUTDOWN_GRACE_S),
-                proc.terminate()
+                ("terminate", _WORKER_TERMINATE_GRACE_S),
+                ("kill", _WORKER_KILL_GRACE_S),
+            ):
+                if not alive:
+                    break
+                if action is not None:
+                    for proc in alive:
+                        try:
+                            logger.warning("Calling %s on diffusion worker %s (pid=%s)", action, proc.name, proc.pid)
+                            getattr(proc, action)()
+                        except OSError:
+                            logger.exception("Failed to %s diffusion worker %s (pid=%s)", action, proc.name, proc.pid)
 
-            terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S
+                deadline = time.monotonic() + grace
-            for proc in alive:
+                for proc in alive:
-                proc.join(max(0.0, terminate_deadline - time.monotonic()))
+                    try:
+                        proc.join(max(0.0, deadline - time.monotonic()))
+                    except OSError:
+                        logger.exception("Failed to join diffusion worker %s (pid=%s)", proc.name, proc.pid)
+                alive = [proc for proc in alive if proc.is_alive()]
+
+            self.processes = alive
+            if alive:
+                logger.error(
+                    "Diffusion worker cleanup incomplete after kill: %s; retaining processes for shutdown retry",
+                    [(proc.name, proc.pid) for proc in alive],
+                )
 
 
 class MultiprocDiffusionExecutor(DiffusionExecutor):
@@ -1006,8 +1038,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
     def shutdown(self) -> None:
         self._closed = True
         self._pump_stop.set()
+        cleaner = self._shutdown_cleaner
         try:
-            self._finalizer()
+            if self._finalizer.alive:
+                self._finalizer()
+            elif cleaner is not None:
+                cleaner()
         finally:
             pump_threads = getattr(self, "_result_pump_threads", [])
             for thread in pump_threads:
@@ -1031,5 +1067,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 self._rpc_futures.clear()
                 self._output_futures.clear()
                 self._batch_split_map.clear()
-            self._shutdown_cleaner = None
+            self._processes = (cleaner.processes or []) if cleaner is not None else []
-            self._processes = []
+            if not self._processes:
+                self._shutdown_cleaner = None

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(
+        loader = AutoWeightsLoader(self)
-            self,
+        return loader.load_weights(
-            skip_prefixes=skip_prefixes,
+            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)
+        if img_query.device.type == "cuda":
-        img_sin = torch.imag(vid_freqs).to(img_query.dtype)
+            img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs)
-        txt_cos = torch.real(txt_freqs).to(txt_query.dtype)
+            img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs)
-        txt_sin = torch.imag(txt_freqs).to(txt_query.dtype)
+            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_query = self.rope(img_query, img_cos, img_sin)
-        img_key = self.rope(img_key, 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_query = self.rope(txt_query, txt_cos, txt_sin)
-        txt_key = self.rope(txt_key, 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:
+    def add_cli_args(cls, parser: argparse.ArgumentParser, async_args_only: bool = False) -> argparse.ArgumentParser:
-        parser = AsyncEngineArgs.add_cli_args(parser)
+        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(
+        renderer = _build_token_only_renderer(stage_vllm_config)
-            vllm_config=stage_vllm_config,
-            renderer=_build_token_only_renderer(stage_vllm_config),
-        )
     else:
-        input_processor = InputProcessor(vllm_config=stage_vllm_config)
+        renderer = build_omni_renderer(stage_vllm_config)
-    input_processor.input_preprocessor = OmniInputPreprocessor(
+    return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer)
-        vllm_config=stage_vllm_config,
-        renderer=input_processor.renderer,
-    )
-    return input_processor
 
 
 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
-
+from vllm.entrypoints.serve import create_error_response
-# vLLM < 0.28 keeps create_error_response under serve.utils (it is not
+from vllm.entrypoints.serve.engine.protocol import (
-# re-exported from vllm.entrypoints.serve); 0.28+ moved it under
+    ErrorResponse,
-# serve.exception_handling and re-exports it from the package root.
+    ModelCard,
-try:
+    ModelList,
-    from vllm.entrypoints.serve import create_error_response
+    ModelPermission,
-except ImportError:
+)
-    from vllm.entrypoints.serve.utils.error_response import create_error_response
 
 # 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.launchers.launcher import terminate_if_errored
-from vllm.entrypoints.openai.engine.protocol import (
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
-    ErrorResponse,
-    RequestResponseMetadata,
-)
 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
+``OmniRenderer`` is a mixin. ``omni_renderer_cls`` puts it in front of a
-from vllm.inputs import EmbedsInput, MultiModalInput, SingletonInput
+concrete ``BaseRenderer`` subclass, so the Omni instance *is* an upstream
-from vllm.inputs.preprocess import InputPreprocessor
+renderer (tokenizer, executors, multimodal cache and every public entry point
-from vllm.logger import init_logger
+come from upstream unchanged) and only ``_process_singleton`` /
-from vllm.renderers.inputs import SingletonDictPrompt
+``_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 (
+from __future__ import annotations
-    OmniEmbedsPrompt,
-    OmniTextPrompt,
-    OmniTokenInputs,
-    OmniTokensPrompt,
-    token_inputs_omni,
-)
 
-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):
+class OmniRenderer:
-    """Input preprocessor for omni models.
+    """Mixin that routes Omni prompts through the upstream renderer.
 
-    Extends the base InputPreprocessor to handle omni-specific input
+    It has no base class of its own; ``omni_renderer_cls`` prepends it to a
-    types including prompt embeddings and additional information payloads.
+    concrete ``BaseRenderer`` subclass, so ``self`` and ``super()`` are that
-    Supports processing tokens, embeddings, text, and multimodal inputs.
+    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(
+    @staticmethod
-        self,
+    def _with_omni_extras(inputs, prompt):
-        parsed_content: OmniTextPrompt,
+        for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"):
-        tokenization_kwargs: dict[str, Any] | None = None,
+            if key in prompt:
-        *,
+                inputs[key] = prompt[key]
-        mm_uuids: Any | None = None,
+        return inputs
-    ) -> OmniTokenInputs | MultiModalInput:
-        """Process text prompts with support for mm_processor_kwargs.
 
-        Extends base class to support mm_processor_kwargs without multi_modal_data.
+    @staticmethod
-        This is needed for models like GLM-Image where text-to-image generation
+    def _routes_no_media_kwargs(prompt) -> bool:
-        requires processor kwargs (target_h, target_w) to format the prompt.
+        return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data")
-        """
-        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")
 
-        inputs: OmniTokenInputs | MultiModalInput
+    def _process_singleton(self, prompt, *, skip_mm_cache: bool = False):
-        if multi_modal_data := parsed_content.get("multi_modal_data"):
+        renderer = cast(BaseRenderer, self)
-            inputs = self._process_multimodal(
+        if self._routes_no_media_kwargs(prompt):
-                prompt_text,
+            inputs = renderer._process_multimodal(
-                multi_modal_data,
+                prompt["prompt_token_ids"],
-                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,
                 {},
-                mm_processor_kwargs,
+                mm_processor_kwargs=prompt["mm_processor_kwargs"],
-                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=prompt.get("multi_modal_uuids"),
-                mm_uuids=effective_mm_uuids,
+                skip_mm_cache=skip_mm_cache,
             )
         else:
-            prompt_token_ids = self._tokenize_prompt(
+            inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
-                prompt_text,
+        return self._with_omni_extras(inputs, prompt)
-                tokenization_kwargs=tokenization_kwargs,
+
+    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(
+            inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache)
-                prompt_token_ids=prompt_token_ids,
+        return self._with_omni_extras(inputs, prompt)
-                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
 
-        return inputs
 
-    def _process_embeds(
+_OMNI_RENDERER_CLASSES: dict[type[BaseRenderer], type[BaseRenderer]] = {}
-        self,
-        parsed_content: OmniEmbedsPrompt,
-    ) -> EmbedsInput:
-        """Process embeddings prompt with omni-specific extensions.
 
-        Extends base _process_embeds to handle additional_information payload
-        for direct transfer between pipeline stages.
-        """
-        # Call parent implementation for base embeds processing
-        inputs = super()._process_embeds(parsed_content)
 
-        # Add omni-specific additional_information if present
+def omni_renderer_cls(base_cls: type[BaseRenderer]) -> type[BaseRenderer]:
-        additional_information = parsed_content.get("additional_information")
+    """Return the Omni subclass of ``base_cls`` (created once per base class)."""
-        if additional_information is not None:
+    if issubclass(base_cls, OmniRenderer):
-            inputs["additional_information"] = additional_information  # type: ignore[typeddict-unknown-key]
+        return base_cls
-        model_intermediate_buffer = parsed_content.get("model_intermediate_buffer")
+    cls = _OMNI_RENDERER_CLASSES.get(base_cls)
-        if model_intermediate_buffer is not None:
+    if cls is None:
-            inputs["model_intermediate_buffer"] = model_intermediate_buffer  # type: ignore[typeddict-unknown-key]
+        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(
+def build_omni_renderer(
-        self,
+    vllm_config: VllmConfig,
-        prompt: SingletonDictPrompt,
+    *,
-        tokenization_kwargs: dict[str, Any] | None = None,
+    renderer_cls: type[BaseRenderer] | None = None,
-        *,
+    tokenizer: TokenizerLike | None = None,
-        mm_uuids: Any | None = None,
+) -> BaseRenderer:
-    ) -> SingletonInput:
+    """Build the Omni renderer for ``vllm_config``.
-        """
-        Extract the singleton inputs from a prompt.
 
-        Arguments:
+    With ``renderer_cls`` unset this mirrors upstream ``renderer_from_config``:
-
+    the tokenizer and renderer mode come from the model config and the class
-        * prompt: single encoder or decoder input prompt
+    from ``RENDERER_REGISTRY``. Pass ``renderer_cls`` (and optionally
-
+    ``tokenizer``) to skip that resolution, e.g. for tokenizer-less stages.
-        Returns:
+    """
-
+    model_config = vllm_config.model_config
-        * [`SingletonInput`][vllm.inputs.engine.SingletonInput] instance
+    if renderer_cls is None:
-        """
+        if tokenizer is None:
-        if "prompt_embeds" in prompt:
+            tokenizer = cached_tokenizer_from_config(model_config)
-            return self._process_embeds(prompt)  # type: ignore[arg-type]
+        renderer_mode, *_ = tokenizer_args_from_config(model_config)
-
+        renderer_cls = RENDERER_REGISTRY.load_renderer_cls(renderer_mode)
-        if "prompt_token_ids" in prompt:
+    return omni_renderer_cls(renderer_cls)(vllm_config, tokenizer)
-            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]

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],
+    ) -> BatchFeature:
-    ) -> bool:
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
-        return False
+        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,
+            self, ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."]
-            skip_prefixes=["vit_pos_embed.pos_embed"],
+        )
-            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]):
+class CovoAudioMultiModalProcessor(OmniMultiModalProcessor[CovoAudioProcessingInfo]):
-    def _hf_processor_applies_updates(self, prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs) -> bool:
-        return False
-
     def _call_hf_processor(
         self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], tok_kwargs: Mapping[str, object]
     ) -> BatchFeature: