Synchronize CPU ring workers before group teardown

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
(cherry picked from commit 4ad71d8270)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/common/scripts/upload_pipeline.py

@@ -102,7 +102,6 @@ CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_depe
 SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback"
 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'
-CUDA_RUNTIME_CHECK = "python3 .buildkite/cuda/scripts/check_vllm_runtime.py"
 
 # Bootstrap Buildkite ``if`` expressions.
 # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
@@ -621,11 +620,11 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
     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, CUDA_RUNTIME_CHECK, *commands]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands]
         elif commands is None:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT]
         else:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, commands]
+            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/scripts/check_vllm_runtime.py

@@ -1,38 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Verify the pinned vLLM and compiled CUDA imports on a GPU CI worker."""
-
-from __future__ import annotations
-
-import importlib
-import re
-from pathlib import Path
-
-
-def main() -> None:
-    dockerfile = Path(__file__).resolve().parents[3] / "docker" / "Dockerfile.ci"
-    match = re.search(
-        r"^ARG VLLM_VERSION=([0-9]+\.[0-9]+\.[0-9]+)$",
-        dockerfile.read_text(),
-        re.MULTILINE,
-    )
-    if match is None:
-        raise RuntimeError(f"Missing vLLM release version in {dockerfile}")
-    expected_version = match.group(1)
-
-    import torch
-    import vllm
-
-    assert vllm.__version__ == expected_version, vllm.__version__
-    assert torch.cuda.is_available(), "vLLM runtime check requires a GPU worker"
-    for module in ("flashinfer", "vllm._C_stable_libtorch"):
-        importlib.import_module(module)
-    print(
-        f"Verified vLLM CUDA runtime: vllm={vllm.__version__} "
-        f"expected={expected_version} torch={torch.__version__} "
-        f"cuda={torch.version.cuda} device={torch.cuda.get_device_name()}"
-    )
-
-
-if __name__ == "__main__":
-    main()

.buildkite/cuda/test-nightly.yml

@@ -48,8 +48,6 @@ steps:
       - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
         source_file_dependencies: omni_qwen3_omni_doc
         timeout_in_minutes: 90
-        artifact_paths:
-          - audio_*.wav
         commands:
           - pytest -s -v tests/examples/ -m "full_model and omni and H100 and B200 and cards_2" --run-level "full_model"
 
@@ -158,11 +156,6 @@ steps:
       - label: ":full_moon: TTS · Function Test with L4"
         source_file_dependencies: tts_qwen3_tts_function
         timeout_in_minutes: 120
-        env:
-          VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR: .
-        artifact_paths:
-          - failed_speech_*.wav
-          - failed_speech_*.json
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"

docker/Dockerfile.ci

@@ -1,7 +1,10 @@
-# Use the official release image with its matching CUDA/PyTorch stack.
+# Target vLLM version: v0.30.0 (tag @ ced6857afa0e).
-ARG VLLM_VERSION=0.31.0
+# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer
+# runtime. Do not reinstall a commit wheel or repair dependencies on top of
+# an older base image at a release boundary. For a future unreleased target,
+# restore an explicit wheel pin and validate its complete runtime stack.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v${VLLM_VERSION}
+ARG VLLM_BASE_TAG=v0.30.0
 
 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents
@@ -9,7 +12,6 @@ ARG VLLM_BASE_TAG=v${VLLM_VERSION}
 # everything except the actual project dir, since that should not
 # change much across branches.
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} AS deps
-ARG VLLM_VERSION
 # Keep APP_DIR outside /workspace so k8s Buildkite workspace mounts cannot
 # shadow the editable install target (#5364). Docker WORKDIR creates /opt if
 # needed.
@@ -34,17 +36,12 @@ RUN apt-get update && \
 # for CI
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-# Scope the override to Omni; dependency source builds need their own versions.
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
-ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
 # Editable install: deps + project metadata are cached in this layer, while
 # the package itself is linked to APP_DIR (populated later by COPY --link).
 # Do NOT use a non-editable install here — that copies an empty stub into
 # site-packages and breaks ``import vllm_omni.entrypoints``.
 RUN uv pip install --system -e ".[dev]"
-# Fail the image build if dependency resolution replaced the selected vLLM.
-# The CPU builder has no NVIDIA driver; CUDA imports are checked at the start
-# of each GPU test job by .buildkite/cuda/scripts/check_vllm_runtime.py.
-RUN python3 -c "import vllm; assert vllm.__version__ == '${VLLM_VERSION}', vllm.__version__"
 
 # The final stage just copies the working dir on top of the
 # dependencies image. This uses --link so that we can skip

docker/Dockerfile.ci.dockerignore

@@ -1,6 +1,6 @@
 # NOTE: this is largely taken from vLLM
 
-# Exclude .git for CI builds — version is set via VLLM_OMNI_VERSION_OVERRIDE
+# Exclude .git for CI builds — version is set via SETUPTOOLS_SCM_PRETEND_VERSION
 .git
 .venv
 /build

docker/Dockerfile.cuda

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

docker/Dockerfile.rocm

@@ -1,7 +1,7 @@
 # Keep the ROCm runtime aligned with the vLLM release targeted by
 # docker/Dockerfile.ci. A stale base image can build successfully while failing
 # at runtime when vLLM-Omni imports APIs added by the newer vLLM release.
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0
 
 # Keep the dependency-heavy portion independent from the source tree. The
 # registry cache is keyed on these inputs, while the final source copy can be
@@ -12,11 +12,11 @@ FROM ${BASE_IMAGE} AS deps
 # BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0.
 # VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when
 # USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's
-# release version. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
+# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
 # intentional custom source pin must update that regression in the same change,
 # and should return to the stable release as soon as one is available.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=v0.31.0
+ARG VLLM_VERSION_OR_COMMIT_HASH=v0.30.0
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 
@@ -75,8 +75,7 @@ WORKDIR ${COMMON_WORKDIR}/vllm-omni
 RUN mkdir -p vllm_omni && touch vllm_omni/__init__.py
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-# Scope the override to Omni; dependency source builds need their own versions.
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
-ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
 
 RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
         --no-cache-dir -e ".[dev]" --no-build-isolation && \

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.31 configuration guide](https://docs.vllm.ai/en/v0.31.0/configuration/index.html).
+For options within a vLLM Engine, please refer to the [vLLM 0.30 configuration guide](https://docs.vllm.ai/en/v0.30.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

@@ -186,7 +186,7 @@ their keys only.
 ## Inherited vLLM variables
 
 vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
-the [vLLM 0.31 environment-variable reference](https://docs.vllm.ai/en/v0.31.0/configuration/env_vars.html)
+the [vLLM 0.30 environment-variable reference](https://docs.vllm.ai/en/v0.30.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. This source checkout targets vLLM 0.31.0. The published vLLM-Omni 0.30.0 wheel instructions use vLLM 0.30.0; choose the installation path that matches your Omni revision.
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x; the source-install instructions pin vLLM 0.30.0.
 
 vLLM-Omni supports the following hardware platforms:
 

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

@@ -1,13 +1,11 @@
 # --8<-- [start:requirements]
 
-- GPU: compute capability 7.5 or higher (e.g., T4, RTX20xx, A100, L4, H100, etc.)
+- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.)
 
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
 
-vLLM-Omni depends on the matching major/minor release of vLLM. This source
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x.
-checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
-the published 0.30.0 release with its matching vLLM dependency.
 
 !!! 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.
@@ -43,17 +41,13 @@ uv pip install 'vllm-omni[demo]==0.30.0'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, install the final 0.31.0
+If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library
-release wheel from its commit-pinned index:
 
 ```bash
-uv pip install vllm==0.31.0 --torch-backend=auto \
+uv pip install vllm==0.30.0 --torch-backend=auto
-  --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
 ```
 
-The default 0.31.0 wheel contains CUDA 12.9 binaries. If you need a different
+The 0.30.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.
-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
@@ -74,16 +68,16 @@ 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.31.0
+git checkout v0.30.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
-# Default CUDA 12.9 wheel for v0.31.0 on x86_64
+#For CUDA 13.0 (the default for v0.30.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.30.0/vllm-0.30.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
-VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto
+uv pip install --editable .
 ```
 Install vllm with command below (If you already have PyTorch).
 ```bash
@@ -134,13 +128,12 @@ docker run --runtime nvidia --gpus 2 \
 DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.cuda -t vllm-omni-cuda .
 ```
 
-The source build defaults to `vllm/vllm-openai:v0.31.0`. To select that
+If you want to specify the base vLLM version:
-upstream release image explicitly:
 
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.31.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.30.0 \
   -t vllm-omni-cuda .
 ```
 

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

@@ -10,9 +10,8 @@
 
 For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images.
 
-vLLM-Omni depends on the matching major/minor release of vLLM. This source
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni
-checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
+0.30.x release line uses vLLM 0.30.x.
-the published 0.30.0 release with its matching vLLM dependency.
 
 The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The
 `vllm-omni` package does not install vLLM as a dependency, so non-Docker source
@@ -50,14 +49,13 @@ uv pip install onnxruntime-rocm
 
 #### Installation of vLLM
 
-If you do not need to modify source code of vLLM, install the final 0.31.0
+If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library
-release wheel for ROCm 7.2.3:
 
 ```bash
-uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
+uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
 ```
 
-The pre-built 0.31.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.30.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
 
@@ -77,7 +75,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.31.0
+git checkout v0.30.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -89,7 +87,7 @@ python3 setup.py develop
 #### Build docker image
 
 The source build defaults to the published upstream base image
-`vllm/vllm-openai-rocm:v0.31.0`, aligned with this source checkout.
+`vllm/vllm-openai-rocm:v0.30.0`, aligned with the vLLM release used by CI.
 Older or custom bases must provide the vLLM APIs checked by the Dockerfile's
 image-build canary. This upstream base is distinct from the prebuilt
 `vllm/vllm-omni-rocm` images discussed below; their published-tag availability
@@ -104,7 +102,7 @@ To select the upstream ROCm base explicitly, pass `BASE_IMAGE`:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.rocm \
-  --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0 \
   -t vllm-omni-rocm .
 ```
 

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

@@ -20,16 +20,12 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
 #### Build docker image
 
 ```bash
-export VLLM_VERSION=v0.31.0
+DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
-DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
-  --build-arg VLLM_VERSION="$VLLM_VERSION" .
 ```
 
 This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
-base image, which Docker pulls automatically. The Dockerfile still defaults to
+base image (v0.30.0 by default), which Docker pulls automatically. To target a different vLLM release,
-v0.30.0; the command above explicitly selects v0.31.0 to match this source
+pass `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
-checkout. To target a different vLLM release, pass
-`--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
 the base from upstream's own Dockerfile first and point the build at it:
 
 ```bash

docs/getting_started/installation/npu/npu.inc.md

@@ -71,14 +71,10 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 
 # --8<-- [start:installation-main]
 
-You can also install vLLM-Omni from source if you want to use the latest features
+You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.30.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
-or bug fixes. This source checkout targets vLLM 0.31.0 and needs a matching
-vLLM-Ascend runtime. The v0.30.0 images above are for the published Omni v0.30.0
-release. Check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886)
-for aligned versions and NPU validation before selecting a development image.
 
 ```bash
-# Inside a container with matching vLLM and vLLM-Ascend, install Omni from source
+# Inside the vLLM-Ascend v0.30.0 container, install vLLM-Omni from main
 cd /vllm-workspace
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/getting_started/quickstart.md

@@ -19,20 +19,17 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.31.0 --torch-backend=auto \
+uv pip install vllm==0.30.0 --torch-backend=auto
-  --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
 
 # On ROCm
-uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
+uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .
 ```
 
-These commands install vLLM 0.31.0 for this source checkout. The CUDA index is
+For additional installation methods — please see the [installation guide](installation/README.md).
-pinned to the final release commit. For published vLLM-Omni wheels and additional
-installation methods, see the [installation guide](installation/README.md).
 
 !!! 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.

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -130,7 +130,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- H200 automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when the CUDA MPS control tool is available; other devices and environments use `cosyvoice3.yaml`. See the [serving defaults and tradeoffs](../online_serving/text_to_speech.md#cosyvoice3). Pass `--deploy-config cosyvoice3.yaml` for the previous profile or `--deploy-config <path>` to override. The example uses the resolved profile's sampling parameters.
+- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
 - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
 
 ---

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -155,22 +155,16 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
 ./cosyvoice3/run_server.sh
 ```
 
-Streaming is on by default. On CUDA Hopper GPUs with at least 140 GiB of device memory (H200), serving automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when `nvidia-cuda-mps-control` is available on `PATH`: AR32 / codec32, standard sampling, packed Flow, cached ISTFT, HiFT decode graphs, and local CUDA MPS. No `--hf-overrides` is needed. Other devices and environments without the MPS control tool use `cosyvoice3.yaml` with RAS sampling and AR8 / codec8.
+Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
 
-The optimized profile is not a universal hardware configuration. Its packed kernels are Hopper-only; the talker's Model Runner V2 is unsupported on NPU/XPU, and CUDA MPS requires a local NVIDIA GPU and its control tool. Larger batches, 25-frame chunks, and graph allocations also require per-device validation on smaller or slower GPUs. Automatic selection retains the generic profile outside the validated device class; an explicit deploy configuration takes precedence.
+Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call:
-
-The H200 default targets streaming throughput at C32/C64. It changes sampling from RAS to standard (temperature 0.7, top-p 0.8, top-k 20, repetition penalty 1.21); output lengths and quality can differ. HiFT captures a recurring exact shape on its third use, which adds latency to that request, and its graph pool needs memory headroom beyond the engine reservations. This profile is not a claim of optimal first-audio latency at low concurrency. Pass `--deploy-config cosyvoice3.yaml` to retain the previous configuration, or `--deploy-config <path>` to select another profile.
-
-Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path; use `--deploy-config cosyvoice3.yaml` alongside it when selecting that path on H200.
-
-Cross-request Stage-1 flow batching is enabled in the H200 profile. For the generic configuration, enable it in the server environment when concurrent requests should share a flow-estimator call:
 
 ```bash
 export COSYVOICE3_BATCH_FLOW=1
 vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
 ```
 
-Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Use `--deploy-config cosyvoice3.yaml` and leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
+Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
 
 ### CLI client
 
@@ -200,7 +194,7 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads the device-specific CosyVoice3 default described above. Pass `--deploy-config <path>` to override.
+- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
 - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
 
 ---

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
-import copy
 import functools
 import json
 import os
@@ -11,11 +10,13 @@ from typing import Any
 
 import numpy as np
 import soundfile as sf
+from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -58,7 +59,8 @@ def run_e2e():
         "--deploy-config",
         type=str,
         default=None,
-        help="Override the deploy config path. If unset, selects the bundled CosyVoice3 profile for the device.",
+        help="Override the deploy config path. If unset, auto-loads "
+        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
     )
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
@@ -103,6 +105,8 @@ def run_e2e():
         profiler_config=args.profiler_config,
     )
 
+    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
+
     print("Model initialized. Preparing inputs...")
     ref_audio_path = args.ref_audio or _default_ref_audio()
     if not os.path.exists(ref_audio_path):
@@ -134,7 +138,7 @@ def run_e2e():
 
     print(f"Generating for prompt: {args.text}")
 
-    config = omni.engine.stage_vllm_configs[0].model_config.hf_config
+    config = CosyVoice3Config()
     tokenizer = get_qwen_tokenizer(
         token_path=args.tokenizer,
         skip_special_tokens=config.skip_special_tokens,
@@ -145,10 +149,29 @@ def run_e2e():
     min_len = int(base_len * config.min_token_text_ratio)
     max_len = int(base_len * config.max_token_text_ratio)
 
-    # Keep the deploy profile's sampling policy, penalties and required stops.
+    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
-    sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
+    gpt_sampling = SamplingParams(
-    sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
+        temperature=1.0,
-    sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
+        top_p=sampling_cfg["top_p"],
+        top_k=sampling_cfg["top_k"],
+        repetition_penalty=2.0,
+        min_tokens=min_len,
+        max_tokens=max_len,
+        stop_token_ids=[sampling_cfg["eos_token_id"]],
+        # allowed_token_ids=list(range(6561+3)),
+        detokenize=False,
+    )
+    # Not used
+    s2mel_sampling = SamplingParams(
+        temperature=1.0,
+        top_p=1.0,
+        top_k=-1,
+        repetition_penalty=2.0,
+        max_tokens=256,
+        detokenize=False,
+    )
+
+    sampling_params_list = [gpt_sampling, s2mel_sampling]
 
     profiler_enabled = args.profiler_config is not None
     if profiler_enabled:

examples/online_serving/text_to_speech/cosyvoice3/run_server.sh

@@ -1,15 +1,11 @@
 #!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 # Launch vLLM-Omni server for CosyVoice3 TTS
 #
 # Usage:
 #   ./run_server.sh
 #   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
 #
-# Streaming (async-chunk) is on by default. H200 selects the optimized profile;
+# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
-# other devices use vllm_omni/deploy/cosyvoice3.yaml.
 # Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
 
 set -e

requirements/common.txt

@@ -2,7 +2,7 @@
 # kernels==0.16.1 requires the LayerRepository version=/revision= integration
 # introduced in transformers 5.13. Keep the tested upper bound: 5.15 has
 # a reported model-construction regression; later versions need validation.
-# This range also satisfies vLLM 0.31.0's transformers>=5.10.4,<5.18.0 requirement.
+# This range also satisfies vLLM 0.30.0's transformers>=5.10.4 requirement.
 # MOSS-TTS-Nano's list/set weight-loading compatibility is handled by
 # transformers_keys_to_ignore_compat in model_executor/models/utils.py.
 transformers >= 5.13.0, < 5.15

tests/buildkite/test_rocm_dockerfile.py

@@ -10,8 +10,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 REPO_ROOT = Path(__file__).resolve().parents[2]
 AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2"
 AMD_BUILD_SCRIPT = REPO_ROOT / ".buildkite/amd/scripts/build-ci-image.sh"
-CUDA_RELEASE_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.cuda"
+CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
-CUDA_CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
 ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm"
 ROCM_DOCKERIGNORE = REPO_ROOT / "docker/Dockerfile.rocm.dockerignore"
 
@@ -31,18 +30,14 @@ def _line_index(lines: list[str], prefix: str) -> int:
     return matches[0]
 
 
-def test_rocm_base_tracks_cuda_vllm_release() -> None:
+def test_rocm_base_tracks_ci_vllm_release() -> None:
-    cuda_base = _docker_arg(CUDA_RELEASE_DOCKERFILE, "BASE_IMAGE")
+    ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
-    cuda_image, _, cuda_release = cuda_base.rpartition(":")
     rocm_base = _docker_arg(ROCM_DOCKERFILE, "BASE_IMAGE")
 
     image_ref, separator, image_tag = rocm_base.rpartition(":")
     assert separator, f"expected a tagged ROCm base image, got {rocm_base}"
-    assert cuda_image == "vllm/vllm-openai"
+    assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref
-    assert image_ref == "vllm/vllm-openai-rocm"
+    assert image_tag == ci_release
-    assert image_tag == cuda_release == f"v{_docker_arg(CUDA_CI_DOCKERFILE, 'VLLM_VERSION')}"
-    assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_IMAGE") == cuda_image
-    assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_TAG") == "v${VLLM_VERSION}"
 
 
 def test_rocm_defaults_to_prebuilt_base_image() -> None:
@@ -108,11 +103,11 @@ def test_rocm_build_context_excludes_git_history() -> None:
     assert ".git" in patterns
 
 
-def test_rocm_source_ref_tracks_cuda_vllm_release() -> None:
+def test_rocm_source_ref_tracks_ci_vllm_release() -> None:
-    cuda_release = _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_VERSION")
+    ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
     rocm_source_ref = _docker_arg(ROCM_DOCKERFILE, "VLLM_VERSION_OR_COMMIT_HASH")
 
-    assert rocm_source_ref == f"v{cuda_release}"
+    assert rocm_source_ref == ci_release
 
 
 def test_rocm_dockerfile_contains_vllm_api_canary() -> None:

tests/buildkite/test_upload_pipeline.py

@@ -172,11 +172,7 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest
     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"] == [
+    assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"]
-        CUDA_HF_TOKEN_EXPORT,
-        "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
-        "pytest -sv tests/example",
-    ]
 
 
 def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
@@ -243,10 +239,6 @@ def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None:
         assert "VLLM_CI_HF_TOKEN" in env_names
         assert "HF_TOKEN" not in env_names
         assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT
-        assert step["commands"][1:] == [
-            "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
-            "pytest -sv tests/example",
-        ]
 
 
 def _gpu_limit(step: dict) -> int:

tests/config/test_config_factory.py

@@ -571,56 +571,6 @@ class TestPipelineDiscovery:
         assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
         assert "qwen3_tts" in OMNI_PIPELINES
 
-    @pytest.mark.parametrize(
-        "cuda,major,memory_gib,mps,optimized",
-        [
-            (False, None, 0, True, False),
-            (True, None, 141, True, False),
-            (True, 8, 141, True, False),
-            (True, 9, 80, True, False),
-            (True, 10, 180, True, False),
-            (True, 9, 141, False, False),
-            (True, 9, 141, True, True),
-        ],
-    )
-    def test_cosyvoice3_device_default_and_explicit_overrides(
-        self, monkeypatch, cuda, major, memory_gib, mps, optimized
-    ):
-        from vllm.platforms import current_platform
-        from vllm.platforms.interface import DeviceCapability
-
-        from vllm_omni.platforms import current_omni_platform
-
-        monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
-        monkeypatch.setattr(current_omni_platform, "device_name", "cuda" if cuda else "cpu")
-        capability = DeviceCapability(major, 0) if major is not None else None
-        monkeypatch.setattr(current_platform, "get_device_capability", lambda: capability)
-        monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: memory_gib * 1024**3)
-        monkeypatch.setattr(
-            "vllm_omni.model_executor.models.cosyvoice3.pipeline.shutil.which",
-            lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps else None,
-        )
-        pipeline = resolve_pipeline_config("cosyvoice3")
-        assert pipeline is not None
-        config = VllmOmniConfig.from_pipeline_config(pipeline)
-        talker, codec = config.stage_configs
-        assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == (32 if optimized else 8)
-        assert talker.runtime_config.cuda_mps == codec.runtime_config.cuda_mps == optimized
-        assert talker.model_config.hf_overrides == ({"cosyvoice3_sampling_mode": "standard"} if optimized else None)
-        assert pipeline.stages[0].sampling_constraints["stop_token_ids"] == list(range(6561, 6761))
-        if optimized:
-            assert codec.runtime_config.env["COSYVOICE3_CACHED_ISTFT"] == "1"
-            assert codec.runtime_config.env["COSYVOICE3_HIFT_GRAPH"] == "1"
-            override = VllmOmniConfig.from_pipeline_config(
-                pipeline, cli_overrides={"hf_overrides": {"cosyvoice3_sampling_mode": "ras", "seed": 7}}
-            )
-            assert override.stage_by_id(0).model_config.hf_overrides == {"cosyvoice3_sampling_mode": "ras", "seed": 7}
-        explicit = VllmOmniConfig.from_pipeline_config(
-            pipeline, deploy_config_path=get_deploy_config_path("cosyvoice3.yaml")
-        )
-        assert explicit.stage_by_id(0).scheduler_config.max_num_seqs == 8
-        assert explicit.stage_by_id(0).model_config.hf_overrides is None
-
     def test_registry_resolver_qwen3_omni_all_stages(self):
         """Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
         pipeline = resolve_pipeline_config(

tests/core/sched/test_generation_batch_coalescing.py

@@ -171,7 +171,6 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock
     s, _ = scheduler(monkeypatch, [event])
     s._generation_max_wait_s = 0
     s.waiting = []
-    s.kv_holding_waiting = []
     s.input_coordinator._async_chunk = True
     s.input_coordinator.update_request_metadata = mocker.Mock()
     s.input_coordinator.process_pending_chunks = mocker.Mock()
@@ -179,11 +178,7 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock
     s.input_coordinator.update_request_metadata.assert_called_once_with(
         s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
     )
-    s.input_coordinator.process_pending_chunks.assert_called_once()
+    s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set())
-    waiting, running, ready, finished = s.input_coordinator.process_pending_chunks.call_args.args
-    assert list(waiting) == []
-    assert running is s.running
-    assert ready == {"0"} and finished == set()
 
 
 @pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])

