Compare commits

..
Author SHA1 Message Date
Sy03 118f646c87 [CI] Retrigger CosyVoice3 ready validation
Signed-off-by: Sy03 <1370724210@qq.com>
2026-10-02 22:56:24 +08:00
Sy03 7d5a6bc5cb [Doc] Clarify CosyVoice3 lazy graph latency and memory headroom
Signed-off-by: Sy03 <1370724210@qq.com>
2026-10-02 22:56:24 +08:00
Sy03 8361765b33 [Model] Minimize CosyVoice3 streaming regression tests
Signed-off-by: Sy03 <1370724210@qq.com>
2026-10-02 22:56:24 +08:00
Sy03 b3c25a608d [Model] Optimize CosyVoice3 packed streaming with bounded HiFT graphs
Signed-off-by: Sy03 <1370724210@qq.com>
2026-10-02 22:56:24 +08:00
234 changed files with 1255 additions and 13712 deletions

No files matched your search

+3 -4
View File
@@ -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"]
@@ -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()
-24
View File
@@ -124,30 +124,6 @@ steps:
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_2
- label: "Omni · Qwen3-Omni Async-Chunk SHM Leak"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_omni_shm_leak.py
- tests/helpers/runtime.py
- tests/helpers/stage_config.py
- vllm_omni/core/sched/omni_ar_scheduler.py
- vllm_omni/distributed/omni_connectors/connectors/shm_connector.py
- vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
- vllm_omni/engine/orchestrator.py
- vllm_omni/engine/stage_engine_core_proc.py
- vllm_omni/engine/stage_pool.py
- vllm_omni/entrypoints/openai/serving_chat.py
- vllm_omni/model_executor/models/qwen3_omni/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
timeout_in_minutes: 40
commands:
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Test"
source_file_dependencies: omni_minicpmo_4_5_function
timeout_in_minutes: 50
-7
View File
@@ -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"
+7 -10
View File
@@ -1,7 +1,10 @@
# Use the official release image with its matching CUDA/PyTorch stack.
ARG VLLM_VERSION=0.31.0
# Target vLLM version: v0.30.0 (tag @ ced6857afa0e).
# 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 VLLM_OMNI_VERSION_OVERRIDE=0.0.0
ENV SETUPTOOLS_SCM_PRETEND_VERSION=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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+4 -5
View File
@@ -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 VLLM_OMNI_VERSION_OVERRIDE=0.0.0
ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
--no-cache-dir -e ".[dev]" --no-build-isolation && \
+1 -1
View File
@@ -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.
+4 -4
View File
@@ -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.
@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.
## Model-specific variables
The inventory includes 93 variables read by a single model or pipeline family. They are
The inventory includes 66 variables read by a single model or pipeline family. They are
not listed as public usage options here because doing so would turn implementation
escape hatches into an accidental compatibility contract.
@@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the
| Disposition | Count | Required outcome |
| --- | ---: | --- |
| Promote | 64 | Move a stable setting into typed stage or model configuration. |
| Promote | 38 | Move a stable setting into typed stage or model configuration. |
| Request scope | 5 | Move request-varying behavior into a declared request-option schema. |
| External | 0 | Retain only when a supported third-party library owns the contract. |
| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. |
| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. |
| Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |
The disposition is a migration target, not a statement that the existing
-47
View File
@@ -514,50 +514,3 @@ for that comparison; this switch only controls automatic MPS management.
MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on
the intended GPU, and warm the complete pipeline before measuring performance.
## MOSS-TTS Local 1.5 with Model Runner V2
`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using
the shared runtime introduced for Qwen3-TTS:
```bash
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
```
This profile inherits the batching, codec graph buckets and 1-frame/15-frame
chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local
Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp`
capabilities while retaining its V1 `talker_mtp` implementation, sampling
defaults and explicit request-seed handling. Codec chunks use the native data
plane; the internal Talker retains its final-only orchestrator output policy.
On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration.
The codec retains its own CUDA Graph capture/replay, including batch buckets
through 64, but disables Inductor compilation (`compilation_config.mode: 0`)
of the stateful decoder to avoid lengthy compilation at startup. This is
independent of `enforce_eager`; the profile keeps `enforce_eager: false`.
It explicitly selects `cudagraph_mode: FULL`, which works without Inductor;
the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear
the codec's capture buckets when compilation is disabled.
For sustained C128 serving on a large-memory CUDA GPU, select the
separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128
stream slots per stage, a 32 GiB Talker KV budget, and the codec's
`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs.
The bounded KV budget leaves room for codec state and graphs; it does not
guarantee that 128 maximum-length prompts fit simultaneously. See the
[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention)
for activation, backend comparisons, memory requirements and benchmark commands.
Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1.
NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable
MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio;
set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM.
Event-driven orchestration remains independently selectable with
`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical
when comparing these modes. Model-runner selection does not change the
orchestration default or enable experimental reference encoding, chunk ramps,
generation-output draining or MPS.
+1 -1
View File
@@ -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:
@@ -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
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.
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x.
!!! note
PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -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
release wheel from its commit-pinned index:
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
```bash
uv pip install vllm==0.31.0 --torch-backend=auto \
--extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
uv pip install vllm==0.30.0 --torch-backend=auto
```
The default 0.31.0 wheel contains CUDA 12.9 binaries. If you need a different
CUDA variant or want to reuse an existing PyTorch installation, build vLLM
from source instead.
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.
#### 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
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl
#For CUDA 13.0 (the default for v0.30.0; the wheel filename has no `+cu130` suffix)
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
upstream release image explicitly:
If you want to specify the base vLLM version:
```bash
DOCKER_BUILDKIT=1 docker build \
-f docker/Dockerfile.cuda \
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.31.0 \
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.30.0 \
-t vllm-omni-cuda .
```
@@ -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
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.
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni
0.30.x release line uses vLLM 0.30.x.
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
release wheel for ROCm 7.2.3:
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
```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 .
```
@@ -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 \
--build-arg VLLM_VERSION="$VLLM_VERSION" .
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
```
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
v0.30.0; the command above explicitly selects v0.31.0 to match this source
checkout. To target a different vLLM release, pass
`--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
base image (v0.30.0 by default), which Docker pulls automatically. 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
@@ -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
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.
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.
```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
+3 -6
View File
@@ -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 \
--extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
uv pip install vllm==0.30.0 --torch-backend=auto
# 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
pinned to the final release commit. For published vLLM-Omni wheels and additional
installation methods, see the [installation guide](installation/README.md).
For additional installation methods — please 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.
@@ -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.
---
@@ -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.
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:
Cross-request Stage-1 flow batching is opt-in. 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).
---
@@ -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.
sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
# Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
gpt_sampling = SamplingParams(
temperature=1.0,
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:
@@ -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;
# other devices use vllm_omni/deploy/cosyvoice3.yaml.
# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
# Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
set -e
-297
View File
@@ -35,10 +35,6 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses
audio-reference generation. Reference transcripts must match the reference
audio; they are not style instructions.
For the optional CUDA MRV2 runner, see the
[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2).
The default Local deployment continues to use V1.
## References
- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
@@ -156,296 +152,3 @@ curl -X POST http://localhost:8091/v1/audio/speech \
- No `ref_audio` required or accepted for MOSS-SoundEffect.
- Input field maps to the `ambient_sound` parameter in the upstream processor.
- Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
## Local 1.5 MRV2 and slot attention
`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with
an explicit deploy profile. The default Local profile continues to use V1.
Both profiles below preserve 1-frame initial and 15-frame steady codec chunks
and the model's sampling defaults.
```bash
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
```
This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned
CUDA graphs without compiling the codec with Inductor.
For sustained high concurrency on a large-memory CUDA GPU, use:
```bash
CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \
--stage-init-timeout 1200 --init-timeout 1500
```
The high-concurrency profile places both stages on one GPU, sets each stage's
capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses
mixed FULL CUDA graphs with both its token budget and capture limit set to
512, so new prompt chunks remain within graph coverage. It also enables
`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and
continuation control stay in GPU request slots, gathered through MRV2's
batch mapping. Prefill conditioning and explicit per-request sampling seeds
retain their original behavior. Published audio codes own their storage, so
slot reuse cannot overwrite an in-flight output. Set the override to `false`
to compare with the generic Omni model state. The profile selects
`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels
disabled, and codec CUDA graph buckets through 128. The configuration was
validated on one H200. The earlier capacity-256 variant used about 110 GiB of
sampled peak GPU memory, including loading and warmup. Reduce the stage capacities,
graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec
compilation can take several minutes; subsequent starts can reuse the AOT
cache. Graph capture is still performed at startup.
The high-concurrency profile also coalesces ready codec streams with
`connectors.shm.extra.generation_min_batch_size: 32` and
`generation_max_wait_ms: 12`. The target controls when to stop waiting; it
neither caps execution at 32 nor guarantees every codec group has 32 rows.
First and final chunks share the bounded window at high concurrency. When
fewer than 32 requests remain, dispatch is immediate. Pending outputs are
retired before waiting, with the original deadline preserved, so waiting does
not prevent the previous batch from releasing its in-flight state. Cancellation
and input notifications keep their ordinary scheduler bookkeeping. This option
requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to
`0` to disable coalescing; other profiles keep the immediate-dispatch default.
The high-concurrency profile enables two additional Local model-state
optimizations under the CUDA stage-0 `hf_overrides`, alongside
`mrv2_gpu_slot_state: true`:
- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines
CPU reference-code slices and their destination positions into one pinned
upload. Host staging is reused only after the upload event completes;
GPU-resident references keep the original path.
- `mrv2_direct_tokens: true` returns the already determined text token after
the Local audio/stop decision. It avoids constructing and sampling the
full text vocabulary. Audio-code and binary-stop sampling are unchanged.
Requests requiring distribution metadata or token constraints, and
unsupported execution modes, retain the normal sampler.
Other profiles leave both flags disabled. Set both flags explicitly when
comparing their combined effect. Deployment
inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also
retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as
throughput: faster prompt preparation can change competition between the
Talker and codec sharing the GPU.
For the single-H200 capacity-128 throughput baseline, keep both switches on.
In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch
rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct
tokens alone had no reliable gain. The combined path had 685 ms mean first
audio versus 679 ms across six control rounds, with lower mean completion time.
A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus
3.78% for the control. This sample does not establish speaker similarity or
exclude small quality changes. At concurrency 256, the combined path had
raised mean first-audio latency by about 0.33 s, so other profiles leave the
switches off until their latency and quality are checked.
The optional connector setting `generation_coalescing_policy: idle_wait`
waits for an inbox notification when no stream is runnable, all receivers
are parked and registered, and no output needs retirement. The first ready
arrival starts a fresh coalescing window, preserving the batch budget.
Control messages, including cancellations, can wait up to two windows
instead of one (24 ms at the 12 ms setting). The default policy remains
`fixed`; neither policy changes the global orchestration default.
The codec backends differ in state access:
- `sdpa` uses PyTorch attention after gathering and updating the ring cache.
- `triton` replaces the attention calculation, retaining the ring-cache
gather/copy and explicit mask construction.
- `triton_slot` writes surviving K/V directly into request slots, attends
directly to the ring, and advances the active slot offsets in three
ordered kernels. Graph-padding rows do not advance persistent state.
The slot path preserves the existing chunk-complete ring semantics, including
retaining the final cache-capacity tokens when a chunk exceeds the ring.
It does not introduce another cache owner or change request-slot lifetime.
Zero-length padding rows skip the attention computation and emit zeros.
The high-concurrency profile includes intermediate codec batch buckets 6, 12
and 24. To compare with power-of-two buckets, change the stage-1
`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`.
Keep the maximum bucket equal to the state capacity to retain terminal-tail
coalescing. Smaller buckets reduce padding work but require more graphs;
measure complete serving runs before selecting them for a deployment.
For a codec output-path comparison, set
`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous
codec output path; `1` selects GPU output snapshots on a private codec
stream. Both configurations use MRV2 and keep streaming audio responses.
The option changes input ownership, metadata staging and snapshot handling
as well as transfer scheduling, so its timing difference is not just D2H
copy time.
Those chunk-complete semantics truncate the causal window: every ring holds
exactly `context` entries and a whole chunk is written before attention runs,
so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the
tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/
400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so
its first 80 tokens attend to nothing and no chunk sees the previous one there.
Padding a terminal tail to 15 frames therefore decodes it differently from an
exact-length decode; the difference is deterministic (identical in fp32 and
fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each
ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming
then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails
equal exact tails. It increases state/graph memory and capture time; historical C256 experiments
observed about a 2% throughput cost. It is off by default and its perceptual
quality benefit has not been established. Terminal tails are padded
into the regular 15-frame graph bucket in either mode.
The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths
retain their existing attention implementation. Event-driven orchestration
remains independently selectable and its default is unchanged.
### Optional progressive chunks
Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks
before steady decoding, for example:
```yaml
connectors:
shm:
extra:
codec_chunk_ramp: [1, 4, 15]
```
The ramp's first entry overrides `initial_codec_chunk_frames`. After the last
entry, chunks use `codec_chunk_frames`. Each request advances independently;
the final partial chunk is flushed and an empty terminal carries no codec
tokens. The codec captures every ramp length, sizes the maximum execution
step for the largest entry, and uses the ramp's first entry for its dedicated
first-chunk graph. Additional shapes increase startup and graph memory.
Ramps are opt-in. They can reduce early playback underrun while increasing
codec calls, first-audio latency or total generation time. The measured C128
throughput presets retain 1→15; their published figures do not establish an
additional benefit from enabling ramps. The processor and graph-shape tests
run without GPU or model weights:
```bash
CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \
tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \
-m 'core_model and cpu' --run-level core_model
```
### Low-latency and reference-encoding options
`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and
uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP.
Its codec first-chunk fast path uses dedicated graphs and a stream handoff
before the regular decoder reuses the request slot. A gate limits contention
with regular codec work. Slot waits are bounded at 30 seconds. A failed or
stalled decode raises an error rather than replaying a partially advanced
slot; a failed output enqueue retains owned PCM for regular delivery.
Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch
policy; it does not change the client concurrency or stage capacities.
Reference encoding runs in the API layer, independently of MRV2. Enable
reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default
eager encoder is retained when the option is absent. The graph path uses the
loaded tokenizer's encoder/quantizer, length buckets and optional compilation
(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed
attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention.
Compilation and attention changes need not produce bit-identical codes.
For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared
reference-code storage and, by default, a single encoder host with four
workers. Use a dedicated directory per service, checkpoint and encoding
configuration. Workers have their own graph resources; all graph captures
complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains
separate encoders while sharing codes. More workers/graphs consume memory;
they do not imply more GPUs. Example after selecting an available GPU and
configuring any operator-managed MPS daemon:
```bash
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \
VLLM_OMNI_MOSS_REF_GRAPHS=1 \
VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \
VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \
VLLM_OMNI_MOSS_REF_INFLIGHT=4 \
VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \
VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
--api-server-count 4 \
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \
--stage-init-timeout 1200 --init-timeout 1500
```
This example is a configurable serving profile, not the complete launch
configuration of a historical benchmark. In particular, the low-latency YAML
uses synchronous stage-1 scheduling; the retained source46 experiment used
asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder
environment settings. Do not infer a throughput result from the YAML alone.
### Cold versus hot reference measurements
A cold reference misses the encoded-reference cache: the request pays for
reference parsing/preparation and GPU encoding before speech generation.
A hot reference reuses reference codes but still synthesizes new speech;
generated audio is not cached. Pre-registering a reference moves encoding
into registration and can make subsequent synthesis hot, but registration
latency is still part of the first-use cost.
Model startup/compilation is separate from reference coldness. EN1088 contains
repeated references, so its first pass is not uniformly cold. For an all-unique
cold pass, use its 666 first-occurrence references, keep order, warm execution
with disjoint references, then repeat the same 666 requests. A second pass
through independent API-local caches is not guaranteed hot; verify cache
coverage before comparing it with a shared-cache result. Report requests/s
alongside audio-s/s and retain failed, empty, long and near-silent outputs.
### Reproduce the serving benchmark
Use the complete Seed-TTS English test set, including reference audio and text.
Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088
entries. Run the native benchmark after the server is ready:
```bash
export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000
export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2
export SEED_TTS_WER_EVAL=0
for concurrency in 128; do
for phase in warm r1 r2; do
vllm bench serve --omni \
--model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
--backend openai-audio-speech --endpoint /v1/audio/speech \
--dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \
--seed-tts-locale en --disable-shuffle --num-prompts 1088 \
--num-warmups 0 --ready-check-timeout-sec 0 \
--output-len 256 --max-concurrency "$concurrency" \
--request-rate inf --seed 42 \
--extra-body '{"task_type":"Base","max_new_tokens":256}' \
--save-result --save-detailed --result-dir results/moss-local \
--result-filename "c${concurrency}-${phase}.json"
done
done
```
Discard the complete `warm` pass and combine measured runs as total generated
audio seconds divided by total benchmark duration. Require 1088 successes,
zero failures and 1088 nonempty-audio metric samples in each run. The output
cap matches the benchmark protocol; success alone does not establish speech
quality or that every sentence ended before the cap. Dataset seed 42 does
not fix an independent sampling seed for every request.
For an attention-only comparison, copy the high-concurrency YAML beside the
original to preserve relative `base_config` resolution. Change its
codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`,
and compare `codec_attention_backend: triton` against `triton_slot`. Keep all
other settings, warmup and client concurrency identical. Restart the server
between configurations and use separate result directories. Comparing the
compiled slot profile to uncompiled Triton includes both changes.
Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages,
and no model weights:
```bash
python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \
tests/model_executor/models/moss_tts/test_streaming_attention.py \
-m 'core_model and cuda' --run-level=core_model
```
+1 -1
View File
@@ -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
+8 -13
View File
@@ -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"
CUDA_CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
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:
cuda_base = _docker_arg(CUDA_RELEASE_DOCKERFILE, "BASE_IMAGE")
cuda_image, _, cuda_release = cuda_base.rpartition(":")
def test_rocm_base_tracks_ci_vllm_release() -> None:
ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
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 == "vllm/vllm-openai-rocm"
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}"
assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref
assert image_tag == ci_release
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:
cuda_release = _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_VERSION")
def test_rocm_source_ref_tracks_ci_vllm_release() -> None:
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:
+1 -9
View File
@@ -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"] == [
CUDA_HF_TOKEN_EXPORT,
"python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
"pytest -sv tests/example",
]
assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "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:
+3 -115
View File
@@ -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(
@@ -2978,17 +2928,9 @@ class TestPlatformOverrides:
assert deploy.stages[0].engine_extras["block_size"] == 128
@pytest.mark.parametrize(
"filename,pipeline_key",
[
("qwen3_tts_mrv2.yaml", "qwen3_tts"),
("moss_tts_local_mrv2.yaml", "moss_tts_local"),
("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
],
)
@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key):
@pytest.mark.parametrize("platform", ["cuda", "npu"])
def test_recommended_native_runner_platform_defaults(self, platform):
filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
deploy = _apply_platform_overrides(deploy, platform=platform)
expect_v2 = platform == "cuda"
@@ -3004,60 +2946,6 @@ class TestPlatformOverrides:
# v2 only engages the native plane on stages declaring support.
assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
pipeline = resolve_pipeline_config("moss_tts_local")
base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
base = load_deploy_config(base_path)
candidate = load_deploy_config(candidate_path)
assert base.model_runner == "v1"
assert candidate.stages == base.stages
assert candidate.connectors == base.connectors
assert candidate.async_chunk == base.async_chunk
cuda = _apply_platform_overrides(candidate, platform="cuda")
stages = merge_pipeline_deploy(pipeline, cuda)
assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512
codec_args = stages[1].yaml_engine_args
assert codec_args["max_num_seqs"] == 64
assert codec_args["enforce_eager"] is False
assert codec_args["compilation_config"] == {
"mode": 0,
"cudagraph_mode": "FULL",
"cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64],
"cudagraph_num_of_warmups": 1,
}
for platform in ("npu", "xpu", "rocm", "musa"):
# Platform resolution mutates the deployment; each server loads
# a fresh profile before applying its own platform overrides.
fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform)
original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
assert fallback.stages == original.stages
for variant in ("moss_tts_delay", "moss_tts_realtime"):
assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages)
def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self):
pipeline = resolve_pipeline_config("moss_tts_local")
path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml"))
deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda")
talker, codec = merge_pipeline_deploy(pipeline, deploy)
assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128
assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True
assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True
assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True
assert talker.yaml_engine_args["max_num_batched_tokens"] == 512
assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3
assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot"
compilation = codec.yaml_engine_args["compilation_config"]
assert compilation["mode"] == 3 and compilation["backend"] == "inductor"
assert compilation["cudagraph_mode"] == "FULL"
assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
for platform in ("npu", "xpu", "rocm", "musa"):
actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
assert actual.stages == expected.stages
@pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
from unittest.mock import Mock
+2 -2
View File
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
EnvironmentVariableCategory.PUBLIC_OMNI: 32,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
EnvironmentVariableCategory.MODEL_SPECIFIC: 77,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
EnvironmentVariableCategory.INTERNAL: 7,
}
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
)
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
ModelEnvironmentVariableDisposition.PROMOTE: 66,
ModelEnvironmentVariableDisposition.PROMOTE: 48,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 17,
@@ -1,258 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Bounded readiness coalescing must not allocate or consume request state."""
import queue
from types import SimpleNamespace
import pytest
from vllm.v1.core.sched.interface import PauseState
from vllm.v1.request import RequestStatus
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
from vllm_omni.outputs import OmniConnectorOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch):
events = [notification("0"), notification("1", "2", "3")]
s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])])
s._generation_coalescing_policy = "idle_wait"
for request in s.requests.values():
request.status = RequestStatus.WAITING_FOR_CHUNK
assert s._drain_omni_connector_outputs() == events
assert clock[0] == pytest.approx(10.006)
assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012])
@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"])
def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending):
s, clock = scheduler(monkeypatch)
s._generation_coalescing_policy = "idle_wait"
for request in s.requests.values():
request.status = RequestStatus.WAITING_FOR_CHUNK
if pending == "in_flight":
s.requests["0"].num_in_flight_tokens = 1
elif pending == "admission":
s.requests["0"].status = RequestStatus.WAITING
elif pending == "registration":
s.input_coordinator.pending_chunk_registrations = [object()]
elif pending == "terminal":
s._pending_data_plane_terminal_req_ids = {"0"}
else:
s.requests.clear()
assert s._drain_omni_connector_outputs() == []
assert clock[0] == 10
assert not s._omni_connector_output_inbox.waits
def test_idle_wait_without_arrivals_is_bounded(monkeypatch):
s, clock = scheduler(monkeypatch)
s._generation_coalescing_policy = "idle_wait"
for request in s.requests.values():
request.status = RequestStatus.WAITING_FOR_CHUNK
assert s._drain_omni_connector_outputs() == []
assert clock[0] == pytest.approx(10.012)
def notification(*ids, terminal=()):
return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal))
class Inbox:
def __init__(self, clock, initial=(), delayed=()):
self.clock = clock
self.initial = list(initial)
self.delayed = list(delayed)
self.waits = []
def get_nowait(self):
if self.initial:
return self.initial.pop(0)
raise queue.Empty
def get(self, timeout):
self.waits.append(timeout)
if not self.delayed or self.delayed[0][0] > timeout:
self.clock[0] += timeout
raise queue.Empty
delay, value = self.delayed.pop(0)
self.clock[0] += delay
return value
def scheduler(monkeypatch, initial=(), delayed=()):
clock = [10.0]
monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0])
s = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
s._generation_max_wait_s = 0.012
s._generation_min_batch_size = 4
s._pause_state = PauseState.UNPAUSED
s.requests = {
str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8)
}
s.running = list(s.requests.values())
s.input_coordinator = SimpleNamespace(
requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set()
)
s._latest_omni_connector_output = None
s._omni_connector_output_inbox = Inbox(clock, initial, delayed)
return s, clock
def test_arrivals_form_target_batch_in_order(monkeypatch):
events = [notification("0"), notification("1", "2"), notification("3")]
s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])])
assert s._drain_omni_connector_outputs() == events
assert clock[0] == pytest.approx(10.005)
assert not s.input_coordinator.requests_with_ready_chunks
assert all(r.num_in_flight_tokens == 0 for r in s.requests.values())
def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch):
events = [notification("0", "gone", "1", "2"), notification("0")]
s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])])
s.requests["1"].is_finished = lambda: True
s.requests["2"].is_finished = lambda: True
assert s._drain_omni_connector_outputs() == events
assert clock[0] == pytest.approx(10.012)
assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009])
@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"])
def test_immediate_dispatch(monkeypatch, reason):
event = notification("0")
s, clock = scheduler(monkeypatch, [event])
if reason == "enough":
event.chunk_ready_req_ids.update(["1", "2", "3"])
elif reason == "empty":
event.chunk_ready_req_ids.clear()
elif reason == "low":
s.requests = {"0": s.requests["0"]}
elif reason == "paused":
s._pause_state = PauseState.PAUSED_ALL
elif reason == "off":
s._generation_max_wait_s = 0
assert s._drain_omni_connector_outputs() == [event]
assert not s._omni_connector_output_inbox.waits
assert clock[0] == 10
def test_terminal_arrival_keeps_original_deadline(monkeypatch):
events = [notification("0"), notification("1", terminal=["1"])]
s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])])
assert s._drain_omni_connector_outputs() == events
assert clock[0] == pytest.approx(10.012)
def test_pending_ready_and_legacy_notification_preserved(monkeypatch):
s, _ = scheduler(monkeypatch)
s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"])
last = notification("3")
s._latest_omni_connector_output = last
assert s._drain_omni_connector_outputs() == [last]
assert s._latest_omni_connector_output is None
assert not s._omni_connector_output_inbox.waits
def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch):
s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))])
assert len(s._drain_omni_connector_outputs()) == 2
assert clock[0] == pytest.approx(10.012)
s._omni_connector_output_inbox.initial.append(notification("3"))
assert len(s._drain_omni_connector_outputs()) == 2
assert clock[0] == pytest.approx(10.024)
def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker):
event = notification("0", "cancelled")
event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}}
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()
s._consume_pending_connector_output("generation")
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()
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"))])
def test_invalid_config_rejected(monkeypatch, target, wait):
from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler
def initialize(s, *args, **kwargs):
s.max_num_running_reqs = 256
s.vllm_config = SimpleNamespace(
model_config=SimpleNamespace(
retains_state_across_chunks=True,
stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}},
)
)
monkeypatch.setattr(VLLMScheduler, "__init__", initialize)
monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None)
with pytest.raises(ValueError):
OmniGenerationScheduler()
@pytest.mark.parametrize("terminal", [False, True])
def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal):
event = notification("0", terminal=["0"] if terminal else [])
s, clock = scheduler(monkeypatch, [event])
s.running = s.running[1:]
assert s._drain_omni_connector_outputs() == [event]
assert clock[0] == pytest.approx(10.012)
def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch):
events = [notification("0"), notification("1")]
s, clock = scheduler(monkeypatch, events[:1])
s.requests["7"].num_in_flight_tokens = 1
assert s._drain_omni_connector_outputs() == events[:1]
assert s._generation_defer_batch
assert not s._omni_connector_output_inbox.waits
s.input_coordinator.requests_with_ready_chunks.add("0")
# Simulate EngineCore retiring the previous batch after five milliseconds.
clock[0] += 0.005
s.requests["7"].num_in_flight_tokens = 0
s._omni_connector_output_inbox.initial.append(events[1])
assert s._drain_omni_connector_outputs() == events[1:]
assert not s._generation_defer_batch
assert s._omni_connector_output_inbox.waits == pytest.approx([0.007])
assert clock[0] == pytest.approx(10.012)
assert s._generation_batch_deadline is None
def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch):
s, clock = scheduler(monkeypatch, [notification("0")])
s.requests["7"].num_in_flight_tokens = 1
s._generation_batch_deadline = clock[0] - 0.001
assert len(s._drain_omni_connector_outputs()) == 1
assert not s._generation_defer_batch
assert not s._omni_connector_output_inbox.waits
def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker):
from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler
request = _chunk_request("pending")
s = _make_generation_scheduler(request)
s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True)
allocate = mocker.spy(s.kv_cache_manager, "allocate_slots")
output = s.schedule()
assert output.total_num_scheduled_tokens == 0
assert not output.num_scheduled_tokens
allocate.assert_not_called()
assert list(s.waiting) == [request]
assert request.num_computed_tokens == request.num_in_flight_tokens == 0
assert s.chunk_transfer_adapter.restore_called
@@ -1,107 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Exercise generation constructor policy after actual deploy/platform resolution."""
from types import SimpleNamespace
import pytest
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config.pipeline_registry import resolve_pipeline_config
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture
def construct_scheduler(monkeypatch, mocker):
# Only upstream cache allocation and connector construction are omitted.
# Keep the actual Omni I/O initialization, native decision and constructor.
mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter")
def upstream_init(self, config):
self.vllm_config = config
self.max_num_running_reqs = config.scheduler_config.max_num_seqs
monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init)
return OmniGenerationScheduler
def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128):
return SimpleNamespace(
model_config=SimpleNamespace(
use_v2_model_runner=native,
supports_native_mrv2_data_plane=True,
retains_state_across_chunks=stateful,
async_chunk=True,
stage_id=1,
stage_connector_config={"extra": extras or {}},
),
parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp),
scheduler_config=SimpleNamespace(max_num_seqs=capacity),
)
@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
def test_moss_profile_generation_constructor_after_platform_resolution(
construct_scheduler, monkeypatch, mocker, profile, platform
):
# merge_pipeline_deploy owns platform resolution. Mock the detected device
# instead of applying an override first and then resolving the worker's
# actual platform a second time.
monkeypatch.setattr(current_omni_platform, "device_name", platform)
deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml"))
pipeline = resolve_pipeline_config("moss_tts_local")
codec = merge_pipeline_deploy(pipeline, deploy)[1]
extras = deploy.connectors["shm"]["extra"]
config = _config(
native=codec.yaml_engine_args["use_v2_model_runner"],
stateful=pipeline.stages[1].retains_state_across_chunks,
capacity=codec.yaml_engine_args["max_num_seqs"],
extras=extras,
)
warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
scheduler = construct_scheduler(config)
# The raw extras persist through fallback; the live constructor must cope.
assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
if platform == "cuda":
assert scheduler._native_data_plane and scheduler.input_coordinator is not None
assert scheduler._generation_min_batch_size == 32
assert scheduler._generation_max_wait_s == 0.012
assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
warning.assert_not_called()
if profile == "low_latency":
assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1
else:
assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None
assert scheduler._generation_min_batch_size == 1
assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0
warning.assert_called_once()
@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)])
def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp):
config = _config(
stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12}
)
with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"):
construct_scheduler(config)
@pytest.mark.parametrize(
"extras,message",
[
({"generation_min_batch_size": 0}, "generation_min_batch_size"),
({"generation_min_batch_size": 129}, "generation_min_batch_size"),
({"generation_max_wait_ms": -1}, "generation_max_wait_ms"),
({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"),
({"generation_max_regular_batch": -1}, "generation_max_regular_batch"),
({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"),
],
)
def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message):
with pytest.raises(ValueError, match=message):
construct_scheduler(_config(native=False, extras=extras))
@@ -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.deferred_waiting = set()
scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
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)
class _HashableChunkRequest(SimpleNamespace):
# Deferred queues require identity hashing; SimpleNamespace types it as None.
__hash__ = object.__hash__ # type: ignore[assignment]
__eq__ = object.__eq__
return SimpleNamespace(**defaults)
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.deferred_waiting = set()
scheduler.skipped_waiting = create_request_queue(scheduler.policy)
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
@@ -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.deferred_waiting = set()
for req in kv_holding_waiting:
scheduler.kv_holding_waiting.add_request(req)
scheduler.skipped_waiting = create_request_queue(policy)
for req in skipped_waiting:
scheduler.skipped_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)
@@ -1,16 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
from vllm import SamplingParams
from vllm.v1.engine import EngineCoreRequest
from vllm.v1.request import RequestStatus
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
from vllm_omni.engine.serialization import deserialize_additional_information
@@ -31,7 +25,7 @@ def _make_engine_request() -> EngineCoreRequest:
)
def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
tagged = apply_omni_final_stage_metadata(
_make_engine_request(),
final_stage_id=0,
@@ -43,187 +37,21 @@ def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
request_id="req",
additional_information=tagged.additional_information,
)
omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request)
omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request)
result = scheduler._request_omits_kv_transfer_to_next_stage(request)
metadata = deserialize_additional_information(tagged.additional_information)
return omits_kv, omits_chunk, metadata
return result, metadata
def test_stage_zero_request_omits_kv_transfer():
omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False)
omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
assert omits_kv
assert omits_chunk
assert omits_transfer
assert "omni_force_kv_transfer" not in metadata
def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True)
omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
assert not omits_kv
assert omits_chunk
assert not omits_transfer
assert metadata["omni_final_stage_id"] == 0
assert metadata["omni_force_kv_transfer"] is True
class _ChunkRequest(SimpleNamespace):
def __hash__(self):
return hash(self.request_id)
def __eq__(self, other):
return isinstance(other, _ChunkRequest) and other.request_id == self.request_id
def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False):
tagged = apply_omni_final_stage_metadata(
_make_engine_request(),
final_stage_id=final_stage_id,
force_kv_transfer=force_kv_transfer,
)
request = _ChunkRequest(
request_id="req-shm-leak",
additional_information=tagged.additional_information,
num_in_flight_tokens=1,
num_stale_output_tokens=0,
sampling_params=SimpleNamespace(num_logprobs=None),
has_encoder_inputs=False,
pooling_params=None,
status=RequestStatus.RUNNING,
resumable=False,
num_output_placeholders=0,
spec_token_ids=[],
_output_token_ids=[],
num_computed_tokens=1,
stop_reason=None,
num_nans_in_logits=None,
client_index=0,
trace_headers=None,
)
request.is_finished = lambda: False
request.get_finished_reason = lambda: None
request.take_prefill_stats = lambda: None
request.take_events = lambda: None
return request
def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None):
adapter = mocker.MagicMock()
adapter._confirmed_num_computed_tokens.return_value = 1
sched = mocker.MagicMock()
sched.requests = {request.request_id: request}
sched.perf_metrics = None
sched.defer_block_free = False
sched.structured_output_manager.should_advance.return_value = False
sched._update_request_with_output.return_value = ([42], True)
sched._process_kv_transfer_trigger.return_value = False
sched._handle_stopped_request.return_value = True
def _free_request(req, delay_free_blocks=False):
sched._omits_kv_transfer_cache.pop(req.request_id, None)
return None, None
sched._free_request.side_effect = _free_request
sched.chunk_transfer_adapter = adapter
sched.running = [request]
sched.waiting_for_transfer_free = set()
sched.transfer_triggered_requests = set()
sched.active_kv_transfers = set()
sched.pending_stop_after_extraction = set()
sched.connector = None
sched.kv_cache_manager.take_events.return_value = None
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None
bind_omits_transfer_helpers(sched)
scheduler_output = mocker.MagicMock()
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
scheduler_output.total_num_scheduled_tokens = 1
scheduler_output.scheduled_spec_decode_tokens = {}
scheduler_output.num_invalid_spec_tokens = 0
model_runner_output = mocker.MagicMock()
model_runner_output.sampled_token_ids = [[42]]
model_runner_output.logprobs = None
model_runner_output.prompt_logprobs_dict = {}
model_runner_output.pooler_output = None
model_runner_output.num_nans_in_logits = None
model_runner_output.kv_connector_output = None
model_runner_output.cudagraph_stats = None
model_runner_output.req_id_to_index = {request.request_id: 0}
model_runner_output.routed_experts = None
model_runner_output.inter_stage_outputs = inter_stage_outputs
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
return adapter
@pytest.mark.parametrize(
("final_stage_id", "force_kv_transfer", "expect_save"),
[
pytest.param(0, False, False, id="stage-zero-final"),
pytest.param(1, False, True, id="downstream-final"),
],
)
def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save):
request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer)
adapter = _run_finished_save_step(mocker, request)
assert adapter.save_async.called is expect_save
def test_stage_zero_final_skips_inter_stage_output(mocker):
warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning")
request = _make_chunk_request(0)
adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}])
adapter.save_async.assert_not_called()
assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list)
def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
scheduler = OmniARScheduler.__new__(OmniARScheduler)
scheduler._omits_kv_transfer_cache = {}
scheduler.deferred_waiting = set()
scheduler.kv_holding_waiting = []
scheduler.log_stats = False
tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
class Session(SimpleNamespace):
__hash__ = object.__hash__
session = Session(
request_id="req",
additional_information=tagged.additional_information,
status=RequestStatus.WAITING,
)
assert scheduler._request_omits_chunk_transfer_to_next_stage(session)
assert "req" in scheduler._omits_kv_transfer_cache
update = SimpleNamespace(
additional_information={"omni_final_stage_id": 1},
model_intermediate_buffer=None,
arrival_time=0.0,
sampling_params=None,
)
OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update)
assert scheduler._omits_kv_transfer_cache == {}
assert session.additional_information == {"omni_final_stage_id": 1}
assert not scheduler._request_omits_chunk_transfer_to_next_stage(session)
assert not scheduler._request_omits_kv_transfer_to_next_stage(session)
def test_omits_cache_misses_when_payload_object_is_replaced():
scheduler = OmniARScheduler.__new__(OmniARScheduler)
scheduler._omits_kv_transfer_cache = {}
tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
request = SimpleNamespace(
request_id="req",
additional_information=tagged.additional_information,
)
assert scheduler._request_omits_chunk_transfer_to_next_stage(request)
request.additional_information = {"omni_final_stage_id": 1}
assert not scheduler._request_omits_chunk_transfer_to_next_stage(request)
@@ -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(),
deferred_waiting=set(),
skipped_waiting=_RequestQueue(),
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=[[]],
@@ -40,7 +40,6 @@ from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
# isort: on
@@ -73,8 +72,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.deferred_waiting = set()
sched.skipped_waiting = set()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
update = StreamingUpdate(
mm_features=None,
@@ -105,7 +103,6 @@ def _make_drain_sched(session: Request) -> MagicMock:
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None
bind_omits_transfer_helpers(sched)
return sched
@@ -26,7 +26,6 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS
from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
OmniChunkTransferAdapter,
)
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
# isort: on
@@ -42,8 +41,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.deferred_waiting = set()
sched.skipped_waiting = set()
sched._free_request_blocks = MagicMock()
sched.encoder_cache_manager = MagicMock()
sched._inflight_prefills = set()
@@ -177,7 +175,6 @@ def _run_resumable_segment_stop(
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None
bind_omits_transfer_helpers(sched)
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -315,7 +312,6 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None
bind_omits_transfer_helpers(sched)
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -404,7 +400,6 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None
bind_omits_transfer_helpers(sched)
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -1170,8 +1165,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.deferred_waiting = set()
sched.skipped_waiting = queue(session)
sched.num_waiting_for_streaming_input = 1
sched.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=True,
@@ -1183,7 +1177,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 +1204,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.deferred_waiting = set()
sched.skipped_waiting = MagicMock()
sched._enqueue_waiting_request = MagicMock()
sched._resume_downstream_chunk_receiver(session)
@@ -1219,7 +1212,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 +1268,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.deferred_waiting.add(session)
sched.skipped_waiting.add_request(session)
sched.num_waiting_for_streaming_input = 1
@@ -1299,7 +1291,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 +1399,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.deferred_waiting = set()
sched.skipped_waiting = create_request_queue(sched.policy)
sched.running = []
sched.requests = {}
sched._free_request = MagicMock() # needs a KV manager; not what this covers
@@ -1437,7 +1428,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 +1447,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 session not in sched.deferred_waiting
assert len(sched.waiting) == 1
assert session.status == RequestStatus.WAITING
assert not getattr(sched, "_streaming_context_overflow", {})
sched._free_request.assert_not_called()
@@ -1479,7 +1469,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 +1492,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 +1525,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 +1551,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 +1559,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
@@ -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.deferred_waiting = set()
sched.skipped_waiting = _MockQueue()
return sched
@pytest.mark.parametrize("queue_name", ["waiting", "kv_holding_waiting"])
def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready(queue_name) -> None:
def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
scheduler = _make_scheduler()
scheduler.running = [_MockRequest("decode")]
scheduler.waiting = _MockQueue()
setattr(scheduler, queue_name, _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]))
scheduler.waiting = _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]
@@ -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.deferred_waiting = set()
scheduler.skipped_waiting = create_request_queue(scheduler.policy)
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
@@ -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.deferred_waiting: set[Request] = set()
self.skipped_waiting = _SkippedWaitingStub()
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.deferred_waiting = set()
sched.skipped_waiting = MagicMock()
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.deferred_waiting = set()
sched.skipped_waiting = MagicMock()
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:
@@ -123,21 +123,13 @@ 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
# Unset MagicMock attributes are truthy. Keep optional scheduling modes
# disabled so this fixture actually schedules its one-shot request.
sched._generation_defer_batch = False
sched._generation_first_chunk_only = False
sched._generation_max_regular_batch = 0
sched._pause_state = PauseState.UNPAUSED
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)
@@ -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
@@ -1,6 +1,3 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from types import SimpleNamespace
@@ -52,8 +49,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"]
@@ -81,7 +76,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
scheduler = OmniARScheduler.__new__(OmniARScheduler)
scheduler.input_coordinator = coordinator
scheduler.chunk_transfer_adapter = None
scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)}
scheduler._omits_kv_transfer_cache = {"req-free": True}
scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
scheduler.finished_req_ids = set()
scheduler.finished_req_ids_dict = None
@@ -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)
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.waiting = ["waiting"]
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-full", [holding, waiting], scheduler.running),
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
("restore-full", scheduler.waiting, scheduler.running),
]
@@ -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"})
assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"})
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())
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
coordinator.process_pending_chunks.assert_called_with([], [], set(), set())
@@ -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.deferred_waiting = set()
scheduler.skipped_waiting = skipped if skipped is not None else []
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:
"""``get_num_unfinished_requests`` sums ``waiting`` and ``kv_holding_waiting``,
def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
"""``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.deferred_waiting = set()
scheduler.skipped_waiting = []
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.deferred_waiting = set()
scheduler.kv_holding_waiting.add_request(aborted)
scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
scheduler.skipped_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 == []
@@ -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",
id="session_update_missing_session",
("model", "Missing required field"),
id="session_update_missing_model",
),
pytest.param(
json.dumps({"type": "session.update", "session": []}),
"invalid_event",
"Invalid or unrecognized client event",
id="session_update_invalid_session",
_REALTIME_WS_MODEL_MISMATCH,
"model_not_found",
("model", "does not exist"),
id="session_update_model_mismatch",
),
pytest.param(
json.dumps({"type": "input_audio_buffer.commit"}),
"invalid_request_error",
"Input audio buffer is empty",
id="commit_without_audio",
json.dumps({"type": "input_audio_buffer.commit", "final": False}),
"model_not_validated",
("session.update", "validate the model"),
id="commit_without_model_validation",
),
pytest.param(
json.dumps({"type": "unknown.event"}),
"invalid_event",
"Invalid or unrecognized client event",
"unknown_event",
"Unknown event type",
id="unknown_event_type",
),
pytest.param(
json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
"invalid_request_error",
"Invalid base64 audio data",
_REALTIME_WS_INVALID_AUDIO_APPEND,
"invalid_audio",
"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,
+1 -3
View File
@@ -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
# skew without changing the collective deadlock checks below.
timeout=datetime.timedelta(seconds=60),
timeout=datetime.timedelta(seconds=10),
)
try:
fs_groups = [
@@ -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):
@@ -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)
@@ -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):
# Exercise the real upstream selector, replacing only platform resolution.
def select_backend(*, backend, attn_selector_config, num_heads):
# 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")
+4 -15
View File
@@ -14,7 +14,6 @@ from typing import Literal
import pytest
import torch
from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
from vllm_omni.diffusion.distributed.parallel_state import (
@@ -57,8 +56,6 @@ def _init_worker(
"MASTER_PORT": str(master_port),
}
)
if device_kind == "cpu":
configure_cpu_collective_worker()
backend = "gloo" if device_kind == "cpu" else None
init_distributed_environment(backend=backend)
@@ -230,10 +227,6 @@ def _run_ring_p2p(
atol=atol,
msg=f"[Rank {local_rank}] Ring P2P data mismatch",
)
if device_kind == "cpu":
# P2P completion only synchronizes neighbors. Keep fast CPU ranks
# alive until every rank has finished constructing and using groups.
torch.distributed.barrier()
finally:
destroy_distributed_env()
@@ -250,7 +243,7 @@ def _spawn_4d_identity(
device_kind: DeviceKind,
master_port: int,
) -> None:
start_collective_workers(
torch.multiprocessing.spawn(
_run_4d_identity,
args=(
world_size,
@@ -264,7 +257,6 @@ def _spawn_4d_identity(
master_port,
),
nprocs=world_size,
device_kind=device_kind,
)
@@ -280,7 +272,7 @@ def _spawn_5d_identity(
device_kind: DeviceKind,
master_port: int,
) -> None:
start_collective_workers(
torch.multiprocessing.spawn(
_run_5d_identity,
args=(
world_size,
@@ -294,7 +286,6 @@ def _spawn_5d_identity(
master_port,
),
nprocs=world_size,
device_kind=device_kind,
)
@@ -308,11 +299,10 @@ def _spawn_ring_p2p(
device_kind: DeviceKind,
master_port: int,
) -> None:
start_collective_workers(
torch.multiprocessing.spawn(
_run_ring_p2p,
args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
nprocs=world_size,
device_kind=device_kind,
)
@@ -509,9 +499,8 @@ def _run_fused_qkv_matches_three_4d(
@pytest.mark.cpu
@pytest.mark.parametrize("world_size", [2, 4])
def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
start_collective_workers(
torch.multiprocessing.spawn(
_run_fused_qkv_matches_three_4d,
args=(world_size, "cpu", 29613),
nprocs=world_size,
device_kind="cpu",
)
@@ -19,7 +19,6 @@ from vllm.utils.network_utils import get_file_store_init_method
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
from tests.helpers.distributed import collective_context, start_collective_workers
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
@@ -40,25 +39,6 @@ DeviceKind = Literal["cpu", "cuda"]
_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _restore_collective_platforms():
"""CPU baselines run in the parent process; retain its original platforms."""
import vllm.platforms
from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
original_platform = vllm.platforms.current_platform
original_omni_platforms = {
module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state)
}
try:
yield
finally:
vllm.platforms.current_platform = original_platform
for module, platform in original_omni_platforms.items():
module.current_omni_platform = platform
def _cleanup_distributed() -> None:
"""Destroy omni distributed state left dangling after a test.
@@ -429,10 +409,6 @@ def init_dist(
device = _worker_device(local_rank, device_kind)
if device_kind == "cuda":
current_omni_platform.set_device(device)
else:
from tests.helpers.distributed import configure_cpu_collective_worker
configure_cpu_collective_worker()
backend = "gloo" if device_kind == "cpu" else None
init_distributed_environment(
world_size=world_size,
@@ -524,14 +500,13 @@ def isend_irecv_worker(
def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
mp_context = collective_context(device_kind)
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
q = manager.Queue()
start_collective_workers(
torch.multiprocessing.spawn(
isend_irecv_worker,
args=(pp_size, init_method, device_kind, q),
nprocs=pp_size,
device_kind=device_kind,
)
results = {label: tensor for label, tensor in [q.get(), q.get()]}
torch.testing.assert_close(
@@ -681,16 +656,15 @@ def _run_predict_noise(
baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
mp_context = collective_context(device_kind)
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
pp_q = manager.Queue()
world_size = pp_size * cfg_size
start_collective_workers(
torch.multiprocessing.spawn(
predict_noise_worker,
args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
nprocs=world_size,
device_kind=device_kind,
)
pp_out = pp_q.get()
@@ -926,16 +900,15 @@ def _run_scheduler_step(
baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
mp_context = collective_context(device_kind)
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
q = manager.Queue()
world_size = pp_size * cfg_size
start_collective_workers(
torch.multiprocessing.spawn(
scheduler_step_worker,
args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
nprocs=world_size,
device_kind=device_kind,
)
result = q.get()
@@ -61,10 +61,7 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa
# The conditioning is prepared once per request and reused by every step.
assert prepare_spy.call_count == 1
# _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5)
# The time grid reaches prepare, so the adaLN modulations are computed once.
assert prepare_spy.call_args.args[7].numel() == 2
assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
assert step_spy.call_count == 2
capture_spy.assert_not_called()
assert not wrapper._cache
@@ -87,8 +84,7 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
assert padded[4].shape == (1, 100, 4)
assert padded[5].shape == (1, 100)
assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0)
assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
assert wrapper.max_graphs == 32
@@ -157,21 +153,18 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
inputs["ref"],
inputs["ref_mask"],
)
# sample_latents passes the grid, so the key carries its step count (two steps here).
key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2)
key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
assert key in wrapper._cache
# The static context was refreshed for this request's conditioning.
entry = wrapper._cache[key]
grid = torch.tensor([0.0, 0.4], device="cuda")
fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
if want is None:
assert static is None
else:
torch.testing.assert_close(static, want)
torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
assert int(entry.static_step) == 1
assert len(wrapper._cache) == 1
@@ -1,139 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The AuK DiT fast paths keep the reference numerics.
Two paths are checked: the grouped conv position embedding as one batched
GEMM, and the adaLN modulations precomputed for a whole time grid, eager and
under the per-step CUDA graph.
"""
import pytest
import torch
from vllm_omni.diffusion.models.auk import auk_transformer as T
from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid
from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_dit(device: str = "cpu") -> AuKTransformer:
torch.manual_seed(5)
return (
AuKTransformer(
dim=32,
heads=2,
dim_head=16,
ff_mult=2,
latent_dim=4,
text_hidden_dim=8,
num_layers=2,
num_single_layers=2,
)
.eval()
.to(device)
)
def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
"""The channels-first conv stack as the reference implementation writes it."""
keep = None if mask is None else mask.unsqueeze(1)
x = x.transpose(1, 2)
if keep is not None:
x = x.masked_fill(~keep, 0.0)
for layer in module.conv1d:
x = layer(x)
if keep is not None and isinstance(layer, torch.nn.Conv1d):
x = x.masked_fill(~keep, 0.0)
return x.transpose(1, 2)
@torch.inference_mode()
def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None:
torch.manual_seed(0)
conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15)
x = torch.randn(2, 37, 64)
expected = conv(x.transpose(1, 2)).transpose(1, 2)
torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5)
@torch.inference_mode()
@pytest.mark.parametrize("masked", [False, True])
def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None:
torch.manual_seed(1)
module = T.ConvPosEmbedding(64).eval()
x = torch.randn(2, 37, 64)
mask = None
if masked:
mask = torch.ones(2, 37, dtype=torch.bool)
mask[1, 30:] = False
torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5)
@torch.inference_mode()
@pytest.mark.parametrize("cfg_infer", [False, True])
def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None:
dit = _make_dit()
text = torch.randn(1, 5, 8)
ref = torch.randn(1, 3, 4)
x = torch.randn(1, 6, 4)
grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu")
plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer)
tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1])
assert tabled.modulation is not None and tabled.modulation.shape[0] == 4
for i in range(4):
expected = dit.step(x, grid[i], plain)
torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5)
index = torch.tensor(i)
torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5)
def _padded_inputs(device: str) -> dict[str, torch.Tensor]:
generator = torch.Generator(device="cpu").manual_seed(3)
text = torch.randn(1, 7, 8, generator=generator)
c_mask = torch.ones(1, 7, dtype=torch.bool)
c_mask[:, 5:] = False
text[:, 5:] = 0.0
ref = torch.randn(1, 6, 4, generator=generator)
ref_mask = torch.ones(1, 6, dtype=torch.bool)
ref_mask[:, 4:] = False
mask = torch.ones(1, 9, dtype=torch.bool)
mask[:, 7:] = False
inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask}
return {name: value.to(device) for name, value in inputs.items()}
@pytest.mark.cuda
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
@torch.inference_mode()
@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None:
dit = _make_dit("cuda")
inputs = _padded_inputs("cuda")
grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda")
x = torch.randn(1, 9, 4, device="cuda")
wrapper = AuKCUDAGraphWrapper(dit)
common = dict(
text=inputs["text"],
c_mask=inputs["c_mask"],
ref=inputs["ref"],
ref_mask=inputs["ref_mask"],
cfg_strength=cfg_strength,
timesteps=grid[:-1],
)
for i in range(4):
graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common)
ctx = dit.prepare(
inputs["text"],
target_len=9,
c_mask=inputs["c_mask"],
ref=inputs["ref"],
ref_mask=inputs["ref_mask"],
cfg_infer=cfg_strength >= 1e-5,
timesteps=grid[:-1],
)
eager = wrapper._step(x, grid[i], ctx, cfg_strength, i)
torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4)
# One graph serves every step of the grid.
assert len(wrapper._cache) == 1
assert next(iter(wrapper._cache))[-1] == 4
@@ -127,9 +127,7 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non
v = dit(x, time=timesteps[i], **inputs)
x = x + (timesteps[i + 1] - timesteps[i]) * v
# sample_latents precomputes the adaLN modulations of the whole grid in one
# GEMM per layer, whose rows round differently from one-row GEMMs in fp32.
torch.testing.assert_close(got, x, rtol=0, atol=1e-6)
torch.testing.assert_close(got, x, rtol=0, atol=0)
assert dit.text_cond is None and dit.text_uncond is None
@@ -12,8 +12,7 @@ import pytest
import torch
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.models.auk import auk_vae
from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample
from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
from vllm_omni.platforms import current_omni_platform
@@ -247,72 +246,6 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
AuKVAEDecodeGraph(vae, tile_frames=60)
@pytest.mark.cpu
@torch.inference_mode()
def test_fused_activation_is_skipped_off_cuda() -> None:
activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True))
x = torch.randn(1, 4, 16)
assert not activation._fused_eligible(x)
assert activation(x).shape == (1, 4, 16)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
@torch.inference_mode()
def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None:
activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda")
x = torch.randn(1, 4, 16, device="cuda")
assert activation._fused_eligible(x)
monkeypatch.setattr(auk_vae, "HAS_TRITON", False)
assert not activation._fused_eligible(x)
assert activation(x).shape == (1, 4, 16)
def _activation_input(layout: str, length: int) -> torch.Tensor:
"""A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer."""
if layout == "contiguous":
x = torch.randn(2, 6, length, device="cuda")
elif layout == "transposed":
x = torch.randn(2, length, 6, device="cuda").transpose(1, 2)
elif layout == "sliced":
x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2]
else:
raise ValueError(layout)
assert x.shape == (2, 6, length)
if length > 1: # contiguity of a single-sample view depends on the layout
assert x.is_contiguous() == (layout == "contiguous")
return x * 3
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"])
@pytest.mark.parametrize("causal", [False, True])
# The kernels tile the output in blocks of 1024 samples: the 2x oversampled
# signal crosses a block boundary at 512 input samples and the decimated one at
# 1024, so both sides of each boundary are covered, plus short and odd lengths.
@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025])
@torch.inference_mode()
def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None:
torch.manual_seed(0)
activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda")
with torch.no_grad():
activation.act.alpha.normal_(0.0, 0.3)
activation.act.beta.normal_(0.0, 0.3)
x = _activation_input(layout, length)
original = x.clone()
eager = activation.downsample(activation.act(activation.upsample(x)))
assert activation._fused_eligible(x)
fused = activation(x)
# Same fp32 arithmetic, only the order in which the taps are summed differs.
assert fused.shape == eager.shape
torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5)
# Neither path writes into its (possibly strided) input.
torch.testing.assert_close(x, original, atol=0, rtol=0)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
@torch.inference_mode()
@@ -291,18 +291,6 @@ class TestRequestParsing:
with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"):
build_pipeline(model_config={"max_dit_graphs": max_dit_graphs})
@pytest.mark.parametrize(
"model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)]
)
def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected):
pipeline, _ = build_pipeline(model_config=model_config)
assert pipeline._ref_cache_size == expected
@pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None])
def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size):
with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"):
build_pipeline(model_config={"auk_ref_cache_size": size})
def test_pipeline_declares_audio_output(self, build_pipeline):
pipeline, _ = build_pipeline()
@@ -504,40 +492,6 @@ class TestRequestParsing:
assert call["sample"] is True
assert isinstance(call["generator"], torch.Generator)
def test_reference_latents_are_cached_by_content(self, build_pipeline):
pipeline, calls = build_pipeline()
clip = _silence(1.0)
for _ in range(3):
pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
# A different clip is a new entry, the first one still hits.
other = (np.full_like(clip[0], 0.25), clip[1])
pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1))
pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
assert len(pipeline.vae.encode_calls) == 2
assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"]
def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline):
pipeline, _ = build_pipeline()
for _ in range(2):
pipeline.forward(
_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)
)
assert len(pipeline.vae.encode_calls) == 2
pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0})
for _ in range(2):
pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1))
assert len(pipeline.vae.encode_calls) == 2
def test_reference_cache_is_bounded(self, build_pipeline):
pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2})
for level in (0.1, 0.2, 0.3):
clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE)
pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
assert len(pipeline._ref_cache) == 2
def test_one_request_per_forward(self, build_pipeline):
pipeline, _ = build_pipeline()
batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)
@@ -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_omni.quantization.fp8_config import DiffusionFp8Config
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
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(
@@ -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,
@@ -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,
+16 -81
View File
@@ -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_omni.quantization.fp8_config import DiffusionFp8Config
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
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_omni.quantization.fp8_config import DiffusionFp8Config
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
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_omni.quantization.fp8_config import DiffusionFp8Config
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
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"
@@ -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()
@@ -11,12 +11,10 @@ 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
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
@@ -97,30 +95,6 @@ def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatc
assert not adapter.code_prompt_token_ids
def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
adapter, connector = shm_sender
ext_id = f"release_{uuid.uuid4().hex}"
request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id)
sibling_key = f"{ext_id}_1_0_0"
assert connector.put("0", "1", sibling_key, "sibling")[0]
for _ in range(3):
adapter.save_async(None, request)
adapter._send_single_request(adapter._pending_save_reqs.popleft())
# The consumer read only the first chunk before its own request ended.
assert connector.get("0", "1", f"{ext_id}_0_0") is not None
adapter.release_shm_resources(ext_id)
# The release is queued behind sends; nothing is unlinked on the caller.
assert f"{ext_id}_0_1" in connector._pending_keys
while adapter._pending_save_reqs:
adapter._send_single_request(adapter._pending_save_reqs.popleft())
assert adapter._reclaimed_shm_total == 2
assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
assert connector.get("0", "1", f"{ext_id}_0_1") is None
assert connector.get("0", "1", sibling_key)[0] == "sibling"
@pytest.mark.parametrize("fail_terminal", [False, True])
def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal):
adapter, connector = shm_sender
@@ -636,45 +610,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 +2420,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.deferred_waiting = set()
scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
freed = []
@@ -2510,7 +2444,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 +2467,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.deferred_waiting = set()
scheduler.skipped_waiting = DummyWaitingQueue()
scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
assert (
@@ -2608,7 +2541,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
@@ -2688,10 +2620,10 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
scheduler = mocker.MagicMock()
scheduler.aux_output_connector = None
scheduler.chunk_transfer_adapter = adapter_mock
bind_omits_transfer_helpers(scheduler)
scheduler.connector = None
scheduler.perf_metrics = None
scheduler.log_stats = False
@@ -2788,8 +2720,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 +2735,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.deferred_waiting = set()
scheduler.skipped_waiting = DummyWaitingQueue([session])
scheduler.num_waiting_for_streaming_input = 1
scheduler.finished_req_ids_dict = {}
@@ -213,7 +213,7 @@ class TestCleanup:
connector.put("s0", "s1", "cleanup_req_42", data)
assert "cleanup_req_42" in connector._pending_keys
assert connector.cleanup("cleanup_req_42") is True
connector.cleanup("cleanup_req_42")
assert "cleanup_req_42" not in connector._pending_keys
result = connector.get("s0", "s1", "cleanup_req_42", metadata=None)
@@ -224,7 +224,7 @@ class TestCleanup:
connector.put("s0", "s1", "consumed_req_99", data)
connector.get("s0", "s1", "consumed_req_99", metadata=None)
assert connector.cleanup("consumed_req_99") is False
connector.cleanup("consumed_req_99")
assert "consumed_req_99" not in connector._pending_keys
def test_close_cleans_all_pending(self, connector):
@@ -296,24 +296,6 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch):
assert key not in connector._pending_keys
@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_0", "_suffix"])
def test_cleanup_prefix_only_matches_own_chunk_keys(connector, suffix):
request_id = f"prefix_{uuid.uuid4().hex}"
own_keys = [f"{request_id}_0_{i}" for i in range(3)]
sibling_key = f"{request_id}{suffix}_0_0"
for key in own_keys:
assert connector.put("0", "1", key, key)[0]
assert connector.put("0", "1", sibling_key, "sibling")[0]
# A consumed chunk is no longer counted as reclaimed.
assert connector.get("0", "1", own_keys[0])[0] == own_keys[0]
assert connector.cleanup_prefix(f"{request_id}_0_") == 2
for key in own_keys:
assert key not in connector._pending_keys
assert connector.get("0", "1", key) is None
assert connector.get("0", "1", sibling_key)[0] == "sibling"
# ── Arrival wakeups ───────────────────────────────────────────────────
-26
View File
@@ -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()
@@ -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
@@ -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.
@@ -1,175 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""AuK encoder stage: FULL CUDA graph replay matches eager, request by request.
``auk.yaml`` captures the encoder prefill as FULL CUDA graphs padded to a few
token buckets. A replay reuses static input buffers and the graph pool, so a
request could in principle read conditioning left behind by an earlier, longer
prompt. This test runs the same request sequence through an eager encoder and
through the graphed one and compares the text condition stage 0 hands to
stage 1 (the fused thinker hidden states) for every request. The sequence
alternates prompts of different lengths inside one capture bucket, crosses
into a larger bucket, and then returns to the first prompt.
Like ``test_auk.py`` it needs an assembled checkpoint directory:
VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk_encoder_cudagraph.py
"""
from __future__ import annotations
import bisect
import os
from pathlib import Path
from typing import Any
import pytest
import torch
import yaml
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.model_executor.stage_input_processors import auk as auk_input_processor
from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params
MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR"
# Text-only instruct prompts: the token count, and so the capture bucket,
# follows the text length. SHORT_A and SHORT_B differ in length but share a
# bucket; LONG lands in a larger one. The test checks these properties on the
# recorded condition lengths, so a tokenizer change fails loudly instead of
# silently shrinking the coverage.
_TEMPLATE = 'Generate speech based on the following description: "{desc}". The content to speak is: "{text}".'
SHORT_A = _TEMPLATE.format(
desc="A calm young woman speaking warmly.",
text="Welcome back, how was your day? I kept some soup warm for you, and the kettle is on if you want tea.",
)
SHORT_B = _TEMPLATE.format(
desc="An older man with a deep, slightly raspy voice, speaking slowly.",
text=(
"The train to the coast leaves at seven, so please be on time. Bring a warm coat, "
"some cash for the ferry, and the blue umbrella from the hall, because the forecast says rain."
),
)
LONG = _TEMPLATE.format(
desc=(
"A cheerful radio host in her thirties, bright and energetic, with clear articulation, "
"a slight smile in her voice, and a brisk but relaxed pace, recorded close to the microphone"
),
text=(
"Good morning and welcome to the weekend edition. Today we are talking about small gardens, "
"the plants that thrive on a sunny balcony, and the neighbours who turned an empty lot into "
"a place where the whole street now meets on Saturday afternoons to trade seedlings and stories. "
"Stay with us, because after the news we will open the phone lines and hear from you."
),
)
SEQUENCE = [SHORT_A, SHORT_B, SHORT_A, LONG, SHORT_B, SHORT_A]
# The DiT output is not inspected; keep stage 1 as cheap as it goes (it also
# runs eager in both configs, see _deploy).
GEN_SECONDS = 1.0
NFE = 1
# bf16 hidden states. Graph replay pads the prefill to the bucket size but runs
# the same kernels on the real tokens, so the two paths agree to bf16 rounding.
# Measured on H20: cosine >= 0.99998 and max-abs error <= 2.9e-3 of the
# condition's max magnitude for every request; the bounds leave about 3x
# headroom for other GPUs.
MIN_COSINE = 0.9999
MAX_REL_ABS = 1e-2
_model_dir = os.environ.get(MODEL_DIR_ENV)
pytestmark = [
pytest.mark.slow,
pytest.mark.tts,
pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"),
]
def _deploy(tmp_path: Path, *, graph_encoder: bool) -> str:
"""``auk.yaml`` with stage 1 eager, and the encoder stage eager or graphed as shipped.
Only the encoder differs between the two runs. The conditions compared here
are recorded before stage 1 runs, so its DiT compile, graph capture and
codec bucket warm-up would only add startup time; eager stage 1 skips them
and still takes the real stage handoff.
"""
with open(get_deploy_config_path("auk.yaml")) as f:
deploy = yaml.safe_load(f)
for stage in deploy["stages"]:
if stage["stage_id"] == 1:
stage["enforce_eager"] = True
elif not graph_encoder:
stage["enforce_eager"] = True
stage.pop("compilation_config", None)
path = tmp_path / f"auk_{'graph' if graph_encoder else 'eager'}_encoder.yaml"
path.write_text(yaml.safe_dump(deploy))
return str(path)
def _graph_capture_sizes() -> list[int]:
with open(get_deploy_config_path("auk.yaml")) as f:
deploy = yaml.safe_load(f)
stage = next(s for s in deploy["stages"] if s["stage_id"] == 0)
assert stage.get("enforce_eager") is False, "auk.yaml no longer graphs the encoder; update this test"
return sorted(stage["compilation_config"]["cudagraph_capture_sizes"])
def _run_sequence(deploy_config: str, monkeypatch: pytest.MonkeyPatch) -> list[torch.Tensor]:
"""Run SEQUENCE one request at a time and return each request's text condition."""
conditions: list[torch.Tensor] = []
original = auk_input_processor.encoder2dit
def recording_encoder2dit(*args: Any, **kwargs: Any) -> dict[str, Any] | None:
result = original(*args, **kwargs)
if result is not None:
conditions.append(result["prompt_embeds"].detach().to("cpu", torch.float32).clone())
return result
# The orchestrator runs in this process and resolves the stage input
# function by module attribute at startup, so the patch must precede Omni.
with monkeypatch.context() as patch:
patch.setattr(auk_input_processor, "encoder2dit", recording_encoder2dit)
with OmniRunner(str(Path(_model_dir).resolve()), deploy_config=deploy_config) as runner:
for text in SEQUENCE:
prompt = auk_prompt(text, None, gen_seconds=GEN_SECONDS)
outputs = runner.omni.generate(prompt, auk_sampling_params(nfe=NFE, cfg=2.0, seed=0))
assert outputs, "no outputs returned"
assert len(conditions) == len(SEQUENCE), f"recorded {len(conditions)} conditions for {len(SEQUENCE)} requests"
return conditions
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_encoder_graph_replay_matches_eager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
capture_sizes = _graph_capture_sizes()
eager = _run_sequence(_deploy(tmp_path, graph_encoder=False), monkeypatch)
graphed = _run_sequence(_deploy(tmp_path, graph_encoder=True), monkeypatch)
lengths = {text: cond.shape[0] for text, cond in zip(SEQUENCE, eager)}
for i, (ref, got) in enumerate(zip(eager, graphed)):
assert got.shape == ref.shape, f"request {i}: shape {tuple(got.shape)} != eager {tuple(ref.shape)}"
assert torch.isfinite(got).all(), f"request {i}: non-finite condition"
cosine = torch.nn.functional.cosine_similarity(got.flatten(), ref.flatten(), dim=0).item()
rel_abs = ((got - ref).abs().max() / ref.abs().max()).item()
assert cosine >= MIN_COSINE and rel_abs <= MAX_REL_ABS, (
f"request {i} ({lengths[SEQUENCE[i]]} tokens): cosine {cosine:.6f}, max-abs error {rel_abs:.2e} "
"of the condition's magnitude"
)
# The sequence must exercise what it claims to: two different lengths in
# one bucket, a second bucket, and every prompt inside the graphed range.
assert max(lengths.values()) <= capture_sizes[-1], f"{lengths} exceeds the largest capture size"
bucket = {text: capture_sizes[bisect.bisect_left(capture_sizes, n)] for text, n in lengths.items()}
assert lengths[SHORT_A] != lengths[SHORT_B] and bucket[SHORT_A] == bucket[SHORT_B], (lengths, bucket)
assert bucket[LONG] > bucket[SHORT_A], (lengths, bucket)
# Repeats of a prompt reproduce its first condition exactly, eager or graphed:
# nothing from the requests in between leaks into the replay.
for run in (eager, graphed):
first: dict[str, torch.Tensor] = {}
for text, cond in zip(SEQUENCE, run):
if text in first:
torch.testing.assert_close(cond, first[text], atol=0, rtol=0)
else:
first[text] = cond
@@ -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, server_args=["--enable-layerwise-offload"]),
OmniServerParams(model=model),
id="base",
marks=BASE_FEATURE_MARKS,
),
@@ -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(
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}}},
),
stage_config_path=get_deploy_config_path("moss_tts_realtime.yaml"),
server_args=["--disable-log-stats"],
),
id="moss_tts_realtime",
@@ -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 = [
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"}}},
)
]
stage_configs = [modify_stage_config(get_deploy_config_path("ci/qwen2_5_omni.yaml"))]
# Create parameter combinations for model and stage config
test_params = [
@@ -1,155 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Live /dev/shm leak check for Qwen3-Omni ``async_chunk`` (stage-0-final).
The default L2 Qwen3-Omni job uses ``ci/qwen3_omni_moe.yaml`` with
``async_chunk: false`` plus ``--no-async-chunk``, so it never exercises the
thinker→talker SharedMemoryConnector path.
A text-only request is tagged ``omni_final_stage_id=0``, so nothing downstream
ever ``get()``s what stage-0 ``put()``s for it. Before #7245 the scheduler still
``put()`` a finished marker on finish and nothing unlinked it, so every such
request left a ``SharedMemoryConnector`` segment plus lockfile in ``/dev/shm``.
Stage-0-final finishes now skip the ``put()`` and the orchestrator reclaims any
undrained segment on request cleanup.
This module starts the production async-chunk deploy and asserts the three
text-only request ids leave no connector lockfile or POSIX segment in
``/dev/shm``.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
_MODEL = os.environ.get("VLLM_OMNI_TEST_MODEL", "Qwen/Qwen3-Omni-30B-A3B-Instruct")
_SHM_DIR = Path("/dev/shm")
_NUM_TEXT_REQUESTS = 3
_SETTLE_S = 15.0
# CI overlay keeps seq/token limits small; flip async_chunk back on so stage-0
# still uses SharedMemoryConnector. Production yaml is already async_chunk:true.
_ASYNC_CHUNK_DEPLOY = modify_stage_config(
get_deploy_config_path("ci/qwen3_omni_moe.yaml"),
updates={"async_chunk": True},
)
test_params = [
pytest.param(
OmniServerParams(
model=_MODEL,
stage_config_path=_ASYNC_CHUNK_DEPLOY,
use_stage_cli=True,
),
id="async_chunk",
)
]
def _system_prompt() -> dict:
return {
"role": "system",
"content": [
{
"type": "text",
"text": (
"You are Qwen, a virtual human developed by the Qwen Team, "
"Alibaba Group, capable of perceiving auditory and visual inputs, "
"as well as generating text and speech."
),
}
],
}
def _shm_used_mib() -> str:
try:
st = os.statvfs(_SHM_DIR)
used = (st.f_blocks - st.f_bavail) * st.f_frsize
return f"{used / (1024 * 1024):.1f}MiB"
except OSError:
return "unknown"
def _is_entry_for_request(name: str, request_id: str) -> bool:
"""Match a segment or ``shm_`` lockfile keyed by the ``{request_id}-{suffix}`` orchestrator id."""
return name.removeprefix("shm_").startswith((f"{request_id}-", f"{request_id}_"))
def _entries_for_request_ids(request_ids: list[str]) -> list[str]:
if not _SHM_DIR.is_dir():
pytest.skip("/dev/shm is not available")
found: list[str] = []
with os.scandir(_SHM_DIR) as entries:
for entry in entries:
if any(_is_entry_for_request(entry.name, request_id) for request_id in request_ids):
found.append(entry.name)
return found
def _wait_no_entries_for_request_ids(request_ids: list[str], timeout_s: float = _SETTLE_S) -> list[str]:
"""Reclaim is fire-and-forget; poll until these ids are gone or *timeout_s*."""
deadline = time.monotonic() + timeout_s
time.sleep(1.0)
leftover: list[str] = _entries_for_request_ids(request_ids)
while leftover:
if time.monotonic() >= deadline:
return leftover
time.sleep(2.0)
leftover = _entries_for_request_ids(request_ids)
return leftover
@pytest.mark.advanced_model
@pytest.mark.omni
@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_only_async_chunk_does_not_leak_shm(omni_server, online_client) -> None:
"""Stage-0-final (text-only) requests must not leave connector SHM behind."""
print(f"[shm-leak] before requests: /dev/shm_used={_shm_used_mib()}", flush=True)
messages = dummy_messages_from_mix_data(
system_prompt=_system_prompt(),
content_text="What is the capital of China? Answer in 20 words.",
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": False,
"modalities": ["text"],
"key_words": {"text": ["beijing"]},
}
responses = online_client.send_omni_request(request_config, request_num=_NUM_TEXT_REQUESTS)
produced = [resp.text_content for resp in responses if getattr(resp, "text_content", None)]
assert len(produced) == _NUM_TEXT_REQUESTS, (
f"Need {_NUM_TEXT_REQUESTS} text completions to exercise stage-0 finish; "
f"got {len(produced)} non-empty texts. A preprocess/engine miss would "
f"leave /dev/shm unchanged and hide a leak."
)
request_ids = [resp.request_id for resp in responses]
assert all(request_ids) and len(request_ids) == _NUM_TEXT_REQUESTS, (
f"Need chat completion ids to assert per-request SHM cleanup; got {request_ids}"
)
leftover = _wait_no_entries_for_request_ids(request_ids)
print(
f"[shm-leak] after {_NUM_TEXT_REQUESTS} text-only requests: "
f"/dev/shm_used={_shm_used_mib()} leftover={leftover} request_ids={request_ids}",
flush=True,
)
assert not leftover, (
f"SharedMemoryConnector leaked /dev/shm entries for {leftover} after stage-0-final request(s) {request_ids}"
)
-34
View File
@@ -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()
@@ -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)
@@ -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)
@@ -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
# 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)
importlib.reload(client_mod)
assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
"StageEngineCoreClient | DPLBStageEngineCoreClient"
)
finally:
client_mod.__dict__.clear()
client_mod.__dict__.update(saved)
assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
"StageEngineCoreClient | DPLBStageEngineCoreClient"
)
def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):
+1 -4
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Deterministic tests for the CFG-companion output lifecycle.
Regression suite for the companion-output race behind the nightly Bagel
@@ -94,9 +94,6 @@ class _FakePool:
def release_bindings(self, request_ids):
pass
async def release_request_resources(self, request_ids):
pass
def _make_orchestrator(num_stages: int = 2) -> Orchestrator:
orch = Orchestrator.__new__(Orchestrator)
+1 -27
View File
@@ -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"
-13
View File
@@ -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.
@@ -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),
@@ -7,9 +7,8 @@ from unittest.mock import MagicMock
import pytest
import torch
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.v1.engine import EngineCoreRequest, FinishReason
from vllm.v1.engine.logprobs import LogprobsProcessor
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine import FinishReason
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.
+1
View File
@@ -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",
+1 -52
View File
@@ -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 (
DPLBStageEngineCoreClient,
StageEngineCoreClient,
StageEngineCoreClientBase,
)
from vllm_omni.engine.stage_init_utils import StageMetadata
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
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,
)
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from typing import Any
from unittest.mock import patch
import pytest
@@ -26,15 +25,11 @@ def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
expected,
):
executor = UniProcExecutor.__new__(UniProcExecutor)
executor.vllm_config = SimpleNamespace(
parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)
)
executor.driver_worker = mocker.Mock()
runner = executor.driver_worker.worker.model_runner
runner.model.first_frame_decoder = object() if first_decoder else None
runner.model.stream_decoder = object() if stream_decoder else None
runner.model.stream_first_audio = stream_first_audio
runner._omni_data_plane = None
assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
assert runner.model_state.set_first_audio_sink.call_count == int(expected)
@@ -59,64 +54,3 @@ def test_preprocess_add_request_preserves_omni_fields():
assert current_wave == 3
assert result.external_req_id == "external"
assert result.additional_information == {"conditioning": "payload"}
def test_codec_uses_generic_first_audio_binding(monkeypatch):
import queue
import torch
from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
from vllm_omni.engine import stage_engine_core_proc as module
calls: dict[str, Any] = {}
def hook(*args):
return False
def bind(sink):
calls["sink"] = sink
return hook
plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value))
model = SimpleNamespace(bind_first_chunk_fast_path=bind)
runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane)
executor = SimpleNamespace(
vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
)
monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
outputs: queue.Queue = queue.Queue()
scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)})
assert module._bind_first_audio_sink(executor, outputs, scheduler)
assert calls["hook"] is hook
calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
client, batch = outputs.get_nowait()
assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
def test_first_audio_binding_preserves_talker_marker(monkeypatch):
import queue
import torch
from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
from vllm_omni.engine import stage_engine_core_proc as module
calls: dict[str, Any] = {}
runner = SimpleNamespace(
model=SimpleNamespace(first_frame_decoder=object()),
model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)),
)
executor = SimpleNamespace(
vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
)
monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
outputs: queue.Queue = queue.Queue()
scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)})
assert module._bind_first_audio_sink(executor, outputs, scheduler)
calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY]
executor.vllm_config.parallel_config.tensor_parallel_size = 2
assert not module._bind_first_audio_sink(executor, outputs, scheduler)
@@ -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:
@@ -175,46 +154,6 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
asyncio.run(run())
@pytest.mark.cpu
def test_release_request_resources_skips_without_async_chunk():
async def run() -> None:
call = AsyncMock()
pool = StagePool(
0,
[SimpleNamespace(call_utility_async=call)], # type: ignore[list-item]
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)),
)
await pool.release_request_resources(["req-1"])
call.assert_not_awaited()
asyncio.run(run())
@pytest.mark.cpu
def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch):
async def run() -> None:
async def hang(*_args):
await asyncio.Event().wait()
live = AsyncMock()
pool = StagePool(
0,
[SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)], # type: ignore[list-item]
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)),
)
monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2)
release = asyncio.create_task(pool.release_request_resources(["req-1"]))
await asyncio.sleep(0.05)
live.assert_awaited_once_with("omni_release_request_resources", ["req-1"])
assert not release.done()
await asyncio.wait_for(release, timeout=1.0)
asyncio.run(run())
@pytest.mark.cpu
@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
@pytest.mark.parametrize("failure", ["error", "timeout", "missing"])
@@ -1,85 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _Processor:
"""Segment API of the model's processor with a deterministic tokenizer."""
def __init__(self):
self.model_config = SimpleNamespace(
n_vq=3,
im_start_token_id=1,
im_end_token_id=2,
audio_start_token_id=3,
audio_pad_token_id=1024,
audio_assistant_slot_token_id=5,
)
self.encoded = []
def _encode_text(self, text):
self.encoded.append(text)
return [100 + ord(c) % 50 for c in text]
def _user_prompt_prefix_ids(self):
return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
def _user_prompt_after_reference_ids(self, language, fields):
return self._encode_text(f"lang={fields.get('language', language)}")
def _assistant_prompt_prefix_ids(self):
return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
def reference(self, text, language, codes):
# Mirrors the processor's _build_continuation_codes + _pad (one item).
fields = {} if language is None else {"language": language}
ids = (
self._user_prompt_prefix_ids()
+ self._encode_text("None")
+ self._user_prompt_after_reference_ids(language, fields)
+ self._encode_text(text)
+ self._assistant_prompt_prefix_ids()
+ [3]
)
text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
text_rows[:, 0] = torch.tensor(ids)
audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
audio_rows[:, 0] = 5
audio_rows[:, 1:] = codes
unified = torch.cat([text_rows, audio_rows])
return unified[:, 0].tolist(), unified[:, 1:].contiguous()
@pytest.mark.parametrize("language", [None, "English"])
def test_matches_processor_and_caches_constant_segments(language):
proc = _Processor()
codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
kwargs = {"text": "ref words target words"}
if language is not None:
kwargs["language"] = language
ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
assert ids == expected_ids
assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
proc.encoded.clear()
_local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
assert proc.encoded == ["other"]
def test_unsupported_layouts_fall_back():
proc = _Processor()
codes = torch.zeros((2, 3), dtype=torch.long)
assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None
@@ -691,45 +691,6 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker)
audio_tokenizer.to.assert_called_once()
@pytest.mark.parametrize(
"variant,graphs_env,prepares",
[("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)],
)
def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch):
if graphs_env is None:
monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False)
else:
monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env)
adapter = _moss_adapter_with_stage_configs([], mocker)
adapter._moss_variant = variant
encoder = mocker.Mock()
mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
asyncio.run(adapter.warmup())
assert encoder.prepare.called is prepares
def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch):
monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
adapter = _moss_adapter_with_stage_configs([], mocker)
encoder = mocker.Mock()
encoder.prepare.side_effect = RuntimeError("capture failed")
mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
asyncio.run(adapter.warmup())
encoder.prepare.assert_called_once_with()
def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch):
monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
adapter = _moss_adapter_with_stage_configs([], mocker)
encoder = mocker.Mock()
encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling")
mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
with pytest.raises(SharedReferenceEncoderStartupError):
asyncio.run(adapter.warmup())
def test_qwen3_tts_metadata():
assert Qwen3TTSAdapter.backend == "ar"
assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
-34
View File
@@ -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",
@@ -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:
@@ -288,9 +288,6 @@ def _setup_log_mocks(monkeypatch):
def build_and_log_summary(self):
return "Fake summary"
def log_timing_summary(self):
pass
monkeypatch.setattr(
"vllm_omni.entrypoints.omni.OrchestratorAggregator",
_FakeOrchestratorAggregator,
+7 -41
View File
@@ -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(
getattr(response, "audio_bytes", None), request_config
_assert_transcript_matches(
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"),
-4
View File
@@ -97,7 +97,6 @@ def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[
class OmniResponse:
"""Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
request_id: str | None = None
text_content: str | None = None
audio_data: list[str] | None = None
audio_content: str | None = None
@@ -422,8 +421,6 @@ class OnlineOmniClient:
text_content = ""
audio_data = []
for chunk in chat_completion:
if result.request_id is None:
result.request_id = getattr(chunk, "id", None)
for choice in chunk.choices:
content = getattr(getattr(choice, "delta", None), "content", None)
modality = getattr(chunk, "modality", None)
@@ -456,7 +453,6 @@ class OnlineOmniClient:
"""Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
result = OmniResponse()
try:
result.request_id = getattr(chat_completion, "id", None)
audio_data = None
text_content = None
for choice in chat_completion.choices:
-19
View File
@@ -1,19 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Import preload for the CPU-only collective forkserver.
This module is imported in the forkserver process, before it starts workers.
Its CUDA visibility does not alter the pytest process or CUDA parity workers.
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import torch # noqa: E402
# Retain the expensive worker-module imports across CPU process groups.
from tests.diffusion.distributed import test_comm, test_pipeline_parallel # noqa: E402, F401
if torch.cuda.is_initialized():
raise RuntimeError("CPU collective forkserver must not initialize CUDA")
-62
View File
@@ -1,62 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Process and platform setup for isolated collective test workers."""
from collections.abc import Callable
from multiprocessing.context import BaseContext
import torch
def collective_context(device_kind: str) -> BaseContext:
"""Reuse imports in a clean CPU forkserver; CUDA workers use spawn."""
if device_kind == "cpu" and "forkserver" in torch.multiprocessing.get_all_start_methods():
torch.multiprocessing.set_forkserver_preload(["tests.helpers.cpu_collective_server"])
return torch.multiprocessing.get_context("forkserver")
return torch.multiprocessing.get_context("spawn")
def start_collective_workers(
worker: Callable[..., None], *, args: tuple[object, ...], nprocs: int, device_kind: str
) -> None:
context = collective_context(device_kind)
if context.get_start_method() == "spawn":
torch.multiprocessing.spawn(worker, args=args, nprocs=nprocs)
else:
torch.multiprocessing.start_processes(worker, args=args, nprocs=nprocs, start_method=context.get_start_method())
def configure_cpu_collective_worker() -> None:
"""Select CPU platforms inside a spawned worker that only uses Gloo.
Upstream coordinators resolve local ranks through the selected
platform's visible devices. CPU ranks can outnumber visible GPUs, so they
must use the CPU platform even when the test host has CUDA installed.
The worker still creates real process groups and runs real collectives.
"""
import vllm.platforms
from vllm.platforms.cpu import CpuPlatform
from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
from vllm_omni.platforms.interface import UnspecifiedOmniPlatform
class CpuCollectivePlatform(UnspecifiedOmniPlatform):
dist_backend = "gloo"
@classmethod
def get_device_count(cls) -> int:
# CPU tensors share one logical device regardless of process rank.
return 1
@classmethod
def set_device(cls, device: torch.device) -> None:
assert device.type == "cpu"
@classmethod
def synchronize(cls) -> None:
pass
vllm.platforms.current_platform = CpuPlatform()
platform = CpuCollectivePlatform()
for module in (comm, group_coordinator, parallel_state):
module.current_omni_platform = platform
-28
View File
@@ -1,28 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Scheduler test helpers shared by AR / chunk-transfer unit tests."""
from __future__ import annotations
from types import MethodType
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
def bind_omits_transfer_helpers(sched) -> None:
"""Bind the real omit helpers so a MagicMock scheduler does not skip save_async.
MagicMock is truthy, so an unbound ``_request_omits_chunk_transfer_to_next_stage``
silently skips the chunk put.
"""
sched._omits_kv_transfer_cache = {}
sched._omni_final_stage_flags = MethodType(OmniARScheduler._omni_final_stage_flags, sched)
sched._request_omits_kv_transfer_to_next_stage = MethodType(
OmniARScheduler._request_omits_kv_transfer_to_next_stage,
sched,
)
sched._request_omits_chunk_transfer_to_next_stage = MethodType(
OmniARScheduler._request_omits_chunk_transfer_to_next_stage,
sched,
)
+12 -58
View File
@@ -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)
if proc.create_time() == created:
print(f"Found owned vllm process: PID={proc.pid}")
cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
name = proc.name().lower()
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()
try:
omni = getattr(self, "omni", None)
if omni is not None and hasattr(omni, "close"):
omni.close()
finally:
self._cleanup_process()
cleanup_test_environment()
omni = getattr(self, "omni", None)
if omni is not None and hasattr(omni, "close"):
omni.close()
self._cleanup_process()
cleanup_test_environment()
# ---------------------------------------------------------------------------
-81
View File
@@ -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",
[
+1 -117
View File
@@ -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
+6 -88
View File
@@ -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"}))
class AudioItems(AudioProcessorItems):
def get_passthrough_data(self):
return {"embedding": "preserve"}
items = MultiModalDataItems({"audio": AudioItems(["wave"]), "image": ImageProcessorItems([])} if counts else {})
items.select = Mock(wraps=items.select)
items = Mock()
items.get_all_counts.return_value = counts
kwargs = {"sampling_rate": 16000}
if original_text is not None:
kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text
@@ -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)
+1 -32
View File
@@ -8,9 +8,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
import torch
from vllm.exceptions import VLLMValidationError
from vllm.inputs import mm_input, tokens_input
from vllm.multimodal.inputs import PlaceholderRange
from vllm.inputs import tokens_input
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."""
+2 -4
View File
@@ -111,7 +111,7 @@ class TestFailureCounterWiring:
obj.request_states["req-1"] = SimpleNamespace(
metrics=SimpleNamespace(
e2e_done=set(),
log_timing_summary=lambda: None,
build_and_log_summary=lambda: None,
),
failure_recorded=False,
)
@@ -127,7 +127,7 @@ class TestFailureCounterWiring:
obj.request_states["req-1"] = SimpleNamespace(
metrics=SimpleNamespace(
e2e_done=set(),
log_timing_summary=lambda: None,
build_and_log_summary=lambda: None,
),
failure_recorded=False,
)
@@ -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
Loaded 100 of 234 files, more files were not shown because too many files have changed in this diff. Show more