tests/core/sched/test_generation_scheduler_restore.py

@@ -11,9 +11,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm import SamplingParams
-from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
 from vllm.v1.core.sched.interface import PauseState
-from vllm.v1.core.sched.output import CachedRequestData
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus
 
@@ -55,14 +53,12 @@ class FakeAdapter:
 def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
     scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
     scheduler.max_num_scheduled_tokens = 8
-    scheduler.max_num_active_reqs = 1
     scheduler.max_num_running_reqs = 1
     scheduler._pause_state = PauseState.UNPAUSED
     scheduler.running = []
     scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
     scheduler.waiting.add_request(waiting_request)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
     scheduler.requests = {waiting_request.request_id: waiting_request}
     scheduler.policy = SchedulingPolicy.FCFS
     scheduler.chunk_transfer_adapter = FakeAdapter()
@@ -90,7 +86,6 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
         get_manager_metadata=lambda: None,
     )
     scheduler.connector = None
-    scheduler.aux_output_connector = None
     scheduler.ec_connector = None
     scheduler.prev_step_scheduled_req_ids = set()
     scheduler._pending_finish_reqs = []
@@ -110,45 +105,6 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
     return scheduler
 
 
-def test_generation_scheduler_preserves_auxiliary_output_metadata() -> None:
-    request = Request("aux-contract", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
-    request.block_hashes.append(b"block-hash")
-    scheduler = _make_generation_scheduler(request, use_v2_model_runner=True)
-    scheduler.aux_output_connector = AuxOutputSchedulerConnector()
-
-    output = scheduler._build_generation_scheduler_output(
-        new_reqs_data=[],
-        cached_reqs_data=CachedRequestData.make_empty(),
-        num_scheduled_tokens={request.request_id: 1},
-        scheduled_spec_decode_tokens={},
-        scheduled_encoder_inputs={},
-        num_common_prefix_blocks=[0],
-    )
-
-    metadata = output.aux_output_connector_metadata
-    assert metadata is not None
-    assert metadata.requests == {request.request_id: 0}
-    assert list(metadata.block_hashes[request.request_id]) == request.block_hashes
-
-
-@pytest.mark.parametrize(("runner_slots", "active_limit"), [(1, 2), (2, 1)])
-def test_generation_admits_kv_holder_before_fresh_requests(runner_slots, active_limit) -> None:
-    fresh = Request("fresh", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
-    holding = Request("holding", [3, 4], SamplingParams(max_tokens=2), pooling_params=None)
-    holding.num_computed_tokens = 1
-    scheduler = _make_generation_scheduler(fresh)
-    scheduler.max_num_running_reqs = runner_slots
-    scheduler.max_num_active_reqs = active_limit
-    scheduler.requests[holding.request_id] = holding
-    scheduler.kv_holding_waiting.add_request(holding)
-
-    output = scheduler.schedule()
-
-    assert set(output.num_scheduled_tokens) == {holding.request_id}
-    assert list(scheduler.waiting) == [fresh]
-    assert list(scheduler.kv_holding_waiting) == []
-
-
 def _chunk_request(request_id, **kwargs):
     defaults = dict(
         request_id=request_id,
@@ -167,13 +123,7 @@ def _chunk_request(request_id, **kwargs):
         record_event=lambda *args, **kwargs: None,
     )
     defaults.update(kwargs)
-    return _HashableChunkRequest(**defaults)
+    return SimpleNamespace(**defaults)
-
-
-class _HashableChunkRequest(SimpleNamespace):
-    # Deferred queues require identity hashing; SimpleNamespace types it as None.
-    __hash__ = object.__hash__  # type: ignore[assignment]
-    __eq__ = object.__eq__
 
 
 def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None:
@@ -188,8 +138,7 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
     )
     scheduler.running = []
     scheduler.waiting = create_request_queue(scheduler.policy)
-    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
+    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
-    scheduler.deferred_waiting = set()
 
     completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None)
     completed.status = RequestStatus.RUNNING
@@ -343,7 +292,6 @@ def _express_scheduler():
     scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True)
     scheduler.waiting.add_request(first)
     scheduler.requests[first.request_id] = first
-    scheduler.max_num_active_reqs = 4
     scheduler.max_num_running_reqs = 4
     scheduler._native_data_plane = True
     scheduler.chunk_transfer_adapter = None

tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py

@@ -9,7 +9,7 @@ handing control to upstream, because upstream ``Scheduler.schedule()`` raises
 admits in a finished state -- which kills the stage's engine core, not just
 the request.
 
-Regression for the sweep missing ``kv_holding_waiting``, the third queue
+Regression for the sweep missing ``skipped_waiting``, the third queue
 upstream admits from. An aborted duplex session parked there was re-selected
 by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
 stage.
@@ -49,30 +49,28 @@ class _StubRequest:
         return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
 
 
-def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), kv_holding_waiting=(), running=()):
+def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
     """Bare scheduler with just the surface the sweep reads."""
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.waiting = create_request_queue(policy)
     for req in waiting:
         scheduler.waiting.add_request(req)
-    scheduler.kv_holding_waiting = create_request_queue(policy)
+    scheduler.skipped_waiting = create_request_queue(policy)
-    scheduler.deferred_waiting = set()
+    for req in skipped_waiting:
-    for req in kv_holding_waiting:
+        scheduler.skipped_waiting.add_request(req)
-        scheduler.kv_holding_waiting.add_request(req)
     scheduler.running = list(running)
     return scheduler
 
 
 @pytest.mark.parametrize("policy", _POLICIES)
-def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
+def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
     monkeypatch: pytest.MonkeyPatch,
     policy: SchedulingPolicy,
 ) -> None:
     """The public schedule path must remove aborted skipped requests first."""
     aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
-    scheduler = _make_scheduler(policy, kv_holding_waiting=[aborted])
+    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
     scheduler.chunk_transfer_adapter = None
-    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
     scheduler.input_coordinator = None
     scheduler.num_waiting_for_streaming_input = 0
@@ -87,7 +85,7 @@ def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
     sentinel = object()
 
     def upstream_schedule(_self, _throttle_prefills: bool = False):
-        assert list(scheduler.kv_holding_waiting) == []
+        assert list(scheduler.skipped_waiting) == []
         return sentinel
 
     monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -185,15 +185,10 @@ def test_stage_zero_final_skips_inter_stage_output(mocker):
 def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler._omits_kv_transfer_cache = {}
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting = []
-    scheduler.kv_holding_waiting = []
     scheduler.log_stats = False
     tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
-
+    session = SimpleNamespace(
-    class Session(SimpleNamespace):
-        __hash__ = object.__hash__
-
-    session = Session(
         request_id="req",
         additional_information=tagged.additional_information,
         status=RequestStatus.WAITING,

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -11,14 +11,12 @@ from types import MethodType, SimpleNamespace
 
 import numpy as np
 import pytest
-import torch
 from vllm.v1.engine import FinishReason
 from vllm.v1.outputs import LogprobsLists
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.outputs import OmniModelRunnerOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -134,13 +132,11 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
     scheduler = SimpleNamespace(
         perf_metrics=None,
         connector=None,
-        aux_output_connector=None,
         chunk_transfer_adapter=None,
         requests={request.request_id: request for request in requests},
         running=list(requests),
         waiting=_RequestQueue(),
-        kv_holding_waiting=_RequestQueue(),
+        skipped_waiting=_RequestQueue(),
-        deferred_waiting=set(),
         structured_output_manager=SimpleNamespace(accept_tokens=lambda _request, _token_ids: True),
         transfer_triggered_requests=set(),
         active_kv_transfers=set(),
@@ -296,36 +292,6 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
 
 
-@pytest.mark.parametrize("scored", [False, True])
-def test_fixed_token_scores_reach_engine_output(scored) -> None:
-    request = _Request("req")
-    request.sampling_params.num_logprobs = None
-    scheduler = _make_scheduler_stub([request])
-    _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
-    scheduler_output = SimpleNamespace(
-        num_scheduled_tokens={"req": 3},
-        scheduled_spec_decode_tokens={},
-        num_invalid_spec_tokens=0,
-    )
-    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
-    model_runner_output = OmniModelRunnerOutput(
-        req_ids=["req"],
-        req_id_to_index={"req": 0},
-        sampled_token_ids=[[7]],
-        prompt_logprobs_dict={},
-        prompt_token_id_logprobs_dict={"req": scores} if scored else {},
-    )
-
-    outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
-    (output,) = outputs[0].outputs
-
-    assert output.new_token_ids == [7]
-    if scored:
-        torch.testing.assert_close(output.prompt_token_id_logprobs, scores)
-    else:
-        assert output.prompt_token_id_logprobs is None
-
-
 def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
     return SimpleNamespace(
         sampled_token_ids=[[]],

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -73,8 +73,7 @@ def _replace_streaming_session(session: Request) -> None:
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
-    sched.kv_holding_waiting = set()
+    sched.skipped_waiting = set()
-    sched.deferred_waiting = set()
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
     update = StreamingUpdate(
         mm_features=None,

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -42,8 +42,7 @@ def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARS
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
-    sched.kv_holding_waiting = set()
+    sched.skipped_waiting = set()
-    sched.deferred_waiting = set()
     sched._free_request_blocks = MagicMock()
     sched.encoder_cache_manager = MagicMock()
     sched._inflight_prefills = set()
@@ -1170,8 +1169,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
     sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
     sched.running = []
     sched.waiting = queue()
-    sched.kv_holding_waiting = queue(session)
+    sched.skipped_waiting = queue(session)
-    sched.deferred_waiting = set()
     sched.num_waiting_for_streaming_input = 1
     sched.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=True,
@@ -1183,7 +1181,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
 
     assert session.status == RequestStatus.WAITING
     assert session in sched.waiting.requests
-    assert session not in sched.kv_holding_waiting.requests
+    assert session not in sched.skipped_waiting.requests
     assert sched.num_waiting_for_streaming_input == 0
     assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
 
@@ -1210,8 +1208,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
         receives_chunks=receives_chunks,
         segment_finished_requests={session.request_id},
     )
-    sched.kv_holding_waiting = MagicMock()
+    sched.skipped_waiting = MagicMock()
-    sched.deferred_waiting = set()
     sched._enqueue_waiting_request = MagicMock()
 
     sched._resume_downstream_chunk_receiver(session)
@@ -1219,7 +1216,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
     assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
     assert sched.num_waiting_for_streaming_input == 1
     assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
-    sched.kv_holding_waiting.remove_requests.assert_not_called()
+    sched.skipped_waiting.remove_requests.assert_not_called()
     sched._enqueue_waiting_request.assert_not_called()
 
 
@@ -1275,8 +1272,7 @@ def _park_session(sched: OmniARScheduler, session: Request) -> None:
     session.num_computed_tokens = 6
     session.num_output_placeholders = 0
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    sched.kv_holding_waiting.add_request(session)
+    sched.skipped_waiting.add_request(session)
-    sched.deferred_waiting.add(session)
     sched.num_waiting_for_streaming_input = 1
 
 
@@ -1299,7 +1295,7 @@ def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_sessi
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.request_id not in sched.requests
     assert len(sched.waiting) == 0
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
     assert sched.num_waiting_for_streaming_input == 0
     sched._free_request_blocks.assert_called_once_with(session)
     assert sched.finished_req_ids == {session.request_id}
@@ -1407,8 +1403,7 @@ def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
     sched.max_model_len = max_model_len
     sched.policy = SchedulingPolicy.FCFS
     sched.waiting = create_request_queue(sched.policy)
-    sched.kv_holding_waiting = create_request_queue(sched.policy)
+    sched.skipped_waiting = create_request_queue(sched.policy)
-    sched.deferred_waiting = set()
     sched.running = []
     sched.requests = {}
     sched._free_request = MagicMock()  # needs a KV manager; not what this covers
@@ -1437,7 +1432,7 @@ def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -
 
     assert finished is True
     assert len(sched.waiting) == 0
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.resumable is False
     assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
@@ -1456,8 +1451,7 @@ def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
     finished = sched._handle_stopped_request(session)
 
     assert finished is False
-    assert list(sched.kv_holding_waiting) == [session]
+    assert len(sched.waiting) == 1
-    assert session not in sched.deferred_waiting
     assert session.status == RequestStatus.WAITING
     assert not getattr(sched, "_streaming_context_overflow", {})
     sched._free_request.assert_not_called()
@@ -1479,7 +1473,7 @@ def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:
 
     assert finished is True
     assert len(sched.waiting) == 0
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.resumable is False
     sched._free_request.assert_not_called()
@@ -1502,7 +1496,7 @@ def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanc
     finished = sched._handle_stopped_request(session)
 
     assert finished is True
-    assert len(sched.kv_holding_waiting) == 0
+    assert len(sched.skipped_waiting) == 0
     assert sched.num_waiting_for_streaming_input == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     sched._free_request.assert_not_called()
@@ -1535,7 +1529,6 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.running = []
     sched._native_data_plane = native
     sched.use_v2_model_runner = native
-    sched.max_num_active_reqs = 8
     sched.max_num_running_reqs = 8
     sched.input_coordinator = (
         SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
@@ -1562,7 +1555,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     observed_limits: list[int] = []
 
     def fake_schedule(self, _throttle_prefills=False):
-        observed_limits.append(self.max_num_active_reqs)
+        observed_limits.append(self.max_num_running_reqs)
         return SimpleNamespace(scheduled_new_reqs=[])
 
     monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule)
@@ -1570,4 +1563,4 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.schedule()
 
     assert observed_limits == [7 if native else 8]
-    assert sched.max_num_active_reqs == 8
+    assert sched.max_num_running_reqs == 8

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -79,19 +79,16 @@ def _make_scheduler(
         else:
             hf_config.voxcpm2_runtime_config = runtime_config
     sched.vllm_config = SimpleNamespace(model_config=model_config)
-    # Upstream ``Scheduler.__init__`` always creates ``kv_holding_waiting``; the
+    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
     # abort sweep at the top of ``schedule()`` reads it.
-    sched.kv_holding_waiting = _MockQueue()
+    sched.skipped_waiting = _MockQueue()
-    sched.deferred_waiting = set()
     return sched
 
 
-@pytest.mark.parametrize("queue_name", ["waiting", "kv_holding_waiting"])
+def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
-def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready(queue_name) -> None:
     scheduler = _make_scheduler()
     scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue()
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
-    setattr(scheduler, queue_name, _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]))
 
     assert scheduler._should_defer_waiting_for_unified_decode_graph()
 
@@ -147,17 +144,12 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa
 
 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
     scheduler = _make_scheduler()
-    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
     scheduler.running = [_MockRequest("decode")]
     original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
     deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
     original_waiting = _MockQueue([original_waiting_req])
     scheduler.waiting = original_waiting
-    kv_holder = _MockRequest("kv-holder", status=RequestStatus.WAITING)
-    original_kv_holding = _MockQueue([kv_holder])
-    scheduler.kv_holding_waiting = original_kv_holding
-    preempted_by_upstream = _MockRequest("preempted-by-upstream", status=RequestStatus.PREEMPTED)
     scheduler.policy = "fcfs"
     scheduler.chunk_transfer_adapter = None
     scheduler.input_coordinator = None
@@ -168,10 +160,7 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
     def fake_upstream_schedule(self, throttle_prefills: bool = False):
         assert self.waiting is not original_waiting
         assert not self.waiting
-        assert self.kv_holding_waiting is not original_kv_holding
-        assert not self.kv_holding_waiting
         self.waiting.add_request(deferred_by_upstream)
-        self.kv_holding_waiting.add_request(preempted_by_upstream)
         raise RuntimeError("stop before output wrapping")
 
     monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
@@ -180,6 +169,3 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
         scheduler.schedule()
 
     assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]
-    assert scheduler.waiting is original_waiting
-    assert scheduler.kv_holding_waiting is original_kv_holding
-    assert scheduler.kv_holding_waiting._items == [preempted_by_upstream, kv_holder]

tests/core/sched/test_omni_cfg_ar_scheduler.py

@@ -23,12 +23,10 @@ def _scheduler(monkeypatch, waiting, running=()):
     scheduler.waiting = create_request_queue(scheduler.policy)
     for request in waiting:
         scheduler.waiting.add_request(request)
-    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
+    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
-    scheduler.deferred_waiting = set()
     scheduler.running = list(running)
     scheduler.requests = {request.request_id: request for request in [*waiting, *running]}
     scheduler.max_model_len = 128
-    scheduler.max_num_active_reqs = 2
     scheduler.max_num_running_reqs = 2
     scheduler.max_num_scheduled_tokens = 256
     scheduler.block_size = 16

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -52,11 +52,10 @@ class _SchedulerStub(OmniGenerationScheduler):
         self.num_waiting_for_streaming_input = 0
         self.log_stats = log_stats
         self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
-        self.kv_holding_waiting = _SkippedWaitingStub()
+        self.skipped_waiting = _SkippedWaitingStub()
-        self.deferred_waiting: set[Request] = set()
 
     def _enqueue_waiting_request(self, session: Request) -> None:
-        raise AssertionError("unexpected enqueue for kv_holding_waiting miss")
+        raise AssertionError("unexpected enqueue for skipped_waiting miss")
 
 
 class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
@@ -135,8 +134,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     sched._handle_stopped_request.side_effect = rearm_resumable_request
     sched.running = [session]
     sched.waiting = MagicMock()
-    sched.kv_holding_waiting = MagicMock()
+    sched.skipped_waiting = MagicMock()
-    sched.deferred_waiting = set()
     sched.structured_output_manager.accept_tokens.return_value = True
     # Async scheduling can observe the segment boundary in the next schedule()
     # before this model output is applied and defer the same request for finish.
@@ -208,8 +206,7 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete:
         OmniSchedulerMixin._remove_stopped_requests_from_queues(sched, running, preempted)
     )
     sched.waiting = MagicMock()
-    sched.kv_holding_waiting = MagicMock()
+    sched.skipped_waiting = MagicMock()
-    sched.deferred_waiting = set()
     sched.structured_output_manager.accept_tokens.return_value = False
     sched._pending_finish_reqs = []
     sched.recompute_kv_load_failures = False
@@ -353,8 +350,6 @@ class _RealignSchedulerStub(OmniSchedulerMixin):
         self.requests = requests
         self.running = running
         self.waiting = waiting
-        self.kv_holding_waiting: list[Request] = []
-        self.deferred_waiting: set[Request] = set()
 
 
 class TestRealignRequestStatusToQueues:

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -123,7 +123,6 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
 
     sched = MagicMock()
     sched.max_num_scheduled_tokens = 8
-    sched.max_num_active_reqs = 8
     sched.max_num_running_reqs = 8
     sched._native_data_plane = False
     sched._retains_state_across_chunks = False
@@ -136,8 +135,6 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
     sched.requests = {request.request_id: request}
     sched.running = [request]
     sched.waiting = []  # falsy: the waiting loop never runs
-    sched.kv_holding_waiting = []
-    sched.deferred_waiting = set()
     sched.policy = SchedulingPolicy.FCFS
     sched.chunk_transfer_adapter = None
     sched.scheduler_config = MagicMock(enable_chunked_prefill=True)

tests/core/sched/test_omni_scheduler_finish_requests_purge.py

@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 """Integration tests for the post-finish defensive purge of
 ``self.running`` placed in
 ``OmniARScheduler.finish_requests`` /
@@ -55,8 +52,6 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting):
     scheduler.requests = requests
     scheduler.running = running
     scheduler.waiting = waiting
-    scheduler.kv_holding_waiting = []
-    scheduler.deferred_waiting = set()
     return scheduler
 
 

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -52,8 +52,6 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
     scheduler.requests = {}
     scheduler.running = []
     scheduler.waiting = []
-    scheduler.kv_holding_waiting = []
-    scheduler.deferred_waiting = set()
 
     def fake_finish_requests(self, request_ids, finished_status):
         assert request_ids == ["req-a", "req-b"]

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -7,10 +7,7 @@ from typing import Any
 
 import pytest
 from vllm.config import SchedulerConfig, VllmConfig
-from vllm.sampling_params import SamplingParams
-from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.engine import FinishReason
-from vllm.v1.request import Request
 
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.core.sched import omni_scheduler_mixin
@@ -21,7 +18,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _Scheduler(OmniSchedulerMixin):
-    _is_blocked_waiting_status = staticmethod(Scheduler._is_blocked_waiting_status)
+    pass
 
 
 def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
@@ -123,11 +120,7 @@ def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_i
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     calls: list[tuple[Any, ...]] = []
     scheduler = _Scheduler()
-    waiting = Request("waiting", [1], SamplingParams(max_tokens=1), pooling_params=None)
+    scheduler.waiting = ["waiting"]
-    holding = Request("holding", [1, 2], SamplingParams(max_tokens=1), pooling_params=None)
-    holding.num_computed_tokens = 1
-    scheduler.waiting = [waiting]
-    scheduler.kv_holding_waiting = [holding]
     scheduler.running = ["running"]
     scheduler.requests = {"request": object()}
     scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
@@ -144,16 +137,16 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     scheduler.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=True,
         process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
-            ("process", list(waiting), running, scheduler_requests)
+            ("process", waiting, running, scheduler_requests)
         ),
         restore_queues=lambda waiting, running, scheduler_requests: calls.append(
-            ("restore-chunks", list(waiting), running, scheduler_requests)
+            ("restore-chunks", waiting, running, scheduler_requests)
         ),
         collect_timed_out_request_ids=_collect_timed_out,
         collect_failed_send_request_ids=_collect_failed_sends,
     )
     scheduler.input_coordinator = SimpleNamespace(
-        restore_queues=lambda waiting, running: calls.append(("restore-full", list(waiting), running))
+        restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running))
     )
 
     scheduler._process_pending_omni_inputs("ar")
@@ -162,13 +155,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     assert calls == [
         ("consume", "ar"),
         ("timeouts",),
-        ("process", [holding, waiting], scheduler.running, scheduler.requests),
+        ("process", scheduler.waiting, scheduler.running, scheduler.requests),
         # The chunk deadline runs after chunks are applied, so a chunk that
         # arrived this cycle resets the clock before it is measured (R1.1).
         ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
         ("failed-sends",),
-        ("restore-chunks", [holding, waiting], scheduler.running, scheduler.requests),
+        ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
-        ("restore-full", [holding, waiting], scheduler.running),
+        ("restore-full", scheduler.waiting, scheduler.running),
     ]
 
 

tests/core/sched/test_omni_scheduler_ready_inbox.py

@@ -5,13 +5,8 @@
 from types import SimpleNamespace
 
 import pytest
-from vllm import SamplingParams
-from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
-from vllm.v1.request import Request, RequestStatus
 
-from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator
 from vllm_omni.outputs import OmniConnectorOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -20,8 +15,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
     scheduler = OmniSchedulerMixin()
     scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], []
-    scheduler.kv_holding_waiting = []
-    scheduler._is_blocked_waiting_status = OmniARScheduler._is_blocked_waiting_status
     coordinator = SimpleNamespace(
         _async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock()
     )
@@ -39,92 +32,7 @@ def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
     coordinator.update_request_metadata.assert_called_once_with(
         scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar"
     )
-    coordinator.process_pending_chunks.assert_called_once_with(mocker.ANY, [], {"r"}, {"r"})
+    coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"})
-    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
     scheduler._consume_pending_connector_output(model_mode="ar")
     assert coordinator.update_request_metadata.call_count == 1
-    coordinator.process_pending_chunks.assert_called_with(mocker.ANY, [], set(), set())
+    coordinator.process_pending_chunks.assert_called_with([], [], set(), set())
-    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
-
-
-@pytest.mark.parametrize("policy", [SchedulingPolicy.FCFS, SchedulingPolicy.PRIORITY])
-@pytest.mark.parametrize("async_chunk", [False, True])
-def test_native_input_gate_parks_and_restores_kv_holders(policy, async_chunk):
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.waiting = create_request_queue(policy)
-    scheduler.kv_holding_waiting = create_request_queue(policy)
-    scheduler.deferred_waiting = set()
-    scheduler.running = []
-    scheduler.chunk_transfer_adapter = None
-    scheduler.input_coordinator = OmniSchedulingCoordinator(
-        stage_id=1, scheduler_max_num_seqs=2, async_chunk=async_chunk
-    )
-    scheduler._init_omni_connector_output_inbox()
-    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
-    holder.num_computed_tokens = 2
-    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
-    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
-    scheduler._enqueue_waiting_request(fresh)
-    scheduler._enqueue_waiting_request(holder)
-
-    scheduler._consume_pending_connector_output(model_mode="ar")
-
-    assert not scheduler.waiting and not scheduler.kv_holding_waiting
-    expected_status = RequestStatus.WAITING_FOR_CHUNK if async_chunk else RequestStatus.WAITING_FOR_INPUT
-    assert holder.status == fresh.status == expected_status
-    coordinator = scheduler.input_coordinator
-    registrations = coordinator.pending_chunk_registrations if async_chunk else coordinator.pending_input_registrations
-    assert [entry.request_id for entry in registrations] == ["holder", "fresh"]
-
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-    ready_output = (
-        OmniConnectorOutput(chunk_ready_req_ids={"holder"})
-        if async_chunk
-        else OmniConnectorOutput(stage_recv_req_ids={"holder"})
-    )
-    scheduler.enqueue_omni_connector_output(ready_output)
-    scheduler._consume_pending_connector_output(model_mode="ar")
-
-    assert holder.status == RequestStatus.WAITING
-    assert fresh.status == expected_status
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert not scheduler.waiting
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-
-
-@pytest.mark.parametrize("async_chunk", [False, True])
-@pytest.mark.parametrize(
-    "status",
-    [
-        RequestStatus.WAITING_FOR_STREAMING_REQ,
-        RequestStatus.WAITING_FOR_REMOTE_KVS,
-        RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR,
-    ],
-)
-def test_native_input_gate_preserves_upstream_blocked_waits(async_chunk, status):
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
-    scheduler.running = []
-    request = Request("blocked", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
-    request.status = status
-    request.num_computed_tokens = 2
-    scheduler.requests = {request.request_id: request}
-    scheduler._enqueue_waiting_request(request)
-    scheduler.input_coordinator = OmniSchedulingCoordinator(
-        stage_id=1, scheduler_max_num_seqs=1, async_chunk=async_chunk
-    )
-    scheduler._init_omni_connector_output_inbox()
-
-    scheduler._consume_pending_connector_output(model_mode="ar")
-
-    assert request.status == status
-    assert list(scheduler.kv_holding_waiting) == [request]
-    assert scheduler.deferred_waiting == {request}
-    assert not scheduler.input_coordinator.pending_chunk_registrations
-    assert not scheduler.input_coordinator.pending_input_registrations

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -59,8 +59,7 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipp
     scheduler.requests = requests
     scheduler.running = running
     scheduler.waiting = waiting
-    scheduler.kv_holding_waiting = skipped if skipped is not None else []
+    scheduler.skipped_waiting = skipped if skipped is not None else []
-    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = counter
     return scheduler
 
@@ -98,8 +97,8 @@ def test_finish_requests_clears_counter_left_by_a_stomped_status(
     assert scheduler.num_waiting_for_streaming_input == 0
 
 
-def test_resync_counts_parked_requests_in_kv_holding_waiting() -> None:
+def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
-    """``get_num_unfinished_requests`` sums ``waiting`` and ``kv_holding_waiting``,
+    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
     so the resync has to look at both or it under-counts and reintroduces the
     spin the counter exists to prevent."""
     parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
@@ -129,11 +128,9 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk
     scheduler.requests = {parked.request_id: parked}
     scheduler.running = [parked]
     scheduler.waiting = []
-    scheduler.kv_holding_waiting = []
+    scheduler.skipped_waiting = []
-    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
     scheduler.policy = "fcfs"
-    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
 
     seen: dict[str, int] = {}
@@ -164,13 +161,12 @@ def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -
     scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
     scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
     scheduler.waiting.add_request(aborted)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
+    scheduler.skipped_waiting.add_request(aborted)
-    scheduler.kv_holding_waiting.add_request(aborted)
     scheduler.running = [aborted]
 
     scheduler._drop_aborted_queued_requests()
 
     assert list(scheduler.waiting) == []
-    assert list(scheduler.kv_holding_waiting) == []
+    assert list(scheduler.skipped_waiting) == []
     assert scheduler.running == []

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -216,40 +216,44 @@ def test_video_chat_stream_invalid_requests(
 # WS /v1/realtime
 # ─────────────────────────────────────────────────────────────────────────────
 
+# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
+_REALTIME_WS_MODEL_MISMATCH = object()
+_REALTIME_WS_INVALID_AUDIO_APPEND = object()
+
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize(
     "send_frames_spec, ws_error_code, err_message",
     [
-        pytest.param("{not-json", "invalid_event", "Invalid or unrecognized client event", id="invalid_json"),
+        pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
         pytest.param(
             json.dumps({"type": "session.update"}),
             "invalid_event",
-            "Invalid or unrecognized client event",
+            ("model", "Missing required field"),
-            id="session_update_missing_session",
+            id="session_update_missing_model",
         ),
         pytest.param(
-            json.dumps({"type": "session.update", "session": []}),
+            _REALTIME_WS_MODEL_MISMATCH,
-            "invalid_event",
+            "model_not_found",
-            "Invalid or unrecognized client event",
+            ("model", "does not exist"),
-            id="session_update_invalid_session",
+            id="session_update_model_mismatch",
         ),
         pytest.param(
-            json.dumps({"type": "input_audio_buffer.commit"}),
+            json.dumps({"type": "input_audio_buffer.commit", "final": False}),
-            "invalid_request_error",
+            "model_not_validated",
-            "Input audio buffer is empty",
+            ("session.update", "validate the model"),
-            id="commit_without_audio",
+            id="commit_without_model_validation",
         ),
         pytest.param(
             json.dumps({"type": "unknown.event"}),
-            "invalid_event",
+            "unknown_event",
-            "Invalid or unrecognized client event",
+            "Unknown event type",
             id="unknown_event_type",
         ),
         pytest.param(
-            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
+            _REALTIME_WS_INVALID_AUDIO_APPEND,
-            "invalid_request_error",
+            "invalid_audio",
-            "Invalid base64 audio data",
+            "Invalid audio data",
             id="invalid_audio_append",
         ),
     ],
@@ -262,12 +266,27 @@ def test_realtime_invalid_requests(
     ws_error_code: str,
     err_message: str | tuple[str, ...],
 ) -> None:
+    send_frames: Any
+    if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
+        send_frames = json.dumps(
+            {
+                "type": "session.update",
+                "model": "this-model-is-not-served-on-this-instance",
+            }
+        )
+    elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
+        send_frames = [
+            json.dumps({"type": "session.update", "model": omni_server.model}),
+            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
+        ]
+    else:
+        send_frames = send_frames_spec
     online_client.send_realtime_ws_request(
         {
-            "send_frames": send_frames_spec,
+            "send_frames": send_frames,
             "timeout": 120,
             "ws_max_size": None,
-            "ws_skip_types": ["session.created", "conversation.created"],
+            "ws_skip_types": ["session.created"],
             "ws_json_type": "error",
             "ws_error_code": ws_error_code,
             "err_message": err_message,

tests/diffusion/cache/test_seacache.py

@@ -329,9 +329,7 @@ def _uneven_cfg_sharded_worker(
         init_method=init_method,
         rank=rank,
         world_size=4,
-        # Spawned workers import torch/vLLM before joining; allow startup
+        timeout=datetime.timedelta(seconds=10),
-        # skew without changing the collective deadlock checks below.
-        timeout=datetime.timedelta(seconds=60),
     )
     try:
         fs_groups = [

tests/diffusion/diffusion_kv/test_kv_cache_utils.py

@@ -103,7 +103,7 @@ def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
     # Independent expected native extras: a crossing item has a negative
     # relative offset; a block can include the end of one item and the next.
     first, second = (feature.identifier for feature in row.mm_features)
-    extras = [(("mm", first, 2),), (("mm", first, -2), ("mm", second, 3)), (("mm", second, -1),)]
+    extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
     parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
     expected = []
     for start, keys in zip(range(0, 12, 4), extras, strict=True):

tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py

@@ -31,10 +31,6 @@ def producer(monkeypatch):
     # vLLM 0.30 send_kv_to_decode also reads the PP size (ce08bb5b34,
     # #56033), the MLA/Mamba flags and the head-resharding validator.
     worker.pp_size = 1
-    worker.opaque_packed_storages = set()
-    worker._prepared_transfer_regions = {}
-    worker.region_shared_groups = []
-    worker.region_row_offsets = []
     worker.use_mla = False
     worker.kv_cache_config = SimpleNamespace(has_mamba_layers=False)
     worker._validate_head_resharding_layout = Mock(return_value=None)

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -1253,15 +1253,14 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         speculative_config=None,
     )
 
-    def select_backend(*, backend, attn_selector_config, num_heads, _run_kv_cache_dtype):
+    def select_backend(*, backend, attn_selector_config, num_heads):
-        # Exercise the real upstream selector, replacing only platform resolution.
+        # Exercise the real 0.29 selector, replacing only platform resolution.
         # Ulysses has already gathered tokens and sharded heads before the
         # native kernel; the MoE PCP mapping must not request PCP attention.
         assert not attn_selector_config.use_pcp
         assert not attn_selector_config.use_dcp
         assert attn_selector_config.use_non_causal
         assert num_heads == 4
-        assert _run_kv_cache_dtype == config.cache_config.cache_dtype
         selected_backends.append((backend, config.attention_config.backend_per_kind))
         if selection_fails:
             raise ValueError("test backend unavailable")

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -239,11 +239,11 @@ def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mo
 )
 def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
     from transformers import AutoModel
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
     from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym
 
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     od_config = OmniDiffusionConfig(
         model="dummy-boogu",
@@ -287,14 +287,15 @@ def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_co
         mock_dependencies["mllm_config_loader"].assert_not_called()
         mock_dependencies["mllm_builder"].assert_not_called()
         assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
-    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config)
+    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config)
 
 
 def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
-    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    quant_config = DiffusionFp8Config(
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    quant_config = Fp8Config(
         ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
     )
     od_config = OmniDiffusionConfig(

tests/diffusion/models/flux2/test_flux2_quantization.py

@@ -16,13 +16,17 @@ class _RecordingBlock(nn.Module):
 
 
 def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch):
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
     from vllm_omni.diffusion.models.flux2 import flux2_transformer
-    from vllm_omni.quantization import build_quant_config
 
     monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
     monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)
 
-    quant_config = build_quant_config("fp8")
+    quant_config = Fp8Config(
+        is_checkpoint_fp8_serialized=False,
+        activation_scheme="dynamic",
+    )
     model = flux2_transformer.Flux2Transformer2DModel(
         in_channels=4,
         out_channels=4,

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -71,12 +71,12 @@ def _small_od_config(*, class_name=None):
 
 def test_fp8_scope_and_prefix_propagation(monkeypatch):
     from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.utils.quant_utils import (
         is_layer_skipped,
     )
 
     from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
     monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
@@ -92,7 +92,7 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
         "blocks.0.adaln_proj.linear",
         "final_layer.adaln_proj.linear",
     }
-    fp8_config = DiffusionFp8Config(ignored_layers=sorted(ignored_layers))
+    fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
     model = h3.MiniMaxH3DiTModel(
         _small_od_config(),
         quant_config=fp8_config,

tests/diffusion/quantization/test_fp8_config.py

@@ -15,74 +15,6 @@ def test_build_quant_config_fp8():
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "dynamic"
-    assert not config.is_checkpoint_fp8_serialized
-
-
-def test_build_quant_config_fp8_uses_native_online_quantization():
-    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
-    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config("fp8", ignored_layers=["proj_out"])
-    assert isinstance(config, OnlineQuantizationConfig)
-    assert config.args.linear.weight == kFp8StaticTensorSym
-    assert config.args.moe.weight == kFp8StaticTensorSym
-    assert config.ignored_layers == ["proj_out"]
-
-
-def test_build_quant_config_upstream_online_fp8_shorthand():
-    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
-    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config("fp8_per_tensor", ignore=["proj_out"])
-    assert isinstance(config, OnlineQuantizationConfig)
-    assert config.args.linear.weight == kFp8StaticTensorSym
-    assert config.ignored_layers == ["proj_out"]
-
-
-def test_online_fp8_rejects_static_activation():
-    from vllm_omni.quantization import build_quant_config
-
-    with pytest.raises(ValueError, match="activation_scheme='dynamic'"):
-        build_quant_config("fp8", activation_scheme="static")
-
-
-def test_build_online_quant_config_preserves_explicit_native_args():
-    from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
-    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
-
-    from vllm_omni.quantization import build_quant_config
-
-    args = QuantizationConfigArgs(linear=QuantSpec(weight="fp8_per_tensor_static"), ignore=["proj_out"])
-    config = build_quant_config("online", args=args)
-    assert isinstance(config, OnlineQuantizationConfig)
-    assert config.args is args
-    assert config.ignored_layers == ["proj_out"]
-
-
-def test_build_online_quant_config_requires_args():
-    from vllm_omni.quantization import build_quant_config
-
-    with pytest.raises(ValueError, match="requires quantization config arguments"):
-        build_quant_config("online")
-
-
-def test_serialized_checkpoint_replaces_online_fp8_config():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
-    from vllm_omni.quantization import build_quant_config
-    from vllm_omni.quantization.factory import resolve_quant_config_from_disk
-
-    config = resolve_quant_config_from_disk(
-        build_quant_config("fp8"),
-        {"quant_method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
-    )
-    assert isinstance(config, Fp8Config)
-    assert config.is_checkpoint_fp8_serialized
-    assert config.activation_scheme == "static"
 
 
 def test_build_quant_config_none():
@@ -107,7 +39,7 @@ def test_build_quant_config_invalid():
 def test_build_quant_config_dict():
     from vllm_omni.quantization import build_quant_config
 
-    config = build_quant_config({"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"})
+    config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "static"
@@ -116,7 +48,7 @@ def test_build_quant_config_dict():
 def test_build_quant_config_dict_not_mutated():
     from vllm_omni.quantization import build_quant_config
 
-    original = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
+    original = {"method": "fp8", "activation_scheme": "static"}
     copy = original.copy()
     build_quant_config(original)
     assert original == copy
@@ -167,7 +99,7 @@ def test_build_quant_config_per_component_inner_dict_not_mutated():
     """Inner component dicts should not be mutated by build_quant_config."""
     from vllm_omni.quantization import build_quant_config
 
-    inner = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
+    inner = {"method": "fp8", "activation_scheme": "static"}
     original = inner.copy()
     build_quant_config({"transformer": inner, "vae": None})
     assert inner == original
@@ -183,18 +115,20 @@ def test_flat_dict_not_misdetected_as_per_component():
 
 
 def test_build_quant_config_passthrough():
-    from vllm_omni.quantization import build_quant_config
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = DiffusionFp8Config()
+    from vllm_omni.quantization import build_quant_config
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
     assert build_quant_config(fp8) is fp8
 
 
 def test_component_config_routing():
-    from vllm_omni.quantization import ComponentQuantizationConfig
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = DiffusionFp8Config()
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
     config = ComponentQuantizationConfig(
         component_configs={"transformer": fp8, "vae": None},
     )
@@ -206,10 +140,11 @@ def test_component_config_routing():
 
 
 def test_component_config_with_default():
-    from vllm_omni.quantization import ComponentQuantizationConfig
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = DiffusionFp8Config()
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
     config = ComponentQuantizationConfig(
         component_configs={"vae": None},
         default_config=fp8,
@@ -232,7 +167,7 @@ def test_integration_dict():
 
     config = OmniDiffusionConfig(
         model="test",
-        quantization_config={"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
+        quantization_config={"method": "fp8", "activation_scheme": "static"},
     )
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "fp8"

tests/diffusion/quantization/test_online_fp8_gpu.py

@@ -1,72 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Exercise online FP8 loading and inference through the native CUDA path."""
-
-import pytest
-import torch
-from torch.nn import functional as F
-from vllm.config import set_current_vllm_config
-from vllm.distributed import (
-    destroy_distributed_environment,
-    destroy_model_parallel,
-    init_distributed_environment,
-    initialize_model_parallel,
-)
-from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
-from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
-from vllm.utils.network_utils import get_open_port
-from vllm.utils.torch_utils import set_default_torch_dtype
-
-from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
-from vllm_omni.quantization import build_quant_config
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
-
-
-@pytest.mark.parametrize("method", ["fp8", "fp8_per_tensor"])
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_online_fp8_loads_bf16_weights_and_runs_native_kernel(method):
-    ignored_key = "ignored_layers" if method == "fp8" else "ignore"
-    config = build_quant_config(method, **{ignored_key: ["skip"]})
-    od_config = OmniDiffusionConfig(model="test", dtype=torch.bfloat16, quantization_config=config)
-    vllm_config = create_diffusion_vllm_config(torch.device("cuda:0"), od_config)
-
-    try:
-        init_distributed_environment(
-            world_size=1,
-            rank=0,
-            local_rank=0,
-            distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
-            backend="nccl",
-        )
-        with set_current_vllm_config(vllm_config), set_default_torch_dtype(torch.bfloat16), torch.device("cuda:0"):
-            initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
-            layer = ReplicatedLinear(
-                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="linear", disable_tp=True
-            )
-            assert isinstance(layer.quant_method, Fp8PerTensorOnlineLinearMethod)
-            assert layer.weight.is_meta
-
-            generator = torch.Generator(device="cuda").manual_seed(42)
-            weights = torch.randn(128, 128, generator=generator, dtype=torch.bfloat16) / 128**0.5
-            inputs = torch.randn(32, 128, generator=generator, dtype=torch.bfloat16)
-            expected = F.linear(inputs, weights).float()
-
-            # This is the loader installed by the native online quantizer.
-            # It materializes and quantizes the BF16 checkpoint in one pass.
-            layer.weight.weight_loader(layer.weight, weights)
-            assert layer.weight.dtype == torch.float8_e4m3fn
-            actual = layer(inputs)[0].float()
-            relative_mse = (actual - expected).square().mean() / expected.square().mean()
-            assert relative_mse.item() < 0.002
-
-            ignored = ReplicatedLinear(
-                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="skip", disable_tp=True
-            )
-            assert isinstance(ignored.quant_method, UnquantizedLinearMethod)
-            assert ignored.weight.dtype == torch.bfloat16
-    finally:
-        destroy_model_parallel()
-        destroy_distributed_environment()

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -11,7 +11,6 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams
-from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import Request, RequestStatus
@@ -636,45 +635,6 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
     return _build
 
 
-@pytest.mark.parametrize("active_stream_window", [0, 2])
-def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window):
-    adapter, _ = build_adapter(active_stream_window=active_stream_window)
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.deferred_waiting = set()
-    scheduler.running = []
-    scheduler.chunk_transfer_adapter = adapter
-    scheduler.input_coordinator = None
-    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
-    holder.num_computed_tokens = 2
-    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
-    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
-    scheduler._enqueue_waiting_request(fresh)
-    scheduler._enqueue_waiting_request(holder)
-
-    scheduler._process_pending_omni_inputs(model_mode="ar")
-
-    assert not scheduler.waiting and not scheduler.kv_holding_waiting
-    assert list(adapter.waiting_for_chunk_waiting_requests) == [holder, fresh]
-    assert holder.status == fresh.status == RequestStatus.WAITING_FOR_CHUNK
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-
-    adapter._finished_load_reqs.add(holder.request_id)
-    scheduler._process_pending_omni_inputs(model_mode="ar")
-
-    assert holder.status == RequestStatus.WAITING
-    assert fresh.status == RequestStatus.WAITING_FOR_CHUNK
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert not scheduler.waiting
-    assert adapter.requests_with_ready_chunks == {holder.request_id}
-    scheduler._restore_omni_wait_queues()
-    assert list(scheduler.kv_holding_waiting) == [holder]
-    assert list(scheduler.waiting) == [fresh]
-
-
 def _dequeue_load_entry(adapter, request):
     """Register and take the entry as the background recv loop would."""
     adapter.load_async(request)
@@ -2485,8 +2445,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
     scheduler.requests = {request.request_id: request}
     scheduler.running = [request] if placement == "running" else []
     scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
-    scheduler.kv_holding_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
+    scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
-    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
 
     freed = []
@@ -2510,7 +2469,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
     assert not adapter.waiting_for_chunk_running_requests
     assert scheduler.running == []
     assert list(scheduler.waiting) == []
-    assert list(scheduler.kv_holding_waiting) == []
+    assert list(scheduler.skipped_waiting) == []
     assert scheduler.num_waiting_for_streaming_input == 0
 
     fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
@@ -2533,8 +2492,7 @@ def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapt
     scheduler.requests = {request.request_id: request}
     scheduler.running = []
     scheduler.waiting = DummyWaitingQueue()
-    scheduler.kv_holding_waiting = DummyWaitingQueue()
+    scheduler.skipped_waiting = DummyWaitingQueue()
-    scheduler.deferred_waiting = set()
     scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
 
     assert (
@@ -2608,7 +2566,6 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
 
     scheduler = mocker.MagicMock()
     scheduler._first_chunk_express = False
-    scheduler.aux_output_connector = None
     scheduler._express_min_slack_s = 0.0
     scheduler.chunk_transfer_adapter = adapter_mock
     scheduler.connector = None
@@ -2689,7 +2646,6 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
 
     scheduler = mocker.MagicMock()
-    scheduler.aux_output_connector = None
     scheduler.chunk_transfer_adapter = adapter_mock
     bind_omits_transfer_helpers(scheduler)
     scheduler.connector = None
@@ -2788,8 +2744,6 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
     sched.requests = {}
     sched.running = []
     sched.waiting = []
-    sched.kv_holding_waiting = []
-    sched.deferred_waiting = set()
 
     with patch.object(VLLMScheduler, "finish_requests", _super_finish):
         OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
@@ -2805,8 +2759,7 @@ def _parked_sender_scheduler(adapter, session):
     scheduler.requests = {session.request_id: session}
     scheduler.running = []
     scheduler.waiting = DummyWaitingQueue()
-    scheduler.kv_holding_waiting = DummyWaitingQueue([session])
+    scheduler.skipped_waiting = DummyWaitingQueue([session])
-    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = 1
     scheduler.finished_req_ids_dict = {}
 

tests/e2e/accuracy/sana_video/conftest.py

@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import pytest
-
-
-@pytest.fixture
-def sana_model_parallel(tmp_path):
-    """Construct native parallel layers with real single-rank GPU groups."""
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_distributed_env,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method=f"file://{tmp_path / 'sana-distributed'}",
-    )
-    try:
-        initialize_model_parallel()
-        yield
-    finally:
-        destroy_distributed_env()

tests/e2e/accuracy/sana_video/test_sana_video_golden.py

@@ -202,7 +202,6 @@ def test_sana_video_transformer_matches_frozen_golden(
     variant: str,
     _hardware,
     accuracy_artifact_root: Path,
-    sana_model_parallel,
 ) -> None:
     del _hardware
     from vllm.config import LoadConfig

tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py

@@ -189,7 +189,6 @@ def _cfg_prediction(
 def test_sana_video_quick_pipeline_stages_match_diffusers(
     variant: str,
     _hardware,
-    sana_model_parallel,
 ) -> None:
     """Compare the first denoising step without replacing native attention.
 

tests/e2e/online_serving/test_flux_kontext_expansion.py

@@ -23,8 +23,7 @@ BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
 def _get_diffusion_feature_cases(model: str):
     return [
         pytest.param(
-            # The unsharded transformer alone exceeds the L4's 22 GiB.
+            OmniServerParams(model=model),
-            OmniServerParams(model=model, server_args=["--enable-layerwise-offload"]),
             id="base",
             marks=BASE_FEATURE_MARKS,
         ),

tests/e2e/online_serving/test_moss_tts_expansion.py

@@ -28,7 +28,7 @@ import pytest
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from tests.helpers.stage_config import get_deploy_config_path
 
 pytestmark = [
     pytest.mark.advanced_model,
@@ -84,12 +84,7 @@ tts_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=modify_stage_config(
+            stage_config_path=get_deploy_config_path("moss_tts_realtime.yaml"),
-                get_deploy_config_path("moss_tts_realtime.yaml"),
-                # Leave room on the single L4 for the codec weights, reference
-                # encoder, and speech-serving activations in addition to KV.
-                updates={"stages": {0: {"gpu_memory_utilization": 0.35}}},
-            ),
             server_args=["--disable-log-stats"],
         ),
         id="moss_tts_realtime",

tests/e2e/online_serving/test_qwen2_5_omni_expansion.py

@@ -21,14 +21,7 @@ models = ["Qwen/Qwen2.5-Omni-7B"]
 
 # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
 # platforms: section in vllm_omni/deploy/ci/qwen2_5_omni.yaml.
-stage_configs = [
+stage_configs = [modify_stage_config(get_deploy_config_path("ci/qwen2_5_omni.yaml"))]
-    modify_stage_config(
-        get_deploy_config_path("ci/qwen2_5_omni.yaml"),
-        # This suite requests two cards. Share the Talker's card with
-        # Code2Wav and leave room for its whole-sequence DiT workspace.
-        updates={"stages": {1: {"gpu_memory_utilization": 0.3}, 2: {"devices": "1"}}},
-    )
-]
 
 # Create parameter combinations for model and stage config
 test_params = [

tests/engine/test_arg_utils.py

@@ -20,11 +20,7 @@ import pytest
 from omegaconf import OmegaConf
 from pydantic import ValidationError
 from transformers import PretrainedConfig, Qwen3OmniMoeConfig
-from vllm.config import VllmConfig
 from vllm.engine.arg_utils import EngineArgs
-from vllm.exceptions import VLLMValidationError
-from vllm.v1.engine.input_processor import InputProcessor
-from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
 
 from tests.helpers.mock import patch_hf_snapshot_download
 from vllm_omni.config.model import OmniModelConfig
@@ -36,36 +32,6 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorMode
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-@pytest.mark.parametrize("use_v2", [False, True])
-def test_stage_runner_selects_matching_input_processor_validator(monkeypatch, use_v2):
-    from vllm.sampling_params import SamplingParams
-
-    from vllm_omni.model_executor.models.audex.cfg import AudexCFGLogitsProcessor
-
-    # Upstream's process-wide preference must not override a stage's deploy choice.
-    monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", not use_v2)
-    config = VllmConfig()
-    model_config = object.__new__(OmniModelConfig)
-    model_config.use_v2_model_runner = use_v2
-    model_config.runner_type = "generate"
-    model_config.logits_processors = [LogitBiasState if use_v2 else AudexCFGLogitsProcessor]
-    config.model_config = model_config
-    assert config.use_v2_model_runner is use_v2
-    processor = object.__new__(InputProcessor)
-    processor.vllm_config, processor.model_config = config, model_config
-    validator = processor._build_logits_processors_params_validator()
-    validator(SamplingParams(extra_args={"cfg_role": "cond", "cfg_scale": 1.5}))
-    if not use_v2:
-        with pytest.raises(VLLMValidationError, match="cfg_role"):
-            validator(SamplingParams(extra_args={"cfg_role": "invalid"}))
-
-
-@pytest.mark.parametrize("use_v2", [False, True])
-def test_plain_vllm_config_retains_upstream_runner_preference(monkeypatch, use_v2):
-    monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", use_v2)
-    assert VllmConfig().use_v2_model_runner is use_v2
-
-
 def test_sync_config_is_omni():
     """Ensure create_model_config gives the right type."""
     cfg = OmniEngineArgs().create_model_config()

tests/engine/test_arg_utils_shared_fields.py

@@ -8,7 +8,6 @@ from __future__ import annotations
 from dataclasses import fields
 
 import pytest
-from vllm.config import AuxOutputConfig
 
 from vllm_omni.engine.arg_utils import (
     SHARED_FIELDS,
@@ -50,26 +49,3 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
     """
     blacklist = internal_blacklist_keys()
     assert blacklist == orchestrator_field_names() - SHARED_FIELDS
-
-
-@pytest.mark.parametrize(
-    "aux_args",
-    [
-        {"enable_return_routed_experts": True},
-        {"aux_output_config": {"enable_return_routed_experts": True}},
-        {"aux_output_config": AuxOutputConfig(enable_return_routed_experts=True)},
-    ],
-)
-@pytest.mark.parametrize("use_v2_model_runner", [False, True])
-@pytest.mark.parametrize("is_npu", [False, True])
-def test_aux_outputs_require_a_supported_omni_stage_runner(aux_args, use_v2_model_runner, is_npu, monkeypatch):
-    # Upstream runner selection reads this environment variable; Omni workers
-    # select their runner from the separate per-stage flag instead.
-    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
-    monkeypatch.setattr("vllm_omni.engine.arg_utils.current_omni_platform.is_npu", lambda: is_npu)
-    if use_v2_model_runner or is_npu:
-        args = OmniEngineArgs(model="unused", use_v2_model_runner=use_v2_model_runner, **aux_args)
-        assert args.aux_output_config.enabled
-    else:
-        with pytest.raises(ValueError, match="use_v2_model_runner=True for this Omni stage"):
-            OmniEngineArgs(model="unused", **aux_args)

tests/engine/test_async_omni_engine_input.py

@@ -168,37 +168,6 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
     output_processor.add_request.assert_not_called()
 
 
-def test_build_add_request_message_preserves_kv_hints_after_omni_payload_upgrade(mocker):
-    from vllm.v1.engine.input_processor import InputProcessor
-    from vllm.v1.kv_hints import KvHintsEnvelope
-
-    engine = object.__new__(AsyncOmniEngine)
-    params = SamplingParams(max_tokens=8)
-    engine.default_sampling_params_list = [params]
-    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
-    engine.supported_tasks = ("generate",)
-    processor = mocker.create_autospec(InputProcessor, instance=True)
-    hints = KvHintsEnvelope("1", "hint-message", [])
-
-    def process(**kwargs):
-        request = _make_engine_core_request()
-        request.kv_hints = kwargs["kv_hints"]
-        return request
-
-    processor.process_inputs.side_effect = process
-    engine.input_processor = processor
-    msg = engine._build_add_request_message(
-        request_id="req-1",
-        prompt={"prompt_token_ids": [1], "additional_information": {"speaker": "voice"}},
-        sampling_params_list=[params],
-        kv_hints=hints,
-    )
-    assert processor.process_inputs.call_args.kwargs["kv_hints"] is hints
-    assert isinstance(msg.prompt, OmniEngineCoreRequest)
-    assert msg.prompt.kv_hints is hints
-    assert msg.prompt.additional_information is not None
-
-
 def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
     params = SamplingParams(max_tokens=8)

tests/engine/test_async_omni_engine_stage_init.py

@@ -268,21 +268,11 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     """Regression test for forward references in make_async_mp_client."""
     import vllm_omni.engine.stage_engine_core_client as client_mod
 
-    # Snapshot the module namespace so a reload cannot strand other test
+    importlib.reload(client_mod)
-    # modules that captured the pre-reload class objects (isinstance checks
-    # against the stale class would then fail). ``importlib.reload`` mutates
-    # the module dict in place; restoring it leaves the runtime exactly as it
-    # was before this test ran.
-    saved = dict(client_mod.__dict__)
-    try:
-        importlib.reload(client_mod)
 
-        assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
+    assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
-            "StageEngineCoreClient | DPLBStageEngineCoreClient"
+        "StageEngineCoreClient | DPLBStageEngineCoreClient"
-        )
+    )
-    finally:
-        client_mod.__dict__.clear()
-        client_mod.__dict__.update(saved)
 
 
 def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):

tests/engine/test_omni_request.py

@@ -13,16 +13,13 @@ substitutable with the base class — including base-style calls that pass
 
 import inspect
 
-import msgspec
 import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
-from vllm.v1.engine import EngineCoreRequest
-from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope
 from vllm.v1.request import Request
 
-from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload
+from vllm_omni.engine import PromptEmbedsPayload
 from vllm_omni.request import OmniRequest, OmniStreamingUpdate
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -114,26 +111,3 @@ def test_model_intermediate_buffer_round_trips_to_streaming_update():
     assert update is not None
     assert request.model_intermediate_buffer is buffer
     assert update.model_intermediate_buffer is buffer
-
-
-def test_kv_hints_survive_omni_upgrade_wire_transport_and_scheduler_conversion():
-    hints = KvHintsEnvelope("1", "hint-message", [KvHintAction("a", "prefetch", "1", {"blocks": [7]})])
-    core = EngineCoreRequest(
-        request_id="hint-request",
-        prompt_token_ids=[1, 2],
-        mm_features=None,
-        sampling_params=SamplingParams(max_tokens=2),
-        pooling_params=None,
-        arrival_time=0.0,
-        lora_request=None,
-        cache_salt=None,
-        data_parallel_rank=None,
-        session_id="session",
-        kv_hints=hints,
-    )
-    upgraded = OmniEngineCoreRequest.from_request(core)
-    decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(upgraded), type=OmniEngineCoreRequest)
-    scheduled = OmniRequest.from_engine_core_request(decoded, block_hasher=None)
-    assert decoded.session_id == "session"
-    assert scheduled.kv_hints == hints
-    assert scheduled.session_id == "session"

tests/engine/test_omni_request_output.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for OmniRequestOutput class."""
 
-import numpy as np
 import pytest
 from PIL import Image
 from vllm.outputs import CompletionOutput, RequestOutput
@@ -158,7 +157,6 @@ class TestMsgpackRoundTrip:
         )
 
         source = _make_text_request_output(text="round trip")
-        source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
         out = OmniRequestOutput.from_stage_output(
             source,
             request_id="rt-1",
@@ -180,7 +178,6 @@ class TestMsgpackRoundTrip:
         assert decoded.outputs[0].text == "round trip"
         assert decoded.prompt == "test prompt"
         assert decoded.prompt_token_ids == [1, 2, 3]
-        np.testing.assert_array_equal(decoded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
 
     def test_round_trip_diffusion_output(self):
         """A diffusion output round-trips without data loss."""
@@ -305,16 +302,6 @@ class TestRequestOutputFieldParity:
         assert omni.num_cached_tokens is None
         assert omni.num_cache_creation_tokens is None
 
-    def test_fixed_token_scores_survive_stage_output_wrapping(self) -> None:
-        source = _make_text_request_output()
-        source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
-
-        omni = OmniRequestOutput.from_stage_output(source, stage_id=0)
-        forwarded = OmniRequestOutput.from_stage_output(omni, stage_id=1)
-
-        np.testing.assert_array_equal(forwarded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
-        assert OmniRequestOutput(request_id="unscored").prompt_token_id_logprobs is None
-
     def test_transfer_param_fields_are_copied_from_the_stage_output(self) -> None:
         """Declaring the fields is not enough — they must also be copied.
 

tests/engine/test_orchestrator_error_handling.py

@@ -1282,7 +1282,6 @@ async def test_cache_reset_rpc_failure_is_nonfatal_and_visits_all_replicas(metho
     [
         ("reset_mm_cache", {}, True),
         ("sleep", {}, True),
-        ("release_kv_cache_memory", {}, True),
         ("pause_scheduler", {"clear_cache": True}, True),
         ("pause_scheduler", {"clear_cache": False}, False),
         ("reset_prefix_cache", {}, False),

tests/engine/test_output_processor_mrv2_text.py

@@ -7,9 +7,8 @@ from unittest.mock import MagicMock
 
 import pytest
 import torch
-from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import EngineCoreRequest, FinishReason
+from vllm.v1.engine import FinishReason
-from vllm.v1.engine.logprobs import LogprobsProcessor
 
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.output_modality import OutputModality
@@ -17,7 +16,6 @@ from vllm_omni.engine.output_processor import (
     MultimodalOutputProcessor,
     OmniRequestState,
 )
-from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -111,37 +109,6 @@ def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload():
     assert not state.mm_accumulated.is_empty
 
 
-def test_fixed_token_scores_survive_output_processing_and_stage_wrapping():
-    processor, state = _make_processor(OutputModality.TEXT, _Detokenizer())
-    request = EngineCoreRequest(
-        request_id="r",
-        prompt_token_ids=[1, 2, 3],
-        mm_features=None,
-        sampling_params=SamplingParams(prompt_logprob_token_ids=[4, 5]),
-        pooling_params=None,
-        arrival_time=0.0,
-        lora_request=None,
-        cache_salt=None,
-        data_parallel_rank=None,
-    )
-    state.logprobs_processor = LogprobsProcessor.from_new_request(None, request)
-    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
-    processed = processor.process_outputs(
-        [
-            OmniEngineCoreOutput(
-                request_id="r",
-                new_token_ids=[42],
-                prompt_token_id_logprobs=scores,
-                finish_reason=FinishReason.STOP,
-            )
-        ]
-    )
-
-    (raw_output,) = processed.request_outputs
-    final_output = OmniRequestOutput.from_stage_output(raw_output, stage_id=0)
-    torch.testing.assert_close(torch.from_numpy(final_output.prompt_token_id_logprobs), scores)
-
-
 @pytest.mark.parametrize("streaming", [False, True])
 def test_generation_stage_mm_delivered_only_with_terminal(streaming):
     # Non-terminal chunks accumulate; the terminal output carries the audio.

tests/engine/test_stage_engine_args.py

@@ -97,6 +97,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
         "enable_multithread_weight_load",
         "env",
         "has_sampling_extra_args",
+        "log_level",
         "log_stats",
         "model_arch",
         "model_subdir",

tests/engine/test_stage_engine_core_client.py

@@ -5,19 +5,11 @@
 from __future__ import annotations
 
 from types import SimpleNamespace
-from unittest.mock import patch
 
 import pytest
-from vllm import SamplingParams
-from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.engine.stage_engine_core_client import (
+from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
-    DPLBStageEngineCoreClient,
-    StageEngineCoreClient,
-    StageEngineCoreClientBase,
-)
-from vllm_omni.engine.stage_init_utils import StageMetadata
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -38,46 +30,3 @@ def test_check_health_raises_when_resources_engine_dead():
     client = _make_client(engine_dead=True)
     with pytest.raises(EngineDeadError, match="engine core is dead"):
         client.check_health()
-
-
-@pytest.mark.parametrize("dp_size, external_lb", [(1, False), (2, False), (2, True)])
-@pytest.mark.parametrize("renderer", [None, object()])
-def test_stage_client_forwards_renderer_to_the_actual_upstream_constructor(dp_size, external_lb, renderer):
-    config = SimpleNamespace(
-        model_config=None,
-        parallel_config=SimpleNamespace(data_parallel_size=dp_size, data_parallel_external_lb=external_lb),
-    )
-    metadata = StageMetadata(
-        stage_id=0,
-        stage_type="llm",
-        engine_output_type="text",
-        is_comprehension=True,
-        requires_multimodal_data=False,
-        engine_input_source=[],
-        final_output=True,
-        final_output_type="text",
-        default_sampling_params=SamplingParams(),
-        custom_process_input_func=None,
-        model_stage=None,
-        runtime_cfg=None,
-    )
-    internal_lb = dp_size > 1 and not external_lb
-    parent = DPLBAsyncMPClient if internal_lb else AsyncMPClient
-    # Keep Omni construction/factory routing real; intercept only upstream's
-    # subprocess/socket boundary, using its current signature via autospec.
-    with patch.object(parent, "__init__", autospec=True, return_value=None) as init:
-        addresses = {"input_address": "tcp://192.0.2.1:1"}
-        client = StageEngineCoreClientBase.make_async_mp_client(
-            config, object, metadata=metadata, client_addresses=addresses, renderer=renderer
-        )
-        assert isinstance(client, DPLBStageEngineCoreClient if internal_lb else StageEngineCoreClient)
-        init.assert_called_once_with(
-            client,
-            config,
-            object,
-            log_stats=False,
-            client_addresses=addresses,
-            client_count=1,
-            client_index=0,
-            renderer=renderer,
-        )

tests/engine/test_stage_pool_collective_rpc.py

@@ -83,27 +83,6 @@ def test_collective_rpc_control_method_reraises_worker_error():
     asyncio.run(run())
 
 
-@pytest.mark.cpu
-def test_kv_memory_release_uses_engine_core_and_propagates_pause_errors():
-    async def run() -> None:
-        release = AsyncMock(return_value=None)
-        pool, client = _make_pool(release_kv_cache_memory_async=release)
-        await pool.collective_rpc(0, "release_kv_cache_memory", timeout=2.0)
-        release.assert_awaited_once_with()
-        client.collective_rpc_async.assert_not_awaited()
-
-        release.side_effect = RuntimeError("requires a completed pause first")
-        with pytest.raises(RuntimeError, match="completed pause"):
-            await pool.collective_rpc(0, "release_kv_cache_memory")
-
-        del client.release_kv_cache_memory_async
-        result = await pool.collective_rpc(0, "release_kv_cache_memory")
-        assert result["supported"] is False
-        client.collective_rpc_async.assert_not_awaited()
-
-    asyncio.run(run())
-
-
 @pytest.mark.cpu
 def test_collective_rpc_reset_caches_use_control_helpers():
     async def run() -> None:

tests/entrypoints/test_async_omni.py

@@ -10,7 +10,6 @@ import pytest
 from vllm.engine.protocol import StreamingInput
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
-from vllm.v1.kv_hints import KvHintsEnvelope
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.stage_config import get_deploy_config_path
@@ -139,39 +138,6 @@ def test_generate_forwards_lora_request_to_engine():
     asyncio.run(run())
 
 
-@pytest.mark.cpu
-@pytest.mark.parametrize("streaming", [None, "chunk", "empty"])
-def test_generate_forwards_kv_hints_on_initial_stage_submission(streaming):
-    async def run():
-        submissions = []
-
-        async def add(**kwargs):
-            submissions.append(kwargs)
-
-        omni = get_async_omni_instance(fake_add_request=add)
-        omni.engine.add_streaming_update_async = _noop
-        omni.engine.stage_configs = [SimpleNamespace(is_comprehension=True, stage_type="llm")]
-        omni.engine.stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))]
-        hints = KvHintsEnvelope("1", "message", [])
-        params = SamplingParams(max_tokens=2, output_kind=RequestOutputKind.DELTA)
-
-        async def chunks():
-            if streaming == "chunk":
-                yield StreamingInput(prompt={"prompt_token_ids": [1]})
-
-        async for _ in omni.generate(
-            prompt=chunks() if streaming else {"prompt_token_ids": [1]},
-            request_id="hints",
-            sampling_params_list=[params],
-            kv_hints=hints,
-        ):
-            pass
-        assert len(submissions) == 1
-        assert submissions[0]["kv_hints"] is hints
-
-    asyncio.run(run())
-
-
 @pytest.mark.cpu
 @pytest.mark.parametrize(
     "req_ids,cancel_prefix,expected_cancel_count",

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -194,56 +194,6 @@ def test_reset_encoder_cache_forwards_to_ar_stages():
     asyncio.run(run())
 
 
-@pytest.mark.cpu
-def test_release_kv_memory_requires_pause_and_wakes_only_ar_kv_allocations():
-    async def run() -> None:
-        omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
-        with pytest.raises(RuntimeError, match="completed pause"):
-            await omni.release_kv_cache_memory()
-        omni.collective_rpc.assert_not_awaited()
-
-        omni._paused = True
-        await omni.release_kv_cache_memory()
-        omni._clear_frontend_mm_cache.assert_awaited_once_with()
-        omni.collective_rpc.assert_awaited_once_with(
-            method="release_kv_cache_memory",
-            args=(),
-            kwargs=None,
-            stage_ids=[0, 2],
-            timeout=CACHE_RESET_TIMEOUT_S,
-        )
-        assert omni._stage_sleeping_tags == {0: {"kv_cache"}, 2: {"kv_cache"}}
-        assert not omni._level2_sleeping
-
-        omni.collective_rpc.reset_mock()
-        await omni.wake_up(stage_ids=[0, 2], tags=["kv_cache"])
-        omni.collective_rpc.assert_awaited_once_with(
-            method="wake_up", args=(), kwargs={"tags": ["kv_cache"]}, stage_ids=[0, 2]
-        )
-        assert not omni._sleeping_tags
-        assert omni._paused
-
-    asyncio.run(run())
-
-
-@pytest.mark.cpu
-def test_release_kv_memory_skips_diffusion_and_reports_engine_failures():
-    async def run() -> None:
-        diffusion = _make_omni(stage_types=["diffusion"])
-        await diffusion.release_kv_cache_memory()
-        diffusion.collective_rpc.assert_not_awaited()
-        assert not diffusion._paused
-
-        omni = _make_omni(stage_types=["llm"])
-        omni._paused = True
-        omni.collective_rpc.return_value = [{"supported": False, "error": "requires all executor memory resident"}]
-        with pytest.raises(RuntimeError, match="all executor memory resident"):
-            await omni.release_kv_cache_memory()
-        assert not omni._sleeping_tags
-
-    asyncio.run(run())
-
-
 @pytest.mark.cpu
 def test_reset_mm_cache_clears_frontend_then_forwards_to_ar_stages():
     async def run() -> None:

tests/helpers/assertions.py

@@ -4,11 +4,9 @@
 """Assertion and response validation helpers for tests."""
 
 import base64
-import hashlib
 import io
 import json
 import math
-import os
 import tempfile
 import threading
 import wave
@@ -1001,31 +999,6 @@ def _assert_transcript_matches(
     )
 
 
-def _retain_failed_speech_audio(
-    audio_bytes: bytes | None, transcript: str, expected_text: Any, error: AssertionError
-) -> None:
-    """Retain the exact ASR input when a CI speech-content assertion fails."""
-    directory = os.environ.get("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR")
-    if not directory or not audio_bytes:
-        return
-    stem = f"failed_speech_{hashlib.sha256(audio_bytes).hexdigest()[:16]}"
-    root = Path(directory)
-    try:
-        root.mkdir(parents=True, exist_ok=True)
-        (root / f"{stem}.wav").write_bytes(audio_bytes)
-        (root / f"{stem}.json").write_text(
-            json.dumps(
-                {"transcript": transcript, "expected_text": str(expected_text), "assertion": str(error)},
-                ensure_ascii=False,
-                indent=2,
-            )
-        )
-        print(f"Retained failed speech ASR input: {root / f'{stem}.wav'}")
-    except OSError as exc:
-        # Diagnostic I/O must never replace or suppress the content failure.
-        print(f"Could not retain failed speech ASR input: {exc}")
-
-
 def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str | None = None) -> None:
     """Validate speech API results from :class:`~tests.helpers.client.OmniResponse`.
 
@@ -1090,21 +1063,14 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
             if expected_text:
                 print(f"audio content is: {transcript}")
                 print(f"input text is: {expected_text}")
-                transcription_audio = _speech_audio_for_transcription(
+                _assert_transcript_matches(
-                    getattr(response, "audio_bytes", None), request_config
+                    transcript,
+                    _speech_audio_for_transcription(getattr(response, "audio_bytes", None), request_config),
+                    expected_text,
+                    threshold=0.9,
+                    escalation_model=request_config.get("transcript_escalation_model"),
+                    language=request_config.get("transcript_language"),
                 )
-                try:
-                    _assert_transcript_matches(
-                        transcript,
-                        transcription_audio,
-                        expected_text,
-                        threshold=0.9,
-                        escalation_model=request_config.get("transcript_escalation_model"),
-                        language=request_config.get("transcript_language"),
-                    )
-                except AssertionError as exc:
-                    _retain_failed_speech_audio(transcription_audio, transcript, expected_text, exc)
-                    raise
         _assert_preset_voice_gender_from_audio(
             response.audio_bytes,
             request_config.get("voice"),

tests/helpers/runtime.py

@@ -659,14 +659,6 @@ class OmniRunner:
         self.seed = seed
         self._prompt_len_estimate_cache: dict[str, Any] = {}
         self.omni: Any = None
-        self._runner_process = psutil.Process()
-        self._initial_child_ids: set[tuple[int, float]] = set()
-        self._owned_engine_processes: dict[int, float] = {}
-        for child in self._runner_process.children(recursive=True):
-            try:
-                self._initial_child_ids.add((child.pid, child.create_time()))
-            except psutil.Error:
-                pass
         try:
             from vllm_omni.entrypoints.omni import Omni
 
@@ -678,14 +670,12 @@ class OmniRunner:
                 deploy_config=deploy_config,
                 **kwargs,
             )
-            self._remember_engine_processes(startup=True)
             startup_s = time.perf_counter() - startup_t0
             if log_stats:
                 print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True)
         except BaseException:
             # ``with OmniRunner(...)`` never reaches ``__enter__``/``__exit__``
             # when construction fails after worker processes have started.
-            self._remember_engine_processes(startup=True)
             self.__exit__(None, None, None)
             raise
 
@@ -909,51 +899,18 @@ class OmniRunner:
     def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
         return self.omni.stop_profile(stages=stages)
 
-    def _remember_engine_processes(self, *, startup: bool = False) -> None:
-        """Track startup workers, then only descendants of those workers.
-
-        Existing children may belong to another runner or a shared transcriber.
-        Record identities before close() can reparent surviving workers.
-        """
-        try:
-            if startup:
-                candidates = self._runner_process.children(recursive=True)
-            else:
-                candidates = []
-                for pid, created in list(self._owned_engine_processes.items()):
-                    try:
-                        proc = psutil.Process(pid)
-                        if proc.create_time() == created:
-                            candidates.extend(proc.children(recursive=True))
-                    except psutil.Error:
-                        pass
-            for proc in candidates:
-                try:
-                    identity = (proc.pid, proc.create_time())
-                    if identity in self._initial_child_ids:
-                        continue
-                    if startup and any(
-                        (parent.pid, parent.create_time()) in self._initial_child_ids for parent in proc.parents()
-                    ):
-                        continue
-                    cmdline = " ".join(proc.cmdline()).lower()
-                    if "enginecore" in cmdline or "enginecore" in proc.name().lower():
-                        self._owned_engine_processes[identity[0]] = identity[1]
-                except psutil.Error:
-                    pass
-        except psutil.Error as error:
-            print(f"Warning: could not inspect runner workers: {error}")
-
     def _cleanup_process(self):
         try:
+            keywords = ["enginecore"]
             matched = []
-            for pid, created in self._owned_engine_processes.items():
+            for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
                 try:
-                    proc = psutil.Process(pid)
+                    cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
-                    if proc.create_time() == created:
+                    name = proc.name().lower()
-                        print(f"Found owned vllm process: PID={proc.pid}")
+                    if any(k in cmdline for k in keywords) or any(k in name for k in keywords):
+                        print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
                         matched.append(proc)
-                except psutil.Error:
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
                     pass
             for proc in matched:
                 try:
@@ -981,14 +938,11 @@ class OmniRunner:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self._remember_engine_processes()
+        omni = getattr(self, "omni", None)
-        try:
+        if omni is not None and hasattr(omni, "close"):
-            omni = getattr(self, "omni", None)
+            omni.close()
-            if omni is not None and hasattr(omni, "close"):
+        self._cleanup_process()
-                omni.close()
+        cleanup_test_environment()
-        finally:
-            self._cleanup_process()
-            cleanup_test_environment()
 
 
 # ---------------------------------------------------------------------------

tests/helpers/tests/test_assertions.py

@@ -2,9 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import io
-import json
 import wave
-from contextlib import nullcontext
 from types import SimpleNamespace
 
 import numpy as np
@@ -251,85 +249,6 @@ def test_pcm_transcript_escalation_uses_same_valid_wav_and_rejects_two_wrong_tra
         _assert_wav_wraps_pcm(wav_bytes, pcm_bytes, 24_000)
 
 
-@pytest.mark.parametrize("response_format", ["wav", "pcm"])
-@pytest.mark.parametrize("escalation_model", [None, "large-v3"])
-def test_failed_speech_capture_retains_exact_asr_input_and_original_failure(
-    monkeypatch, tmp_path, response_format, escalation_model
-):
-    pcm_bytes = _pcm_sine()
-    wav = io.BytesIO()
-    with wave.open(wav, "wb") as writer:
-        writer.setnchannels(1)
-        writer.setsampwidth(2)
-        writer.setframerate(24_000)
-        writer.writeframes(pcm_bytes)
-    audio_bytes = wav.getvalue() if response_format == "wav" else pcm_bytes
-    transcribed = []
-
-    def transcribe(audio_bytes, **kwargs):
-        transcribed.append(audio_bytes)
-        return "this is unrelated speech"
-
-    directory = tmp_path / "failed_audio"
-    monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
-    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", transcribe)
-    config = {
-        "input": "please read this requested sentence accurately",
-        "response_format": response_format,
-        "transcript_pcm_sample_rate": 24_000,
-        "transcript_escalation_model": escalation_model,
-    }
-    response = OmniResponse(success=True, audio_bytes=audio_bytes, audio_format=f"audio/{response_format}")
-    with pytest.raises(AssertionError, match="Transcript doesn't match input") as error:
-        assert_audio_speech_response(response, config, "full_model")
-
-    clips = list(directory.glob("*.wav"))
-    assert len(clips) == 1
-    assert all(clip == clips[0].read_bytes() for clip in transcribed)
-    _assert_wav_wraps_pcm(clips[0].read_bytes(), pcm_bytes, 24_000)
-    metadata = json.loads(clips[0].with_suffix(".json").read_text())
-    assert metadata["expected_text"] == config["input"]
-    assert metadata["transcript"] == "this is unrelated speech"
-    assert metadata["assertion"] == str(error.value)
-
-
-@pytest.mark.parametrize("capture_enabled", [False, True])
-@pytest.mark.parametrize("content_failure", [False, True])
-def test_speech_capture_is_opt_in_and_does_not_write_successes(monkeypatch, tmp_path, capture_enabled, content_failure):
-    directory = tmp_path / "failed_audio"
-    if capture_enabled:
-        monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
-    else:
-        monkeypatch.delenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", raising=False)
-    expected = "please read this requested sentence accurately"
-    monkeypatch.setattr(
-        assertions,
-        "convert_audio_bytes_to_text",
-        lambda *_args, **_kwargs: "this is unrelated speech" if content_failure else expected,
-    )
-    with pytest.raises(AssertionError, match="Transcript doesn't match input") if content_failure else nullcontext():
-        assert_audio_speech_response(
-            OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
-            {"input": expected, "response_format": "wav"},
-            "full_model",
-        )
-    assert directory.exists() is (capture_enabled and content_failure)
-    assert len(list(tmp_path.rglob("*.wav"))) == int(capture_enabled and content_failure)
-
-
-def test_speech_capture_io_failure_does_not_mask_content_failure(monkeypatch, tmp_path):
-    directory = tmp_path / "blocked"
-    directory.write_text("a file cannot be an output directory")
-    monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
-    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "this is unrelated speech")
-    with pytest.raises(AssertionError, match="Transcript doesn't match input"):
-        assert_audio_speech_response(
-            OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
-            {"input": "please read this requested sentence accurately", "response_format": "wav"},
-            "full_model",
-        )
-
-
 @pytest.mark.parametrize(
     "pcm_bytes",
     [

tests/helpers/tests/test_runtime_omni_runner.py

@@ -1,15 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Tests for OmniRunner startup rollback and worker ownership."""
+"""Tests for OmniRunner startup rollback."""
 
 from __future__ import annotations
 
-import subprocess
-import sys
-from collections.abc import Generator
-
-import psutil
 import pytest
 
 from tests.helpers.runtime import OmniRunner
@@ -43,114 +38,3 @@ def test_runner_init_rolls_back_on_omni_startup_failure(monkeypatch: pytest.Monk
     assert cleaned.count("process") == 1
     # Once at the start of ``__init__``, once from the constructor rollback.
     assert cleaned.count("env") == 2
-
-
-@pytest.fixture
-def live_workers() -> Generator[list[subprocess.Popen[bytes]], None, None]:
-    workers: list[subprocess.Popen[bytes]] = []
-    try:
-        yield workers
-    finally:
-        for worker in workers:
-            if worker.poll() is None:
-                worker.kill()
-            worker.wait(timeout=10)
-            if worker.stdout is not None:
-                worker.stdout.close()
-
-
-def _start_engine_worker(
-    workers: list[subprocess.Popen[bytes]], *, ignore_terminate: bool = False
-) -> subprocess.Popen[bytes]:
-    script = "import signal, time; "
-    if ignore_terminate:
-        script += "signal.signal(signal.SIGTERM, signal.SIG_IGN); "
-    script += "print('ready', flush=True); time.sleep(90)"
-    worker = subprocess.Popen([sys.executable, "-c", script, "enginecore"], stdout=subprocess.PIPE)
-    workers.append(worker)
-    assert worker.stdout is not None
-    assert worker.stdout.readline() == b"ready\n"
-    return worker
-
-
-@pytest.mark.parametrize("close_raises", [False, True])
-def test_runner_cleanup_preserves_unrelated_engine_workers(
-    monkeypatch: pytest.MonkeyPatch,
-    live_workers: list[subprocess.Popen[bytes]],
-    close_raises: bool,
-) -> None:
-    existing = _start_engine_worker(live_workers)
-    owned: list[subprocess.Popen[bytes]] = []
-    late_unrelated: list[subprocess.Popen[bytes]] = []
-    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
-
-    class _OwnedOmni:
-        def __init__(self, *args: object, **kwargs: object) -> None:
-            owned.append(_start_engine_worker(live_workers))
-
-        def close(self) -> None:
-            late_unrelated.append(_start_engine_worker(live_workers))
-            if close_raises:
-                raise RuntimeError("close failed")
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _OwnedOmni)
-    if close_raises:
-        with pytest.raises(RuntimeError, match="close failed"), OmniRunner("fake-model"):
-            pass
-    else:
-        with OmniRunner("fake-model"):
-            pass
-
-    assert owned[0].poll() is not None
-    assert existing.poll() is None
-    assert late_unrelated[0].poll() is None
-
-
-def test_runner_startup_failure_cleans_only_new_workers(
-    monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
-) -> None:
-    unrelated = _start_engine_worker(live_workers)
-    owned: list[subprocess.Popen[bytes]] = []
-    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
-
-    class _FailingOmni:
-        def __init__(self, *args: object, **kwargs: object) -> None:
-            owned.append(_start_engine_worker(live_workers))
-            raise RuntimeError("startup failed")
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _FailingOmni)
-    with pytest.raises(RuntimeError, match="startup failed"):
-        OmniRunner("fake-model")
-
-    assert owned[0].poll() is not None
-    assert unrelated.poll() is None
-
-
-def test_runner_cleanup_rejects_stale_process_identity(live_workers: list[subprocess.Popen[bytes]]) -> None:
-    unrelated = _start_engine_worker(live_workers)
-    runner = object.__new__(OmniRunner)
-    runner._owned_engine_processes = {unrelated.pid: psutil.Process(unrelated.pid).create_time() + 1}
-
-    runner._cleanup_process()
-
-    assert unrelated.poll() is None
-
-
-def test_runner_cleanup_kills_owned_worker_ignoring_terminate(
-    monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
-) -> None:
-    owned: list[subprocess.Popen[bytes]] = []
-    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
-
-    class _StubbornOmni:
-        def __init__(self, *args: object, **kwargs: object) -> None:
-            owned.append(_start_engine_worker(live_workers, ignore_terminate=True))
-
-        def close(self) -> None:
-            pass
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _StubbornOmni)
-    with OmniRunner("fake-model"):
-        assert owned[0].poll() is None
-
-    assert owned[0].poll() is not None

tests/inputs/test_mm_processor.py

@@ -1,17 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-from importlib import import_module
 from types import SimpleNamespace
 from unittest.mock import Mock
 
 import pytest
 from transformers import BatchFeature
-from vllm.config.multimodal import MultiModalDummyOptions
-from vllm.multimodal.parse import AudioProcessorItems, ImageProcessorItems, MultiModalDataItems
-from vllm.multimodal.processing import BaseMultiModalProcessor
 
-from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor
+from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -28,14 +24,10 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts,
 
     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()
-    class AudioItems(AudioProcessorItems):
+    items.get_all_counts.return_value = counts
-        def get_passthrough_data(self):
-            return {"embedding": "preserve"}
-
-    items = MultiModalDataItems({"audio": AudioItems(["wave"]), "image": ImageProcessorItems([])} if counts else {})
-    items.select = Mock(wraps=items.select)
     kwargs = {"sampling_rate": 16000}
     if original_text is not None:
         kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text
@@ -44,83 +36,9 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts,
     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"]} if counts else {},
+        {"audios": ["wave"]},
         {"sampling_rate": 16000},
         {},
     )
-    assert dict(result) == ({"features": "processed", "embedding": "preserve"} if counts else {"features": "processed"})
+    assert dict(result) == {"features": "processed", "embedding": "preserve"}
     assert kwargs == before
-
-
-@pytest.mark.parametrize(
-    "module_name,builder_name,processor_name,expected_kwargs",
-    [
-        (
-            "cosyvoice3",
-            "CosyVoice3DummyInputsBuilder",
-            "CosyVoice3MultiModalProcessor",
-            {"prompt_text": "Testing my voices. Why should I not?"},
-        ),
-        (
-            "glm_tts",
-            "GLMTTSDummyInputsBuilder",
-            "GLMTTSMultiModalProcessor",
-            {"prompt_text": "This is the reference voice."},
-        ),
-        (
-            "omnivoice",
-            "OmniVoiceDummyInputsBuilder",
-            "OmniVoiceMultiModalProcessor",
-            {"ref_text": "Testing voice cloning."},
-        ),
-        ("mimo_audio", "MiMoAudioLLMDummyInputsBuilder", "MiMoAudioLLMMultiModalProcessor", {}),
-    ],
-)
-def test_dummy_budget_preserves_model_conditioning(module_name, builder_name, processor_name, expected_kwargs):
-    """The v0.31 budget path must reach each model's custom builder hook."""
-    module = import_module(f"vllm_omni.model_executor.models.{module_name}.{module_name}")
-    builder = object.__new__(getattr(module, builder_name))
-    assert isinstance(builder, OmniDummyInputsBuilder)
-    items = MultiModalDataItems({})
-    builder.info = SimpleNamespace(ctx=SimpleNamespace(tokenizer=None), parse_mm_data=Mock(return_value=items))
-    builder.get_dummy_mm_data = Mock(return_value={})
-    processor = object.__new__(getattr(module, processor_name))
-    processor.dummy_inputs = builder
-
-    inputs = processor.get_dummy_inputs(128, {"audio": 1}, MultiModalDummyOptions())
-
-    assert inputs.mm_data_items is items
-    assert inputs.hf_processor_mm_kwargs == expected_kwargs
-    assert inputs.prompt == ("<|empty|>" if module_name == "mimo_audio" else [])
-
-
-def test_dummy_budget_keeps_upstream_path_for_standard_builders(monkeypatch):
-    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()
-    expected = object()
-    upstream = Mock(return_value=expected)
-    monkeypatch.setattr(BaseMultiModalProcessor, "get_dummy_inputs", upstream)
-    options = MultiModalDummyOptions()
-
-    assert processor.get_dummy_inputs(128, {"audio": 1}, options) is expected
-    upstream.assert_called_once_with(128, {"audio": 1}, options)
-
-
-def test_voxtral_dummy_budget_uses_audio_token_prompt():
-    from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import VoxtralTTSMultiModalProcessor
-
-    processor = object.__new__(VoxtralTTSMultiModalProcessor)
-    expected = SimpleNamespace(prompt=[1, 25, 24, 35])
-    builder = SimpleNamespace(get_dummy_processor_inputs=Mock(return_value=expected))
-    processor.dummy_inputs = builder
-    options = MultiModalDummyOptions()
-
-    assert processor.get_dummy_inputs(128, {"audio": 1}, options) is expected
-    builder.get_dummy_processor_inputs.assert_called_once_with(128, {"audio": 1}, options)

tests/inputs/test_preprocess.py

@@ -8,9 +8,7 @@ from unittest.mock import AsyncMock, Mock
 
 import pytest
 import torch
-from vllm.exceptions import VLLMValidationError
+from vllm.inputs import tokens_input
-from vllm.inputs import mm_input, tokens_input
-from vllm.multimodal.inputs import PlaceholderRange
 from vllm.pooling_params import PoolingParams
 from vllm.renderers import BaseRenderer
 from vllm.renderers.params import TokenizeParams
@@ -117,35 +115,6 @@ def test_omni_renderer_is_a_real_subclass_created_once():
     assert not hasattr(instance, "_renderer")
 
 
-@pytest.mark.parametrize("async_mode", [False, True])
-@pytest.mark.parametrize("no_media_kwargs", [False, True])
-@pytest.mark.parametrize("limit", [2, 4])
-def test_renderer_applies_expanded_prompt_bound_without_splitting_media(renderer, async_mode, no_media_kwargs, limit):
-    renderer._process_multimodal.return_value = mm_input(
-        list(range(6)), {}, {"image": ["hash"]}, {"image": [PlaceholderRange(offset=1, length=2)]}
-    )
-    prompt = {"prompt_token_ids": [1, 2, 3], "additional_information": {"speaker": 1}}
-    if no_media_kwargs:
-        prompt["mm_processor_kwargs"] = {"target_h": 512}
-    else:
-        prompt["multi_modal_data"] = {"image": "image"}
-    params = TokenizeParams(max_total_tokens=128, truncate_prompt_tokens=limit, truncation_side="right")
-
-    def render():
-        if async_mode:
-            return asyncio.run(renderer.render_cmpl_async([prompt], params))
-        return renderer.render_cmpl([prompt], params)
-
-    if limit == 2:
-        with pytest.raises(VLLMValidationError, match="would split or drop"):
-            render()
-    else:
-        (result,) = render()
-        assert result["prompt_token_ids"] == [0, 1, 2, 3]
-        assert result["mm_placeholders"]["image"] == [PlaceholderRange(offset=1, length=2)]
-        assert result["additional_information"] == {"speaker": 1}
-
-
 def test_concrete_renderer_overrides_are_not_suppressed():
     """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides."""
 

tests/metrics/test_emit_calls.py

@@ -426,8 +426,6 @@ def _make_scheduler_shell() -> OmniARScheduler:
     obj.requests = {}
     obj.running = []
     obj.waiting = []
-    obj.kv_holding_waiting = []
-    obj.deferred_waiting = set()
     obj._kv_wait_start_ts = {}
     obj._omni_kv_config = None
     return obj

tests/metrics/test_stat_logger.py

@@ -477,7 +477,6 @@ def stats_config():
         observability_config=ObservabilityConfig(kv_cache_metrics=True),
         speculative_config=Mock(spec=SpeculativeConfig, num_speculative_tokens=3),
         kv_transfer_config=None,
-        ec_transfer_config=None,
         lora_config=None,
     )
     yield config

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -928,51 +928,4 @@ def test_talker_output_contract(monkeypatch, mrv2, packed):
         model._sampling_eps = 1e-6
         model.mrv2_custom_sampler(SimpleNamespace(penalties_state=None))
     output = model.forward(torch.ones(2, dtype=torch.long), torch.arange(2), inputs_embeds=hidden)
-    assert output is hidden
+    assert output is (hidden if mrv2 or packed else sentinel)
-
-
-def test_request_conditioning_normalizes_singletons_without_changing_schema():
-    from vllm_omni.data_entry_keys import to_struct
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
-
-    token = torch.tensor([[7, 8]], dtype=torch.int32)
-    feat = torch.randn(1, 4, 80)
-    speaker = torch.randn(1, 192)
-    length = torch.tensor([[2]], dtype=torch.int32)
-    raw = {"embed": {"speech_token": token, "speech_feat": feat, "embedding": speaker, "speech_token_len": [length]}}
-    result = to_struct(_normalize_request_conditioning(raw))
-    assert result.embed.speech_token is token
-    assert result.embed.speech_feat is feat
-    assert result.embed.embedding is speaker
-    assert result.embed.speech_token_len is length
-    assert raw["embed"]["speech_token_len"] == [length]
-
-
-def test_request_conditioning_rejects_unsplit_voices():
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
-
-    with pytest.raises(ValueError, match="unsplit batch"):
-        _normalize_request_conditioning({"embed": {"embedding": [torch.ones(1, 192), torch.zeros(1, 192)]}})
-
-
-def test_request_conditioning_reaches_code2wav_with_true_prompt_length():
-    model = _make_code2wav_model()
-    output = model.forward(
-        torch.tensor([1, 2]),
-        torch.arange(2),
-        model_intermediate_buffer=[
-            {
-                "embed": {
-                    "speech_token": torch.tensor([[3, 2, 0, 0]], dtype=torch.int32),
-                    "speech_feat": torch.randn(1, 8, 80),
-                    "embedding": torch.randn(1, 192),
-                    "speech_token_len": [torch.tensor([[2]], dtype=torch.int32)],
-                },
-                "meta": {"req_id": ["voice"]},
-            }
-        ],
-    )
-    call = model.code2wav.forward_calls[0]
-    assert call["prompt_token"].shape == (1, 2)
-    assert call["prompt_feat"].shape == (1, 4, 80)
-    assert output.multimodal_outputs

tests/model_executor/models/moss_tts/test_local_determined_tokens.py

@@ -7,11 +7,7 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
-from vllm.config import VllmConfig
 from vllm.sampling_params import SamplingParams
-from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
-from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
-from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 from vllm.v1.worker.gpu.states import RequestState
 
@@ -31,11 +27,11 @@ def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed,
         pytest.skip("CUDA required")
     device = torch.device("cuda")
     reqs = RequestState(5, 32, 16, 0, 8, device)
-    sampler = Sampler(VllmConfig(), 5, 8, device, reqs)
+    sampler = Sampler(5, 8, device, reqs)
     for i, length in enumerate([4, 1, 3]):
         reqs.add_request(str(i), length, [2] * length, 0, 8)
         slot = reqs.req_id_to_index[str(i)]
-        sampler.add_request(slot, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
+        sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
     reqs.apply_staged_writes()
     sampler.apply_staged_writes()
     state = _state(MossLocalModelState, device)
@@ -43,7 +39,6 @@ def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed,
     state.model._batch_state = None
     batch = _batch(device, [2, 4, 3], [1, 2, 1])
     batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32)
-    batch.seq_lens_cpu_upper_bound = torch.tensor([5, 2, 2], dtype=torch.int32)
     batch.cu_num_logits_np = np.arange(4, dtype=np.int32)
     batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device)
     batch.expanded_idx_mapping = batch.idx_mapping
@@ -78,10 +73,10 @@ def test_real_sampler_constraints_disable_determined_tokens(constraint):
         pytest.skip("CUDA required")
     device = torch.device("cuda")
     reqs = RequestState(5, 32, 16, 0, 8, device)
-    sampler = Sampler(VllmConfig(), 5, 8, device, reqs)
+    sampler = Sampler(5, 8, device, reqs)
     reqs.add_request("constrained", 1, [2], 0, 8)
     slot = reqs.req_id_to_index["constrained"]
-    sampler.add_request(slot, SamplingParams(**constraint))
+    sampler.add_request(slot, 1, SamplingParams(**constraint))
     state = _state(MossLocalModelState, device)
     state._direct_tokens = True
     # Use the actual upstream request registration, including its combined
@@ -90,22 +85,16 @@ def test_real_sampler_constraints_disable_determined_tokens(constraint):
 
 
 @pytest.mark.cpu
-@pytest.mark.parametrize(
+@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"])
-    "feature", ["logprobs", "bias", "penalty", "bad_words", "custom", "thinking", "trace", "mask", "nans"]
-)
 def test_distribution_features_fall_back_before_using_gpu(feature):
     sampler = object.__new__(Sampler)
     sampler.compute_nans = feature == "nans"
     sampler.return_sampling_mask = feature == "mask"
     sampler.trace_replay_state = object() if feature == "trace" else None
     sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None
-    bias = object.__new__(LogitBiasState)
+    sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"]))
-    bias.use_logit_bias = np.array([feature == "bias"])
+    sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"]))
-    penalties = object.__new__(PenaltiesState)
+    sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"])))
-    penalties.use_penalty = np.array([feature == "penalty"])
-    bad_words = object.__new__(BadWordsState)
-    bad_words.num_bad_words = SimpleNamespace(np=np.array([feature == "bad_words"]))
-    sampler.logits_processors = [bias, penalties, bad_words] + ([object()] if feature == "custom" else [])
     sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"]))
     state = object.__new__(MossLocalModelState)
     state._direct_tokens = True

tests/model_executor/models/moss_tts/test_local_eager_mtp.py

@@ -132,10 +132,10 @@ def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatc
         hidden_states=torch.zeros(2, 4),
         finished_req_ids=set(),
         ec_connector_output=None,
+        routed_experts=None,
     )
     runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
     runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
-    runner.aux_output_connector = None
     runner.model_config = SimpleNamespace(async_chunk=False)
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
     runner.model_state, runner.model = state, model
@@ -154,10 +154,7 @@ def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatc
             torch.zeros(1),
         )
     )
-    runner.prompt_logprobs_worker = SimpleNamespace(
+    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={}))
-        compute_prompt_logprobs=mocker.Mock(return_value={}),
-        compute_prompt_token_id_logprobs=mocker.Mock(return_value={}),
-    )
     runner.postprocess_sampled = mocker.Mock()
     runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None))
     mock_out = SimpleNamespace(copy_event=None)

tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py

@@ -42,26 +42,10 @@ def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProce
     return processor
 
 
-class _FakeProcessorItem:
-    def __init__(self, processor_data: dict[str, object]) -> None:
-        self._processor_data = processor_data
-
-    def get_processor_data(self) -> dict[str, object]:
-        return dict(self._processor_data)
-
-    def get_passthrough_data(self) -> dict[str, object]:
-        return {}
-
-
 def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace:
-    class _Selected:
-        @staticmethod
-        def values():
-            return [_FakeProcessorItem(mm_data)] if mm_data else []
-
     return SimpleNamespace(
         get_all_counts=lambda: counts,
-        select=lambda _keys: _Selected(),
+        select=lambda _keys: mm_data,
     )
 
 
@@ -91,8 +75,9 @@ def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> Non
 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}, {"audios": [audio]}), {"ref_text": "reference"})
+    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]
@@ -104,6 +89,7 @@ def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.Monke
 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"})
 
@@ -113,6 +99,4 @@ def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.Monke
 
 
 def test_legacy_call_hf_processor_is_gone() -> None:
-    # Upstream reintroduced a 2-arg `_call_hf_processor` on the base class; the
+    assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor")
-    # omni-specific 4-arg hook must not be defined on the omni class itself.
-    assert "_call_hf_processor" not in OmniVoiceMultiModalProcessor.__dict__

tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py

@@ -262,7 +262,6 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
         "audio": 1,
     }[modality]
     mm_data_items.__getitem__ = MagicMock(side_effect=lambda modality: [f"{modality}-0"])
-    mm_data_items.values = MagicMock(return_value=[])
 
     inputs = MagicMock()
     inputs.mm_data_items = mm_data_items
@@ -276,7 +275,6 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     cache = MagicMock()
     # Video miss, audio hit under pair-aware keys — expansion must force audio.
     cache.is_cached.side_effect = lambda keys: [False] if keys[0].startswith("v") else [True]
-    inputs.cache = cache
 
     captured = {}
 
@@ -311,8 +309,6 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     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
-    fake_self._postprocess_prompt = lambda prompt: prompt
-    inputs.prompt = [0, 1, 2]
 
     with patch(
         "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs",
@@ -553,6 +549,24 @@ def test_qwen2_5_prompt_updates_validate_mask_against_request_video_count():
         )
 
 
+def test_qwen2_5_mm_only_dummy_counts_subtract_only_videos_using_audio():
+    captured_counts = {}
+    fake_self = SimpleNamespace(
+        dummy_inputs=SimpleNamespace(get_dummy_text=lambda counts: captured_counts.update(counts) or "dummy"),
+        _apply_hf_processor_text_mm=lambda **_kwargs: ([], {}, False),
+    )
+    mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2, "audio": 1})
+
+    Qwen2_5OmniThinkerMultiModalProcessor._apply_hf_processor_mm_only(
+        fake_self,
+        mm_items,
+        {"use_audio_in_video": [True, False]},
+        {},
+    )
+
+    assert captured_counts == {"video": 2, "audio": 0}
+
+
 def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hits():
     fake_self = _fake_processor()
     update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID]))
@@ -647,12 +661,7 @@ def test_qwen3_processor_inherits_vllm_omni_per_video_helpers():
         Qwen3OmniMoeThinkerMultiModalProcessor,
         Qwen2_5OmniThinkerMultiModalProcessor,
     )
-    # Keep Qwen3's native HF preprocessing while inheriting Omni pair handling.
+    assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data
-    assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_inputs is UpstreamQwen3Processor._get_hf_mm_inputs
-    assert (
-        Qwen3OmniMoeThinkerMultiModalProcessor._get_audio_in_video_pairs
-        is Qwen2_5OmniThinkerMultiModalProcessor._get_audio_in_video_pairs
-    )
     mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__
     assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor)
 

tests/model_executor/models/qwen3_tts/test_qwen3_tts_prompt_text_batch.py

@@ -1,140 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from typing import Any
-
-import pytest
-import torch
-
-from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
-    Qwen3TTSConfig,
-    Qwen3TTSTalkerConfig,
-)
-from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
-    PRECOMPUTED_TEXT_IDS_KEY,
-    Qwen3TTSPromptEmbedsBuilder,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_batch_preprocess_projects_new_non_streaming_texts_once():
-    torch.manual_seed(0)
-    text_embedding = torch.nn.Embedding(64, 4)
-    projection = torch.nn.Linear(4, 4)
-    projected_shapes: list[tuple[int, ...]] = []
-
-    def text_projection(x: torch.Tensor) -> torch.Tensor:
-        projected_shapes.append(tuple(x.shape))
-        return projection(x)
-
-    builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder)
-    builder._text_embedding = text_embedding
-    builder._text_projection = text_projection
-    builder._batched_text_embeds = {}
-
-    ids_a = list(range(1, 13))
-    ids_b = list(range(20, 31))
-    # Serving stores the assistant-template ids wrapped in a one-element list.
-    serving_ids = [ids_a]
-    buf: dict[str, dict[str, Any]] = {
-        "a": {"req_id": "a", "task_type": ["CustomVoice"], PRECOMPUTED_TEXT_IDS_KEY: [ids_a]},
-        "b": {"req_id": "b", "task_type": ["VoiceDesign"], PRECOMPUTED_TEXT_IDS_KEY: [torch.tensor(ids_b)]},
-        "streaming": {
-            "req_id": "streaming",
-            "task_type": ["CustomVoice"],
-            "non_streaming_mode": [False],
-            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
-        },
-        "base": {"req_id": "base", "task_type": ["Base"], PRECOMPUTED_TEXT_IDS_KEY: serving_ids},
-        "built": {
-            "req_id": "built",
-            "task_type": ["CustomVoice"],
-            "embed": {"prefill": torch.zeros(1)},
-            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
-        },
-    }
-
-    builder.preprocess_infos_batch(req_infos=list(buf.values()), device=torch.device("cpu"))
-
-    # One embedding + projection over both requests' text tokens (template stripped).
-    assert projected_shapes == [(1, (len(ids_a) - 8) + (len(ids_b) - 8), 4)]
-    assert set(builder._batched_text_embeds) == {"a", "b"}
-    with torch.no_grad():
-        for req_id, ids in (("a", ids_a), ("b", ids_b)):
-            expected = projection(text_embedding(torch.tensor([ids[3:-5]])))
-            torch.testing.assert_close(builder._batched_text_embeds[req_id], expected)
-            assert buf[req_id][PRECOMPUTED_TEXT_IDS_KEY].tolist() == [ids]
-    for skipped in ("streaming", "base", "built"):
-        assert buf[skipped][PRECOMPUTED_TEXT_IDS_KEY] is serving_ids
-
-
-class _Tokenizer:
-    def __call__(self, text: str, **kwargs: Any) -> dict[str, torch.Tensor]:
-        return {"input_ids": torch.tensor([[2, 3, 4, 5]])}
-
-
-def _prompt_builder() -> Qwen3TTSPromptEmbedsBuilder:
-    """Use the real constructor and prompt assembly with small embedding tables."""
-    torch.manual_seed(42)
-    config = Qwen3TTSConfig(tts_bos_token_id=50, tts_eos_token_id=51, tts_pad_token_id=52)
-    talker_config = Qwen3TTSTalkerConfig(
-        codec_nothink_id=10,
-        codec_think_id=11,
-        codec_think_bos_id=12,
-        codec_think_eos_id=13,
-        codec_pad_id=14,
-        codec_bos_id=15,
-        codec_language_id={"english": 16},
-        spk_id={"vivian": 17, "serena": 18},
-    )
-    builder = Qwen3TTSPromptEmbedsBuilder(
-        config=config,
-        talker_config=talker_config,
-        model_path="",
-        text_embedding=torch.nn.Embedding(64, 4),
-        text_projection=torch.nn.Linear(4, 4),
-        codec_embed=torch.nn.Embedding(64, 4),
-        residual_code_embeddings=lambda: [],
-        speaker_encoder=torch.nn.Identity(),
-        tts_pad_embed=torch.zeros(1, 4),
-        encode_ref_audio_batch=lambda *args, **kwargs: [],
-    )
-    builder._text_tokenizer = _Tokenizer()
-    return builder
-
-
-@pytest.mark.parametrize("task_type", ["CustomVoice", "VoiceDesign"])
-@pytest.mark.parametrize("non_streaming", [False, True])
-def test_batched_prompt_consumption_matches_serial_and_clears_cache(task_type: str, non_streaming: bool):
-    batched, serial = _prompt_builder(), _prompt_builder()
-
-    def request(req_id: str, ids: list[int], speaker: str) -> dict[str, Any]:
-        return {
-            "req_id": req_id,
-            "text": ["hello"],
-            "task_type": [task_type],
-            "language": ["English"],
-            "speaker": [speaker],
-            "instruct": ["Speak calmly"],
-            "non_streaming_mode": [non_streaming],
-            PRECOMPUTED_TEXT_IDS_KEY: [ids],
-        }
-
-    with torch.inference_mode():
-        # The same request id can be used again after its earlier prompt was consumed.
-        for ids_a, ids_b in [(list(range(12)), list(range(20, 31))), (list(range(10, 24)), list(range(30, 46)))]:
-            pending = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
-            expected = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
-            batched.preprocess_infos_batch(req_infos=pending, device=torch.device("cpu"))
-            assert set(batched._batched_text_embeds) == ({"a", "b"} if non_streaming else set())
-            for actual_info, expected_info in zip(pending, expected, strict=True):
-                actual = batched.build_prompt_embeds(task_type=task_type, info_dict=actual_info)
-                reference = serial.build_prompt_embeds(task_type=task_type, info_dict=expected_info)
-                torch.testing.assert_close(actual[0], reference[0])
-                torch.testing.assert_close(actual[1], reference[1])
-                assert actual[2:] == reference[2:]
-                assert PRECOMPUTED_TEXT_IDS_KEY not in actual_info
-            assert not batched._batched_text_embeds
-        batched.preprocess_infos_batch(req_infos=[], device=torch.device("cpu"))
-        assert not batched._batched_text_embeds

tests/model_executor/models/test_omni_processing.py

@@ -79,16 +79,31 @@ def get_text_token_prompts(
     """Return ``(text_prompt, token_prompt)`` for the given multimodal data."""
     dummy_inputs = processor.dummy_inputs
     tokenizer: TokenizerLike = processor.info.get_tokenizer()
+    model_config = processor.info.ctx.model_config
 
     parsed_data = processor.info.parse_mm_data(mm_data)
     mm_counts = {k: len(vs) for k, vs in parsed_data.items()}
 
-    text_prompt = dummy_inputs.get_dummy_text(mm_counts)
+    inputs = dummy_inputs.get_dummy_processor_inputs(
-    token_prompt = tokenizer.encode(
+        model_config.max_model_len,
-        text_prompt,
+        mm_counts,
-        **processor.info.get_default_tok_params().get_encode_kwargs(),
+        mm_options={},
     )
 
+    text_prompt: str | None
+    token_prompt: list[int]
+    if isinstance(inputs.prompt, list):
+        text_prompt = None
+        token_prompt = inputs.prompt
+    elif isinstance(inputs.prompt, str):
+        text_prompt = inputs.prompt
+        token_prompt = tokenizer.encode(
+            text_prompt,
+            **processor.info.get_default_tok_params().get_encode_kwargs(),
+        )
+    else:
+        raise TypeError(type(inputs.prompt))
+
     return text_prompt, token_prompt
 
 
@@ -184,8 +199,8 @@ def _test_processing_correctness(
         modality: _to_dummy_options(modality, count) for modality, count in limit_mm_per_prompt_ints.items()
     }
 
-    baseline_processor = factories.build_processor(ctx)
+    baseline_processor = factories.build_processor(ctx, cache=None)
-    cached_processor = factories.build_processor(ctx)
+    cached_processor = factories.build_processor(ctx, cache=cache)
 
     rng = np.random.RandomState(0)
 
@@ -223,7 +238,6 @@ def _test_processing_correctness(
             baseline_processor,
             cached_processor,
             batch_idx,
-            cache,
         )
 
 
@@ -233,7 +247,6 @@ def _test_processing_correctness_one(
     baseline_processor: BaseMultiModalProcessor,
     cached_processor: BaseMultiModalProcessor,
     batch_idx: int,
-    cache: MultiModalProcessorOnlyCache,
 ):
     model_type = model_config.hf_config.model_type
 
@@ -245,14 +258,12 @@ def _test_processing_correctness_one(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
-        cache=None,
     )
 
     cached_tokenized_result = cached_processor(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
-        cache=cache,
     )
 
     _assert_inputs_equal(
@@ -267,13 +278,11 @@ def _test_processing_correctness_one(
             text_prompt,
             mm_items=mm_items,
             hf_processor_mm_kwargs={},
-            cache=None,
         )
         cached_text_result = cached_processor(
             text_prompt,
             mm_items=mm_items,
             hf_processor_mm_kwargs={},
-            cache=cache,
         )
 
         _assert_inputs_equal(

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -13,7 +13,6 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     text2flow_full_payload,
     text2flow_token_only,
 )
-from vllm_omni.outputs.mm_outputs import MultimodalPayload
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -122,24 +121,6 @@ def test_text2flow_full_payload_does_not_send_codec_ids():
     assert torch.equal(payload["embed"]["speech_token"], torch.tensor([[1, 2]], dtype=torch.long))
 
 
-def test_text2flow_full_payload_accepts_structured_multimodal_payload():
-    speech_token = torch.tensor([[1, 2]], dtype=torch.long)
-    payload = text2flow_full_payload(
-        None,
-        MultimodalPayload(
-            tensors={
-                "embed.speech_token": speech_token,
-                "codes.audio": torch.tensor([7, 8, 9], dtype=torch.long),
-            }
-        ),
-        SimpleNamespace(external_req_id="req-structured"),
-    )
-
-    assert payload is not None
-    assert "codes" not in payload
-    assert torch.equal(payload["embed"]["speech_token"], speech_token)
-
-
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     transfer_manager = _transfer_manager()
     request = SimpleNamespace(
@@ -374,36 +355,3 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     assert payload_final is not None
     assert payload_final.meta.finished.item() is True
     assert payload_final.codes.audio.tolist() == []
-
-
-def test_full_payload_retains_true_length_for_padded_conditioning():
-    from vllm_omni.data_entry_keys import to_struct
-
-    token = torch.tensor([[7, 8, 0, 0]], dtype=torch.int32)
-    feat = torch.randn(1, 8, 80)
-    speaker = torch.randn(1, 192)
-    length = torch.tensor([[2]], dtype=torch.int32)
-    result = text2flow_full_payload(
-        None,
-        {
-            "embed.speech_token": token,
-            "embed.speech_feat": feat,
-            "embed.embedding": speaker,
-            "embed.speech_token_len": length,
-        },
-        SimpleNamespace(request_id="padded"),
-    )
-    payload = to_struct(result)
-    assert payload.embed.speech_token_len is length
-
-
-def test_reference_features_are_snapshots_under_output_accumulation():
-    from vllm_omni.outputs.mm_outputs import MultimodalPayload
-    from vllm_omni.outputs.output_modality import OutputModality
-
-    first = torch.ones(1, 174, 80)
-    latest = torch.zeros(1, 1500, 80)
-    payload = MultimodalPayload.from_dict({"embed.speech_feat": first})
-    payload = payload.merged_with(MultimodalPayload.from_dict({"embed.speech_feat": latest}))
-    payload.consolidate_tensors(OutputModality.LATENT)
-    assert payload["embed.speech_feat"] is latest

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -971,14 +971,12 @@ def test_cosyvoice3_text2flow_full_payload_nested_fallback() -> None:
 
 
 def test_cosyvoice3_full_payload_replace_keys_present() -> None:
-    """Reference tensors and their length travel as complete snapshots."""
+    """Confirm _FULL_PAYLOAD_REPLACE_KEYS lists the three embed.* keys."""
     from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
         _FULL_PAYLOAD_REPLACE_KEYS,
     )
 
-    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset(
+    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"})
-        {"embed.speech_token", "embed.speech_token_len", "embed.speech_feat", "embed.embedding"}
-    )
 
 
 def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:

tests/model_executor/test_cfg_pairing.py

@@ -66,8 +66,6 @@ def _make_scheduler_cls() -> type:
     class FakeScheduler:
         def __init__(self) -> None:
             self.waiting = _FakeWaitingQueue()
-            self.kv_holding_waiting = _FakeWaitingQueue()
-            self.deferred_waiting: set[_FakeRequest] = set()
             self.running: list[_FakeRequest] = []
             self.requests: dict[str, _FakeRequest] = {}
             self.finished: list[set[str]] = []
@@ -198,44 +196,6 @@ def test_lone_member_is_held_then_released(scheduler: Any) -> None:
     assert scheduler.running == []
 
 
-def test_lone_kv_holding_member_is_held(scheduler: Any) -> None:
-    request = _FakeRequest("req-1", {"cfg_role": "cond"})
-    scheduler.add_request(request)
-    scheduler.waiting.remove_requests([request])
-    scheduler.kv_holding_waiting.append(request)
-
-    held = cfg_pairing._hold_incomplete_pairs(scheduler)
-
-    assert list(scheduler.kv_holding_waiting) == []
-    cfg_pairing._release_held(held)
-    assert list(scheduler.kv_holding_waiting) == [request]
-
-
-def test_deferred_kv_partner_holds_the_fresh_member(scheduler: Any) -> None:
-    cond, uncond = _add_pair(scheduler)
-    scheduler.waiting.remove_requests([uncond])
-    scheduler.kv_holding_waiting.append(uncond)
-    scheduler.deferred_waiting.add(uncond)
-
-    held = cfg_pairing._hold_incomplete_pairs(scheduler)
-
-    assert list(scheduler.waiting) == []
-    assert list(scheduler.kv_holding_waiting) == [uncond]
-    cfg_pairing._release_held(held)
-    assert list(scheduler.waiting) == [cond]
-
-
-def test_kv_holding_queue_keeps_pair_partners_adjacent(scheduler: Any) -> None:
-    cond, uncond = _add_pair(scheduler)
-    other = _FakeRequest("other")
-    scheduler.kv_holding_waiting.extend([cond, other, uncond])
-    scheduler.waiting.clear()
-
-    cfg_pairing._reorder_waiting_for_cfg(scheduler)
-
-    assert list(scheduler.kv_holding_waiting) == [cond, uncond, other]
-
-
 def test_complete_pair_is_admitted_together(scheduler: Any) -> None:
     scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
     scheduler.schedule()

tests/model_executor/test_seeded_exponential_platform.py

@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from unittest.mock import Mock
-
-import pytest
-import torch
-
-from tests.helpers.mark import hardware_test
-from vllm_omni.utils import seeded_exponential
-
-pytestmark = [pytest.mark.core_model]
-
-
-@pytest.mark.cpu
-@pytest.mark.parametrize("is_nvidia", [False, True], ids=["rocm", "cuda"])
-def test_cuda_tensor_requires_nvidia_distribution_contract(monkeypatch, is_nvidia: bool) -> None:
-    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: is_nvidia)
-    # HIP tensors expose the same Torch device predicate as NVIDIA tensors.
-    q = Mock(spec=torch.Tensor)
-    q.is_cuda = True
-    q.dtype = torch.float32
-    q.dim.return_value = 2
-    q.is_contiguous.return_value = True
-    generators = {0: object(), 1: object()}
-
-    assert seeded_exponential.batched_seeded_exponential_supported(q, generators) is is_nvidia
-
-
-@pytest.mark.cpu
-def test_shared_generator_keeps_ordered_torch_draws(monkeypatch) -> None:
-    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: True)
-    q = Mock(spec=torch.Tensor)
-    q.is_cuda = True
-    q.dtype = torch.float32
-    q.dim.return_value = 2
-    q.is_contiguous.return_value = True
-    shared = object()
-
-    assert not seeded_exponential.batched_seeded_exponential_supported(q, {0: shared, 1: shared})
-
-
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
-@pytest.mark.parametrize("seeded_rows", [(0, 1, 2), (0, 2)], ids=["all_seeded", "mixed"])
-@pytest.mark.parametrize("use_fp64_gumbel", [False, True])
-def test_torch_fallback_preserves_samples_and_generator_state(monkeypatch, seeded_rows, use_fp64_gumbel) -> None:
-    from vllm.v1.sample.ops import topk_topp_sampler as sampler_ops
-
-    import vllm_omni.patch  # noqa: F401
-
-    original = sampler_ops.random_sample.__wrapped__
-    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
-
-    def reject_kernel(*args, **kwargs):
-        raise AssertionError("the CUDA distribution kernel must not run on the fallback path")
-
-    monkeypatch.setattr(seeded_exponential, "fill_exponential_rows", reject_kernel)
-    probs = torch.arange(1, 258, device="cuda", dtype=torch.float32).expand(3, -1).contiguous()
-    probs /= probs.sum(dim=-1, keepdim=True)
-    expected_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
-    actual_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
-    expected_default = torch.cuda.get_rng_state()
-    actual_default = expected_default.clone()
-    try:
-        for _ in range(5):
-            torch.cuda.set_rng_state(expected_default)
-            # FP32 sampling divides probabilities in place. Give both paths
-            # the same untouched input on every step.
-            expected = original(probs.clone(), expected_gens, use_fp64_gumbel)
-            expected_default = torch.cuda.get_rng_state()
-            torch.cuda.set_rng_state(actual_default)
-            actual = sampler_ops.random_sample(probs.clone(), actual_gens, use_fp64_gumbel)
-            actual_default = torch.cuda.get_rng_state()
-            assert torch.equal(actual, expected)
-            assert torch.equal(actual_default, expected_default)
-            for row in seeded_rows:
-                assert torch.equal(actual_gens[row].get_state(), expected_gens[row].get_state())
-    finally:
-        torch.cuda.set_rng_state(expected_default)
-
-
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
-@pytest.mark.parametrize("rows", [None, [0, 2, 4]], ids=["grouped_codebooks", "selected_rows"])
-def test_direct_fill_fallback_preserves_ordered_draws(monkeypatch, rows) -> None:
-    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
-
-    def reject_kernel(*args, **kwargs):
-        raise AssertionError("the NVIDIA kernel must not run for direct fallback callers")
-
-    monkeypatch.setattr(seeded_exponential, "torch_exponential_policy", reject_kernel)
-    expected = torch.ones(6, 257, device="cuda")
-    actual = expected.clone()
-    expected_gen = torch.Generator(device="cuda").manual_seed(71)
-    actual_gen = torch.Generator(device="cuda").manual_seed(71)
-    # A shared seeded generator and an unseeded request make loop order and
-    # default-state advancement observable, including grouped codebooks.
-    expected_gens = [expected_gen, None, expected_gen]
-    actual_gens = [actual_gen, None, actual_gen]
-    original_default = torch.cuda.get_rng_state()
-    expected_default = original_default.clone()
-    actual_default = original_default.clone()
-    try:
-        for _ in range(5):
-            torch.cuda.set_rng_state(expected_default)
-            targets = expected.reshape(3, -1) if rows is None else [expected[row] for row in rows]
-            for target, generator in zip(targets, expected_gens):
-                target.exponential_(generator=generator)
-            expected_default = torch.cuda.get_rng_state()
-            torch.cuda.set_rng_state(actual_default)
-            assert seeded_exponential.fill_exponential_rows(actual, actual_gens, rows) is actual
-            actual_default = torch.cuda.get_rng_state()
-            assert torch.equal(actual, expected)
-            assert torch.equal(actual_gen.get_state(), expected_gen.get_state())
-            assert torch.equal(actual_default, expected_default)
-    finally:
-        torch.cuda.set_rng_state(original_default)

tests/platforms/npu/test_aux_output.py

@@ -1,161 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""CPU contract coverage for NPU routing capture and the native data plane."""
-
-from types import SimpleNamespace
-
-import numpy as np
-import pytest
-import torch
-from vllm.config import AuxOutputConfig
-from vllm.distributed.aux_output_connector import worker
-from vllm.distributed.aux_output_connector.connector import (
-    AuxOutputConnectorMetadata,
-    AuxOutputSchedulerConnector,
-    PackedBlockHashes,
-)
-from vllm.model_executor.layers.fused_moe import routed_experts_capturer
-from vllm.v1.kv_cache_interface import KVCacheConfig
-
-from vllm_omni.platforms.npu.worker.aux_output import NPUAuxOutputMixin
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _Router:
-    def _select_experts(self, logical_ids):
-        return torch.ones_like(logical_ids, dtype=torch.float32), logical_ids
-
-
-class _AscendMoE(torch.nn.Module):
-    __module__ = "vllm_ascend.ops.fused_moe"
-
-    def __init__(self):
-        super().__init__()
-        self.layer_id = 0
-        self.router = _Router()
-
-
-class _Runner(NPUAuxOutputMixin):
-    def _sync_device(self):
-        self.synchronized += 1
-
-
-@pytest.fixture
-def runner(monkeypatch):
-    # Exercise the real capturer, block store, worker and scheduler connectors
-    # with CPU tensors; importing vllm-ascend or owning NPU hardware is unnecessary.
-    monkeypatch.setattr(routed_experts_capturer, "current_platform", SimpleNamespace(device_type="cpu"))
-    monkeypatch.setattr(routed_experts_capturer, "get_forward_context", lambda: SimpleNamespace(dp_metadata=None))
-    monkeypatch.setattr(worker, "get_tp_group", lambda: SimpleNamespace(is_first_rank=True))
-    runner = _Runner()
-    runner.model = torch.nn.Sequential(_AscendMoE())
-    runner.vllm_config = SimpleNamespace(
-        aux_output_config=AuxOutputConfig(enable_return_routed_experts=True),
-        model_config=SimpleNamespace(
-            get_total_num_hidden_layers=lambda: 1,
-            get_num_experts=lambda: 16,
-            get_num_experts_per_tok=lambda: 2,
-            hf_text_config=SimpleNamespace(model_type="test_moe"),
-        ),
-        scheduler_config=SimpleNamespace(max_num_batched_tokens=16, max_num_seqs=4),
-        parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1, decode_context_parallel_size=1),
-        cache_config=SimpleNamespace(block_size=4, prefix_match_unit=None),
-        max_concurrent_batches=2,
-    )
-    runner.input_batch = SimpleNamespace(req_ids=["r"], num_computed_tokens_cpu=np.array([0]))
-    runner.query_start_loc = SimpleNamespace(np=np.array([0, 4]))
-    runner.synchronized = 0
-    runner._init_omni_aux_output(KVCacheConfig(8, [], []))
-    try:
-        yield runner
-    finally:
-        runner._close_omni_aux_output()
-
-
-def _step(runner, *, emit_start=0, hashes=None, finished=()):
-    scheduler_output = SimpleNamespace(
-        aux_output_connector_metadata=AuxOutputConnectorMetadata(
-            generation=0,
-            requests={} if finished else {"r": emit_start},
-            block_hashes=hashes or {},
-            finished_requests=finished,
-        ),
-        scheduled_spec_decode_tokens={},
-    )
-    runner._begin_omni_aux_output_step(scheduler_output)
-    return scheduler_output
-
-
-def _capture(runner, ids):
-    logical_ids = torch.tensor(ids, dtype=torch.int32)
-    weights, returned = runner.model[0].router._select_experts(logical_ids)
-    assert torch.equal(returned, logical_ids)
-    assert weights.shape == logical_ids.shape
-    return logical_ids.numpy()[:, None, :].astype(np.uint8)
-
-
-def test_routing_reaches_scheduler_and_owns_legacy_batch_snapshot(runner):
-    step = _step(runner)
-    expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]])
-    pending = runner._prepare_omni_aux_output()
-    # Simulate bookkeeping and the next batch reusing the runner's buffers.
-    runner.input_batch.num_computed_tokens_cpu[:] = 99
-    runner.query_start_loc.np[:] = 99
-    _capture(runner, [[7, 8]] * 4)
-    output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]))
-    request = SimpleNamespace(request_id="r", num_tokens=5)
-    rows = AuxOutputSchedulerConnector().take_output(request, output)
-    np.testing.assert_array_equal(rows, expected)
-    assert output["r"].token_start == 0
-    assert runner.synchronized == 1
-
-
-def test_speculative_rejection_drops_unaccepted_routing_suffix(runner):
-    step = _step(runner)
-    step.scheduled_spec_decode_tokens = {"r": [7, 8, 9]}
-    expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]])
-    pending = runner._prepare_omni_aux_output()
-    output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5, 6, -1, -1]]))
-    np.testing.assert_array_equal(output["r"].rows, expected[:2])
-
-
-def test_partial_prefill_accumulates_rows_and_no_forward_step_tears_down(runner):
-    step = _step(runner)
-    prefix = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]])
-    pending = runner._prepare_omni_aux_output()
-    assert runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]), [0]) == {}
-
-    # Both completed four-token blocks need their scheduler-provided hashes.
-    step = _step(runner, hashes={"r": PackedBlockHashes(b"h0h1", 2)})
-    runner.input_batch.num_computed_tokens_cpu[:] = 4
-    suffix = _capture(runner, [[4, 5], [5, 6], [6, 7], [7, 8]])
-    pending = runner._prepare_omni_aux_output()
-    output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]))
-    np.testing.assert_array_equal(output["r"].rows, np.concatenate((prefix, suffix)))
-    _step(runner, finished=("r",))
-    assert not runner.aux_output_connector._requests
-    store = runner.aux_output_connector._store
-    runner._close_omni_aux_output()
-    assert not store._thread.is_alive()
-    assert runner.model[0].capture_fn is None
-    runner._close_omni_aux_output()  # shutdown is idempotent
-
-
-def test_reinitialization_does_not_wrap_router_twice(runner):
-    selection = runner.model[0].router._select_experts
-    runner._init_omni_aux_output(KVCacheConfig(8, [], []))
-    assert runner.model[0].router._select_experts is selection
-    step = _step(runner)
-    expected = _capture(runner, [[0, 1]] * 4)
-    output = runner._finish_omni_aux_output(runner._prepare_omni_aux_output(), step, torch.tensor([[5]]))
-    np.testing.assert_array_equal(output["r"].rows, expected)
-
-
-def test_disabled_capture_preserves_normal_runner_path(runner):
-    runner._close_omni_aux_output()
-    runner.vllm_config.aux_output_config.enable_return_routed_experts = False
-    runner._init_omni_aux_output(KVCacheConfig(8, [], []))
-    runner._begin_omni_aux_output_step(SimpleNamespace())
-    assert runner._prepare_omni_aux_output() is None
-    assert runner._finish_omni_aux_output(None, SimpleNamespace()) is None

tests/worker/test_gpu_generation_model_runner.py

@@ -3,7 +3,6 @@
 
 import contextlib
 from types import SimpleNamespace
-from unittest.mock import Mock
 
 import pytest
 import torch
@@ -23,23 +22,6 @@ class _StopExecutionError(Exception):
     """Halts execute_model at the model boundary."""
 
 
-@pytest.mark.parametrize("randomize_inputs", [False, True])
-def test_generation_profile_accepts_upstream_randomize_inputs(randomize_inputs):
-    """Release workers can pass the profiling flag to codec-stage runners."""
-    runner = object.__new__(GPUGenerationModelRunner)
-    runner.supports_mm_inputs = False
-    runner.max_num_tokens = 128
-    runner.encoder_cache = {"temporary": object()}
-    runner._dummy_run = Mock(return_value=(torch.zeros(1), None))
-    runner._sync_device = Mock()
-
-    runner.profile_run(randomize_inputs=randomize_inputs)
-
-    runner._dummy_run.assert_called_once_with(128, is_profile=True, randomize_inputs=randomize_inputs)
-    runner._sync_device.assert_called_once_with()
-    assert runner.encoder_cache == {}
-
-
 class _DummyInputBatch:
     def __init__(self, num_reqs: int = 1):
         self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)]

tests/worker/test_omni_v2_profile_dispatch.py

@@ -15,8 +15,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
     runner = object.__new__(OmniGPUModelRunner)
     runner.cudagraph_manager = None
-    runner.parallel_config = SimpleNamespace(data_parallel_size=dp_size, data_parallel_rank=1)
+    runner.dp_size = dp_size
-    runner.dp_size, runner.dp_rank = dp_size, 1
+    runner.dp_rank = 1
 
     import vllm.v1.worker.gpu.dp_utils as dp_utils
 
@@ -52,7 +52,6 @@ def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
         return expected
 
     runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch)
-    runner.parallel_config.data_parallel_size = 1
     runner.dp_size = 1
     batch_desc, _ = runner._dispatch_batch_descriptor(
         num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8

tests/worker_v2/test_capture_query_bound_gpu.py

@@ -44,8 +44,8 @@ def _check_prefill():
     spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16)
     kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
     state = OmniModelState.__new__(OmniModelState)
-    state.model_config = SimpleNamespace(max_model_len=256)
     state.vllm_config = c
+    state.max_model_len = 256
     state.supports_mm_inputs = False
     with set_current_vllm_config(c):
         group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0)

tests/worker_v2/test_decode_query_bound_gpu.py

@@ -54,7 +54,7 @@ def test_full_decode_graph_matches_fa3_aot_schedule(graph_mode, tmp_path):
     kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
     state = OmniModelState.__new__(OmniModelState)
     state.vllm_config = c
-    state.model_config = c.model_config
+    state.max_model_len = 2048
     state.supports_mm_inputs = False
     with set_current_vllm_config(c):
         group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0)

tests/worker_v2/test_omni_ar_model_runner.py

@@ -11,13 +11,10 @@ from unittest.mock import MagicMock
 import numpy as np
 import pytest
 import torch
-from vllm.distributed.aux_output_connector.connector import AuxRequestOutput
+from vllm.v1.outputs import RoutedExpertsTensors
-from vllm.sampling_params import SamplingParams
 from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
-from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
 
 import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner
-from tests.helpers.mark import hardware_test
 from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
 from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput
 from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot
@@ -44,7 +41,7 @@ class _FakeEvent:
 def _async_output(req_ids=("req-0",), **overrides) -> OmniAsyncOutput:
     rid2idx = {rid: i for i, rid in enumerate(req_ids)}
     mro = omni_ar_model_runner.OmniModelRunnerOutput(list(req_ids), rid2idx, None, prompt_logprobs_dict={})
-    sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, torch.tensor([1]), torch.tensor([0]))
+    sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, None)
     kwargs = dict(model_runner_output=mro, sampler_output=sampler_output)
     kwargs.update(num_sampled_tokens=torch.tensor([1] * len(req_ids)), copy_event=_FakeEvent())
     kwargs.update(main_stream=_FakeStream(), copy_stream=_FakeStream())
@@ -67,38 +64,23 @@ def test_async_output_blocking_event_and_routing_masks(monkeypatch, compact_widt
         counts=torch.tensor([2, 0]),
         vocab_size=4,
     )
-    sampler_output = SamplerOutput(
+    sampler_output = SamplerOutput(torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), None, masks)
-        torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), torch.tensor([0, 0]), masks
+    routed = RoutedExpertsTensors(torch.tensor([[[2, 3]], [[4, 5]]]), torch.tensor([7, 9]))
-    )
-    copied: list[tuple[np.ndarray, np.ndarray]] = []
 
-    class _FakePendingAuxOutput:
-        def enqueue_cpu_copy(self, *, num_sampled, num_rejected) -> None:
-            copied.append((num_sampled, num_rejected))
-
-        def process_output(self) -> dict[str, AuxRequestOutput]:
-            return {
-                "decode": AuxRequestOutput(token_start=0, rows=np.array([[2, 3]], dtype=np.uint8)),
-                "prefill": AuxRequestOutput(token_start=0, rows=np.array([[4, 5]], dtype=np.uint8)),
-            }
-
-    pending = _FakePendingAuxOutput()
     output = _async_output(
         req_ids=["decode", "prefill"],
         sampler_output=sampler_output,
         num_sampled_tokens=torch.tensor([1, 0]),
         copy_event=None,  # None → constructor builds the default blocking event
-        pending_aux_output=pending,
+        routed_experts=routed,
     ).get_output()
     assert event_kwargs == [{"blocking": True}]  # blocking event by default
     assert output.sampled_token_ids == [[2], []]
     assert output.sampling_masks.to_nested_list() == [[0, 2], []]
     np.testing.assert_array_equal(output.sampling_masks.offsets, [0, 2, 2])
-    np.testing.assert_array_equal(output.aux_output_connector_output["decode"].rows, [[2, 3]])
+    np.testing.assert_array_equal(output.routed_experts.routing_data, [[[2, 3]], [[4, 5]]])
-    np.testing.assert_array_equal(output.aux_output_connector_output["prefill"].rows, [[4, 5]])
+    np.testing.assert_array_equal(output.routed_experts.slot_mapping, [7, 9])
-    assert output.aux_output_connector_output["decode"].token_start == 0
     np.testing.assert_array_equal(output.sampling_masks.token_ids, [0, 2])
-    assert len(copied) == 1 and copied[0][0].tolist() == [1, 0]
 
 
 @pytest.mark.parametrize("needs_history", [False, True])
@@ -106,26 +88,19 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
     runner = OmniARModelRunner.__new__(OmniARModelRunner)
     input_batch = SimpleNamespace(req_ids=["req"], num_reqs=1, seq_lens=torch.tensor([3]))
     input_batch.idx_mapping, input_batch.query_start_loc = torch.tensor([0]), torch.tensor([0, 1])
-    input_batch.idx_mapping_np = np.array([0])
-    input_batch.num_computed_prefill_tokens_np = np.array([0])
-    input_batch.num_scheduled_tokens = np.array([3])
-    input_batch.prefill_len_np = np.array([3])
-    input_batch.query_start_loc_np = np.array([0, 3])
     state = SimpleNamespace(input_batch=input_batch, hidden_states=torch.zeros(1, 2))
-    state.finished_req_ids, state.ec_connector_output = {"finished"}, None
+    state.finished_req_ids, state.ec_connector_output, state.routed_experts = {"finished"}, None, None
     runner.execute_model_state = state
     runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
     runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
-    runner.aux_output_connector = None
     runner.model_config = SimpleNamespace(async_chunk=False)
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
-    text_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+    runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(torch.zeros(1, 2), None)))
-    runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(text_hidden, None)))
     runner.model_state.intermediate_buffer = SimpleNamespace(buffers={0: {"global_request_id": "global-req"}})
     runner.req_states = SimpleNamespace(req_id_to_index={"req": 0})
     runner.req_states.all_token_ids = SimpleNamespace(gpu=torch.tensor([[1]]))
     runner.req_states.num_computed_tokens = SimpleNamespace(gpu=torch.tensor([0]))
-    runner.req_states.prompt_len = SimpleNamespace(np=np.array([3]), gpu=torch.tensor([3]))
+    runner.req_states.prompt_len = SimpleNamespace(np=np.array([1]), gpu=torch.tensor([1]))
     runner.main_stream = runner.output_copy_stream = MagicMock()
     runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = MagicMock()
     sampler_out = (SimpleNamespace(sampled_token_ids=torch.tensor([[2]])), MagicMock(), MagicMock())
@@ -139,17 +114,13 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
         yield
         sampling_active = False
 
-    runner.model = SimpleNamespace(
+    runner.model = SimpleNamespace(compute_logits=None, logitsprocs_need_output_token_ids=needs_history)
-        compute_logits=lambda hidden: hidden, logitsprocs_need_output_token_ids=needs_history
-    )
     runner.model.mrv2_sampling_context = sampling_context
     runner.sampler = None
     runner.sample = MagicMock(return_value=sampler_out)
     runner.sample.side_effect = lambda *_: sampler_out if sampling_active is needs_history else pytest.fail()
     logprobs_mock = MagicMock(side_effect=lambda *_: pytest.fail("inside ctx") if sampling_active else {})
-    runner.prompt_logprobs_worker = PromptLogprobsWorker(1, torch.device("cpu"), logprobs_mode="raw_logits")
+    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=logprobs_mock)
-    runner.prompt_logprobs_worker.add_request("req", 0, SamplingParams(prompt_logprob_token_ids=[1, 0]))
-    runner.prompt_logprobs_worker.compute_prompt_logprobs = logprobs_mock
     runner.postprocess_sampled, connector_output = MagicMock(), object()
 
     def post_forward(finished_req_ids):
@@ -164,67 +135,10 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
     assert runner.sample_tokens(None) is mock_out
     built = omni_ar_model_runner.OmniAsyncOutput.call_args.kwargs["model_runner_output"]
     assert built.kv_connector_output is connector_output
-    torch.testing.assert_close(built.prompt_token_id_logprobs_dict["req"], torch.tensor([[2.0, 1.0], [4.0, 3.0]]))
     assert runner._resolve_global_request_id("req") == "global-req"  # from the intermediate buffer
     assert runner._resolve_global_request_id("unknown") == "unknown"  # fallback to the local id
 
 
-def test_async_output_preserves_fixed_token_scores_on_host(monkeypatch) -> None:
-    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
-    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
-    mro = omni_ar_model_runner.OmniModelRunnerOutput(
-        req_ids=["req-0"],
-        req_id_to_index={"req-0": 0},
-        sampled_token_ids=None,
-        prompt_logprobs_dict={},
-        prompt_token_id_logprobs_dict={"req-0": scores},
-    )
-    output = _async_output(model_runner_output=mro).get_output()
-    assert output.prompt_token_id_logprobs_dict["req-0"].device.type == "cpu"
-    torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req-0"], scores)
-
-
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_chunked_fixed_token_scores_are_copied_from_cuda() -> None:
-    device = torch.device("cuda")
-    hidden = torch.tensor([[1.0, 2.0, 3.0], [3.0, 1.0, 2.0], [2.0, 3.0, 1.0], [4.0, 5.0, 6.0]], device=device)
-    worker = PromptLogprobsWorker(1, device)
-    worker.add_request("req", 0, SamplingParams(prompt_logprob_token_ids=[2, 0], prompt_logprob_start=1))
-    batch = SimpleNamespace(
-        req_ids=["req"],
-        idx_mapping_np=np.array([0]),
-        num_computed_prefill_tokens_np=np.array([0]),
-        num_scheduled_tokens=np.array([2]),
-        prefill_len_np=np.array([4]),
-        query_start_loc_np=np.array([0, 2]),
-    )
-    assert worker.compute_prompt_token_id_logprobs(lambda hidden: hidden, hidden[:2], batch, np.array([4])) == {}
-    batch.num_computed_prefill_tokens_np = np.array([2])
-    scores = worker.compute_prompt_token_id_logprobs(lambda hidden: hidden, hidden[2:], batch, np.array([4]))
-    mro = omni_ar_model_runner.OmniModelRunnerOutput(
-        req_ids=["req"],
-        req_id_to_index={"req": 0},
-        sampled_token_ids=None,
-        prompt_logprobs_dict={},
-        prompt_token_id_logprobs_dict=scores,
-    )
-    num_sampled = torch.tensor([1], device=device, dtype=torch.int32)
-    sampler_output = SamplerOutput(
-        torch.tensor([[1]], device=device), None, None, num_sampled, torch.zeros_like(num_sampled)
-    )
-    output = OmniAsyncOutput(
-        model_runner_output=mro,
-        sampler_output=sampler_output,
-        num_sampled_tokens=num_sampled,
-        main_stream=torch.cuda.current_stream(),
-        copy_stream=torch.cuda.Stream(),
-    ).get_output()
-
-    expected = hidden[1:3].log_softmax(dim=-1)[:, [2, 0]].cpu()
-    assert output.prompt_token_id_logprobs_dict["req"].device.type == "cpu"
-    torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req"], expected)
-
-
 def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     runner = OmniARModelRunner.__new__(OmniARModelRunner)
     runner.model_config = SimpleNamespace(async_chunk=True)
@@ -370,13 +284,7 @@ def test_async_output_slices_request_payloads_with_graph_padding(
     batch.num_tokens_after_padding = padded_total
     output = _async_output(
         req_ids=[f"req-{i}" for i in range(len(lengths))],
-        sampler_output=SamplerOutput(
+        sampler_output=SamplerOutput(torch.ones(len(lengths), 1, dtype=torch.long), None, None, None),
-            torch.ones(len(lengths), 1, dtype=torch.long),
-            None,
-            None,
-            torch.ones(len(lengths), dtype=torch.long),
-            torch.zeros(len(lengths), dtype=torch.long),
-        ),
         text_hidden=hidden,
         multimodal_outputs={"codes": {"audio": codes, "ref": refs}},
         input_batch=batch,
@@ -677,9 +585,7 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch,
 
     result = _async_output(
         req_ids=["audio", "partial"],
-        sampler_output=SamplerOutput(
+        sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None),
-            torch.tensor([[99], [100]]), None, None, torch.tensor([1, 1]), torch.tensor([0, 0])
-        ),
         num_sampled_tokens=torch.tensor([1, 0]),
         multimodal_outputs={"snapshot": snapshot},
         input_batch=batch,
@@ -804,7 +710,7 @@ def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocke
     extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None
     result = _async_output(
         req_ids=["frame", "eos", "prefill"],
-        sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None, None),
+        sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None),
         num_sampled_tokens=torch.tensor([1, 1, 1]),
         multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}},
         input_batch=batch,
@@ -847,9 +753,7 @@ def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming):
 
     result = _async_output(
         req_ids=["audio", "partial"],
-        sampler_output=SamplerOutput(
+        sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None),
-            torch.tensor([[99], [100]]), None, None, torch.tensor([1, 1]), torch.tensor([0, 0])
-        ),
         num_sampled_tokens=torch.tensor([1, 0]),
         multimodal_outputs={"snapshot": snapshot},
         input_batch=batch,

tests/worker_v2/test_omni_generation_model_runner.py

@@ -4,14 +4,12 @@
 """OmniGenerationModelRunner contracts: empty-step lifecycle, output partition,
 CPU-sync vs CUDA-async dispatch, and async-chunk slot recycling."""
 
-from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import numpy as np
 import pytest
 import torch
-from vllm.distributed.aux_output_connector.connector import AuxRequestOutput
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -20,17 +18,6 @@ from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModel
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-@pytest.mark.parametrize("randomize_inputs", [False, True])
-def test_profile_run_accepts_release_flag_without_running_codec(monkeypatch, randomize_inputs):
-    runner = object.__new__(OmniGenerationModelRunner)
-    runner.model = MagicMock(side_effect=AssertionError("profiling ran codec model"))
-    synchronize = MagicMock()
-    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize)
-    runner.profile_run(randomize_inputs=randomize_inputs)
-    synchronize.assert_called_once_with()
-    runner.model.assert_not_called()
-
-
 class _FakeStagedField:
     def __init__(self, data):
         self.np = data
@@ -58,7 +45,6 @@ def _make_runner(model_output, num_reqs=1, prompt_len=10):
     runner.execute_model_state = SimpleNamespace(finished_req_ids={"finished"}, ec_connector_output=None)
     runner.kv_connector = SimpleNamespace(post_forward=MagicMock(return_value=None))
     runner.check_ep_fault = False
-    runner.aux_output_connector = None
     runner.req_states = MagicMock(
         prompt_len=SimpleNamespace(np=np.full(num_reqs, prompt_len, dtype=np.int32)),
         num_computed_tokens=_FakeStagedField(np.zeros(num_reqs, dtype=np.int32)),
@@ -66,8 +52,7 @@ def _make_runner(model_output, num_reqs=1, prompt_len=10):
     return runner
 
 
-@pytest.mark.parametrize("aux_enabled", [False, True])
+def test_control_only_step_keeps_lifecycle_and_skips_input_construction():
-def test_control_only_step_keeps_lifecycle_and_skips_input_construction(aux_enabled):
     runner = object.__new__(OmniGenerationModelRunner)
     order = []
     for name in (
@@ -83,13 +68,8 @@ def test_control_only_step_keeps_lifecycle_and_skips_input_construction(aux_enab
     runner.kv_connector = SimpleNamespace(no_forward=lambda _s: output)
     runner._merge_ec_connector_no_forward = lambda _s, value: value
     runner._attach_native_data_plane_signals = lambda value: value
-    metadata = object()
-    runner.aux_output_connector = SimpleNamespace(begin_step=MagicMock()) if aux_enabled else None
     scheduler_output = SimpleNamespace(
-        total_num_scheduled_tokens=0,
+        total_num_scheduled_tokens=0, scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace(req_ids=[])
-        scheduled_new_reqs=[],
-        scheduled_cached_reqs=SimpleNamespace(req_ids=[]),
-        aux_output_connector_metadata=metadata,
     )
     assert runner.execute_model(scheduler_output) is output
     assert order == [
@@ -98,9 +78,6 @@ def test_control_only_step_keeps_lifecycle_and_skips_input_construction(aux_enab
         "free_states",
         "_apply_block_table_staged_writes_if_available",
     ]
-    if aux_enabled:
-        assert runner.aux_output_connector is not None
-        runner.aux_output_connector.begin_step.assert_called_once_with(metadata)
 
 
 def test_released_chunk_reuses_scheduler_output_and_slot_recycle_clears_state():
@@ -207,72 +184,6 @@ def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypa
     assert output_req_ids == ["req-0"]
 
 
-def test_generation_aux_output_copies_and_materializes_without_sampling(monkeypatch):
-    from vllm_omni.worker_v2 import omni_generation_model_runner as module
-
-    events = []
-
-    class PendingOutput:
-        def enqueue_cpu_copy(self, *, num_sampled, num_rejected):
-            np.testing.assert_array_equal(num_sampled, [0])
-            np.testing.assert_array_equal(num_rejected, [0])
-            events.append("copy")
-
-        def process_output(self):
-            assert events == ["copy", "record", "synchronize"]
-            return {"req-0": AuxRequestOutput(0, np.array([[2, 3]], dtype=np.uint8))}
-
-    class CopyEvent:
-        def record(self, stream):
-            events.append("record")
-
-        def synchronize(self):
-            events.append("synchronize")
-
-    monkeypatch.setattr(torch.cuda, "Event", lambda **kw: CopyEvent())
-    monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
-    monkeypatch.setattr(module, "_async_copy_mm", lambda *a, **kw: {})
-    output = OmniModelRunnerOutput(["req-0"], {"req-0": 0}, [[]])
-    result = module.OmniGenerationAsyncOutput(
-        model_runner_output=output,
-        multimodal_outputs={},
-        num_reqs=1,
-        main_stream=object(),
-        copy_stream=SimpleNamespace(wait_stream=lambda stream: None),
-        pending_aux_output=PendingOutput(),
-    ).get_output()
-
-    assert result.sampled_token_ids == [[]]
-    np.testing.assert_array_equal(result.aux_output_connector_output["req-0"].rows, [[2, 3]])
-
-
-def test_pending_aux_output_forces_async_delivery_without_cuda_payload(monkeypatch):
-    from vllm_omni.worker_v2 import omni_generation_model_runner as module
-
-    runner = _make_runner(OmniOutput(text_hidden_states=torch.empty(0), multimodal_outputs={}))
-    runner.device = SimpleNamespace(type="cuda")
-    runner.main_stream = object()
-    runner.output_copy_stream = object()
-    pending = object()
-    runner.aux_output_connector = SimpleNamespace(prepare_output=MagicMock(return_value=pending))
-    runner._reserve_native_data_plane_outputs = MagicMock()
-    runner._release_generation_slots = MagicMock()
-    runner._finalize_native_data_plane_output = MagicMock()
-    captured = {}
-
-    class AsyncOutput:
-        def __init__(self, **kwargs):
-            captured.update(kwargs)
-
-    monkeypatch.setattr(module, "OmniGenerationAsyncOutput", AsyncOutput)
-    input_batch = runner._gen_input_batch
-    result = runner.sample_tokens()
-
-    assert isinstance(result, AsyncOutput)
-    runner.aux_output_connector.prepare_output.assert_called_once_with(input_batch)
-    assert captured["pending_aux_output"] is pending
-
-
 @pytest.mark.parametrize("exact", [True, False])
 def test_exact_input_shape_models_reject_padded_batches(exact):
     # Qwen3-Omni Code2Wav splits input_ids by seq_token_counts (#6712).

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -13,7 +13,6 @@ from vllm import SamplingParams
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
-from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 
 from vllm_omni.config.model import OmniModelConfig
@@ -96,8 +95,7 @@ def test_full_payload_receive_is_polled_without_scheduled_tokens(mocker):
 
 
 @pytest.mark.parametrize("output_form", ["tuple", "omni"])
-@pytest.mark.parametrize("profile_only", [False, True])
+def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form):
-def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form, profile_only):
     runner = object.__new__(OmniGPUModelRunner)
     hidden = torch.ones(1, 2)
 
@@ -119,21 +117,16 @@ def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form, profile
     runner.model_state = SimpleNamespace(capture_mtp_graphs=MagicMock())
     runner._dispatch_mtp_batch_descriptor = MagicMock(return_value="desc")
 
-    def assert_unwrapped(_self, *, profile_only: bool = False):
+    def assert_unwrapped(_self):
         assert torch.equal(runner.model.forward(), hidden)  # forward unwrapped during capture
-        assert profile_only is expected_profile_only
         return 3
 
-    expected_profile_only = profile_only
     with patch.object(GPUModelRunner, "capture_model", assert_unwrapped):
-        assert runner.capture_model(profile_only=profile_only) == 3
+        assert runner.capture_model() == 3
 
     assert runner.model.forward is original_forward  # restored after capture
     assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.PIECEWISE: [piecewise]}
-    if profile_only:
+    runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor)
-        runner.model_state.capture_mtp_graphs.assert_not_called()
-    else:
-        runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor)
 
 
 @pytest.mark.parametrize(
@@ -170,28 +163,6 @@ def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect
         assert state is upstream.return_value
 
 
-@pytest.mark.parametrize("omni_stage", [False, True])
-@pytest.mark.parametrize("specialized", [False, True])
-def test_auxiliary_default_state_has_omni_lifecycle_but_preserves_upstream_states(monkeypatch, omni_stage, specialized):
-    class SpecializedState(DefaultModelState):
-        pass
-
-    state_cls = SpecializedState if specialized else DefaultModelState
-    upstream_state = object.__new__(state_cls)
-    monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", lambda *_args: upstream_state)
-    monkeypatch.setattr(OmniModelState, "__init__", lambda *_args: None)
-    model_config = object.__new__(OmniModelConfig) if omni_stage else SimpleNamespace()
-    state = init_omni_model_state(
-        SimpleNamespace(model_config=model_config), SimpleNamespace(), None, torch.device("cpu")
-    )
-    if omni_stage and not specialized:
-        assert isinstance(state, OmniModelState)
-        assert callable(state.run_preprocess)
-        assert callable(state.postprocess_model_output)
-    else:
-        assert state is upstream_state
-
-
 def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch):
     runner = _make_runner()
     calls = []
@@ -238,7 +209,7 @@ def test_full_graph_aux_contract_keeps_full_and_round_trips_leaves():
     runner.model_state = SimpleNamespace()
     captured = {}
 
-    def capture(_self, *, profile_only=False):
+    def capture(_self):
         # The graph manager stores (hidden, leaves) in its aux buffers.
         assert runner.use_aux_hidden_state_outputs is True
         hidden, leaves = runner.model.forward(4)
@@ -291,9 +262,7 @@ def test_full_graph_output_without_contract_is_hidden_only():
 
 
 @pytest.mark.parametrize("runner_kind", ["gpu", "ar", "generation"])
-@pytest.mark.parametrize("randomize_inputs", [False, True])
+def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch):
-@pytest.mark.parametrize("is_profile", [False, True])
-def test_dummy_forward_uses_upstream_execution_state(runner_kind, randomize_inputs, is_profile, monkeypatch):
     from vllm.v1.worker.gpu.input_batch import InputBatch
     from vllm.v1.worker.gpu.model_runner import ExecuteModelState
 
@@ -319,7 +288,6 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, randomize_inpu
     runner.model_state.prepare_inputs.return_value = {}
     runner._omni_data_plane = object()
     runner.supports_mm_inputs = False
-    runner.vocab_size = 32
     runner.lora_config = None
     runner.is_encoder_decoder = False
     runner.eplb = MagicMock()
@@ -340,15 +308,7 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, randomize_inpu
         cg_mode=CUDAGraphMode.NONE, num_reqs=1, num_tokens=1, num_active_loras=0, max_query_len=1
     )
     runner._dispatch_batch_descriptor = MagicMock(return_value=(batch_desc, None))
-    make_dummy = MagicMock(return_value=input_batch)
+    monkeypatch.setattr(InputBatch, "make_dummy", lambda *args, **kwargs: input_batch)
-    monkeypatch.setattr(InputBatch, "make_dummy", make_dummy)
-    randomized = []
-
-    def randomize(tensor, low, high):
-        randomized.append((low, high))
-        return tensor.fill_(7)
-
-    monkeypatch.setattr(torch.Tensor, "random_", randomize)
     monkeypatch.setattr("vllm_omni.worker_v2.omni_model_runner.build_slot_mappings_by_layer", lambda *args: {})
     for module in ("omni_model_runner", "omni_generation_model_runner"):
         monkeypatch.setattr(f"vllm_omni.worker_v2.{module}.set_forward_context", lambda *args, **kwargs: nullcontext())
@@ -358,20 +318,7 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, randomize_inpu
 
     # Upstream _dummy_run always supplies valid_dummy_state_slots, and the
     # result must use the real upstream constructor rather than a mocked state.
-    assert (
+    assert runner.execute_model(scheduled, dummy_run=True, valid_dummy_state_slots=True) is None
-        runner.execute_model(
-            scheduled,
-            dummy_run=True,
-            valid_dummy_state_slots=True,
-            randomize_inputs=randomize_inputs,
-            is_profile=is_profile,
-        )
-        is None
-    )
-    assert make_dummy.call_args.kwargs["is_padding"] is not is_profile
-    should_randomize = randomize_inputs and runner_kind != "generation"
-    assert randomized == ([(0, 32)] if should_randomize else [])
-    assert input_batch.input_ids.tolist() == ([7] if should_randomize else [1])
     assert isinstance(runner.execute_model_state, ExecuteModelState)
     assert runner.execute_model_state.input_batch is input_batch
     assert runner.execute_model_state.cudagraph_stats is None

tests/worker_v2/test_omni_model_state.py

@@ -22,19 +22,8 @@ from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _m
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def test_prepare_attn_forwards_release_model_specific_metadata():
-    state = object.__new__(OmniModelState)
-    metadata = object()
-    output = object()
-    batch = SimpleNamespace()
-    with patch.object(DefaultModelState, "prepare_attn", return_value=output) as prepare:
-        assert state.prepare_attn(batch, None, (), None, [], None, model_specific_attn_metadata=metadata) is output
-    assert prepare.call_args.kwargs["model_specific_attn_metadata"] is metadata
-
-
 class _DummyInputBatch:
     is_prefilling_np: np.ndarray
-    req_ids: list[str]
 
     input_ids: SimpleNamespace
     query_start_loc: torch.Tensor
@@ -712,58 +701,6 @@ def test_publish_sampled_embeddings_is_opt_in() -> None:
     assert state.publish_sampled_embeddings(SimpleNamespace(num_reqs=1), torch.tensor([[1]])) is None
 
 
-def test_split_settled_rows_matches_per_row_check():
-    batch = _DummyInputBatch(np.array([3, 0, 2, 1, 4]))
-    batch.req_ids = ["c", "a", "x", "b", "e"]
-    batch.num_scheduled_tokens = [1, 1, 1, 4, 1]
-    batch.query_start_loc_np = [0, 1, 2, 3, 7]
-    req_indices = batch.idx_mapping_np.tolist()
-    # Slot 2 still names a finished request, slot 1 schedules a prefill
-    # chunk, slot 4 is not settled and slot 7 is not scheduled.
-    settled = {3: "c", 0: "a", 2: "old", 1: "b", 7: "z"}
-    split = OmniModelState._split_settled_rows(batch, req_indices, settled)
-    assert split is not None
-    settled_rows, remaining = split
-    assert settled_rows == [(0, 3, 0, "c"), (1, 0, 1, "a")]
-    assert remaining == [(2, 2), (3, 1), (4, 4)]
-    assert OmniModelState._split_settled_rows(batch, req_indices, {2: "old"}) is None
-
-
-def test_settled_row_split_preserves_mixed_batch_and_pending_replay(monkeypatch):
-    """Exercise run_preprocess, including the scalar fallback during replay."""
-
-    def prepare():
-        state = _make_eager_state()
-        _fill_buffers(state, "prefill", "decode")
-        state._eager_fastpath = True
-        state._eager_settled = {1: "decode"}
-        state._eager_ready = {1: "decode"}
-        state._eager_embeds[1].fill_(7)
-        state.model.eager_settled_text_step = lambda: torch.ones(_EAGER_DIM)
-        state.model.eager_decode_settled = lambda info: False
-        state.model.preprocess = lambda input_ids, input_embeds, **info: (input_ids, input_embeds, {})
-        batch = _EagerBatch([3, 1])
-        batch.req_ids = ["prefill", "decode"]
-        inputs = {"input_ids": torch.zeros(4, dtype=torch.long), "inputs_embeds": torch.zeros(4, _EAGER_DIM)}
-        return state, batch, inputs
-
-    vector, batch, inputs = prepare()
-    vector.run_preprocess(batch, inputs)
-    reference, ref_batch, ref_inputs = prepare()
-    monkeypatch.setattr(reference, "_split_settled_rows", lambda *args: None)
-    reference.run_preprocess(ref_batch, ref_inputs)
-    torch.testing.assert_close(inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
-    assert vector._eager_rows[1] == reference._eager_rows[1]
-    assert inputs["inputs_embeds"][3].tolist() == [8.0] * _EAGER_DIM
-
-    replay, replay_batch, replay_inputs = prepare()
-    replay._eager_state._restore_audio["prefill"] = [torch.ones(1)]
-    assert replay._eager_state.has_pending_replay()
-    monkeypatch.setattr(replay, "_split_settled_rows", MagicMock(side_effect=AssertionError("replay must skip split")))
-    replay.run_preprocess(replay_batch, replay_inputs)
-    torch.testing.assert_close(replay_inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
-
-
 @pytest.mark.parametrize("prefilling", [False, True])
 def test_identity_preprocess_skips_only_decode_rows(prefilling):
     state = _make_state(has_preprocess=True)

vllm_omni/config/pipeline_registry.py

@@ -54,7 +54,7 @@ from vllm_omni.model_executor.models.cosmos3.pipeline import (
     COSMOS3_OMNI_DEPLOY_PIPELINE,
     COSMOS3_POLICY_PIPELINE,
 )
-from vllm_omni.model_executor.models.cosyvoice3.pipeline import resolve_cosyvoice3_pipeline
+from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
 from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
 from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
@@ -150,7 +150,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     # key (see deploy/cosmos3_policy_droid.yaml and deploy/cosmos3_omni.yaml).
     "cosmos3_omni_deploy": COSMOS3_OMNI_DEPLOY_PIPELINE,
     "cosmos3_policy": COSMOS3_POLICY_PIPELINE,
-    "cosyvoice3": resolve_cosyvoice3_pipeline,
+    "cosyvoice3": COSYVOICE3_PIPELINE,
     "covo_audio": COVO_AUDIO_PIPELINE,
     "dreamzero": DREAMZERO_PIPELINE,
     "fish_qwen3_omni": FISH_SPEECH_PIPELINE,

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -25,6 +25,7 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.utils import (
     free_kv_blocks_in_physical_order,
+    omni_routed_experts_for_request,
 )
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import deserialize_additional_information
@@ -98,7 +99,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
     core scheduling logic.
     """
 
-    max_num_active_reqs: int
+    max_num_running_reqs: int
 
     def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool:
         model_config = self.vllm_config.model_config
@@ -379,32 +380,32 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # them. Upstream vllm raises RuntimeError on this status; omni allows
         # async abort (e.g. client disconnect during TTS streaming) to leave
         # requests in the waiting/running queues temporarily.
+        waiting = getattr(self, "waiting")
         self._drop_aborted_queued_requests()
         self._process_pending_omni_inputs(model_mode="ar")
         self._drop_aborted_queued_requests()
         self._resync_streaming_input_counter()
-        original_wait_queues = None
+        original_waiting = None
         if self._should_defer_waiting_admission():
-            original_wait_queues = self.waiting, self.kv_holding_waiting
+            original_waiting = waiting
             self.waiting = create_request_queue(self.policy)
-            self.kv_holding_waiting = create_request_queue(self.policy)
 
-        original_max_num_active_reqs = self.max_num_active_reqs
+        original_max_num_running_reqs = self.max_num_running_reqs
         async_chunk_transport = self._async_chunk_transport_enabled()
         reserved_running_slots = (
             self._get_async_chunk_reserved_running_slots() if async_chunk_transport and self.use_v2_model_runner else 0
         )
         if reserved_running_slots:
-            self.max_num_active_reqs = max(0, original_max_num_active_reqs - reserved_running_slots)
+            self.max_num_running_reqs = max(0, original_max_num_running_reqs - reserved_running_slots)
         try:
             scheduler_output = super().schedule(throttle_prefills)
         finally:
-            self.max_num_active_reqs = original_max_num_active_reqs
+            self.max_num_running_reqs = original_max_num_running_reqs
-            if original_wait_queues is not None:
+            if original_waiting is not None:
-                original_waiting, original_kv_holding = original_wait_queues
+                deferred_waiting = list(self.waiting)
-                original_waiting.prepend_requests(self.waiting)
+                if deferred_waiting:
-                original_kv_holding.prepend_requests(self.kv_holding_waiting)
+                    original_waiting.prepend_requests(deferred_waiting)
-                self.waiting, self.kv_holding_waiting = original_wait_queues
+                self.waiting = original_waiting
             self._restore_omni_wait_queues()
 
         self._postprocess_omni_schedule_output(
@@ -427,7 +428,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         sampled_token_ids = model_runner_output.sampled_token_ids
         logprobs = model_runner_output.logprobs
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
-        prompt_token_id_logprobs_dict = getattr(model_runner_output, "prompt_token_id_logprobs_dict", {})
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
         pooler_outputs = model_runner_output.pooler_output
         mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
@@ -699,11 +699,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if stopped:
                 if self.chunk_transfer_adapter is not None:
                     confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request)
-                if (
+                if model_runner_output.routed_experts is not None:
-                    self.aux_output_connector is not None
+                    routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
-                    and (aux_output := model_runner_output.aux_output_connector_output) is not None
-                ):
-                    routed_experts = self.aux_output_connector.take_output(request, aux_output)
 
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
@@ -776,7 +773,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            prompt_token_id_logprobs = prompt_token_id_logprobs_dict.get(req_id)
             has_stage_output = (
                 bool(new_token_ids)
                 or mm_output is not None
@@ -797,7 +793,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     finish_reason=finish_reason,
                     new_logprobs=new_logprobs,
                     new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                    prompt_token_id_logprobs=prompt_token_id_logprobs,
                     pooling_output=pooling_output_payload,
                     multimodal_output=mm_output,
                     stop_reason=request.stop_reason,
@@ -813,7 +808,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
-                assert prompt_token_id_logprobs is None
 
             if omits_chunk_transfer and inter_stage_output is not None:
                 logger.warning(
@@ -972,9 +966,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
                     self.num_waiting_for_streaming_input -= 1
                 session.status = RequestStatus.WAITING
-                self.deferred_waiting.discard(session)
+                if session in self.skipped_waiting:
-                if session in self.kv_holding_waiting:
+                    self.skipped_waiting.remove_requests((session,))
-                    self.kv_holding_waiting.remove_requests((session,))
                     self._enqueue_waiting_request(session)
 
                 if self.log_stats:
@@ -1398,8 +1391,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         ``_free_blocks``) and skips the caller's input-coordinator cleanup.
         """
         self.waiting.remove_requests((request,))
-        self.kv_holding_waiting.remove_requests((request,))
+        self.skipped_waiting.remove_requests((request,))
-        self.deferred_waiting.discard(request)
         if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
             self.num_waiting_for_streaming_input -= 1
         request.status = RequestStatus.FINISHED_ERROR

vllm_omni/core/sched/omni_cfg_ar_scheduler.py

@@ -32,7 +32,7 @@ class OmniCFGARScheduler(OmniARScheduler):
             raise ValueError("Atomic CFG reservations require one full-attention KV group")
         self._capacity = self.kv_cache_manager.block_pool.num_gpu_blocks - 1 - self.kv_cache_manager.watermark_blocks
         pair_pages = 2 * -(-self.max_model_len // self.block_size)
-        if self._capacity < pair_pages or min(self.max_num_running_reqs, self.max_num_active_reqs) < 2:
+        if self._capacity < pair_pages or self.max_num_running_reqs < 2:
             raise ValueError("Atomic CFG requires KV capacity and max_num_seqs for two maximum-length branches")
         if self.max_num_scheduled_tokens < 2 * self.max_model_len:
             raise ValueError("Atomic CFG requires max_num_batched_tokens >= 2 * max_model_len")
@@ -81,7 +81,7 @@ class OmniCFGARScheduler(OmniARScheduler):
         admitted = create_request_queue(self.policy)
         pending = {request.request_id: request for request in original_waiting}
         pages = self._capacity - sum(self._reserved_pages(request) for request in self.running)
-        slots = self.max_num_active_reqs - len(self.running)
+        slots = self.max_num_running_reqs - len(self.running)
         tokens = self.max_num_scheduled_tokens - len(self.running)
         chosen: set[str] = set()
         for request in original_waiting: