Compare commits

...
Author SHA1 Message Date
hsliuustc0106 fc5bafb2af [Fix] Restore MammothModa2 support
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-27 23:40:02 +08:00
Zhanshenzhuiandamy-why-3459 c6df39e9f7 [Benchmark][MiniCPM-o] Add Omni-DuplexEval support (#6634)
Signed-off-by: Zhanshenzhui <99891002+zyforsure@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 16:31:50 +08:00
Zheng Wengangandamy-why-3459 94b1546ce3 [Model] Qwen3-Omni: register thinker-only pipeline for Instruct serve (#6284)
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 08:21:03 +00:00
19da238246 [Bugfix][NPU] Eager-import mindiesd for every diffusion attention backend (#6054)
Signed-off-by: brandneway <gyuan4892@gmail.com>
Co-authored-by: changetheway <guotaoyuan1@h-partners.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-27 14:41:36 +08:00
Zhou TaichangandClaude Opus 5 651b03225b [Rebase] Rebase to vllm 0.28.0 (#6606)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-27 06:12:55 +00:00
Hongsheng Liuandhsliuustc0106 da4a08b65b [Diffusion] Add FLUX.2-klein BF16 HWR contract (#6651)
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-27 12:51:15 +08:00
SYLARandHongsheng Liu 0e44fc7406 [diffusion][feature] Add offline SVDQuant W4A4 support (#6162)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-27 08:45:53 +08:00
Sy03andsy03 432b2c6ca3 [Model] Add native full-duplex Nemotron VoiceChat serving (#6089)
Signed-off-by: Sysy <1370724210@qq.com>
Signed-off-by: sy03 <137072421@qq.com>
Co-authored-by: sy03 <137072421@qq.com>
2026-08-26 13:19:20 -07:00
c2723979c0 [Bugfix] Validate per-stage device layout before spawning workers (#5003) (#5742)
Signed-off-by: ChoHee15 <cc5281@126.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-26 19:25:22 +00:00
vOv 2d4724eb74 [Bugfix][Diffusion] Fix FLASH_ATTN cross-attention key-padding unpad (#5866)
Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>
2026-08-26 17:20:59 +00:00
Anurag 91dd7fcaec [Bugfix] Reject instructions changes on locked native duplex sessions (#6318)
Signed-off-by: ANURAG KANADE <anuragkanade6@gmail.com>
2026-08-26 16:50:31 +00:00
harley 926e7fbfa2 [Recipe] zai-org/GLM-TTS (2x non-standard RTX 4090 48GB) (#5769)
Signed-off-by: harley <480497125@qq.com>
2026-08-26 22:16:49 +08:00
Yueqian LinandYueqian Lin 8cfede4309 [Perf][Engine] Event-driven orchestration loop (opt-in) — S1 of #4855 (#5221)
Signed-off-by: Yueqian Lin <gao@yueqian.us>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <gao@yueqian.us>
2026-08-26 22:15:25 +08:00
Mu GuanLin d3c990dca8 [Model][Bugfix] Improve LTX audio parity and similarity guards (#6342)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-26 20:59:10 +08:00
psv666andamy-why-3459 5ec576ab4c [Bugfix] Fix missing Realtime audio completion (#6564)
Signed-off-by: psv666 <2693925048@qq.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-26 18:27:25 +08:00
WeiQing Chen 030fccb2bf [Test] Add MiniMax-H3 DLO DP2 T2VA smoke (#6555)
Signed-off-by: david6666666 <530634352@qq.com>
2026-08-26 17:17:10 +08:00
Zhichao Zhang 816335cb46 [Doc] Add MiMo-Audio recipe for RTX 5090/5090D 32GB (#6559)
Signed-off-by: chaosansui <zzc15560846421@163.com>
2026-08-26 16:26:24 +08:00
Mu GuanLinandHongsheng Liu 59cf3add40 [Feature] Support MiniMax-H3 Turbo LoRA with DLO (#6550)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-26 16:15:47 +08:00
Gao Han 7439ed41c2 Revert "[Bugfix] Restore name-based model detection for HF cache snapshot paths" (#6642) 2026-08-26 15:02:47 +08:00
186 changed files with 15142 additions and 709 deletions

No files matched your search

+1 -1
View File
@@ -192,7 +192,7 @@ steps:
grade: Blocking
commands:
- |
timeout 20m bash -c '
timeout 45m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
'
+18
View File
@@ -306,6 +306,24 @@ steps:
"
mirror_hardwares: h100_1
- label: "Diffusion · MiniMax-H3 DLO DP2 T2VA Test"
timeout_in_minutes: 90
source_file_dependencies:
- tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py
- tests/helpers/assertions.py
- tests/helpers/runtime.py
- vllm_omni/diffusion/models/minimax_h3/
- vllm_omni/diffusion/offloader/
- vllm_omni/diffusion/distributed/
- vllm_omni/entrypoints/openai/
commands:
- |
timeout 80m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py -m 'core_model and cuda and diffusion' --run-level 'core_model'
"
mirror_hardwares: h100_2
- label: "Diffusion · Qwen Image Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py
+8 -5
View File
@@ -1,4 +1,4 @@
# Target vllm version: v0.27.0 (releases/v0.27.0 branch).
# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5).
# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
#
# Release-boundary form: there is no precompiled-wheel pin here, because the
@@ -14,11 +14,14 @@
# core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy
# under numba's ceiling. All of it was compensation for a base image and a vLLM
# that disagreed. At a real release boundary they agree, so none of it is
# needed. If you pin an untagged commit again, restore the reinstall block —
# agent/tasks/20_update_wheel_url.sh fails the run rather than letting the
# Dockerfile silently ship the wrong vLLM.
# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels
# rode on the v0.27.0 base image with the full reinstall block until the
# official v0.28.0 image shipped, at which point this file dropped back to the
# form you are reading. If you pin an untagged commit again, restore the
# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather
# than letting the Dockerfile silently ship the wrong vLLM.
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.27.0
ARG VLLM_BASE_TAG=v0.28.0
# We use registry-based layer caching for the dependencies,
# where the tag is determined based on the hash of the contents
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0
ARG BASE_IMAGE=vllm/vllm-openai:v0.28.0
FROM ${BASE_IMAGE}
ARG COMMON_WORKDIR=/app
+2 -2
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
# we should swap over to use stable release ASAP.
# We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+1
View File
@@ -53,6 +53,7 @@ nav:
- BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
- SVDQuant W4A4: user_guide/quantization/svdquant.md
- ModelOpt: user_guide/quantization/modelopt.md
- AutoRound: user_guide/quantization/autoround.md
- msModelSlim: user_guide/quantization/msmodelslim.md
+4
View File
@@ -114,6 +114,10 @@ Configuration classes.
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
@@ -83,6 +83,7 @@ depends on the installed kernels and model path.
| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental |
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
+4 -2
View File
@@ -59,7 +59,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
| `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
| `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. |
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only` / `qwen3_omni_moe_thinker_only`). |
| `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. |
| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
@@ -94,12 +94,14 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
| `max_num_batched_tokens` | int \| null | optional | `null` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
| `devices` | str \| null | optional | `null` | Device list assigned to this stage. |
| `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. |
| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
| `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. |
| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). |
| `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. |
**Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage.
### Connector schema
Each entry under top-level `connectors:` follows this shape:
+9
View File
@@ -322,6 +322,15 @@ Each 1-second window records:
On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima).
`loop_idle` and `loop_active` count loop iterations, so their absolute scale is
tied to which orchestration loop is running. Under the default poll loop an idle
orchestrator records roughly one iteration per millisecond. Under the
event-driven loop (`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, see
[Speech API](../serving/speech_api.md#orchestration-loop-experimental)) an idle
orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still
records one iteration per routed output. Window counts and the `loop_active_pct`
summary are therefore not comparable across the two modes.
### Relationship to other diagnostics
This monitor is intentionally separate from the existing profiling tools:
+11 -11
View File
@@ -14,12 +14,12 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
## Status
The first implementation provides contracts and a CPU local-filesystem store.
The initial diffusion consumer contract additionally defines typed,
The implementation provides contracts and a CPU local-filesystem store. The
diffusion consumer additionally defines typed,
representation-independent final-layout identity/restoration mechanics plus a
concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains
library-only: no loader or DLO path selects, publishes, restores, or transports
that artifact yet.
concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
selects, publishes, restores, and transfers these artifacts.
V1 includes:
@@ -33,13 +33,12 @@ V1 includes:
V1 does not include:
- a public CLI, loader activation, or default-on model integration;
- lease handoff to DLO or another transport consumer;
- FP8, quantized, merged-adaptation, or additional model producers;
- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
- default-on activation or consumers outside no-AllGather DLO;
- online FP8, quantized, merged-adaptation, or additional model producers;
- HWR interaction with DLO AllGather;
- a remote artifact provider or cross-node coordination;
- automatic eviction; or
- a change to DLO AllGather or no-AllGather behavior.
- a change to DLO collective or execution behavior.
## Motivation and use cases
@@ -255,7 +254,8 @@ The contract is intentionally separate from loader activation:
- `FinalLayoutBF16Producer` accepts only the matching identity context and a
finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
coordinate. Its BF16 policy preserves model-declared FP32 parameters and
buffers and revalidates MiniMax H3 mixed-precision invariants.
buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
Other representations reuse source identity, typed parallel identity, tensor
ownership, and exact restoration only when their policy proves those semantics.
@@ -220,9 +220,15 @@ model revision and parallel layout with `required`. TP coordinates have distinct
identities, while equivalent DP replicas share them.
The current producer/consumer boundary is the model-declared final-layout BF16
contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1
and TP2 rank identities plus SP layout identities; online quantization, HSDP,
LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
contract covers both transformer block stacks, constructor-stable packed QKV
mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
buffers. The shared HWR machinery supports ordinary-loader final layouts for
TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
in this change is TP1-only; its TP2/SP layouts rely on the existing
per-coordinate identity mechanics and remain unmeasured. Online quantization,
HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
HWR mode
`disabled`, DLO-disabled, and DLO AllGather configurations stop before source
identity or store construction and retain the existing checkpoint-mmap or
ordinary-loader path.
@@ -66,6 +66,12 @@ the boundary affected by the in-flight stage client/process refactor in
[#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and
responsibilities proposed only by that PR are not current contracts.
The orchestration loop also has an opt-in event-driven mode
(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in
[#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll
cadence only: the routing, ordering, and terminal-state contracts below hold
identically on both loops.
## Ownership boundary
This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,
+1 -1
View File
@@ -1,7 +1,7 @@
# Installation
!!! important
vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0.
vLLM-Omni supports the following hardware platforms:
+1 -1
View File
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.26.0 --torch-backend=auto
uv pip install vllm==0.28.0 --torch-backend=auto
```
#### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
```bash
uv pip install vllm==0.26.0 --torch-backend=auto
uv pip install vllm==0.28.0 --torch-backend=auto
```
The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.26.0
git checkout v0.28.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl
#For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm/vllm-omni:v0.26.0 \
vllm/vllm-omni:v0.28.0 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
```bash
DOCKER_BUILDKIT=1 docker build \
-f docker/Dockerfile.cuda \
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \
-t vllm-omni-cuda .
```
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
```bash
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
```
The pre-built 0.26.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.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.26.0
git checkout v0.28.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.26.0 \
vllm/vllm-omni-rocm:v0.28.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.26.0
vllm/vllm-omni-rocm:v0.28.0
```
# --8<-- [end:pre-built-images]
+3 -3
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.26.0 --torch-backend=auto
uv pip install vllm==0.28.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
!!! note
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
## Offline Inference
+2
View File
@@ -69,6 +69,8 @@ th {
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
+45
View File
@@ -811,6 +811,51 @@ If you encounter OOM errors:
Use `/v1/audio/voices` to list available voices for the loaded model.
## Orchestration Loop (experimental)
Multi-stage omni deployments route stage outputs through a single orchestrator
loop. By default that loop polls every stage replica on a 1 ms cadence. An
opt-in event-driven mode replaces the poll with one reader task per live stage
replica awaiting its client directly, and switches the serving-side
final-output drain to a condition-variable wakeup at the same time.
**Configuration (environment variables):**
| Variable | Default | Description |
|----------|---------|-------------|
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
Set it on the process that runs the orchestrator (stage 0 of an omni
deployment) before starting the server:
```bash
export VLLM_OMNI_EVENT_DRIVEN_ORCH=1
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--omni \
--port 8091
```
The server logs the selected loop mode and its reader/poller counts once at
startup, so you can confirm which loop is live.
Routing, output ordering, and terminal-state behavior are identical on both
loops; only the poll cadence changes. Leaving the variable unset keeps the
legacy poll loop, which is the supported default.
**Known limitations:**
- The measured serving A/B (idle CPU 2.43% to 0.07%; TTFP p99 -32% at
concurrency 8) predates the rebuild on the per-replica fault-isolation work
in [#4285](https://github.com/vllm-project/vllm-omni/pull/4285). That work
changed dead-replica handling and reader/poller lifecycle rather than the
steady-state output path, and the parity suite covers it, but the serving
A/B has not been re-run on the current head.
- The diffusion-poller branch is covered by unit tests only. Deployments whose
stages all run as standard engine cores never exercise it, including GLM-TTS,
which deploys its DiT without `stage_type: diffusion`.
- Concurrency 1 and 32 measured at parity with the legacy loop. At 32 the
latency is admission-bound, which this mode does not address.
## Development
Enable debug logging:
@@ -112,6 +112,12 @@ perform a collective and uses the same rank-local transfer behavior.
HWR is an opt-in startup optimization for models that declare the final-layout
BF16 restore contract. Use it only with no-AllGather DLO:
The validated BF16 model contracts currently cover MiniMax H3 and
`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
class but has not been validated against this contract. FLUX.2-dev, online
FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
the validated scope.
```bash
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
+1
View File
@@ -129,6 +129,7 @@ The following tables show which models support each feature:
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
@@ -28,6 +28,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/deploy_config_file
```
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
```
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
the same single-stage pipeline without `--deploy-config`.
### Launch individual stages (stage-based CLI)
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
+50
View File
@@ -0,0 +1,50 @@
# SVDQuant W4A4
## Overview
[SVDQuant](https://arxiv.org/abs/2411.05007) combines four-bit weights
and activations with a small low-rank branch that corrects part of the
quantization error. vLLM-Omni consumes an offline-quantized checkpoint;
it does not calibrate the model while loading.
## Checkpoint contract
Place the following entry in the diffusion transformer's `config.json`:
```json
{
"quantization_config": {
"quant_method": "svdquant",
"rank": 32,
"precision": "nvfp4",
"act_unsigned": false,
"modules_to_not_convert": []
}
}
```
For a quantized linear with input size `K`, output size `N`, and correction
rank `R`, the checkpoint stores:
| Suffix | Shape | dtype |
| --- | --- | --- |
| `qweight` | `(N, K / 2)` | `int8` (two packed FP4 values per byte) |
| `wscales` | `(K / 16, N)` | `float8_e4m3fn` |
| `proj_down` | `(K, R)` | `bfloat16` |
| `proj_up` | `(N, R)` | `bfloat16` |
| `smooth_factor` | `(K,)` | `bfloat16` |
| `wcscales` | `(N,)` | `bfloat16` |
| `wtscale` | `(1,)` | `bfloat16` |
`K` must be divisible by 16 on every tensor-parallel rank. Set `wcscales` to
ones when no per-output correction is needed. Modules listed in
`modules_to_not_convert` keep their checkpoint precision.
## Runtime support
The compatibility path accepts BF16 inputs and executes an NVFP4 GEMM followed
by the BF16 rank correction. It supports vLLM's FlashInfer, CUTLASS, and FBGEMM
NVFP4 tensor layouts; incompatible forced backends fail during model loading.
SM103 is the currently validated and enabled hardware target. Native fusion of
the NVFP4 GEMM and rank correction is separate from this checkpoint-loading
contract.
+9
View File
@@ -0,0 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""vLLM-Omni example scripts, importable as a package.
Keep this ``__init__.py``: with the upstream vLLM installed editable (which also
places an ``examples`` package on ``sys.path``), a regular package at this repo
root wins over the upstream ``examples`` namespace during ``import examples``.
"""
@@ -0,0 +1,117 @@
"""
Offline inference example: MammothModa2 image summarization (single AR stage).
Example:
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
--model path/to/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
--image /path/to/input.jpg \
--question "Please summarize the content of this image."
"""
from __future__ import annotations
import argparse
import os
from PIL import Image
from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
parser.add_argument(
"--max-tokens",
type=int,
default=512,
help="Max new tokens to generate.",
)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--top-p", type=float, default=0.9)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
return parser.parse_args()
def build_prompt(system: str, question: str) -> str:
return (
f"<|im_start|>system\n{system}<|im_end|>\n"
"<|im_start|>user\n"
"<|vision_start|><|image_pad|><|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
def main() -> None:
args = parse_args()
if not os.path.exists(args.image):
raise FileNotFoundError(f"Image file not found: {args.image}")
pil_image = Image.open(args.image)
image_data = convert_image_mode(pil_image, "RGB")
prompt = build_prompt(args.system, args.question)
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
trust_remote_code=args.trust_remote_code,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
try:
sp = SamplingParams(
temperature=float(args.temperature),
top_p=float(args.top_p),
top_k=-1,
max_tokens=int(args.max_tokens),
seed=int(args.seed),
detokenize=True,
)
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
# Consume it inside the try block so the worker isn't closed early.
outputs = list(
omni.generate(
[
{
"prompt": prompt,
"multi_modal_data": {"image": image_data},
"additional_information": {"omni_task": ["chat"]},
}
],
[sp],
)
)
finally:
omni.close()
lines: list[str] = []
for stage_outputs in outputs:
ro = stage_outputs
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
print("\n".join(lines))
if __name__ == "__main__":
main()
@@ -599,6 +599,17 @@ def main():
if args.seed is not None and hasattr(params, "seed"):
params.seed = args.seed
# MammothModa2's AR stage emits one visual token per grid cell,
# one EOL token per row, and one final look-ahead token whose hidden
# state is unavailable. Size the first stage from the prompt metadata
# instead of SamplingParams' default of 16 tokens.
prompt_info = prompt_dict.get("additional_information", {})
if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
ar_width = int(prompt_info.get("ar_width", [0])[0])
ar_height = int(prompt_info.get("ar_height", [0])[0])
if ar_width > 0 and ar_height > 0:
params.max_tokens = ar_height * (ar_width + 1) + 1
if not diffusion_replaced and len(sampling_params_list) == 1:
sampling_params_list = [diffusion_params]
+31 -3
View File
@@ -10,10 +10,12 @@ Generate text from text or image inputs with vLLM-Omni's shared offline entrypoi
| --- | --- | --- | --- |
| `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` |
| `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` |
| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
| `bytedance-research/MammothModa2-Dev` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
The script recognizes these model families from `config.json`, applies the
model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3.
BAGEL uses its registered default deploy.
model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3
and MammothModa2. BAGEL uses its registered default deploy.
## Text-To-Text
@@ -33,6 +35,22 @@ python examples/offline_inference/x_to_text/x_to_text.py \
--prompt "Explain multimodal inference in three concise sentences."
```
### MammothModa2
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Preview \
--prompt "Explain multimodal inference in three concise sentences."
```
The same command supports the Dev checkpoint:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Dev \
--prompt "Explain multimodal inference in three concise sentences."
```
## Image-To-Text
Pass one image with `--image`. The default question can be replaced with any
@@ -54,6 +72,16 @@ python examples/offline_inference/x_to_text/x_to_text.py \
--prompt "Describe this image in detail."
```
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Preview \
--image image.png \
--prompt "Describe this image in detail."
```
Use `bytedance-research/MammothModa2-Dev` in the same command to run I2T with
the Dev checkpoint.
## Key Arguments
| Argument | Default | Description |
@@ -74,4 +102,4 @@ python examples/offline_inference/x_to_text/x_to_text.py \
The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a
different topology, pass a custom AR-only YAML through `--deploy-config`.
BAGEL uses its repository deploy default unless overridden.
BAGEL and MammothModa2 use their repository deploy defaults unless overridden.
@@ -15,6 +15,7 @@ from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_f
_REPO_ROOT = Path(__file__).resolve().parents[3]
_HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml"
_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml"
def parse_args() -> argparse.Namespace:
@@ -61,7 +62,7 @@ def main() -> None:
omni_kwargs: dict[str, Any] = {
"model": args.model,
"mode": "image-to-text" if image is not None else "text-to-text",
"trust_remote_code": args.trust_remote_code or family == "hunyuan_image3",
"trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"},
"log_stats": args.log_stats,
}
if args.enforce_eager is not None:
@@ -70,6 +71,8 @@ def main() -> None:
omni_kwargs["deploy_config"] = args.deploy_config
elif family == "hunyuan_image3":
omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG)
elif family == "mammoth_moda2":
omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG)
omni = Omni(**omni_kwargs)
try:
@@ -192,6 +192,64 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built
two-response WAV, its `--input-sha256`, and an
`--expect-second-response-substring` value.
## Run Omni-DuplexEval
Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex
endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge
served from the same vLLM-Omni environment. This validates the vLLM-Omni
duplex client and local judge integration; it does not reproduce the paper's
original MiniCPM-o inference implementation.
Start the duplex generation server:
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \
--deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
--served-model-name openbmb/MiniCPM-o-4_5 \
--host 0.0.0.0 --port 8099
```
Start a separate multimodal judge. The allowed path must contain any local
videos passed with `--judge-video-mode video_url`:
```bash
vllm serve Qwen/Qwen2.5-VL-7B-Instruct \
--served-model-name Qwen/Qwen2.5-VL-7B-Instruct \
--allowed-local-media-path /data/omni-duplex-eval \
--host 0.0.0.0 --port 8000
```
Generate, evaluate, and summarize:
```bash
vllm bench omni-duplex-eval --omni generate \
--url ws://127.0.0.1:8099/v1/realtime?duplex=1 \
--model openbmb/MiniCPM-o-4_5 \
--ref-audio /data/ref.wav \
--dataset Hothan/Omni-DuplexEval \
--concurrency 2 \
--response-root /data/omni-duplex-eval/responses
vllm bench omni-duplex-eval --omni evaluate \
--dataset Hothan/Omni-DuplexEval \
--response-root /data/omni-duplex-eval/responses \
--score-root /data/omni-duplex-eval/scores \
--judge-base-url http://127.0.0.1:8000 \
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
--judge-video-mode video_url \
--eval-workers 4
vllm bench omni-duplex-eval --omni summarize \
--score-root /data/omni-duplex-eval/scores
```
The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face
splits, not dataset configurations. Pass one with `--split`, or omit it to
run all nine splits. `--limit 1` is useful for a smoke test. Realtime
generation records `clock=media`; artifacts generated with
`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects
them unless `--allow-invalid-clock` is explicit.
## Related examples
- [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)
@@ -24,6 +24,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/your_deploy_config.yaml
```
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
```
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
the same single-stage pipeline without `--deploy-config`.
For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
use `--stage-overrides` on top of the default config:
+212
View File
@@ -0,0 +1,212 @@
# MammothModa2
> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation
## Summary
- Vendor: ByteDance Research
- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
- Mode: Offline inference
- Maintainer: Community
## When to use this recipe
Use this recipe to run MammothModa2-Preview through shared task-oriented
offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script.
The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
and forwards MammothModa2-specific generation parameters through the
pipeline-declared `extra_body` contract.
MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs
interface (not `OmniDiffusionRequest`), so its generation knobs
(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via
`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale`
flags. Image size uses the standard `--height` / `--width` flags.
## References
- Upstream model:
[`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
- Dev model:
[`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
- Related offline example:
[`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
- Related T2T/I2T example:
[`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
- Declared parameters:
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
- Deploy config:
[`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
## Hardware Support
The default deploy config runs both the AR and DiT stages on a single GPU
(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
after rebalancing the split so the AR weights (~23 GB) leave room for the KV
cache — see the note under *1x L40S 48GB*.
## GPU
### 1x L40S 48GB
> **48 GB config adjustment:** the committed
> `vllm_omni/deploy/mammoth_moda2.yaml` uses
> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
> before running. (On an ~80 GB GPU, leave the defaults unchanged.)
#### Environment
- OS: Linux
- Python: Match the repository requirements for your checkout
- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an
~80 GB GPU for the default config
- vLLM version: Match the repository requirements for your checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Offline Commands
Download the model:
```bash
hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
```
Run text-to-image with the shared offline example from the repository root. The
deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
MammothModa2 generation parameters as a JSON object through `--extra-body`:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model ./MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
--output mammoth_t2i.png
```
The `--extra-body` JSON forwards MammothModa2-specific parameters into
`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's
declared `extra_body_params` (see
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
so unknown keys for MammothModa2 are silently dropped:
- `text_guidance_scale` — classifier-free guidance scale for the DiT stage
(default `9.0`; CFG is active only when `> 1.0`).
- `cfg_range` — relative step range `[start, end]` over which CFG is applied
(default `[0.0, 1.0]`).
- `num_inference_steps` — number of DiT denoising steps (default `50`).
`--height` and `--width` must be multiples of 16.
Run text-to-text through the shared understanding example. It recognizes the
MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Preview \
--prompt "Explain multimodal generation in three sentences."
```
Add an image for image-to-text or image summarization. The shared example
uses MammothModa2's chat and vision-token template:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Preview \
--image /path/to/input.jpg \
--prompt "Please summarize the content of this image."
```
#### Verification
The example writes the generated image to the `--output` path. Confirm the file
exists and is a valid image:
```bash
ls -lh mammoth_t2i.png
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
```
### 1x AMD MI300X, MammothModa2 Preview
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
#### Offline Commands
The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
```bash
python3 examples/offline_inference/text_to_image/text_to_image.py \
--model bytedance-research/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--seed 42 \
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
--enable-diffusion-pipeline-profiler \
--log-stats \
--output mammoth_t2i.png
```
#### Verification
The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
The output was a valid 1024 by 1024 RGB PNG.
## MammothModa2-Dev unified inference
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
required.
Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
then sends the generated visual tokens and hidden states to the DiT stage.
Download the checkpoint:
```bash
hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
```
Run text-to-text through the shared understanding example. It recognizes the
Dev checkpoint as MammothModa2 and automatically selects
`mammoth_moda2_ar.yaml`:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Dev \
--prompt "Explain multimodal generation in three sentences."
```
Add an image for image-to-text or image summarization:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Dev \
--image ./image.png \
--prompt "Please summarize the content of this image."
```
The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
and encoder caches. Allow additional GPU memory for those caches and the input
image.
+5 -4
View File
@@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
```
Start from a non-offloaded FL2VA server command and add
Start from a non-offloaded or DLO FL2VA server command and add
`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
`--lora-path` preloads the adapter; each request still activates it and uses
the published sampling settings:
@@ -853,9 +853,10 @@ the published sampling settings:
For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
dynamic-only and does not support prefusion, DLO, or LoRA composition.
It also rejects model-level, layerwise, and distributed layerwise offload;
the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
dynamic-only and does not support prefusion or LoRA composition. DLO is
supported by keeping the request-switchable LoRA A/B buffers resident on the
accelerator while DLO streams only the base blocks; budget for this additional
fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
The five requested sigma points produce the four denoiser evaluations expected
by the Turbo artifact.
+70 -7
View File
@@ -1,7 +1,7 @@
# NemotronLabs VoiceChat 11B
> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
> Speech-to-speech on a 3-stage vLLM-Omni pipeline: offline single-turn
> inference and experimental model-native, frame-locked Realtime serving.
## Summary
@@ -11,8 +11,8 @@
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
function calling, and batch>1 are documented follow-ups.
- Mode: Offline single-turn inference (batch=1), plus experimental native
duplex Realtime serving at one 80 ms microphone frame per scheduler wake.
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
## When to use this recipe
@@ -33,6 +33,11 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
- Staged pipeline config:
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
- Native duplex config and E2E probe:
[`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml),
[`tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py`](../../tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py)
- Nightly model-level gate:
[`tests/e2e/online_serving/test_nemotron_voicechat_duplex.py`](../../tests/e2e/online_serving/test_nemotron_voicechat_duplex.py)
- Upstream model card:
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
@@ -60,7 +65,7 @@ output stayed within one word of the reference in testing.
- OS: Linux
- Python: 3.12
- vLLM version: 0.26.0
- vLLM version: 0.27.0
- vLLM-Omni version or commit: this PR / current `main`
#### Command
@@ -88,6 +93,55 @@ The reply text should read as a coherent spoken-style answer to the question in
the input WAV, and the WAV should transcribe to (approximately) the same text
with any ASR model.
#### Native duplex serving
The checkpoint does not contain the underlying Nemotron text tokenizer. Set
`NEMOTRON_VOICECHAT_LLM_PATH` to a local snapshot of
`nvidia/NVIDIA-Nemotron-Nano-9B-v2` before starting an offline deployment.
```bash
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export NEMOTRON_VOICECHAT_LLM_PATH=/path/to/NVIDIA-Nemotron-Nano-9B-v2
vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--omni \
--served-model-name nemotron-voicechat \
--deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
```
The single-GPU profile runs three eager fp32 engine processes on one device.
In another shell, stream only the user channel of one of NVIDIA's bundled
stereo conversation fixtures:
```bash
python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
--model nemotron-voicechat \
--input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/turn_taking.wav \
--input-channel 0 \
--output-dir results/nemotron_voicechat_duplex \
--timeout-s 600
```
The probe requires a completed response, non-silent 22.05 kHz output, and the
advertised native 80 ms append capabilities. Add `--expect-function-call` and
use `tool_call.wav` to validate the function-call channel. To validate the
Realtime tool round trip, also provide the expected arguments and return a
tool result:
```bash
python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
--model nemotron-voicechat \
--input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/tool_call.wav \
--output-dir results/nemotron_voicechat_tool_call \
--expect-function-call \
--expected-function-name generate_random_number \
--expected-function-arguments '{"min":1,"max":50}' \
--function-output 20 \
--expected-post-tool-text "random number"
```
#### Notes
- Memory usage: the shipped yaml runs all three stages on one GPU
@@ -109,5 +163,14 @@ with any ASR model.
instead of pausing.
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
token, and the stage context is 16384).
- Known limitations: batch=1 only; offline single turn (no online/duplex
serving yet); the function-call channel is decoded but not yet acted on.
- The 80 ms frame period is a model/protocol contract, not a throughput
guarantee. In an eager fp32 H200 profile, Stage 2 emitted 80 ms packets at
roughly 0.87--0.88 s intervals (RTF about 10.8--10.9); placing the three
stages on separate GPUs did not remove the serial 8-step EAR-TTS bottleneck.
The current native duplex path is functionally streaming but not wall-clock
realtime.
- Known limitations: batch=1 only. The native duplex deployment allows one
active session and does not support barge-in. Tool execution remains
client-owned: the server emits function-call events, accepts a validated
`function_call_output`, and resumes the live model with the returned result;
it does not execute arbitrary tools itself.
+2 -1
View File
@@ -85,8 +85,9 @@ recipes/
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB / 2x non-standard RTX 4090 48GB (49,140 MiB each) |
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
| [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |
+161
View File
@@ -0,0 +1,161 @@
# MiMo-Audio for omni audio generation and understanding
> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue
## Summary
- Vendor: Xiaomi MiMo
- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct`
- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue
- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a known-good starting point for running
MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D).
The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids
FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce
Stage-1 GPU pressure.
## References
- Upstream model card: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Instruct>
- Audio tokenizer: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer>
- Offline example:
[`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py)
- Online example:
[`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md)
- Default deploy config (generic single GPU):
[`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml)
- 5090D-tuned deploy config:
[`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
- Pipeline / deploy docs:
[Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/)
- Related issue or discussion:
[RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
- Tokenizer device alignment fix (required for CPU tokenizer on stage 1):
[PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539)
## Hardware Support
## GPU
### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB
Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested
offline with `tts_sft`.
#### Environment
- OS: Ubuntu (Linux x86_64)
- Python: 3.12
- PyTorch: 2.11.0+cu129
- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x)
- GPU: NVIDIA GeForce RTX 5090 D, 32 GB
- vLLM: 0.26.0+cu129
- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support)
Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack.
On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using
PyTorch 2.11+cu129.
Download or cache:
- `XiaomiMiMo/MiMo-Audio-7B-Instruct`
- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo)
Export the runtime knobs below before offline or online runs:
```bash
export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}"
export MIMO_AUDIO_TOKENIZER_DEVICE=cpu
export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0
export VLLM_USE_FLASHINFER_SAMPLER=0
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
```
Optional (mirror / download stability):
```bash
export HF_HUB_DISABLE_XET=1
export HF_ENDPOINT=https://hf-mirror.com # if you use a Hugging Face mirror
```
Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure
`attention_backend: TRITON_ATTN` in the deploy YAML instead.
#### Command
Offline TTS validation from the repository root:
```bash
cd examples/offline_inference/mimo_audio
python3 -u end2end.py \
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft \
--text "The weather is so nice today."
```
Online serving (requires MiMo chat template):
```bash
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
--served-model-name "MiMo-Audio-7B-Instruct" \
--port 18091 \
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
```
Pass the same environment exports to the server process when using a CPU
tokenizer or disabling FlashInfer sampling.
#### Verification
Offline:
```bash
ls -lh output_audio/tts_sft/*.wav
```
Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt`
with the decoded text. A successful run completes both stages without a
Stage-1 `decode_vq` device mismatch.
Online:
```bash
curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health
```
Then use the client under
[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/).
#### Notes
- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB
with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1.
- Key flags:
- `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory
and `TRITON_ATTN`.
- `MIMO_AUDIO_TOKENIZER_PATH` is mandatory.
- `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires
[#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent
fix) so code2wav decode uses the tokenizer device.
- `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler
issues on consumer GPUs.
- Known limitations:
- Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional
tooling; the tested path uses TRITON attention and cu129 wheels.
- Without a working `flash-attn` build, audio quality may be slightly
metallic; install FlashAttention when a compatible wheel is available.
- First request triggers Triton JIT for rotary / attention kernels (latency
spike); subsequent requests are faster.
- For longer outputs or batching, you may need to raise `max_model_len` in
the deploy YAML and match `max_position_embeddings` in the model config
(see the offline example README).
+135 -1
View File
@@ -1,4 +1,4 @@
# GLM-TTS for Chinese/English TTS on 1x GPU
# GLM-TTS for Chinese/English TTS
## Summary
@@ -103,3 +103,137 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
- Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback.
- Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side.
- Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40.
### 2x non-standard RTX 4090 48GB
This configuration was personally validated on two non-standard RTX 4090 cards
reporting 49,140 MiB each in `nvidia-smi`. Stock 24GB RTX 4090 cards were not
validated.
#### Environment
- OS: Ubuntu 22.04.1 LTS
- NVIDIA driver: 580.76.05
- Python: 3.12.13
- PyTorch: 2.11.0+cu130
- CUDA runtime and JIT toolkit: 13.0
- vLLM: 0.26.0
- vLLM-Omni: commit `66a9b0c8`, version
`0.26.0rc2.dev5+g66a9b0c84`
- FlashInfer: 0.6.14
The host's default `nvcc` was 11.8. The working CUDA toolkit came from the
Python environment, so `CUDA_HOME` was set to
`$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13`.
#### Command
From the repository root, start the async-chunk server with:
```bash
export CUDA_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13"
export PATH="$CUDA_HOME/bin:$PATH"
export LD_LIBRARY_PATH="$CUDA_HOME/lib:${LD_LIBRARY_PATH:-}"
CUDA_VISIBLE_DEVICES=0,1 vllm-omni serve zai-org/GLM-TTS \
--omni \
--trust-remote-code \
--port 8091 \
--deploy-config vllm_omni/deploy/glm_tts.yaml \
--stage-overrides '{"1":{"devices":"1"}}' \
--allowed-local-media-path "${PWD}/tests/assets"
```
Stage 0 runs on GPU 0, while the override places stage 1 on GPU 1. For the sync
variant, add `--no-async-chunk` to the same command.
#### Verification
Run the Chinese voice-cloning check:
```bash
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
--model zai-org/GLM-TTS \
--text "今天天气很好,我们一起学习语音模型推理。" \
--ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
--ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
--response-format wav \
--output glm_tts_zh.wav
```
Run the English voice-cloning check with the same Chinese reference:
```bash
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
--model zai-org/GLM-TTS \
--text "Hello, this is a voice cloning test with GLM-TTS." \
--ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
--ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
--response-format wav \
--output glm_tts_en.wav
```
Inspect either output:
```bash
ffprobe -v error -select_streams a:0 \
-show_entries stream=codec_name,sample_rate,channels \
-of default=noprint_wrappers=1 glm_tts_zh.wav
```
Expected output:
```text
codec_name=pcm_s16le
sample_rate=24000
channels=1
```
The English smoke test produced a valid WAV with an actual duration of 5.020
seconds.
#### Observed warm-cache results
These observations are from one controlled run, not a hardware or performance
guarantee. Cold-start and JIT costs were excluded. Values are mean +/- standard
deviation.
##### Measurement procedure
- For the single-GPU server, use the documented command with
`CUDA_VISIBLE_DEVICES=0` and omit `--stage-overrides`; both bundled stages
remain on GPU 0. For the dual-GPU server, use the command as shown.
- Async chunking is the default. Add `--no-async-chunk` for sync mode.
- For non-streaming measurements, use the exact Chinese client command above.
For streaming, add `--stream` to that command and change the output filename
to `glm_tts_zh.pcm`; streaming output is raw PCM, not WAV.
- The fixed Chinese prompt and reference were exactly those shown in
Verification and yielded 3.900 seconds of audio. Each case used one warm-up
followed by three measured requests at concurrency 1.
- Wall time starts immediately before the HTTP POST and stops after the complete
response body. TTFA runs from immediately before the POST to the first
non-empty audio bytes. RTF is wall time divided by audio duration.
- Peak total device memory used was sampled every 0.2 seconds and reported by
NVML via `nvmlDeviceGetMemoryInfo.used`. The readings include background and
baseline use and are not baseline-subtracted.
| GPUs | Mode | Response | Wall time (s) | TTFA (s) | RTF |
|---:|---|---|---:|---:|---:|
| 1 | Async | Non-streaming | 2.144 +/- 0.076 | -- | 0.550 +/- 0.019 |
| 2 | Async | Non-streaming | 1.969 +/- 0.130 | -- | 0.505 +/- 0.033 |
| 1 | Async | Streaming | 2.135 +/- 0.093 | 1.245 +/- 0.178 | 0.548 +/- 0.024 |
| 2 | Async | Streaming | 1.891 +/- 0.006 | 1.170 +/- 0.042 | 0.485 +/- 0.002 |
| 1 | Sync | Non-streaming | 1.822 +/- 0.016 | -- | 0.467 +/- 0.004 |
| 2 | Sync | Non-streaming | 1.792 +/- 0.023 | -- | 0.459 +/- 0.006 |
Moving stage 1 to GPU 1 reduced async non-streaming wall time by 8.2%. It
reduced async streaming wall time by 11.4% and TTFA by 6.0%, and reduced sync
wall time by 1.6%. Single- and dual-GPU outputs were bitwise identical only
within matching execution and response modes.
| Configuration | GPU 0 | GPU 1 |
|---|---:|---:|
| Single-GPU async | 31.52 GiB | -- |
| Dual-GPU async | 29.23 GiB | 2.91 GiB |
| Single-GPU sync | 31.47 GiB | -- |
| Dual-GPU sync | 29.23 GiB | 2.91 GiB |
@@ -0,0 +1,143 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import hashlib
import json
import sys
from types import SimpleNamespace
import pytest
from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
extract_timed_sentences,
normalize_response_items,
split_text,
validate_clock,
)
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
DuplexSample,
canonical_task_type,
family_for_split,
load_samples,
task_type_for_split,
)
from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
PROTOCOL_PIN,
build_content_prompt,
build_reminder_prompt,
build_temporal_prompt,
parse_judge_json,
reminder_window,
summarize_pr_results,
temporal_window,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_dataset_mapping_and_manifest(tmp_path):
manifest = tmp_path / "samples.json"
manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
sample = load_samples(manifest, media_root=tmp_path)[0]
assert sample.family == "rtd"
assert sample.video == str(tmp_path / "clip.mp4")
assert family_for_split("PR_event_reminder") == "pr"
assert canonical_task_type("post-event-reminder") == "post_event_reminder"
assert task_type_for_split("PR_correction") == "correction"
assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
def test_hugging_face_subset_names_are_splits(monkeypatch):
calls = []
def fake_load_dataset(name, **kwargs):
calls.append((name, kwargs))
return [{"id": "pr", "question_text": "Correct this."}]
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
assert sample.split == "PR_correction"
assert sample.task_type == "correction"
def test_response_aliases_and_clock_guard():
assert split_text("One. Two!") == ["One.", "Two!"]
assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
{"sentence": "x", "start": 800.0, "end": 800.0}
]
with pytest.raises(ValueError, match="clock=invalid"):
validate_clock({"clock": "invalid"})
validate_clock({"clock": "invalid"}, allow_invalid=True)
timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
with pytest.raises(ValueError, match="only clock=media"):
extract_timed_sentences([], clock="wall")
def test_protocol_windows_and_parsing():
assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
assert temporal_window(4, 5, 10) == (2.0, 3.0)
assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
assert reminder_window(5) == (5.0, 15.0)
assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
"content_score": 2.35,
"is_relevant": 1,
}
assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
def test_protocol_prompt_wording_is_pinned():
prompts = [
build_temporal_prompt(1, 2, "response", "question"),
build_content_prompt("response", "question", ["one", "two"]),
build_reminder_prompt("instruction", "response", "proactive_reminder"),
build_reminder_prompt("instruction", "response", "post_event_reminder"),
build_reminder_prompt("instruction", "response", "correction", "answer"),
]
assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
"5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
"1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
"d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
"08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
"3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
]
def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
response = tmp_path / "response.json"
response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
video = tmp_path / "video.mp4"
video.write_bytes(b"video")
monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
class Judge:
def temporal(self, *args, **kwargs):
return '{"temporal_score": 3, "is_relevant": 1}'
def content(self, prompt, video, frames=None, *, mode="video_url"):
assert mode == "frame-sample"
assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
return '{"content_score": 3}'
sample = DuplexSample(
id="rtd",
split="RTD_OCR",
family="rtd",
task_type=None,
video=video,
video_duration=4,
)
score = eval_module.evaluate_sample(
sample,
response,
tmp_path / "score.json",
Judge(),
judge_video_mode="frame-sample",
)
assert score["content"]["frame_count"] == 2
+28 -49
View File
@@ -15,11 +15,7 @@ from transformers import PretrainedConfig, Qwen3OmniMoeConfig
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
from vllm_omni.config import config_factory as config_factory_module
from vllm_omni.config.config_factory import (
StageConfigFactory,
_materialize_object_storage_configs,
_name_match_candidate,
)
from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
@@ -538,6 +534,7 @@ class TestPipelineDiscovery:
"""Check that specific models are in OMNI_PIPELINES."""
assert "qwen2_5_omni" in OMNI_PIPELINES
assert "qwen3_omni_moe" in OMNI_PIPELINES
assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
assert "qwen3_tts" in OMNI_PIPELINES
def test_registry_resolver_qwen3_omni_all_stages(self):
@@ -562,6 +559,18 @@ class TestPipelineDiscovery:
assert pipeline.default_deploy_config_name is None
assert len(pipeline.stages) == 1 # thinker only
def test_registry_qwen3_omni_thinker_only_key_is_static(self):
"""Explicit registry key must not re-enter the HF-config resolver."""
pipeline = resolve_pipeline_config(
"qwen3_omni_moe_thinker_only",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
assert len(pipeline.stages) == 1
assert pipeline.stages[0].engine_output_type == "text"
assert pipeline.stages[0].final_output_type == "text"
@pytest.mark.parametrize(
"pipeline",
[registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
@@ -1229,14 +1238,22 @@ stages:
assert deploy.connectors is not None
assert deploy.platforms is not None
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
@pytest.mark.parametrize(
("hf_config", "model"),
[
(Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
(Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
],
ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
)
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
with patch(
"vllm_omni.config.config_factory.get_config",
return_value=Q3_OMNI_THINKER_HF_CONFIG,
return_value=hf_config,
):
pipeline = StageConfigFactory.get_pipeline_config(
"Qwen/Qwen3-Omni-30B-A3B-Thinking",
model,
trust_remote_code=False,
deploy_config_path=str(deploy_path),
)
@@ -1245,6 +1262,7 @@ stages:
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
deploy = load_deploy_config(deploy_path)
assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 1
@@ -1296,6 +1314,8 @@ stages:
@pytest.mark.parametrize(
("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
[
("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
("step_audio_2.yaml", "step_audio_2", 2, "audio", None),
@@ -2936,44 +2956,3 @@ class TestObjectStorageConfigResolution:
matching = "s3://any-bucket/my-cosyvoice3-model"
assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
class TestNameMatchCandidateSnapshotPaths:
"""Name-based matching must survive resolved HF cache snapshot paths."""
_SNAPSHOT_REV = "29e01c4e8d000f4bcd70751be16fa94bf3d85a18"
def test_candidate_recovers_repo_name_from_hf_cache_layout(self):
path = f"/data/hub/models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512/snapshots/{self._SNAPSHOT_REV}"
assert _name_match_candidate(path) == "Fun-CosyVoice3-0.5B-2512"
assert _name_match_candidate(path + "/") == "Fun-CosyVoice3-0.5B-2512"
def test_candidate_keeps_basename_for_non_cache_paths(self):
assert _name_match_candidate("org/My-Model") == "My-Model"
assert _name_match_candidate("/data/snapshots-of/My-Model") == "My-Model"
# A ``snapshots`` parent without the ``models--<org>--<name>`` encoding
# falls back to the basename, same as before.
assert _name_match_candidate(f"/data/notcache/snapshots/{self._SNAPSHOT_REV}") == self._SNAPSHOT_REV
def test_candidate_recovers_bare_repo_name_from_hf_cache_layout(self):
# Legacy un-namespaced repos (e.g. ``gpt2``) cache as ``models--<name>``.
path = f"/data/hub/models--my-cosyvoice3-model/snapshots/{self._SNAPSHOT_REV}"
assert _name_match_candidate(path) == "my-cosyvoice3-model"
def test_candidate_ignores_malformed_cache_repo_dirs(self):
# Hub validation rejects ``--`` inside repo ids, so more than three
# ``--``-separated segments cannot be a real cache dir; keep basename.
assert _name_match_candidate("/hub/models--a--b--c/snapshots/abc123") == "abc123"
assert _name_match_candidate("/hub/models--/snapshots/abc123") == "abc123"
def test_try_infer_model_type_matches_cosyvoice3_snapshot_dir(self, tmp_path):
"""Regression for the HF-cache path shape: CosyVoice3 ships an empty
config.json and relies on name matching, but a resolved snapshot dir
basename is a bare revision hash, so basename-only matching (#5036)
misclassified it as a diffusion pipeline and stage init crashed with
"Diffusers pipeline index not found"."""
snapshot = tmp_path / "hub" / "models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512" / "snapshots" / self._SNAPSHOT_REV
snapshot.mkdir(parents=True)
(snapshot / "config.json").write_text("{}")
assert StageConfigFactory.try_infer_model_type(model=str(snapshot), trust_remote_code=False) == "cosyvoice3"
+1 -1
View File
@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
"""Make an inventory expansion an explicit review decision."""
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
assert category_counts == {
EnvironmentVariableCategory.PUBLIC_OMNI: 23,
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
@@ -100,6 +100,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
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
return sched
@@ -93,6 +93,7 @@ def _run_resumable_segment_stop(
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
@@ -198,6 +199,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
inter_stage_output,
session,
True,
new_token_ids=[42],
confirmed_num_computed_tokens=26,
)
@@ -226,6 +228,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
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
@@ -427,3 +430,33 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
"""A WAITING_FOR_CHUNK stop must not delete its newly parked session."""
session = _make_request()
session.resumable = True
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
def queue(*requests):
result = MagicMock()
result.requests = set(requests)
result.add_request.side_effect = result.requests.add
result.remove_requests.side_effect = result.requests.difference_update
return result
sched = OmniARScheduler.__new__(OmniARScheduler)
sched.running = []
sched.waiting = queue()
sched.skipped_waiting = queue(session)
sched.num_waiting_for_streaming_input = 1
sched.chunk_transfer_adapter = SimpleNamespace(segment_finished_requests={session.request_id})
sched._resume_downstream_chunk_receiver(session)
sched._remove_stopped_requests_from_queues(set(), {session})
assert session.status == RequestStatus.WAITING
assert session in sched.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
@@ -138,6 +138,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
sched.recompute_kv_load_failures = False
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
@@ -0,0 +1,168 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review).
Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every
non-empty scheduled step and, at the top of ``update_from_output()``,
advances ``processed_step_seq`` and drains the deferred frees. The Omni
schedulers each carried only half of that contract:
- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence
advances) but its reimplemented ``update_from_output()`` never advanced
``processed_step_seq`` nor drained -> deferred blocks accumulate without
bound under async/PP KV-consumer configurations.
- ``OmniGenerationScheduler`` reimplements both: the fast-path
``schedule()`` never advanced ``sched_step_seq`` (so finished requests
freed blocks an in-flight GPU step could still be writing) and its
``update_from_output()`` never drained.
These tests pin the restored contract at the call boundary; the drain and
free logic itself is upstream-tested.
"""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
# before Request / RequestStatus are bound in this module.
# isort: off
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.interface import PauseState
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.core.sched.request_queue import SchedulingPolicy
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
# isort: on
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_update_sched(scheduler_cls) -> MagicMock:
"""MagicMock scheduler driving the real ``update_from_output`` with an
empty batch; only the fence state carries real values."""
sched = MagicMock()
sched.requests = {}
sched.perf_metrics = None
sched.chunk_transfer_adapter = None
sched.defer_block_free = True
sched.processed_step_seq = 0
return sched
def _empty_frame(*, total_num_scheduled_tokens: int):
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {}
scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
model_runner_output = MagicMock(spec=ModelRunnerOutput)
model_runner_output.sampled_token_ids = []
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
return scheduler_output, model_runner_output
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None:
sched = _make_update_sched(scheduler_cls)
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
assert sched.processed_step_seq == 1
sched._drain_deferred_frees.assert_called_once_with()
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None:
"""0-token steps do not advance the fence upstream (their seq was never
advanced in schedule()); draining on them would run ahead of the GPU."""
sched = _make_update_sched(scheduler_cls)
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0)
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
assert sched.processed_step_seq == 0
sched._drain_deferred_frees.assert_not_called()
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None:
sched = _make_update_sched(scheduler_cls)
sched.defer_block_free = False
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
assert sched.processed_step_seq == 0
sched._drain_deferred_frees.assert_not_called()
def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]:
"""MagicMock scheduler that drives the real generation fast-path
``schedule()`` over one running one-shot request."""
request = Request(
request_id="req-fence",
prompt_token_ids=[1, 2, 3],
sampling_params=SamplingParams(max_tokens=4),
pooling_params=None,
arrival_time=100.0,
block_hasher=None,
)
request.status = RequestStatus.RUNNING
sched = MagicMock()
sched.max_num_scheduled_tokens = 8
sched._pause_state = PauseState.UNPAUSED
sched.requests = {request.request_id: request}
sched.running = [request]
sched.waiting = [] # falsy: the waiting loop never runs
sched.policy = SchedulingPolicy.FCFS
sched.chunk_transfer_adapter = None
sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
sched.num_lookahead_tokens = 0
sched.log_stats = False
sched.kv_cache_config = MagicMock(kv_cache_groups=[])
sched.use_v2_model_runner = False
sched.prev_step_scheduled_req_ids = set()
sched.needs_kv_cache_zeroing = False
sched.connector = None
sched.ec_connector = None
sched.defer_block_free = defer_block_free
sched.sched_step_seq = 0
# last_sched_seq stamping order: _update_after_schedule must observe the
# already-advanced fence value, as upstream schedule() guarantees.
sched.seq_at_update_after_schedule = None
def _capture(_scheduler_output):
sched.seq_at_update_after_schedule = sched.sched_step_seq
sched._update_after_schedule.side_effect = _capture
return sched, request
def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None:
sched, request = _make_fast_path_sched(defer_block_free=True)
OmniGenerationScheduler.schedule(sched)
assert sched.sched_step_seq == 1
assert sched.seq_at_update_after_schedule == 1
def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
sched, request = _make_fast_path_sched(defer_block_free=False)
OmniGenerationScheduler.schedule(sched)
assert sched.sched_step_seq == 0
@@ -0,0 +1,180 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review).
Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()``
BEFORE the encoder cache is freed (so the connector can inspect per-request
state), honors its delayed-free result, and returns the generated
``ec_transfer_params``; ``update_from_output()`` then carries them into the
``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and
always returned ``None``; both Omni schedulers discarded the second tuple
item, so a remote encoder-cache handle could never reach the frontend.
"""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
# before Request / RequestStatus are bound in this module.
# isort: off
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
# isort: on
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"}
def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]:
"""Minimal real-`_free_request` harness; `order` records the EC hook
firing relative to the encoder-cache free."""
order = MagicMock()
sched = OmniARScheduler.__new__(OmniARScheduler)
sched._omits_kv_transfer_cache = {}
sched._connector_finished = lambda request: (False, None)
sched.ec_connector = ec_connector
if ec_connector is not None:
order.attach_mock(ec_connector.request_finished, "request_finished")
sched.encoder_cache_manager = MagicMock()
order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free")
sched.finished_req_ids = set()
sched._new_prompt_len_snapshot = {}
sched.finished_req_ids_dict = None
sched._should_transfer_kv_for_request = lambda req_id: False
sched._free_blocks = MagicMock()
sched._free_input_coordinator_request = MagicMock()
sched.chunk_transfer_adapter = None
return sched, order
class _FakeFinishedRequest:
def __init__(self, request_id: str) -> None:
self.request_id = request_id
def is_finished(self) -> bool:
return True
def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None:
ec_connector = MagicMock()
ec_connector.request_finished.return_value = (False, EC_PARAMS)
sched, order = _make_free_request_sched(ec_connector=ec_connector)
request = _FakeFinishedRequest("req-ec")
kv_params, ec_params = sched._free_request(request)
assert ec_params == EC_PARAMS
ec_connector.request_finished.assert_called_once_with(request)
hook_calls = [name for name, *_ in order.mock_calls]
assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free")
sched._free_blocks.assert_called_once_with(request)
def test_free_request_honors_ec_delayed_free() -> None:
ec_connector = MagicMock()
ec_connector.request_finished.return_value = (True, EC_PARAMS)
sched, _ = _make_free_request_sched(ec_connector=ec_connector)
request = _FakeFinishedRequest("req-ec-delay")
kv_params, ec_params = sched._free_request(request)
assert ec_params == EC_PARAMS
sched._free_blocks.assert_not_called()
def test_free_request_without_ec_connector_returns_none() -> None:
sched, _ = _make_free_request_sched(ec_connector=None)
request = _FakeFinishedRequest("req-no-ec")
kv_params, ec_params = sched._free_request(request)
assert ec_params is None
sched._free_blocks.assert_called_once_with(request)
def _make_finishing_request() -> Request:
request = Request(
request_id="req-ec-finish",
prompt_token_ids=[1, 2, 3],
sampling_params=SamplingParams(max_tokens=4),
pooling_params=None,
arrival_time=100.0,
block_hasher=None,
)
request.status = RequestStatus.RUNNING
return request
def _make_finish_sched(request: Request) -> MagicMock:
"""MagicMock scheduler driving the real ``update_from_output`` through a
finishing request; ``_free_request`` reports EC transfer params."""
sched = 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._process_kv_transfer_trigger.return_value = False
sched._maybe_decode_pooling_output.return_value = None
sched._handle_stopped_request.return_value = True
sched._free_request.return_value = (None, EC_PARAMS)
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.chunk_transfer_adapter = None
sched.finished_req_ids_dict = None
sched._new_prompt_len_snapshot = {}
return sched
def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None:
request = _make_finishing_request()
sched = _make_finish_sched(request)
sched._update_request_with_output.return_value = ([42], True)
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.ec_transfer_params == EC_PARAMS
sched._free_request.assert_called_once_with(request)
def _finish_frame_parts(request: Request):
scheduler_output = MagicMock(spec=SchedulerOutput)
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 = MagicMock(spec=ModelRunnerOutput)
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
return scheduler_output, model_runner_output
def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None:
request = _make_finishing_request()
# One-shot generation finish: the whole prompt is computed.
request.num_computed_tokens = len(request.prompt_token_ids)
sched = _make_finish_sched(request)
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.ec_transfer_params == EC_PARAMS
sched._free_request.assert_called_once_with(request)
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review).
Upstream v0.28 finalizes prefill statistics with
``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))``
before emitting them, which computes ``num_cache_creation_tokens`` (tokens
computed and written to the prefix cache). Both Omni schedulers previously
emitted ``take_prefill_stats()`` raw, so creation usage was always 0.
"""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
# before Request / RequestStatus are bound in this module.
# isort: off
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
# isort: on
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_request() -> Request:
request = Request(
request_id="req-prefill-stats",
prompt_token_ids=[1, 2, 3],
sampling_params=SamplingParams(max_tokens=4),
pooling_params=None,
arrival_time=100.0,
block_hasher=None,
)
request.status = RequestStatus.RUNNING
# A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache.
request.prefill_stats.set(
num_prompt_tokens=3,
num_local_cached_tokens=1,
num_external_cached_tokens=0,
)
return request
def _make_sched(request: Request) -> MagicMock:
sched = 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._process_kv_transfer_trigger.return_value = False
sched._maybe_decode_pooling_output.return_value = None
sched._handle_stopped_request.return_value = True
# The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created.
sched.kv_cache_manager.estimate_cached_tokens.return_value = 3
def _free_and_invalidate(req):
# Lifecycle guard: freeing releases the KV blocks, after which the
# manager reports 0 cached tokens. If finalize ran after
# _free_request (the reviewed bug), creation count would compute
# against 0 and the assertions below would see 0, not 2.
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
return (None, None)
sched._free_request.side_effect = _free_and_invalidate
sched.chunk_transfer_adapter = None
sched.finished_req_ids_dict = None
sched._new_prompt_len_snapshot = {}
return sched
def _frame(request: Request):
scheduler_output = MagicMock(spec=SchedulerOutput)
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 = MagicMock(spec=ModelRunnerOutput)
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
return scheduler_output, model_runner_output
def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None:
request = _make_request()
sched = _make_sched(request)
sched._update_request_with_output.return_value = ([42], True)
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.prefill_stats is not None
assert eco.prefill_stats.num_cached_tokens == 1
assert eco.prefill_stats.num_cache_creation_tokens == 2
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None:
request = _make_request()
request.num_computed_tokens = len(request.prompt_token_ids)
sched = _make_sched(request)
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.prefill_stats is not None
assert eco.prefill_stats.num_cached_tokens == 1
assert eco.prefill_stats.num_cache_creation_tokens == 2
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
@@ -308,6 +308,91 @@ def test_flash_attn_func_preferred_over_varlen():
print("✓ flash_attn_func forward works correctly!")
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
@pytest.mark.parametrize("k_len", [40, 256])
def test_cross_attn_key_padding_vs_sdpa(k_len):
"""
Case 3: Cross-attention with a key-padding mask.
The mask covers only the key sequence; every query position must attend
to the valid keys. FA and SDPA outputs should be very close. k_len=256
covers equal Q/K lengths, which must not be mistaken for self-attention,
and valid key lengths differ per batch element.
"""
device = torch.device(current_omni_platform.device_type)
dtype = torch.bfloat16
batch_size = 2
q_len = 256
valid_lens = (17, 29)
num_heads = 8
head_dim = 64
fa_impl = FlashAttentionImpl(
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross"
)
sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
torch.manual_seed(7)
query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype)
key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device)
for i, valid_len in enumerate(valid_lens):
attn_mask[i, :valid_len] = True
attn_metadata = AttentionMetadata(attn_mask=attn_mask)
output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item()
mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item()
print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===")
print(f"Max absolute difference: {max_diff:.6f}")
print(f"Mean absolute difference: {mean_diff:.6f}")
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
def test_varlen_masked_routing_by_role(monkeypatch):
"""The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
calls = []
def fake_varlen_func(q, k, v, **kwargs):
calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"]))
return q
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8
query = torch.randn(batch_size, seq_len, num_heads, head_dim)
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]]))
cross_impl = FlashAttentionImpl(
num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross"
)
output = cross_impl.forward_cuda(query, query, query, metadata)
assert output.shape == query.shape
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
assert num_q_rows == batch_size * seq_len
assert cu_seqlens_q.tolist() == [0, 4, 8]
assert cu_seqlens_k.tolist() == [0, 2, 5]
self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False)
output = self_impl.forward_cuda(query, query, query, metadata)
assert output.shape == query.shape
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
assert num_q_rows == 5
assert cu_seqlens_q.tolist() == [0, 2, 5]
assert cu_seqlens_k.tolist() == [0, 2, 5]
def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
calls = []
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# ruff: noqa: N803, E741
"""Layer-level correctness & performance test for MoT parallel linear layers.
@@ -237,6 +240,53 @@ def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
return output
def _assert_gen_expert_parallel_metadata(layer):
"""The secondary expert must satisfy vLLM's parallel-linear contract."""
for attr in (
"tp_rank",
"tp_size",
"input_size",
"input_size_per_partition",
"output_size",
"output_size_per_partition",
):
assert getattr(layer.gen_exp, attr) == getattr(layer, attr)
def test_mot_gen_expert_parallel_metadata():
"""Keep secondary experts compatible with vLLM online quantizers."""
from vllm.model_executor.layers.quantization.online.fp8 import (
_is_tp_sharded,
)
with set_current_vllm_config(VllmConfig()):
layers = (
MoTQKVParallelLinear(
hidden_size=16,
head_size=4,
total_num_heads=4,
total_num_kv_heads=2,
bias=False,
vae_bias=False,
params_dtype=torch.bfloat16,
disable_tp=True,
),
MoTRowParallelLinear(
16,
8,
bias=False,
vae_bias=False,
input_is_parallel=True,
params_dtype=torch.bfloat16,
disable_tp=True,
),
)
for layer in layers:
_assert_gen_expert_parallel_metadata(layer)
assert _is_tp_sharded(layer.gen_exp) is False
def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
"""Compare outputs, print metrics, and assert correctness.
@@ -335,6 +385,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool
disable_tp=True,
).cuda()
_assert_gen_expert_parallel_metadata(mot_linear)
assert text_linear.output_size_per_partition == N, (
f"Expected output_size_per_partition={N}, "
f"got {text_linear.output_size_per_partition}. "
@@ -423,6 +475,8 @@ def test_mot_o_proj(
disable_tp=True,
).cuda()
_assert_gen_expert_parallel_metadata(mot_linear)
_sync_weights(text_linear, vae_linear, mot_linear)
text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
@@ -501,6 +555,13 @@ def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool):
).cuda()
with torch.no_grad():
# vLLM linear layers leave weights uninitialized (torch.empty), so
# a freshly zeroed GPU page would make both outputs zero and the
# cosine-similarity gate degenerate. Seed deterministic values so
# und-mode is compared against a meaningful reference.
ref_linear.weight.normal_(mean=0.0, std=0.02)
if ref_linear.bias is not None:
ref_linear.bias.normal_(mean=0.0, std=0.02)
mot_linear.weight.copy_(ref_linear.weight)
if bias and ref_linear.bias is not None:
mot_linear.bias.copy_(ref_linear.bias)
@@ -549,6 +610,13 @@ def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool):
).cuda()
with torch.no_grad():
# vLLM linear layers leave weights uninitialized (torch.empty), so
# a freshly zeroed GPU page would make both outputs zero and the
# cosine-similarity gate degenerate. Seed deterministic values so
# und-mode is compared against a meaningful reference.
ref_linear.weight.normal_(mean=0.0, std=0.02)
if ref_linear.bias is not None:
ref_linear.bias.normal_(mean=0.0, std=0.02)
mot_linear.weight.copy_(ref_linear.weight)
if bias and ref_linear.bias is not None:
mot_linear.bias.copy_(ref_linear.bias)
+24 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice():
assert torch.allclose(out, expected)
def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
torch.nn.Module.__init__(layer)
layer.base_layer = torch.nn.Module()
layer.n_slices = 1
layer.lora_a_stacked = (torch.zeros(2),)
layer.lora_b_stacked = (torch.zeros(3),)
layer._set_diffusion_lora_buffer_device(torch.device("meta"))
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
def recreate_on_cpu(self, max_loras, lora_config, model_config):
self.lora_a_stacked = (torch.zeros(4),)
self.lora_b_stacked = (torch.zeros(5),)
monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
layer.create_lora_weights(max_loras=1, lora_config=object())
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
# Verify that after reset_lora(), apply() skips LoRA matmuls even if the
# LoRA tensors are still allocated and non-empty.
+33
View File
@@ -3,6 +3,8 @@
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from vllm.lora.lora_weights import LoRALayerWeights
@@ -143,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
assert replace_calls == ["foo"]
def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
import vllm_omni.diffusion.lora.manager as manager_mod
class _DLOLoRALayer(DummyBaseLayerWithLoRA):
lora_buffer_device = None
def _set_diffusion_lora_buffer_device(self, device):
self.lora_buffer_device = device
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
monkeypatch.setattr(
manager_mod,
"from_layer_diffusion",
lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
)
pipeline = torch.nn.Module()
pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
pipeline.transformer = torch.nn.Module()
pipeline.transformer.foo = _FakeLinearBase()
manager = DiffusionLoRAManager(
pipeline=pipeline,
device=torch.device("cpu"),
dtype=torch.bfloat16,
)
manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
import vllm_omni.diffusion.lora.manager as manager_mod
@@ -249,6 +249,30 @@ def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_p
assert loader.take_host_weight_plan() is None
def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
tmp_path: Path,
monkeypatch,
):
canonical_root = tmp_path / "canonical"
canonical_root.mkdir()
save_file(
{"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
str(canonical_root / "model.safetensors"),
)
loader = DiffusersPipelineLoader(
LoadConfig(),
_hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
)
pipeline = _HWRPipeline(canonical_root)
monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
loader.load_model(load_device="cpu", device=torch.device("cpu"))
assert pipeline.load_count == 0
assert loader.take_host_weight_plan() is None
def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
return OmniDiffusionConfig(
model="",
@@ -261,10 +285,11 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
@pytest.mark.parametrize(
("dist_offload", "use_allgather"),
("dist_offload", "use_allgather", "mode"),
[
(False, False),
(True, True),
(False, False, "preferred"),
(True, True, "preferred"),
(True, False, "disabled"),
],
)
def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
@@ -272,12 +297,13 @@ def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
tmp_path,
dist_offload,
use_allgather,
mode,
):
"""Disabled and AllGather paths must never construct or probe HWR."""
from vllm_omni.host_weight_runtime import HostWeightRuntime
root = tmp_path / "must-not-be-touched"
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root))
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
model = _DummyPipelineModel(source_prefix="transformer.")
modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
@@ -303,6 +303,83 @@ class _FakeH3Attention(nn.Module):
super().__init__()
class _FakeFluxLinear(_FakeH3Linear):
def __init__(
self,
*args: object,
bias: bool = False,
params_dtype: torch.dtype | None = None,
total_num_heads: int | None = None,
total_num_kv_heads: int | None = None,
**kwargs: object,
) -> None:
super().__init__(
*args,
bias=bias,
params_dtype=params_dtype,
total_num_heads=total_num_heads,
total_num_kv_heads=total_num_kv_heads,
**kwargs,
)
self.weight.weight_loader = self._load_weight
@staticmethod
def _load_weight(
parameter: nn.Parameter,
loaded_weight: torch.Tensor,
shard_id: str | None = None,
) -> None:
del shard_id
parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
class _FakeFluxAttention(nn.Module):
def __init__(self, **kwargs: object) -> None:
del kwargs
super().__init__()
self.register_buffer("beta", torch.tensor([0.5]))
self.register_buffer("eps", torch.tensor([-1e-6]))
class _Flux2KleinPipeline(nn.Module):
def __init__(self, transformer: nn.Module) -> None:
super().__init__()
self.transformer = transformer
self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
self.vae = nn.Module()
self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
def _reduced_flux2_klein(
monkeypatch: pytest.MonkeyPatch,
*,
device: torch.device | str = "cpu",
) -> nn.Module:
from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
transformer = flux2.Flux2Transformer2DModel(
patch_size=1,
in_channels=8,
out_channels=8,
num_layers=1,
num_single_layers=1,
attention_head_dim=8,
num_attention_heads=1,
joint_attention_dim=8,
timestep_guidance_channels=8,
mlp_ratio=1.0,
axes_dims_rope=(2, 2, 2, 2),
guidance_embeds=False,
)
return transformer.to(device=device, dtype=torch.bfloat16)
class _H3Pipeline(nn.Module):
def __init__(self, transformer: nn.Module) -> None:
super().__init__()
@@ -1200,3 +1277,152 @@ def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
assert contract.implementation_id == "minimax-h3-dit"
assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
def test_reduced_flux2_klein_declares_complete_final_layout_contract(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
Flux2Transformer2DModel,
)
transformer = _reduced_flux2_klein(monkeypatch)
pipeline = _Flux2KleinPipeline(transformer)
source = _prepared_source(tmp_path, directory="flux2-klein-source")
context = _identity(pipeline, source)
targets = collect_final_layout_targets(
pipeline,
(("transformer", transformer),),
policy=FINAL_LAYOUT_BF16_POLICY,
require_materialized=True,
)
target_names = {target.name for target in targets}
contract = Flux2Transformer2DModel.host_weight_restore_contract
assert isinstance(contract, FinalLayoutModelContract)
assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
assert contract.implementation_id == "flux2-klein-dit"
assert contract.version == "1"
assert context.model_contracts == (contract,)
assert transformer.stacked_params_mapping == [
(".to_qkv.", ".to_q.", "q"),
(".to_qkv.", ".to_k.", "k"),
(".to_qkv.", ".to_v.", "v"),
(".add_kv_proj", ".add_q_proj", "q"),
(".add_kv_proj", ".add_k_proj", "k"),
(".add_kv_proj", ".add_v_proj", "v"),
]
loaded = transformer.load_weights(
[
("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
(
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
torch.tensor([3.0], dtype=torch.bfloat16),
),
]
)
assert loaded == {
"transformer_blocks.0.attn.to_qkv.weight",
"transformer_blocks.0.attn.add_kv_proj.weight",
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
}
assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
assert all("text_encoder" not in name and "vae" not in name for name in target_names)
transformer.validate_restored_host_weights()
def test_flux2_klein_restore_validator_rejects_non_ready_state(
monkeypatch: pytest.MonkeyPatch,
) -> None:
transformer = _reduced_flux2_klein(monkeypatch)
mapping = transformer.stacked_params_mapping
transformer.stacked_params_mapping = None
with pytest.raises(ValueError, match="packed QKV mapping"):
transformer.validate_restored_host_weights()
transformer.stacked_params_mapping = mapping
weight = transformer.x_embedder.weight
weight.data = weight.data.float()
with pytest.raises(ValueError, match="must stay bf16"):
transformer.validate_restored_host_weights()
weight.data = weight.data.to(torch.bfloat16)
weight.data = weight.data.T
with pytest.raises(ValueError, match="contiguous strided layout"):
transformer.validate_restored_host_weights()
weight.data = weight.data.contiguous()
attention = transformer.transformer_blocks[0].attn.attn
attention._non_persistent_buffers_set.add("beta")
with pytest.raises(ValueError, match="must be persistent"):
transformer.validate_restored_host_weights()
attention._non_persistent_buffers_set.remove("beta")
attention.beta.data = attention.beta.data.float()
with pytest.raises(ValueError, match="scalar bf16"):
transformer.validate_restored_host_weights()
def test_flux2_klein_publication_restore_fallback_and_teardown(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
source = _prepared_source(tmp_path, directory="flux2-klein-source")
cold_transformer = _reduced_flux2_klein(monkeypatch)
cold_model = _Flux2KleinPipeline(cold_transformer)
with torch.no_grad():
for index, parameter in enumerate(cold_transformer.parameters(), start=1):
parameter.fill_(index)
for name, buffer in cold_transformer.named_buffers():
if name.endswith((".beta", ".eps")):
buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
context = _identity(cold_model, source)
runtime = _runtime(tmp_path / "flux2-klein-store")
miss = runtime.resolve(context.identity)
assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
producer = FinalLayoutBF16Producer(
context,
cold_model,
(("transformer", cold_transformer),),
max_shard_bytes=64,
)
publication = runtime.publish_after_load(context.identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
warm_model = _Flux2KleinPipeline(warm_transformer)
warm_context = _identity(warm_model, source)
assert warm_context.identity == context.identity
hit = runtime.resolve(warm_context.identity)
assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
assert hit.lease is not None
plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
plan.commit()
warm_transformer.validate_restored_host_weights()
cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
for name in warm_context.tensor_names:
assert torch.equal(warm_tensors[name], cold_tensors[name])
assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
del warm_tensors, warm_model, warm_transformer, plan
gc.collect()
hit.lease.close()
assert hit.lease.closed
assert isinstance(runtime.store, FilesystemHostWeightStore)
assert runtime.store.cleanup(context.identity) is None
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -245,7 +245,7 @@ else:
[sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
capture_output=True,
text=True,
timeout=3,
timeout=30,
check=False,
)
@@ -11,6 +11,31 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
class FakeBaseVocoder(torch.nn.Module):
def __init__(self):
super().__init__()
self.input_dtype = None
def forward(self, value):
self.input_dtype = value.dtype
return value
monkeypatch.setattr(
ltx_runtime.torch,
"autocast",
lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")),
)
vocoder = FakeBaseVocoder()
output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16))
assert vocoder.input_dtype == torch.bfloat16
assert output.dtype == torch.bfloat16
class TestLTXOutputRank:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
@@ -0,0 +1,48 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CUDA regression tests for LTX-2 BWE vocoder precision."""
import pytest
import torch
from tests.helpers.mark import hardware_test
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
class _BWEConvVocoder(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv1d(
1,
1,
kernel_size=1,
bias=False,
device="cuda",
dtype=torch.bfloat16,
)
self.bwe_generator = torch.nn.Identity()
self.input_dtype = None
self.conv_output_dtype = None
def forward(self, value):
self.input_dtype = value.dtype
output = self.conv(value)
self.conv_output_dtype = output.dtype
return output
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder
vocoder = _BWEConvVocoder()
generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
output = _run_ltx_vocoder(vocoder, generated_mel)
assert next(vocoder.parameters()).dtype == torch.bfloat16
assert vocoder.input_dtype == torch.float32
assert vocoder.conv_output_dtype == torch.float32
assert output.dtype == torch.bfloat16
@@ -197,7 +197,6 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
[
"model-level CPU offload (--enable-cpu-offload)",
"layerwise offload (--enable-layerwise-offload)",
"distributed layerwise offload (--enable-distributed-layerwise-offload)",
],
)
def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
@@ -214,6 +213,38 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
)
def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import (
pipeline_minimax_h3 as pipeline_module,
)
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._turbo_lora_adapter_ids = set()
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=True,
)
captured = {}
def load_turbo(**kwargs):
captured.update(kwargs)
return object(), object()
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=_request("turbo"),
lora_path="turbo",
dtype=torch.bfloat16,
)
assert loaded is not None
assert captured["unsupported_offload_mode"] is None
def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
_write_tiny_turbo(unsupported)
@@ -35,7 +35,10 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
DistributedLayerwiseOffloadHook,
PinnedResidentLayerGroup,
)
from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
from vllm_omni.diffusion.offloader.host_registration import (
HostRegistrationCleanupError,
HostRegistrationError,
)
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
from vllm_omni.diffusion.offloader.offload_plan import (
OffloadPlan,
@@ -807,7 +810,10 @@ def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_ru
assert torch.equal(block.weight, expected)
def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch):
def test_hwr_registration_uses_transport_budget(
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
plan, _, lease = _fake_hwr_plan()
backend = DistributedLayerwiseOffloadBackend(
@@ -1125,6 +1131,51 @@ class TestMmapWeightLoading:
backend.disable()
assert lease.closed
def test_hwr_registration_failure_falls_back_to_bounded_staging(
self,
monkeypatch,
patched_offload_runtime,
):
plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dp_size=1,
dlo_use_allgather=False,
),
torch.device("cpu"),
host_weight_plan=plan,
)
pipeline = _HWRPipeline()
original_staging_allocator = backend._allocate_shared_cpu_staging_buffers
def fail_registration(*args, **kwargs):
del args, kwargs
raise HostRegistrationError("registration unavailable in CPU test")
def allocate_unpinned_staging(hooks, resident_group=None):
# Registration selection still observes pin_cpu_memory=True. The
# CPU-only test avoids asking the host for CUDA-pinned allocations.
for hook in hooks:
hook.pin_memory = False
return original_staging_allocator(hooks, resident_group)
monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)
backend.enable(pipeline)
assert carrier.taken
assert backend._using_rank_local_mmap
assert not backend._using_registered_mmap
assert backend._host_registration is None
assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
assert not lease.closed
backend.disable()
assert lease.closed
def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
self,
monkeypatch,
@@ -2132,8 +2183,9 @@ class TestDynamicSlotTracking:
prev_idx = (i - 1) % num_blocks
# Dynamic slot tracking: read from prev's prefetched slot
if prefetched_slots[prev_idx] is not None:
current_slots[i] = prefetched_slots[prev_idx]
prefetched_slot = prefetched_slots[prev_idx]
if prefetched_slot is not None:
current_slots[i] = prefetched_slot
read_slot = current_slots[i]
@@ -0,0 +1,146 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from unittest.mock import Mock
import pytest
import torch
from vllm.model_executor.layers.linear import (
LinearBase,
UnquantizedLinearMethod,
)
from vllm_omni.quantization import svdquant_config
from vllm_omni.quantization.factory import (
SUPPORTED_QUANTIZATION_METHODS,
build_quant_config,
)
from vllm_omni.quantization.svdquant_config import (
DiffusionSVDQuantConfig,
DiffusionSVDQuantLinearMethod,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_config_from_checkpoint_dict() -> None:
config = DiffusionSVDQuantConfig.from_config(
{
"rank": 32,
"precision": "nvfp4",
"act_unsigned": False,
"modules_to_not_convert": ["blocks.0.adaln_proj.linear"],
}
)
assert config.rank == 32
assert config.precision == "nvfp4"
assert config.modules_to_not_convert == ["blocks.0.adaln_proj.linear"]
assert config.get_supported_act_dtypes() == [torch.bfloat16]
@pytest.mark.parametrize(
("kwargs", "message"),
[
({"rank": 0}, "rank must be positive"),
({"precision": "int4"}, "NVFP4 checkpoints only"),
({"act_unsigned": True}, "unsigned activations"),
],
)
def test_config_rejects_unsupported_phase1_options(
kwargs: dict[str, object],
message: str,
) -> None:
with pytest.raises(ValueError, match=message):
DiffusionSVDQuantConfig(**kwargs)
def test_factory_builds_svdquant_with_extra_checkpoint_fields() -> None:
config = build_quant_config(
{
"quant_method": "svdquant",
"rank": 16,
"precision": "nvfp4",
"w4a16_modules": ["adaln_proj.linear"],
"w4a16_group_size": 64,
}
)
assert isinstance(config, DiffusionSVDQuantConfig)
assert config.rank == 16
assert "svdquant" in SUPPORTED_QUANTIZATION_METHODS
def test_skipped_and_non_linear_modules_remain_unquantized() -> None:
config = DiffusionSVDQuantConfig(modules_to_not_convert=["blocks.0.adaln_proj.linear"])
linear = Mock(spec=LinearBase)
assert isinstance(
config.get_quant_method(linear, "blocks.0.adaln_proj.linear"),
UnquantizedLinearMethod,
)
assert config.get_quant_method(torch.nn.ReLU(), "blocks.0.act") is None
def test_linear_method_creates_canonical_partitioned_weights(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig(rank=32))
layer = torch.nn.Module()
method.create_weights(
layer,
input_size_per_partition=128,
output_partition_sizes=[64, 64],
input_size=128,
output_size=128,
params_dtype=torch.bfloat16,
)
assert layer.qweight.shape == (128, 64)
assert layer.qweight.dtype == torch.int8
assert layer.wscales.shape == (8, 128)
assert layer.wscales.dtype == torch.float8_e4m3fn
assert layer.proj_down.shape == (128, 32)
assert layer.proj_up.shape == (128, 32)
assert layer.smooth_factor.shape == (128,)
assert layer.wcscales.shape == (128,)
assert layer.wtscale.shape == (1,)
assert layer.qweight.input_dim == 1
assert layer.qweight.output_dim == 0
assert layer.wscales.input_dim == 0
assert layer.wscales.output_dim == 1
assert layer.proj_down.input_dim == 0
assert layer.proj_up.output_dim == 0
assert layer.wcscales.output_dim == 0
def test_linear_method_rejects_misaligned_input_partition(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig())
with pytest.raises(ValueError, match="divisible by the block size 16"):
method.create_weights(
torch.nn.Module(),
input_size_per_partition=24,
output_partition_sizes=[32],
input_size=24,
output_size=32,
params_dtype=torch.bfloat16,
)
def test_active_linear_uses_svdquant_method(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
config = DiffusionSVDQuantConfig()
linear = Mock(spec=LinearBase)
method = config.get_quant_method(linear, "blocks.0.attn.out_proj")
assert isinstance(method, DiffusionSVDQuantLinearMethod)
@@ -0,0 +1,168 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.quantization import svdquant_config as svdquant
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
class _FakeNvfp4Kernel:
def __init__(self, base_weight: torch.Tensor | None = None) -> None:
self.base_weight = base_weight
self.processed = False
self.last_input: torch.Tensor | None = None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
del layer
self.processed = True
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
del layer
assert bias is None
assert self.base_weight is not None
self.last_input = x.clone()
return torch.mm(x, self.base_weight)
def _register_parameter(
layer: torch.nn.Module,
name: str,
value: torch.Tensor,
) -> None:
layer.register_parameter(
name,
torch.nn.Parameter(value, requires_grad=False),
)
def test_supports_only_validated_datacenter_blackwell() -> None:
assert not svdquant._supports_capability(SimpleNamespace(major=10, minor=0))
assert svdquant._supports_capability(SimpleNamespace(major=10, minor=3))
assert not svdquant._supports_capability(SimpleNamespace(major=11, minor=0))
assert not svdquant._supports_capability(SimpleNamespace(major=12, minor=0))
assert not svdquant._supports_capability(None)
def test_rejects_nvfp4_backend_with_incompatible_layer_contract(
monkeypatch: pytest.MonkeyPatch,
) -> None:
svdquant._nvfp4_kernel.cache_clear()
monkeypatch.setattr(
svdquant,
"init_nvfp4_linear_kernel",
lambda: _FakeNvfp4Kernel(),
)
with pytest.raises(RuntimeError, match="incompatible with the SVDQuant checkpoint"):
svdquant._nvfp4_kernel()
svdquant._nvfp4_kernel.cache_clear()
def test_prepare_weights_uses_existing_nvfp4_layout_contract(
monkeypatch: pytest.MonkeyPatch,
) -> None:
kernel = _FakeNvfp4Kernel()
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
layer = torch.nn.Module()
qweight = torch.arange(16, dtype=torch.int8).reshape(4, 4)
wscales = torch.arange(8, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(2, 4)
_register_parameter(layer, "qweight", qweight)
_register_parameter(layer, "wscales", wscales)
_register_parameter(layer, "wtscale", torch.tensor([2.0], dtype=torch.bfloat16))
_register_parameter(
layer,
"wcscales",
torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16),
)
method.process_weights_after_loading(layer)
assert kernel.processed
assert not hasattr(layer, "qweight")
assert not hasattr(layer, "wscales")
assert layer.weight.dtype == torch.uint8
assert torch.equal(layer.weight.view(torch.int8), qweight)
assert torch.equal(layer.weight_scale, wscales.transpose(0, 1))
torch.testing.assert_close(layer.alpha, torch.tensor([2.0]))
assert layer.input_global_scale_inv.dtype == torch.float32
assert layer.output_channel_scale.shape == (4,)
def test_prepare_weights_drops_identity_output_scale(
monkeypatch: pytest.MonkeyPatch,
) -> None:
kernel = _FakeNvfp4Kernel()
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
layer = torch.nn.Module()
_register_parameter(layer, "qweight", torch.zeros(4, 4, dtype=torch.int8))
_register_parameter(
layer,
"wscales",
torch.ones(2, 4, dtype=torch.float8_e4m3fn),
)
_register_parameter(layer, "wtscale", torch.ones(1, dtype=torch.bfloat16))
_register_parameter(layer, "wcscales", torch.ones(4, dtype=torch.bfloat16))
method.process_weights_after_loading(layer)
assert layer.output_channel_scale is None
def test_compatibility_path_applies_qkv_scale_and_rank_correction(
monkeypatch: pytest.MonkeyPatch,
) -> None:
base_weight = torch.tensor(
[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
dtype=torch.bfloat16,
)
kernel = _FakeNvfp4Kernel(base_weight)
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
layer = torch.nn.Module()
layer.smooth_factor = torch.tensor([2.0, 4.0], dtype=torch.bfloat16)
layer.proj_down = torch.tensor([[1.0], [2.0]], dtype=torch.bfloat16)
layer.proj_up = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16)
layer.output_channel_scale = torch.tensor([1.0, 2.0, 4.0], dtype=torch.bfloat16)
layer.output_size_per_partition = 3
x = torch.tensor([[2.0, 4.0], [4.0, 8.0]], dtype=torch.bfloat16)
bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16)
actual = method.apply(layer, x, bias)
smoothed = x / layer.smooth_factor
expected = torch.mm(smoothed, base_weight)
expected = expected * layer.output_channel_scale
expected = expected + torch.mm(torch.mm(x, layer.proj_down), layer.proj_up.T)
expected = expected + bias
torch.testing.assert_close(actual, expected)
torch.testing.assert_close(kernel.last_input, smoothed)
def test_compatibility_path_rejects_fp16_activations(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
layer = torch.nn.Module()
with pytest.raises(ValueError, match="requires BF16 activations"):
method.apply(
layer,
torch.zeros(1, 16, dtype=torch.float16),
)
@@ -0,0 +1,130 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
from vllm_omni.quantization import svdquant_config as svdquant
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _linear_without_init(linear_type: type[torch.nn.Module], tp_rank: int):
layer = object.__new__(linear_type)
torch.nn.Module.__init__(layer)
layer.tp_rank = tp_rank
return layer
def _values(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
values = torch.arange(torch.Size(shape).numel(), dtype=torch.int64).remainder(127)
return values.reshape(shape).to(dtype)
def _method(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
return svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig(rank=8))
def test_column_parallel_loader_shards_svdquant_outputs(
monkeypatch: pytest.MonkeyPatch,
) -> None:
layer = _linear_without_init(ColumnParallelLinear, tp_rank=1)
method = _method(monkeypatch)
method.create_weights(
layer,
input_size_per_partition=128,
output_partition_sizes=[32],
input_size=128,
output_size=64,
params_dtype=torch.bfloat16,
weight_loader=layer.weight_loader,
)
loaded_tensors = {
"qweight": _values((64, 64), torch.int8),
"wscales": _values((8, 64), torch.float8_e4m3fn),
"proj_up": _values((64, 8), torch.bfloat16),
"wcscales": _values((64,), torch.bfloat16),
}
for name, loaded in loaded_tensors.items():
parameter = getattr(layer, name)
parameter.weight_loader(parameter, loaded)
expected = loaded.narrow(parameter.output_dim, 32, 32)
torch.testing.assert_close(parameter, expected)
def test_row_parallel_loader_shards_svdquant_inputs(
monkeypatch: pytest.MonkeyPatch,
) -> None:
layer = _linear_without_init(RowParallelLinear, tp_rank=1)
method = _method(monkeypatch)
method.create_weights(
layer,
input_size_per_partition=64,
output_partition_sizes=[64],
input_size=128,
output_size=64,
params_dtype=torch.bfloat16,
weight_loader=layer.weight_loader,
)
loaded_tensors = {
"qweight": _values((64, 64), torch.int8),
"wscales": _values((8, 64), torch.float8_e4m3fn),
"proj_down": _values((128, 8), torch.bfloat16),
"smooth_factor": _values((128,), torch.bfloat16),
}
for name, loaded in loaded_tensors.items():
parameter = getattr(layer, name)
parameter.weight_loader(parameter, loaded)
shard_size = parameter.shape[parameter.input_dim]
expected = loaded.narrow(parameter.input_dim, shard_size, shard_size)
torch.testing.assert_close(parameter, expected)
def test_qkv_parallel_loader_shards_fused_svdquant_outputs(
monkeypatch: pytest.MonkeyPatch,
) -> None:
layer = _linear_without_init(QKVParallelLinear, tp_rank=1)
layer.total_num_heads = 4
layer.total_num_kv_heads = 2
layer.num_heads = 2
layer.num_kv_heads = 1
layer.num_kv_head_replicas = 1
layer.head_size = 8
layer.v_head_size = 8
method = _method(monkeypatch)
method.create_weights(
layer,
input_size_per_partition=128,
output_partition_sizes=[16, 8, 8],
input_size=128,
output_size=64,
params_dtype=torch.bfloat16,
weight_loader=layer.weight_loader,
)
loaded_tensors = {
"qweight": _values((64, 64), torch.int8),
"wscales": _values((8, 64), torch.float8_e4m3fn),
"proj_up": _values((64, 8), torch.bfloat16),
"wcscales": _values((64,), torch.bfloat16),
}
expected_global_outputs = [slice(16, 32), slice(40, 48), slice(56, 64)]
for name, loaded in loaded_tensors.items():
parameter = getattr(layer, name)
parameter.weight_loader(parameter, loaded)
expected = torch.cat(
[
loaded.narrow(parameter.output_dim, part.start, part.stop - part.start)
for part in expected_global_outputs
],
dim=parameter.output_dim,
)
torch.testing.assert_close(parameter, expected)
@@ -172,12 +172,29 @@ class TestMetricKeys:
assert ".cancel(" not in abort_branch
def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None:
"""A metrics-only sentinel contributes its queue depth and is not routed.
Absorption lives in ``_absorb_diffusion_metrics`` rather than inline in a
loop, so both orchestration loops share one implementation. The ordering
is asserted where it now lives: the snapshot inside the helper, and the
absorb-before-route guard inside every loop that polls diffusion output.
"""
source = _read_source(_ORCHESTRATOR_PATH)
loop_source = _get_function_source(source, "Orchestrator", "_orchestration_loop")
snapshot_pos = loop_source.index("_update_stage_replica_waiting(")
sentinel_pos = loop_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
route_pos = loop_source.index("pool.record_output_timestamps([diffusion_output])")
assert snapshot_pos < sentinel_pos < route_pos
# Snapshot before the verdict, so a sentinel still reports its waiting
# depth on the way out instead of being dropped unaccounted.
absorb_source = _get_function_source(source, "Orchestrator", "_absorb_diffusion_metrics")
snapshot_pos = absorb_source.index("_update_stage_replica_waiting(")
sentinel_pos = absorb_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
assert snapshot_pos < sentinel_pos
# Absorb before routing, or a sentinel reaches the downstream consumer
# as if it were a real request output.
for loop_name in ("_orchestration_loop", "_orchestration_loop_event_driven"):
loop_source = _get_function_source(source, "Orchestrator", loop_name)
absorb_pos = loop_source.index("self._absorb_diffusion_metrics(")
route_pos = loop_source.index("record_output_timestamps(")
assert absorb_pos < route_pos, loop_name
class TestVaeDecodeEmit:
@@ -1209,6 +1209,67 @@ def test_cleanup_preserves_pending_save(build_adapter):
assert len(adapter._pending_save_reqs) == 1
def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter):
adapter, connector = build_adapter(stage_id=1)
request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
)
adapter.save_async(multimodal_output=None, request=request)
stale_task = adapter._pending_save_reqs.popleft()
adapter.cleanup(request.request_id, request.external_req_id)
adapter._send_single_request(stale_task)
connector.put.assert_not_called()
assert request.external_req_id not in adapter.put_req_chunk
replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id)
adapter.save_async(multimodal_output=None, request=replacement)
current_task = adapter._pending_save_reqs.popleft()
adapter._send_single_request(current_task)
connector.put.assert_called_once()
assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0"
def test_finish_requests_does_not_wait_for_inflight_send(build_adapter):
adapter, connector = build_adapter(stage_id=1)
first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
put_started = threading.Event()
release_put = threading.Event()
def blocking_put(**kwargs):
put_started.set()
release_put.wait(timeout=2)
return True, 1, {}
connector.put.side_effect = blocking_put
adapter.save_async(multimodal_output=None, request=first)
task = adapter._pending_save_reqs.popleft()
sender = threading.Thread(target=adapter._send_single_request, args=(task,))
sender.start()
assert put_started.wait(timeout=1)
cleaner = threading.Thread(
target=adapter.finish_requests,
args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}),
)
cleaner.start()
try:
cleaner.join(timeout=0.5)
assert not cleaner.is_alive()
assert task["sender_token"].cancelled
assert first.external_req_id in adapter._sender_tokens
finally:
release_put.set()
sender.join(timeout=1)
cleaner.join(timeout=1)
assert not sender.is_alive()
assert first.external_req_id not in adapter._sender_tokens
assert first.external_req_id not in adapter.put_req_chunk
def test_cleanup_only_affects_target_request(build_adapter):
"""Cleanup for one request must not affect another request's state."""
adapter, _ = build_adapter(stage_id=1)
@@ -1507,6 +1568,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
scheduler.structured_output_manager.should_advance.return_value = False
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
request = _HashableRequest(
@@ -1587,6 +1649,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager = mocker.MagicMock()
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
scheduler.waiting_for_transfer_free = set()
scheduler.transfer_triggered_requests = set()
@@ -1742,6 +1805,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
scheduler.structured_output_manager.should_advance.return_value = False
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
scheduler._pending_finish_reqs = list(pending_finish_reqs)
@@ -2216,3 +2280,43 @@ def test_expiry_is_per_request(build_adapter):
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
assert "healthy" in adapter._waiting_since
def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter):
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
talker2code2wav_async_chunk,
)
adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1})
external_req_id = "ext-native-abort"
first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
first.resumable = True
first.model_intermediate_buffer = None
first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1}
first_frame = torch.full((1, 31), 101, dtype=torch.long)
talker2code2wav_async_chunk(
adapter,
{"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}},
first,
is_finished=False,
)
adapter.requests_num_chunks_sent[external_req_id] = 1
adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first})
assert external_req_id not in adapter.request_payload
assert external_req_id not in adapter.requests_num_chunks_sent
replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
replacement.resumable = True
replacement.model_intermediate_buffer = None
replacement.additional_information = first.additional_information
replacement_payload = talker2code2wav_async_chunk(
adapter,
{"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}},
replacement,
is_finished=False,
)
assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))
@@ -22,6 +22,7 @@ import shutil
import subprocess
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path
import numpy as np
@@ -83,35 +84,27 @@ NEGATIVE_PROMPT = (
# Release parity exercises the accelerated BF16 cuDNN path in both runtimes.
# Degenerate one-token attention falls back to PyTorch SDPA MATH.
ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16"
VIDEO_SSIM_MEAN_THRESHOLD = 0.95
VIDEO_SSIM_MIN_THRESHOLD = 0.90
VIDEO_PSNR_MEAN_THRESHOLD = 30.0
AUDIO_RELATIVE_L2_THRESHOLD = 0.2
AUDIO_COSINE_THRESHOLD = 0.95
# Decoded-output gates for the pinned official/Omni BF16 cuDNN comparison.
# The metrics artifact preserves the exact observed values.
LTX25_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
LTX25_VIDEO_SSIM_MIN_THRESHOLD = 0.99
LTX25_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
LTX25_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
LTX25_AUDIO_RELATIVE_L2_THRESHOLD = 0.32
LTX25_AUDIO_COSINE_THRESHOLD = 0.95
# Full/SFT uses independently converted dev weights, but the release contract
# still requires decoded-video parity at the same 0.99 SSIM floor as the
# distilled paths when both runtimes use BF16 cuDNN attention.
LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD = 0.99
LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD = 0.3
LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD = 0.02
LTX25_FULL_AUDIO_COSINE_THRESHOLD = AUDIO_COSINE_THRESHOLD
LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
@dataclass(frozen=True)
class LTXAccuracyThresholds:
video_ssim_mean: float
video_ssim_min: float
video_psnr_mean_db: float
audio_relative_l2: float
audio_cosine_similarity: float
STRICT_THRESHOLDS = LTXAccuracyThresholds(
video_ssim_mean=0.99,
video_ssim_min=0.99,
video_psnr_mean_db=40.0,
audio_relative_l2=0.10,
audio_cosine_similarity=0.99,
)
def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None:
from vllm_omni.outputs import OmniRequestOutput
@@ -419,8 +412,6 @@ def test_ltx25_i2v_requests_pin_official_crf() -> None:
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
from benchmarks.diffusion.quantization_quality import compute_lpips_video
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
ssim_scores: list[float] = []
@@ -435,7 +426,6 @@ def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
"lpips_alex_mean": compute_lpips_video(reference, prediction),
"max_abs": float(difference.max()),
"mean_abs": float(difference.mean()),
}
@@ -463,6 +453,17 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
}
def _assert_strict_similarity(
video_metrics: dict[str, float],
audio_metrics: dict[str, float | bool],
) -> None:
assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
@pytest.mark.slow
@pytest.mark.benchmark
@pytest.mark.diffusion
@@ -569,6 +570,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
"omni_model": omni_model_source,
"omni_model_revision": omni_model_revision,
"resolved_omni_model_path": str(omni_model),
"thresholds": asdict(STRICT_THRESHOLDS),
"peak_memory_mb": {
"official_allocated": official_metadata["peak_memory_allocated_mb"],
"official_reserved": official_metadata["peak_memory_reserved_mb"],
@@ -582,13 +584,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
ssim_mean_threshold = LTX25_VIDEO_SSIM_MEAN_THRESHOLD if task == "t2v" else LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_mean"] >= ssim_mean_threshold
assert video_metrics["ssim_min"] >= LTX25_VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= LTX25_VIDEO_PSNR_MEAN_THRESHOLD
assert video_metrics["lpips_alex_mean"] <= LTX25_VIDEO_LPIPS_MEAN_THRESHOLD
assert audio_metrics["cosine_similarity"] >= LTX25_AUDIO_COSINE_THRESHOLD
assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
_assert_strict_similarity(video_metrics, audio_metrics)
@pytest.mark.slow
@@ -729,6 +725,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
"omni_model": omni_model_source,
"omni_model_revision": omni_model_revision,
"resolved_omni_model_path": str(omni_model),
"thresholds": asdict(STRICT_THRESHOLDS),
"peak_memory_mb": {
"official_allocated": official_metadata["peak_memory_allocated_mb"],
"official_reserved": official_metadata["peak_memory_reserved_mb"],
@@ -742,12 +739,4 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
assert video_metrics["ssim_mean"] >= LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_min"] >= LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD
lpips_threshold = (
LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD if task == "t2v" else LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD
)
assert audio_metrics["relative_l2"] <= LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD
assert video_metrics["lpips_alex_mean"] <= lpips_threshold
assert audio_metrics["cosine_similarity"] >= LTX25_FULL_AUDIO_COSINE_THRESHOLD
_assert_strict_similarity(video_metrics, audio_metrics)
@@ -3,9 +3,8 @@
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
The original reduced one-stage guards remain unchanged. Four complementary
default-shape cases cover both model versions, both two-stage weight families,
and T2V/I2V without expanding every combination into a duplicate golden.
Four complementary default-shape cases cover both model versions, both
two-stage weight families, and T2V/I2V without duplicating one-stage coverage.
"""
from __future__ import annotations
@@ -69,25 +68,23 @@ class LTXAccuracyThresholds:
video_ssim_mean: float
video_ssim_min: float
video_psnr_mean_db: float
# Zero disables gating for waveform-sensitive audio that is not expected
# to align with the official output.
audio_relative_l2: float
audio_cosine_similarity: float
STRICT_THRESHOLDS = LTXAccuracyThresholds(
video_ssim_mean=0.95,
video_ssim_min=0.90,
video_psnr_mean_db=30.0,
audio_relative_l2=0.20,
audio_cosine_similarity=0.95,
video_ssim_mean=0.99,
video_ssim_min=0.99,
video_psnr_mean_db=40.0,
audio_relative_l2=0.10,
audio_cosine_similarity=0.99,
)
@dataclass(frozen=True)
class LTXAccuracyCase:
name: str
pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
pipeline_kind: Literal["distilled", "two_stage"]
model_id: str
model_revision: str
model_env: str
@@ -99,8 +96,6 @@ class LTXAccuracyCase:
num_inference_steps: int
seed: int
stg_block: int | None
thresholds: LTXAccuracyThresholds
prompt: str = PROMPT
gemma_model_id: str | None = None
gemma_model_revision: str | None = None
gemma_model_env: str | None = None
@@ -168,59 +163,6 @@ I2V_IMAGE = LTXArtifact(
)
LEGACY_CASES = (
LTXAccuracyCase(
name="ltx2",
pipeline_kind="one_stage",
model_id="Lightricks/LTX-2",
model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2Pipeline",
checkpoint=LTX2_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=29,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3_i2v",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
TWO_STAGE_CASES = (
LTXAccuracyCase(
name="ltx2_distilled_t2v",
@@ -240,7 +182,6 @@ TWO_STAGE_CASES = (
gemma_model_id="Lightricks/LTX-2",
gemma_model_revision=LTX2_REVISION,
gemma_model_env="VLLM_TEST_LTX2_MODEL",
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_distilled_i2v",
@@ -260,7 +201,6 @@ TWO_STAGE_CASES = (
gemma_model_id="diffusers/LTX-2.3-Diffusers",
gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
gemma_model_env="VLLM_TEST_LTX23_MODEL",
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
LTXAccuracyCase(
@@ -280,7 +220,6 @@ TWO_STAGE_CASES = (
seed=10,
stg_block=29,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_two_stage_layer_fused_i2v",
@@ -299,12 +238,11 @@ TWO_STAGE_CASES = (
seed=10,
stg_block=28,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
WEEKLY_CASES = TWO_STAGE_CASES
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
@@ -460,10 +398,49 @@ def _resolve_image(case: LTXAccuracyCase) -> Path | None:
return _resolve_artifact(case.image)
def _omni_model_with_pinned_artifacts(
model: Path,
output_root: Path,
artifacts: tuple[Path | None, ...],
) -> Path:
"""Expose the already-resolved sidecars to Omni without copying the model."""
resolved_artifacts = tuple(artifact for artifact in artifacts if artifact is not None)
if not resolved_artifacts:
return model
overlay = output_root / "omni-model"
overlay.mkdir()
for source in model.iterdir():
(overlay / source.name).symlink_to(source.resolve(), target_is_directory=source.is_dir())
for artifact in resolved_artifacts:
target = overlay / artifact.name
if target.exists():
assert target.samefile(artifact), f"Model sidecar does not match pinned artifact: {target} != {artifact}"
else:
target.symlink_to(artifact.resolve())
return overlay
def test_omni_model_overlay_pins_artifacts(tmp_path: Path) -> None:
model = tmp_path / "model"
model.mkdir()
component = model / "model_index.json"
component.write_text("{}\n")
artifact = tmp_path / "sidecar.safetensors"
artifact.write_bytes(b"pinned")
output_root = tmp_path / "output"
output_root.mkdir()
overlay = _omni_model_with_pinned_artifacts(model, output_root, (artifact,))
assert (overlay / component.name).samefile(component)
assert (overlay / artifact.name).samefile(artifact)
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
request: dict[str, object] = {
"pipeline_kind": case.pipeline_kind,
"prompt": case.prompt,
"prompt": PROMPT,
"negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
"width": case.width,
"height": case.height,
@@ -473,14 +450,17 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
"seed": case.seed,
}
if case.pipeline_kind != "distilled":
assert case.stg_block is not None
request.update(
video_cfg_scale=3.0,
audio_cfg_scale=7.0,
video_stg_scale=1.0,
audio_stg_scale=1.0,
video_modality_scale=3.0,
audio_modality_scale=3.0,
)
if case.pipeline_kind == "two_stage":
assert case.stg_block is not None
request.update(
video_stg_scale=1.0,
audio_stg_scale=1.0,
video_rescale_scale=0.7,
audio_rescale_scale=0.7,
video_stg_blocks=[case.stg_block],
@@ -543,7 +523,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
}
@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
@pytest.mark.parametrize("case", WEEKLY_CASES, ids=lambda case: case.name)
@pytest.mark.slow
@pytest.mark.benchmark
@pytest.mark.diffusion
@@ -561,9 +541,11 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
checkpoint = _resolve_artifact(case.checkpoint, model)
spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
omni_model = _omni_model_with_pinned_artifacts(model, output_root, (spatial_upsampler, distilled_lora))
image = _resolve_image(case)
request = _request(case, image)
request_path = output_root / "request.json"
request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
request_path.write_text(json.dumps(request, indent=2) + "\n")
runner = Path(__file__).with_name("run_ltx_reference.py")
runner_args = [
@@ -619,7 +601,7 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
"--output-dir",
str(omni_output),
"--model",
str(model),
str(omni_model),
"--model-class-name",
case.model_class_name,
]
@@ -651,18 +633,16 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
"spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
"distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
"enable_layerwise_offload": enable_layerwise_offload,
"thresholds": asdict(case.thresholds),
"request": _request(case, image),
"thresholds": asdict(STRICT_THRESHOLDS),
"request": request,
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
if case.thresholds.audio_relative_l2 > 0:
assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
if case.thresholds.audio_cosine_similarity > 0:
assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
@@ -0,0 +1,94 @@
from __future__ import annotations
from types import SimpleNamespace
import numpy as np
import pytest
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import (
NemotronVoiceChatDataPlaneContext,
NemotronVoiceChatDataPlaneSession,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_RUNTIME = {
"nvc_text_bos_id": 0,
"nvc_text_eos_id": 1,
"nvc_text_pad_id": 12,
"nvc_function_sotc_id": 20,
"nvc_function_eotc_id": 21,
"nvc_tokenizer_ref": "test-tokenizer",
}
def _projector(encode_audio=lambda *_: None) -> NemotronVoiceChatDataPlaneSession:
projector = NemotronVoiceChatDataPlaneSession(encode_audio)
projector.configure_runtime(_RUNTIME, tokenizer=SimpleNamespace(decode=lambda *_args, **_kwargs: "text"))
return projector
def _stage0_output(request_id: str = "req-0", *, text_token: int = 12, function_token: int | None = None) -> object:
metadata: dict[str, object] = {"nvc_text_token_ids": [text_token]}
if function_token is not None:
metadata["nvc_function_token"] = [function_token]
completion = SimpleNamespace(multimodal_output=metadata)
return SimpleNamespace(
stage_id=0,
request_output=SimpleNamespace(request_id=request_id, outputs=[completion]),
)
def test_function_channel_projects_completed_call_without_ending_speech() -> None:
projector = _projector()
projector._decode = lambda token_ids: (
'[{"name":"weather","arguments":{"city":"Shanghai"}}]' if token_ids == [99] else ""
)
context = NemotronVoiceChatDataPlaneContext(epoch=0)
events = []
for function_token in (20, 99, 21):
events.extend(projector.project_output(_stage0_output(function_token=function_token), context=context))
listen = [event for event in events if event.get("is_listen") is True]
function = [event for event in events if event.get("function_call") is True]
assert len(listen) == 3
assert len(function) == 1
assert function[0]["name"] == "weather"
assert function[0]["arguments"] == '{"city":"Shanghai"}'
projector._decode = lambda _token_ids: "not-json"
bad_events = []
for function_token in (20, 98, 21):
bad_events.extend(projector.project_output(_stage0_output(function_token=function_token)))
assert any(event.get("error_code") == "nemotron_function_call_parse_error" for event in bad_events)
def test_new_epoch_request_does_not_inherit_partial_function_or_frame_state() -> None:
projector = _projector(lambda *_: "audio")
old_request = duplex_resource_request_id(DuplexFence("sid", epoch=1), "stage0")
new_request = duplex_resource_request_id(DuplexFence("sid", epoch=2), "stage0")
projector.begin_request(old_request)
list(projector.project_output(_stage0_output(old_request, text_token=42, function_token=20)))
projector.begin_request(new_request)
# EOTC in the new epoch must not close the old epoch's partial function.
new_events = list(projector.project_output(_stage0_output(new_request, text_token=1, function_token=21)))
assert not [event for event in new_events if event.get("function_call") is True]
assert not [event for event in new_events if event.get("end_of_turn") is True]
audio = SimpleNamespace(
stage_id=2,
request_output=SimpleNamespace(
request_id=new_request,
outputs=[
SimpleNamespace(multimodal_output={"model_outputs": [np.ones(1764, dtype=np.float32)], "sr": [22050]})
],
),
)
audio_events = list(projector.project_output(audio))
assert sum(event.get("end_of_turn") is True for event in audio_events) == 1
@@ -0,0 +1,73 @@
from __future__ import annotations
import base64
import numpy as np
import pytest
from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import (
NemotronVoiceChatPcmAppendBuffer,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _payload(samples: np.ndarray, *, sample_rate_hz: int = 16000, fmt: str = "pcm_f32le") -> dict[str, object]:
return {
"type": "audio",
"audio": base64.b64encode(np.asarray(samples, dtype=np.float32).tobytes()).decode("ascii"),
"format": fmt,
"sample_rate_hz": sample_rate_hz,
}
def _append(buffer, samples, operation_id):
return buffer.prepare_append(
_payload(np.asarray(samples, dtype=np.float32)),
operation_id=operation_id,
chunk_period_ms=80,
allow_emit=True,
)
def test_irregular_browser_packets_emit_exact_ordered_80ms_frame() -> None:
buffer = NemotronVoiceChatPcmAppendBuffer()
assert _append(buffer, np.arange(1000), "packet-1") is None
reservation = _append(buffer, np.arange(1000, 1280), "packet-2")
assert reservation is not None
assert reservation.operation_id == "packet-2"
decoded = np.frombuffer(base64.b64decode(reservation.payload["audio"]), dtype=np.float32)
np.testing.assert_array_equal(decoded, np.arange(1280, dtype=np.float32))
assert buffer.pending_byte_count == 0
def test_append_rejects_invalid_frame_contract() -> None:
cases = (
(1281, 16000, "pcm_f32le", 0.0, "at most 1280"),
(1, 8000, "pcm_f32le", 0.0, "sample_rate_hz"),
(1, 16000, "audio/pcm", 0.0, "format"),
(1, 16000, "pcm_f32le", np.nan, "finite"),
)
for samples, sample_rate_hz, fmt, value, message in cases:
with pytest.raises(ValueError, match=message):
NemotronVoiceChatPcmAppendBuffer().prepare_append(
_payload(np.full(samples, value, dtype=np.float32), sample_rate_hz=sample_rate_hz, fmt=fmt),
operation_id="invalid",
chunk_period_ms=80,
allow_emit=True,
)
def test_terminal_commit_rejects_more_than_one_frame_tail() -> None:
buffer = NemotronVoiceChatPcmAppendBuffer()
for index, samples in enumerate((1000, 281)):
buffer.prepare_append(
_payload(np.zeros(samples, dtype=np.float32)),
operation_id=f"tail-{index}",
chunk_period_ms=80,
allow_emit=False,
)
with pytest.raises(ValueError, match="more than one 80 ms frame"):
buffer.prepare_commit(operation_id="commit", chunk_period_ms=80)
@@ -0,0 +1,115 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import base64
from types import SimpleNamespace
import numpy as np
import pytest
from vllm.sampling_params import SamplingParams
from vllm_omni.experimental.fullduplex.engine.contracts import DuplexInputMode
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
from vllm_omni.experimental.fullduplex.nemotron_voicechat.runtime import (
NemotronVoiceChatDuplexRuntimeExtension,
)
from vllm_omni.experimental.fullduplex.nemotron_voicechat.serving_adapter import (
NemotronVoiceChatServingRuntimeAdapter,
_render_tool_response,
)
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
NemotronVoiceChatThinkerForConditionalGeneration,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _frame() -> dict[str, object]:
raw = np.zeros(1280, dtype=np.float32).tobytes()
return {
"type": "audio",
"audio": base64.b64encode(raw).decode("ascii"),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
def _plan(extension, *, input_seq: int):
return extension.plan_append(
request_id="req",
fence=DuplexFence("sid", incarnation=2, epoch=3),
session_config={},
runtime_config={
"nvc_prompt_token_ids": [0, 42, 1],
"nvc_text_pad_id": 12,
"nvc_max_model_len": 8192,
},
seq=input_seq,
turn_seq=input_seq,
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
payload=_frame(),
final=False,
sampling_params=SamplingParams(),
)
def test_first_append_prefills_prompt_then_each_append_consumes_one_frame() -> None:
extension = NemotronVoiceChatDuplexRuntimeExtension()
params = extension.configure_sampling_params(
runtime_config={}, defaults=(SamplingParams(temperature=0.7, max_tokens=4),)
)[0]
assert params.temperature == 0.0 and params.max_tokens == 1 and params.top_p == 1.0 and params.top_k == 0
first = _plan(extension, input_seq=1)
later = _plan(extension, input_seq=2)
assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12]
assert later.prompt["prompt_token_ids"] == [12]
assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1
assert later.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2
def test_append_rejects_stage0_context_overflow() -> None:
with pytest.raises(ValueError, match="max_model_len"):
_plan(NemotronVoiceChatDuplexRuntimeExtension(), input_seq=8190)
def test_function_output_becomes_versioned_nvidia_channel_tokens() -> None:
encoded: list[str] = []
adapter = NemotronVoiceChatServingRuntimeAdapter(lambda *_: None)
adapter._tokenizer = SimpleNamespace(
encode=lambda text, **_kwargs: encoded.append(text) or [31, 32, 33],
)
first = adapter.runtime_config_for_function_output(
{"type": "function_call_output", "call_id": "call-1", "output": '{"result":20}'},
{},
)
second = adapter.runtime_config_for_function_output(
{"type": "function_call_output", "call_id": "call-2", "output": "plain text"},
first,
)
assert _render_tool_response('{"result":20}') == '<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>'
assert encoded == [
'<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>',
'<TOOL_RESPONSE>["plain text"]</TOOL_RESPONSE>',
]
assert first["nvc_function_response_generation"] == 1
assert second["nvc_function_response_generation"] == 2
assert second["nvc_function_response_token_ids"] == [31, 32, 33]
assert second["nvc_function_response_call_id"] == "call-2"
assert second["nvc_function_response_batches"] == [
{"generation": 1, "call_id": "call-1", "token_ids": [31, 32, 33]},
{"generation": 2, "call_id": "call-2", "token_ids": [31, 32, 33]},
]
session = {"function_response_generation": 0, "forced_function_tokens": []}
NemotronVoiceChatThinkerForConditionalGeneration._sync_forced_function_response(session, second)
assert session["function_response_generation"] == 2
assert session["forced_function_token"] == 31
assert session["forced_function_tokens"] == [31, 32, 33, 31, 32, 33]
@@ -18,6 +18,10 @@ from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
DuplexInputMode,
)
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
from vllm_omni.experimental.fullduplex.openai.runtime_adapter import (
ServingRuntimeConfigError,
)
from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
from vllm_omni.experimental.fullduplex.personaplex.data_plane import (
PersonaPlexDataPlaneContext,
PersonaPlexDataPlaneSession,
@@ -242,6 +246,53 @@ def test_personaplex_capabilities_are_honest() -> None:
assert multi.supports_barge_in is False
def test_personaplex_runtime_config_update_rejects_changed_persona() -> None:
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
short_persona = "You are a helpful assistant."
longer_persona = (
"You are now a swashbuckling pirate captain who speaks entirely in "
"nautical slang, references buried treasure constantly, and never "
"breaks character no matter what the user asks."
)
current = {"personaplex_persona": short_persona, "personaplex_voice_prompt": "NATF2.pt"}
with pytest.raises(ServingRuntimeConfigError, match="persona"):
adapter.runtime_config_for_update(
SimpleNamespace(instructions=longer_persona, extra_body={}),
current,
)
unchanged = adapter.runtime_config_for_update(
SimpleNamespace(instructions=short_persona, extra_body={}),
current,
)
assert unchanged["personaplex_persona"] == short_persona
def test_personaplex_runtime_config_update_rejects_changed_voice() -> None:
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
current = {"personaplex_persona": DEFAULT_PERSONA, "personaplex_voice_prompt": "NATF2.pt"}
with pytest.raises(ServingRuntimeConfigError, match="voice"):
adapter.runtime_config_for_update(
SimpleNamespace(instructions=None, voice="OtherVoice.pt", extra_body={}),
current,
)
unchanged = adapter.runtime_config_for_update(
SimpleNamespace(instructions=None, voice="NATF2.pt", extra_body={}),
current,
)
assert unchanged["personaplex_voice_prompt"] == "NATF2.pt"
def test_personaplex_rejects_minicpmo_only_opt_out_flag() -> None:
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
with pytest.raises(ServingRuntimeConfigError, match="minicpmo45_native_duplex"):
adapter.validate_client_extra_body({"minicpmo45_native_duplex": False})
def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None:
buffer = PersonaPlexPcmAppendBuffer()
reservation = buffer.prepare_append(
@@ -0,0 +1,194 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import argparse
import base64
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import pytest
from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample
from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge
from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli
from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys):
manifest = tmp_path / "manifest.json"
manifest.write_text(
json.dumps(
[
{
"id": "sample-1",
"split": "PR_correction",
"question_text": "What changed?",
"answer1": "The object moved.",
}
]
),
encoding="utf-8",
)
response_root = tmp_path / "responses"
score_root = tmp_path / "scores"
async def fake_generate(sample, *, output_root, **kwargs):
output = Path(output_root) / sample.split / f"{sample.id}.json"
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8")
output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
return output
class FakeJudge:
def __init__(self, *args, **kwargs):
pass
def chat(self, *args, **kwargs):
return '{"success_score": 1, "is_relevant": 1}'
monkeypatch.setattr(cli, "generate_sample", fake_generate)
monkeypatch.setattr(cli, "DuplexJudge", FakeJudge)
parser = argparse.ArgumentParser()
OmniDuplexEvalSubcommand.add_cli_args(parser)
common = ["--dataset", str(manifest), "--family", "pr"]
generate = parser.parse_args(
[
"generate",
*common,
"--model",
"mock",
"--ref-audio",
str(manifest),
"--response-root",
str(response_root),
"--concurrency",
"2",
]
)
OmniDuplexEvalSubcommand.cmd(generate)
evaluate = parser.parse_args(
[
"evaluate",
*common,
"--response-root",
str(response_root),
"--score-root",
str(score_root),
"--judge-model",
"mock-judge",
"--eval-workers",
"2",
]
)
OmniDuplexEvalSubcommand.cmd(evaluate)
OmniDuplexEvalSubcommand.cmd(parser.parse_args(["summarize", "--score-root", str(score_root)]))
summary = json.loads(capsys.readouterr().out)
assert summary["samples"] == 1
assert summary["pr"]["mean_all_success"] == 1.0
@pytest.mark.asyncio
async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
import websockets
received = []
async def handler(websocket):
response_started = False
async for raw in websocket:
event = json.loads(raw)
received.append(event)
if event["type"] == "session.update":
await websocket.send(json.dumps({"type": "session.created"}))
elif event["type"] == "input_audio_buffer.append" and not response_started:
response_started = True
await websocket.send(json.dumps({"type": "response.created", "response": {"id": "r1"}}))
await websocket.send(
json.dumps({"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."})
)
await websocket.send(
json.dumps(
{
"type": "response.audio.delta",
"response_id": "r1",
"delta": base64.b64encode(b"\0\0" * 240).decode(),
"sample_rate_hz": 24_000,
}
)
)
elif event["type"] == "input_audio_buffer.commit":
await websocket.send(json.dumps({"type": "response.done", "response": {"id": "r1"}}))
elif event["type"] == "session.close":
await websocket.send(json.dumps({"type": "session.closed"}))
return
monkeypatch.setattr(runner, "read_audio_pcm16", lambda path: b"\0\0" * (16_000 * 8 // 10))
monkeypatch.setattr(runner, "video_duration", lambda path: 0.8)
monkeypatch.setattr(runner, "iter_jpegs", lambda *args, **kwargs: iter([(0.0, b"jpeg")]))
audio = tmp_path / "question.wav"
video = tmp_path / "video.mp4"
ref = tmp_path / "ref.wav"
for path in (audio, video, ref):
path.write_bytes(b"data")
sample = DuplexSample("sample", "PR_correction", "pr", "correction", video, audio)
async with websockets.serve(handler, "127.0.0.1", 0) as server:
port = server.sockets[0].getsockname()[1]
output = await runner.generate_sample(
sample,
url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1",
model="mock",
ref_audio=ref,
output_root=tmp_path / "responses",
)
assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8"))
assert meta["response_done"] is True
assert meta["drain_timeout"] is None
assert any(event["type"] == "playback.ack" for event in received)
def test_judge_exercises_openai_http_schema():
requests = []
class Handler(BaseHTTPRequestHandler):
def do_POST(self):
length = int(self.headers["Content-Length"])
requests.append((self.path, self.headers["Authorization"], json.loads(self.rfile.read(length))))
payload = json.dumps({"choices": [{"message": {"content": '{"success_score": 1}'}}]}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, *args):
pass
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
judge = DuplexJudge(f"http://127.0.0.1:{server.server_port}", "judge", api_key="token")
assert judge.chat("prompt") == '{"success_score": 1}'
finally:
server.shutdown()
thread.join()
server.server_close()
path, authorization, payload = requests[0]
assert path == "/v1/chat/completions"
assert authorization == "Bearer token"
assert payload["model"] == "judge"
assert payload["messages"] == [{"role": "user", "content": "prompt"}]
@@ -76,6 +76,7 @@ def _send_inpaint_with_generator(
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient):
input_image, mask_image = _create_test_inputs()
@@ -97,6 +98,7 @@ def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient):
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_flux2_klein_inpaint_different_seeds_different_output(offline_client: OfflineOmniClient):
input_image, mask_image = _create_test_inputs()
@@ -0,0 +1,204 @@
"""
End-to-end test for MammothModa2 text-to-image generation.
Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
match a golden reference.
Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
"""
from __future__ import annotations
import json
import os
from pathlib import Path
import pytest
import torch
from huggingface_hub import snapshot_download
from vllm.sampling_params import SamplingParams
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
_IMAGE_TOKEN_ID = 151655 # "<|image_pad|>"
_VIDEO_TOKEN_ID = 151656 # "<|video_pad|>"
_VISION_START_TOKEN_ID = 151652 # "<|vision_start|>"
_VISION_END_TOKEN_ID = 151653 # "<|vision_end|>"
_AR_PATCH_SIZE = 16
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
# Fixed sampling coordinates: (channel, row_fraction, col_fraction)
# Covers corners, centre, and mid-edges across all 3 channels.
_PIXEL_SAMPLE_COORDS = [
(0, 0.0, 0.0),
(0, 0.5, 0.5),
(0, 1.0, 1.0),
(0, 0.25, 0.75),
(1, 0.0, 1.0),
(1, 0.5, 0.0),
(1, 0.75, 0.25),
(1, 1.0, 0.5),
(2, 0.0, 0.5),
(2, 0.5, 1.0),
(2, 0.75, 0.75),
(2, 1.0, 0.0),
]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _load_t2i_gen_config(repo_id: str) -> dict:
weights_dir = Path(snapshot_download(repo_id))
cfg_path = weights_dir / "t2i_generation_config.json"
if not cfg_path.exists():
pytest.skip(f"t2i_generation_config.json not found at {cfg_path}")
with cfg_path.open() as f:
return json.load(f)
def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
return (
"<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
"<|im_start|>assistant\n"
f"<|image start|>{ar_width}*{ar_height}<|image token|>"
)
def _sample_pixels(img_tensor: torch.Tensor) -> list[float]:
"""Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor."""
t = img_tensor.float().clamp(0.0, 1.0)
if t.ndim == 4:
t = t[0] # unbatch
C, H, W = t.shape
values = []
for c, rh, rw in _PIXEL_SAMPLE_COORDS:
ri = min(int(rh * (H - 1)), H - 1)
ci = min(int(rw * (W - 1)), W - 1)
values.append(round(float(t[c, ri, ci]), 6))
return values
# ---------------------------------------------------------------------------
# End-to-end test
# ---------------------------------------------------------------------------
pytestmark = [
pytest.mark.slow,
pytest.mark.diffusion,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3201")
@hardware_test(res={"cuda": "H100"})
def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner):
"""
End-to-end text-to-image generation with MammothModa2 (AR -> DiT).
Verifies:
- Omni pipeline initialises with the two-stage YAML config.
- DiT stage outputs an image tensor with the correct shape.
- A fixed set of pixel values matches a golden reference
(regenerate with ``UPDATE_GOLDEN=1``).
"""
gen_cfg = _load_t2i_gen_config(MODEL_PATH)
eol_token_id = int(gen_cfg["eol_token_id"])
visual_start = int(gen_cfg["visual_token_start_id"])
visual_end = int(gen_cfg["visual_token_end_id"])
height, width = 256, 256 # small for CI speed
ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE
expected_grid_tokens = ar_height * (ar_width + 1)
prompt_text = "A cat sitting on a laptop keyboard"
formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height)
omni = omni_runner.omni
ar_sampling = SamplingParams(
temperature=0.0,
top_k=1,
max_tokens=max(1, expected_grid_tokens + 1),
detokenize=False,
)
dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False)
outputs = list(
omni.generate(
[
{
"prompt": formatted_prompt,
"additional_information": {
"omni_task": ["t2i"],
"ar_width": [ar_width],
"ar_height": [ar_height],
"eol_token_id": [eol_token_id],
"visual_token_start_id": [visual_start],
"visual_token_end_id": [visual_end],
"image_height": [height],
"image_width": [width],
"num_inference_steps": [2],
"text_guidance_scale": [1.0],
"cfg_range": [0.0, 1.0],
"visual_ids": [
_IMAGE_TOKEN_ID,
_VIDEO_TOKEN_ID,
_VISION_START_TOKEN_ID,
_VISION_END_TOKEN_ID,
],
},
}
],
[ar_sampling, dit_sampling],
)
)
assert len(outputs) > 0, "Pipeline produced no outputs"
found_image = False
for out in outputs:
ro_list = out
if not isinstance(ro_list, list):
ro_list = [ro_list]
for ro in ro_list:
completion_outputs = getattr(ro, "outputs", None)
if not isinstance(completion_outputs, list):
continue
for completion in completion_outputs:
mm = getattr(completion, "multimodal_output", None)
if not (isinstance(mm, dict) and "image" in mm):
continue
img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]]
for img_tensor in img_list:
assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}"
assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D"
sampled = _sample_pixels(img_tensor)
if os.environ.get("UPDATE_GOLDEN"):
_GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True)
_GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2))
print(f"\nGolden file written to {_GOLDEN_T2I_PATH}")
elif _GOLDEN_T2I_PATH.exists():
golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"]
for i, (got, exp) in enumerate(zip(sampled, golden)):
assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}"
found_image = True
assert found_image, "No image tensor found in pipeline output"
@@ -0,0 +1,364 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import argparse
import asyncio
import base64
import json
import math
import uuid
import wave
from collections.abc import Sequence
from contextlib import asynccontextmanager, suppress
from pathlib import Path
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import numpy as np
from scipy.signal import resample_poly
from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, wait_for, write_pcm16_wav
INPUT_SAMPLE_RATE_HZ = 16_000
OUTPUT_SAMPLE_RATE_HZ = 22_050
FRAME_SAMPLES = 1_280
FRAME_PERIOD_S = FRAME_SAMPLES / INPUT_SAMPLE_RATE_HZ
DEFAULT_FUNCTION_TOOLS = [
{
"type": "function",
"name": "generate_random_number",
"description": "Generate a random integer between min and max (inclusive).",
"parameters": {
"type": "object",
"properties": {
"min": {"type": "integer", "description": "Minimum value (inclusive)"},
"max": {"type": "integer", "description": "Maximum value (inclusive)"},
},
"required": ["min", "max"],
},
}
]
DEFAULT_INSTRUCTIONS = "You are NVIDIA Voice Chat. Answer briefly. Start by greeting the user."
DEFAULT_FUNCTION_INSTRUCTIONS = (
"You are NVIDIA Voice Chat. If the user's request matches an available tool, "
"you MUST call that tool instead of answering from your own knowledge. "
"Use only argument values spoken by the user and never invent missing values."
)
def _url(base_url: str, model: str, session_id: str) -> str:
parts = urlsplit(base_url)
query = dict(parse_qsl(parts.query, keep_blank_values=True))
query.update(duplex="1", model=model, autostart="0", session_id=session_id)
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
def _read_wav(path: Path, *, input_channel: int = 0) -> np.ndarray:
with wave.open(str(path), "rb") as wav_file:
channels = wav_file.getnchannels()
width = wav_file.getsampwidth()
source_rate = wav_file.getframerate()
raw = wav_file.readframes(wav_file.getnframes())
dtypes = {1: np.uint8, 2: np.dtype("<i2"), 4: np.dtype("<i4")}
if width not in dtypes:
raise ValueError(f"unsupported input sample width: {width}")
pcm = np.frombuffer(raw, dtype=dtypes[width]).astype(np.float32)
pcm = (pcm - 128.0) / 128.0 if width == 1 else pcm / float(1 << (width * 8 - 1))
if not 0 <= input_channel < channels:
raise ValueError(f"input channel {input_channel} is outside WAV channel count {channels}")
if channels > 1:
pcm = pcm.reshape(-1, channels)[:, input_channel]
if source_rate != INPUT_SAMPLE_RATE_HZ:
divisor = math.gcd(source_rate, INPUT_SAMPLE_RATE_HZ)
pcm = resample_poly(pcm, up=INPUT_SAMPLE_RATE_HZ // divisor, down=source_rate // divisor)
return np.ascontiguousarray(pcm, dtype="<f4")
async def _stream(client: RealtimeDuplexClient, pcm: np.ndarray, *, max_frames: int | None, realtime: bool) -> int:
count = math.ceil(pcm.size / FRAME_SAMPLES)
if max_frames is not None:
count = min(count, max_frames)
for seq in range(count):
frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size)).astype("<f4")
await client.send(
{
"type": "input_audio_buffer.append",
"audio": base64.b64encode(frame).decode("ascii"),
"format": "pcm_f32le",
"sample_rate_hz": INPUT_SAMPLE_RATE_HZ,
"duration_ms": 80,
"audio_end_ms": (seq + 1) * 80,
}
)
if realtime:
await asyncio.sleep(FRAME_PERIOD_S)
return count
def _events(client: RealtimeDuplexClient, event_type: str) -> list[dict[str, object]]:
return [event for event in client.events.events if event.get("type") == event_type]
async def _return_function_output_when_ready(
client: RealtimeDuplexClient,
*,
output: str,
timeout_s: float,
) -> tuple[str, int]:
await wait_for(
lambda: bool(client.events.errors()) or client.events.count("response.function_call_arguments.done") > 0,
timeout_s=timeout_s,
label="function call to execute",
)
if client.events.errors():
raise AssertionError(f"function call failed before tool execution: {client.events.errors()}")
function_done = _events(client, "response.function_call_arguments.done")[-1]
call_id = function_done.get("call_id")
if not isinstance(call_id, str) or not call_id:
raise AssertionError(f"completed function call has no call_id: {function_done}")
event_count_before_output = len(client.events.events)
await client.send(
{
"type": "conversation.item.create",
"item": {
"type": "function_call_output",
"call_id": call_id,
"output": output,
},
}
)
return call_id, event_count_before_output
def _write_events(path: Path, client: RealtimeDuplexClient) -> None:
path.write_text(
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
encoding="utf-8",
)
@asynccontextmanager
async def _managed_client(client: RealtimeDuplexClient, *, timeout_s: float):
async with client:
try:
yield
finally:
if client.events.count("session.created") and not client.events.count("session.closed"):
with suppress(Exception):
await client.close_session(timeout_s=min(timeout_s, 30.0))
async def run(args: argparse.Namespace) -> dict[str, object]:
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
if args.expect_function_call and args.instructions == DEFAULT_INSTRUCTIONS:
instructions = DEFAULT_FUNCTION_INSTRUCTIONS
else:
instructions = args.instructions
tools = DEFAULT_FUNCTION_TOOLS if args.expect_function_call else None
session_id = f"nemotron-voicechat-{uuid.uuid4().hex}"
client = RealtimeDuplexClient(_url(args.url, args.model, session_id))
async with _managed_client(client, timeout_s=args.timeout_s):
session_payload: dict[str, object] = {
"session_id": session_id,
"model": args.model,
"modalities": ["audio", "text"],
"input_audio_format": "pcm_f32le",
"output_audio_format": "pcm16",
"instructions": instructions,
"idle_timeout_s": args.timeout_s,
"turn_detection": None,
"extra_body": {"auto_response": True},
}
if tools is not None:
session_payload["tools"] = tools
await client.send({"type": "session.update", "session": session_payload})
await wait_for(
lambda: client.events.count("session.created") > 0 or bool(client.events.errors()),
timeout_s=args.timeout_s,
label="session.created",
)
if client.events.errors():
raise AssertionError(f"session setup failed: {client.events.errors()}")
created = _events(client, "session.created")[-1]
session = created.get("session")
capabilities = session.get("capabilities") if isinstance(session, dict) else None
expected = dict(
implementation_level="model_native_duplex",
chunk_period_ms=80,
supports_core_resumable_request=True,
supports_core_kv_lease=False,
supports_multi_session=False,
)
if not isinstance(capabilities, dict) or any(capabilities.get(key) != value for key, value in expected.items()):
raise AssertionError(f"unexpected capabilities: {capabilities}")
function_output_task = (
asyncio.create_task(
_return_function_output_when_ready(
client,
output=args.function_output,
timeout_s=args.timeout_s,
)
)
if args.function_output is not None
else None
)
pcm = _read_wav(Path(args.input_wav), input_channel=args.input_channel)
frame_count = await _stream(client, pcm, max_frames=args.max_frames, realtime=not args.no_realtime)
completed_responses_at_commit = client.events.count("response.done")
await client.send({"type": "input_audio_buffer.commit", "final": True})
await wait_for(
lambda: (
bool(client.events.errors())
or (
client.events.count("response.function_call_arguments.done") > 0
if args.expect_function_call
else (
client.events.count("response.audio.delta") >= args.minimum_audio_chunks
and (client.events.count("response.done") > completed_responses_at_commit)
)
)
),
timeout_s=args.timeout_s,
label="model output",
)
await asyncio.sleep(args.drain_s)
if client.events.errors():
raise AssertionError(f"Realtime session emitted errors: {client.events.errors()}")
done_events = _events(client, "response.done")
if not args.expect_function_call:
response = done_events[-1].get("response") if done_events else None
status = response.get("status") if isinstance(response, dict) else None
if status != "completed":
raise AssertionError(f"response did not complete successfully: {done_events[-1:]}")
function_events = [
event for event in client.events.events if str(event.get("type", "")).startswith("response.function_call")
]
function_items = [
event
for event in _events(client, "response.output_item.done")
if isinstance(event.get("item"), dict) and event["item"].get("type") == "function_call"
]
if args.expect_function_call and not any(
event.get("type") == "response.function_call_arguments.done" for event in function_events
):
raise AssertionError(f"no completed function call: {function_events}")
if args.expect_function_call:
matching_items = [
event for event in function_items if event["item"].get("name") == args.expected_function_name
]
if not matching_items:
raise AssertionError(f"expected {args.expected_function_name!r}, got {function_items}")
function_item = matching_items[-1]["item"]
try:
function_arguments = json.loads(str(function_item.get("arguments", "")))
except json.JSONDecodeError as exc:
raise AssertionError(f"function arguments are not JSON: {function_item}") from exc
if args.expected_function_arguments is not None:
expected_arguments = json.loads(args.expected_function_arguments)
if function_arguments != expected_arguments:
raise AssertionError(
f"function arguments differ: expected={expected_arguments!r}, actual={function_arguments!r}"
)
if args.function_output is not None:
call_id = function_item.get("call_id")
if not isinstance(call_id, str) or not call_id:
raise AssertionError(f"function item has no call_id: {function_item}")
assert function_output_task is not None
returned_call_id, event_count_before_output = await function_output_task
assert returned_call_id == call_id
def tool_result_completed() -> bool:
later = client.events.events[event_count_before_output:]
transcript = "".join(
str(event.get("delta", ""))
for event in later
if event.get("type") == "response.audio_transcript.delta"
).lower()
return (
any(event.get("type") == "response.audio.delta" for event in later)
and any(event.get("type") == "response.done" for event in later)
and (args.expected_post_tool_text is None or args.expected_post_tool_text.lower() in transcript)
)
await wait_for(
lambda: bool(client.events.errors()) or tool_result_completed(),
timeout_s=args.timeout_s,
label="completed response after function output",
)
if client.events.errors():
raise AssertionError(f"function output failed: {client.events.errors()}")
await asyncio.sleep(args.drain_s)
audio = client.events.audio_bytes()
audio_events = _events(client, "response.audio.delta")
rates = {event.get("sample_rate_hz") for event in audio_events}
if not args.expect_function_call and audio and rates != {OUTPUT_SAMPLE_RATE_HZ}:
raise AssertionError(f"unexpected output sample rates: {rates}")
if not args.expect_function_call and args.minimum_audio_chunks and not audio:
raise AssertionError("model produced no audio")
expected_bytes = 2 * OUTPUT_SAMPLE_RATE_HZ * expected["chunk_period_ms"] // 1000
packet_sizes = [len(base64.b64decode(str(event.get("delta", "")), validate=True)) for event in audio_events]
if not args.expect_function_call and any(size != expected_bytes for size in packet_sizes):
raise AssertionError(f"audio deltas are not fixed 80 ms PCM16 packets: {packet_sizes}")
audio_pcm = np.frombuffer(audio, dtype="<i2").astype(np.float32) / 32768.0
audio_rms = float(np.sqrt(np.mean(np.square(audio_pcm)))) if audio_pcm.size else 0.0
if not args.expect_function_call and audio and audio_rms < args.minimum_audio_rms:
raise AssertionError(
f"model output RMS {audio_rms:.6f} is below {args.minimum_audio_rms:.6f}; "
"received packets contain only silence"
)
_write_events(output_dir / "events.jsonl", client)
if audio:
write_pcm16_wav(output_dir / "output.wav", audio, sample_rate_hz=OUTPUT_SAMPLE_RATE_HZ)
result = {
"ok": True,
"session_id": session_id,
"input_frames": frame_count,
"capabilities": capabilities,
"event_counts": {
event_type: client.events.count(event_type)
for event_type in sorted({str(event.get("type")) for event in client.events.events})
},
"audio_bytes": len(audio),
"audio_rms": audio_rms,
"function_events": function_events,
"function_items": function_items,
"output_dir": str(output_dir),
}
(output_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
return result
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="ws://127.0.0.1:8125/v1/realtime")
parser.add_argument("--model", required=True)
parser.add_argument("--input-wav", required=True)
parser.add_argument("--input-channel", type=int, default=0)
parser.add_argument("--output-dir", default="/tmp/nemotron-voicechat-duplex")
parser.add_argument("--instructions", default=DEFAULT_INSTRUCTIONS)
parser.add_argument("--max-frames", type=int)
parser.add_argument("--minimum-audio-chunks", type=int, default=1)
parser.add_argument("--minimum-audio-rms", type=float, default=1e-4)
parser.add_argument("--expect-function-call", action="store_true")
parser.add_argument("--expected-function-name", default="generate_random_number")
parser.add_argument("--expected-function-arguments")
parser.add_argument("--function-output")
parser.add_argument("--expected-post-tool-text")
parser.add_argument("--no-realtime", action="store_true")
parser.add_argument("--drain-s", type=float, default=2.0)
parser.add_argument("--timeout-s", type=float, default=600.0)
return parser.parse_args(argv)
def main() -> None:
print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
@@ -52,9 +52,8 @@ def get_prompt(prompt_type="zh"):
def _set_rocm_request_timeout(request_config: dict) -> None:
if current_omni_platform.is_rocm():
# A cold AITER MHA kernel build took 301 seconds on MI300X.
# Allow 600 seconds for the first request; later requests reuse the cache.
request_config["timeout"] = 600.0
# Cold AITER MHA builds on loaded MI300 agents can exceed 20 minutes.
request_config["timeout"] = 1800.0
tts_server_params = [
@@ -0,0 +1,128 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax-H3 DLO + DP2 online-serving smoke test."""
from __future__ import annotations
import concurrent.futures
import io
import json
import os
import av
import pytest
import requests
from tests.helpers.assertions import assert_video_valid
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = os.environ.get("VLLM_TEST_MINIMAX_H3_MODEL", "MiniMaxAI/MiniMax-H3")
WIDTH = 1344
HEIGHT = 768
FPS = 24
NUM_INFERENCE_STEPS = 4
REQUEST_TIMEOUT_SECONDS = 1800
H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
SERVER_ARGS = [
"--trust-remote-code",
"--task-type",
"fl2va",
"--num-gpus",
"2",
"--tensor-parallel-size",
"1",
"--data-parallel-size",
"2",
"--request-batch-max-wait-ms",
"500",
"--usp",
"1",
"--ring",
"1",
"--text-encoder-tp-size",
"1",
"--vae-patch-parallel-size",
"1",
"--vae-parallel-mode",
"tile",
"--vae-use-tiling",
"--enable-distributed-layerwise-offload",
]
def _assert_audio_stream_present(video: bytes) -> None:
"""Assert that the generated MP4 contains decodable audio samples."""
with av.open(io.BytesIO(video)) as container:
audio_streams = [stream for stream in container.streams if stream.type == "audio"]
assert audio_streams, "MiniMax-H3 MP4 has no audio stream"
audio_frame = next(container.decode(audio=0), None)
assert audio_frame is not None and audio_frame.samples > 0, "MiniMax-H3 MP4 audio stream is empty"
def _run_t2va_request(client: OpenAIClientHandler, seed: int) -> bytes:
"""Submit one synchronous T2VA request and return its MP4 body."""
request_data = {
"model": MODEL,
"prompt": "In a snowy blue-purple forest, a traveler walks past a sleeping giant; footsteps crunch in the snow while the creature softly breathes.",
"width": str(WIDTH),
"height": str(HEIGHT),
"fps": str(FPS),
"num_inference_steps": str(NUM_INFERENCE_STEPS),
"flow_shift": "12",
"seed": str(seed),
"extra_params": json.dumps(
{
"task": "t2va",
"duration": 4.0,
"aspect_ratio": "16:9",
"audio_flow_shift": 3.0,
},
separators=(",", ":"),
),
}
response = requests.post(
f"{client.base_url.rstrip('/')}/v1/videos/sync",
data=request_data,
headers={"Accept": "video/mp4"},
timeout=REQUEST_TIMEOUT_SECONDS,
)
response.raise_for_status()
assert response.headers.get("content-type", "").startswith("video/mp4")
assert response.content, "MiniMax-H3 returned an empty video body"
return response.content
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.slow
@pytest.mark.parametrize(
"omni_server",
[
pytest.param(
OmniServerParams(
model=MODEL,
server_args=SERVER_ARGS,
stage_init_timeout=1800,
init_timeout=1800,
),
id="minimax_h3_dlo_dp2_t2va",
marks=H100_TWO_CARD_MARKS,
)
],
indirect=True,
)
def test_minimax_h3_dlo_dp2_t2va(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
"""Validate one complete DLO all-gather DP2 wave with two T2VA jobs."""
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
futures = [executor.submit(_run_t2va_request, openai_client, seed) for seed in (1101, 1102)]
videos = [future.result() for future in futures]
for video in videos:
assert_video_valid(video, width=WIDTH, height=HEIGHT, fps=FPS)
_assert_audio_stream_present(video)
@@ -0,0 +1,56 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import asyncio
import os
from pathlib import Path
import pytest
from huggingface_hub import snapshot_download
from tests.e2e.online_serving.nemotron_voicechat_realtime_duplex import parse_args, run
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServerParams, get_model_prefix
from tests.helpers.stage_config import get_deploy_config_path
MODEL = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_MODEL", "nvidia/NVIDIA-NemotronLabs-VoiceChat-11B")
DEPLOY_CONFIG = get_deploy_config_path("nemotron_labs_voicechat_duplex.yaml")
TOKENIZER = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_LLM_PATH")
pytestmark = [pytest.mark.full_model, pytest.mark.omni]
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize(
"omni_server",
[
pytest.param(
OmniServerParams(
model=MODEL,
stage_config_path=DEPLOY_CONFIG,
env_dict={"NEMOTRON_VOICECHAT_LLM_PATH": TOKENIZER} if TOKENIZER else None,
),
id="native-duplex",
)
],
indirect=True,
)
def test_native_duplex_turn_taking_streams_model_audio(omni_server, tmp_path: Path) -> None:
model_prefix = get_model_prefix()
root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
if not root.is_dir():
root = Path(snapshot_download(MODEL, local_files_only=True))
args = parse_args(
(
f"--url ws://{omni_server.host}:{omni_server.port}/v1/realtime --model {omni_server.model} "
f"--input-wav {root / 'turn_taking.wav'} --input-channel 0 --max-frames 190 "
f"--minimum-audio-chunks 48 --minimum-audio-rms 0.001 --no-realtime --timeout-s 300 "
f"--output-dir {tmp_path / 'native_duplex'}"
).split()
)
result = asyncio.run(run(args))
assert result["ok"] is True
assert result["input_frames"] == 190
assert result["event_counts"]["response.speak"] > 0
assert result["audio_bytes"] >= result["input_frames"] * 1764 * 2 // 4
@@ -452,6 +452,10 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
"env": {runtime_env_var: "stage-value"},
}
monkeypatch.setenv(runtime_env_var, "parent-value")
# Pin the logical->physical device mapping so the assertion below is
# deterministic regardless of the ambient CUDA_VISIBLE_DEVICES (which on a
# rebase host can be e.g. "3,4", mapping logical device 0 -> physical 3).
runtime._init_visible_devices_baseline = "0"
captured: dict[str, str | None] = {}
+164
View File
@@ -17,6 +17,7 @@ import janus
import pytest
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason
from vllm.v1.engine.exceptions import EngineDeadError
from vllm.v1.metrics.stats import IterationStats
@@ -106,6 +107,13 @@ class OrchestratorFixture:
result_future: concurrent.futures.Future[None]
@dataclass
class FakePromptRequest:
request_id: str
prompt_token_ids: list[int]
resumable: bool = True
class FakeStageClient:
def __init__(
self,
@@ -291,6 +299,20 @@ def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None)
def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs:
return EngineCoreOutputs(
outputs=[
EngineCoreOutput(
request_id=request_id,
new_token_ids=[],
finish_reason=FinishReason.STOP,
)
],
timestamp=1.0,
finished_requests={request_id},
)
def _build_request_output(
request_id: str,
*,
@@ -764,6 +786,148 @@ async def test_run_async_chunk(orchestrator_factory) -> None:
await _shutdown_orchestrator(orchestrator_fixture)
@pytest.mark.asyncio
async def test_async_chunk_raw_terminal_without_processed_output_finishes_request(orchestrator_factory) -> None:
stage0 = FakeStageClient(stage_type="llm", final_output=False)
stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio")
orchestrator_fixture = orchestrator_factory(
[stage0, stage1],
output_processors=[FakeOutputProcessor(), FakeOutputProcessor()],
async_chunk=True,
)
request = FakePromptRequest(
request_id="req-stream-terminal",
prompt_token_ids=[1, 2, 3, 4],
)
try:
await _enqueue_add_request(
orchestrator_fixture,
request_id=request.request_id,
prompt=request,
original_prompt={"prompt": "stream audio"},
sampling_params_list=[_sampling_params(), _sampling_params()],
final_stage_id=1,
)
await _wait_for(lambda: len(stage1.add_request_calls) == 1)
stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request.request_id))
output_msg = await _get_output_message(orchestrator_fixture)
assert output_msg.request_id == request.request_id
assert output_msg.stage_id == 1
assert output_msg.finished is True
assert output_msg.engine_outputs.finished is True
assert output_msg.engine_outputs.outputs[0].multimodal_output["audio"].numel() == 0
await _wait_for(lambda: request.request_id not in orchestrator_fixture.orchestrator.request_states)
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@pytest.mark.asyncio
async def test_async_chunk_data_then_raw_terminal_finishes_once(orchestrator_factory) -> None:
request_id = "req-stream-data-terminal"
stage0 = FakeStageClient(stage_type="llm", final_output=False)
stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio")
final_data = _build_request_output(
request_id,
token_ids=[20, 21],
finished=False,
text="last audio chunk",
)
orchestrator_fixture = orchestrator_factory(
[stage0, stage1],
output_processors=[
FakeOutputProcessor(),
FakeOutputProcessor(request_outputs=[final_data]),
],
async_chunk=True,
)
request = FakePromptRequest(
request_id=request_id,
prompt_token_ids=[1, 2, 3, 4],
)
try:
await _enqueue_add_request(
orchestrator_fixture,
request_id=request_id,
prompt=request,
original_prompt={"prompt": "stream audio"},
sampling_params_list=[_sampling_params(), _sampling_params()],
final_stage_id=1,
)
await _wait_for(lambda: len(stage1.add_request_calls) == 1)
stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request_id))
data_msg = await _get_output_message(orchestrator_fixture)
terminal_msg = await _get_output_message(orchestrator_fixture)
assert data_msg.engine_outputs is final_data
assert data_msg.finished is False
assert terminal_msg.request_id == request_id
assert terminal_msg.stage_id == 1
assert terminal_msg.finished is True
assert terminal_msg.engine_outputs.finished is True
await _wait_for(lambda: request_id not in orchestrator_fixture.orchestrator.request_states)
await asyncio.sleep(0.05)
with pytest.raises(queue.Empty):
orchestrator_fixture.output_sync_q.get_nowait()
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@pytest.mark.asyncio
async def test_async_chunk_processed_terminal_and_raw_terminal_finishes_once(orchestrator_factory) -> None:
request_id = "req-stream-processed-terminal"
stage0 = FakeStageClient(stage_type="llm", final_output=False)
stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio")
processed_terminal = _build_request_output(
request_id,
token_ids=[20, 21],
finished=True,
text="last audio chunk",
)
orchestrator_fixture = orchestrator_factory(
[stage0, stage1],
output_processors=[
FakeOutputProcessor(),
FakeOutputProcessor(request_outputs=[processed_terminal]),
],
async_chunk=True,
)
request = FakePromptRequest(
request_id=request_id,
prompt_token_ids=[1, 2, 3, 4],
)
try:
await _enqueue_add_request(
orchestrator_fixture,
request_id=request_id,
prompt=request,
original_prompt={"prompt": "stream audio"},
sampling_params_list=[_sampling_params(), _sampling_params()],
final_stage_id=1,
)
await _wait_for(lambda: len(stage1.add_request_calls) == 1)
stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request_id))
terminal_msg = await _get_output_message(orchestrator_fixture)
assert terminal_msg.engine_outputs is processed_terminal
assert terminal_msg.finished is True
await _wait_for(lambda: request_id not in orchestrator_fixture.orchestrator.request_states)
await asyncio.sleep(0.05)
with pytest.raises(queue.Empty):
orchestrator_fixture.output_sync_q.get_nowait()
finally:
await _shutdown_orchestrator(orchestrator_fixture)
@pytest.mark.asyncio
async def test_run_shutdown(orchestrator_factory) -> None:
stages = [
@@ -0,0 +1,247 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Event-driven orchestration loop (``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``) tests.
Parity suite: re-runs the legacy orchestration scenarios from
``test_orchestrator.py`` / ``test_orchestrator_error_handling.py`` with the
event-driven loop selected, so both loops are held to the same behavior. Plus
event-driven-specific coverage: reader reconcile on client swap, the blocking
final-output drain, and flag parsing.
"""
from __future__ import annotations
import asyncio
from types import SimpleNamespace
import janus
import pytest
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.messages import ShutdownRequestMessage
from vllm_omni.engine.orchestrator import _event_driven_orch_enabled
from . import test_orchestrator as legacy
from . import test_orchestrator_error_handling as legacy_errors
from .test_orchestrator import (
FakeOutputProcessor,
FakeStageClient,
OrchestratorFixture,
_build_harness,
_build_request_output,
_engine_core_outputs,
_enqueue_add_request,
_get_output_message,
_sampling_params,
_shutdown_orchestrator,
_wait_for,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture
def orchestrator_factory(monkeypatch):
"""Flag-setting clone of the legacy harness fixture.
Sets ``VLLM_OMNI_EVENT_DRIVEN_ORCH=1`` before any Orchestrator is
constructed and asserts the flag actually took effect, so the parity tests
cannot silently exercise the legacy poll loop.
"""
monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", "1")
fixtures: list[OrchestratorFixture] = []
def _factory(*args, **kwargs) -> OrchestratorFixture:
fixture = _build_harness(*args, **kwargs)
assert fixture.orchestrator._event_driven_orch is True
fixtures.append(fixture)
return fixture
yield _factory
for fixture in fixtures:
if fixture.thread.is_alive():
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
fixture.thread.join(timeout=5)
for q in fixture.queues:
q.close()
# ---------------------------------------------------------------------------
# Parity: the legacy scenario matrix, re-run through the event-driven loop
# ---------------------------------------------------------------------------
_PARITY_TESTS = [
legacy.test_run_two_stage_llm,
legacy.test_run_single_stage_diffusion,
legacy.test_run_single_stage_diffusion_streaming_forwards_intermediate_chunks,
legacy.test_run_llm_to_diffusion,
legacy.test_run_async_chunk,
legacy.test_run_shutdown,
legacy.test_run_abort,
legacy.test_multi_replica_round_robin_distribution,
legacy.test_multi_replica_abort_broadcasts_to_all_replicas,
legacy.test_multi_replica_shutdown_all_replicas,
legacy.test_multi_replica_cfg_companion_inherits_parent_affinity,
# Stats plumbing. The scheduler-stats cases matter most here: a batch with
# no request outputs still carries SchedulerStats on throttled ticks, and a
# reader that drops every output-less batch silently stops reporting
# KV/queue gauges under the event-driven loop.
legacy.test_orchestrator_records_iteration_stats_without_scheduler_stats,
legacy.test_orchestrator_records_scheduler_stats_without_outputs,
legacy.test_orchestrator_does_not_build_iteration_stats_for_finished_only_batch,
legacy.test_orchestrator_does_not_build_iteration_stats_without_stat_logger,
# Per-replica fault isolation (#4285): a dead replica must be evicted and
# the server kept up, on both the LLM reader path and the diffusion poller.
legacy_errors.test_engine_dead_error_evicts_replica_and_keeps_running,
legacy_errors.test_engine_dead_error_fails_only_dead_replica_requests,
legacy_errors.test_forward_to_dead_downstream_stage_fails_request_not_server,
legacy_errors.test_add_request_to_dead_stage_fails_request_not_server,
legacy_errors.test_diffusion_replica_death_on_poll_keeps_server,
legacy_errors.test_diffusion_error_output_routed_as_finished,
legacy_errors.test_diffusion_client_error_output_propagates_status_code,
]
@pytest.mark.asyncio
@pytest.mark.parametrize("legacy_test", _PARITY_TESTS, ids=lambda f: f.__name__)
async def test_event_driven_parity(legacy_test, orchestrator_factory) -> None:
await legacy_test(orchestrator_factory)
# ---------------------------------------------------------------------------
# Event-driven-specific behavior
# ---------------------------------------------------------------------------
def test_flag_parsing(monkeypatch) -> None:
monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False)
assert _event_driven_orch_enabled() is False
for value in ("1", "true", "True", "YES", "on"):
monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", value)
assert _event_driven_orch_enabled() is True
for value in ("0", "false", "off", ""):
monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", value)
assert _event_driven_orch_enabled() is False
def test_default_is_legacy_loop(monkeypatch) -> None:
"""Without the env flag, the harness runs the legacy poll loop."""
monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False)
fixture = _build_harness([FakeStageClient(stage_type="llm", final_output=True)])
try:
assert fixture.orchestrator._event_driven_orch is False
finally:
fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
fixture.thread.join(timeout=5)
for q in fixture.queues:
q.close()
@pytest.mark.asyncio
async def test_reader_reconcile_picks_up_swapped_client(orchestrator_factory) -> None:
"""Outputs from a replica whose client object was replaced still flow.
The event-driven loop binds one reader task per client object; the
periodic reconcile must respawn the reader when ``pool.clients[replica]``
is swapped (replica replacement), otherwise the new client's outputs
would never be drained.
"""
stage0 = FakeStageClient(stage_type="llm", final_output=True)
processor = FakeOutputProcessor(request_outputs=[_build_request_output("req-swap", token_ids=[3], finished=True)])
orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor])
request = SimpleNamespace(request_id="req-swap", prompt_token_ids=[1, 2])
try:
await _enqueue_add_request(
orchestrator_fixture,
request_id="req-swap",
prompt=request,
original_prompt={"prompt": "swap"},
sampling_params_list=[_sampling_params()],
final_stage_id=0,
)
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
# Swap in a fresh client for the same replica slot; keep the pool's
# other wiring intact. The reconcile tick (0.5 s) must respawn the
# reader bound to the new client object.
pool = orchestrator_fixture.orchestrator.stage_pools[0]
replacement = FakeStageClient(stage_type="llm", final_output=True)
replacement.stage_id = stage0.stage_id
replacement.replica_id = stage0.replica_id
pool.clients[0] = replacement
replacement.push_engine_core_outputs(_engine_core_outputs("swapped-raw", 1.0))
output_msg = await _get_output_message(orchestrator_fixture, timeout=5.0)
assert output_msg.request_id == "req-swap"
assert output_msg.finished is True
finally:
await _shutdown_orchestrator(orchestrator_fixture)
# ---------------------------------------------------------------------------
# Blocking final-output drain (AsyncOmniEngine.get_output_blocking_async)
# ---------------------------------------------------------------------------
def _drain_engine(alive: bool = True) -> AsyncOmniEngine:
engine = object.__new__(AsyncOmniEngine)
engine.output_queue = janus.Queue()
engine.orchestrator_thread = SimpleNamespace(is_alive=lambda: alive)
return engine
def _drain_cleanup(engine: AsyncOmniEngine) -> None:
if engine._output_drain_executor is not None:
engine._output_drain_executor.shutdown(wait=False)
engine._output_drain_executor = None
engine.output_queue.close()
@pytest.mark.asyncio
async def test_blocking_drain_returns_queued_message() -> None:
engine = _drain_engine()
try:
engine.output_queue.sync_q.put_nowait("msg-1")
assert await engine.get_output_blocking_async(timeout=1.0) == "msg-1"
finally:
_drain_cleanup(engine)
@pytest.mark.asyncio
async def test_blocking_drain_wakes_on_late_message() -> None:
"""A message put after the wait starts wakes the drain, no polling."""
engine = _drain_engine()
try:
async def _delayed_put() -> None:
await asyncio.sleep(0.05)
engine.output_queue.sync_q.put_nowait("late-msg")
put_task = asyncio.create_task(_delayed_put())
msg = await engine.get_output_blocking_async(timeout=5.0)
await put_task
assert msg == "late-msg"
finally:
_drain_cleanup(engine)
@pytest.mark.asyncio
async def test_blocking_drain_timeout_returns_none_when_alive() -> None:
engine = _drain_engine(alive=True)
try:
assert await engine.get_output_blocking_async(timeout=0.05) is None
finally:
_drain_cleanup(engine)
@pytest.mark.asyncio
async def test_blocking_drain_raises_when_orchestrator_dead() -> None:
engine = _drain_engine(alive=False)
try:
with pytest.raises(RuntimeError, match="Orchestrator died"):
await engine.get_output_blocking_async(timeout=0.05)
finally:
_drain_cleanup(engine)
+44 -7
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
from types import SimpleNamespace
@@ -581,18 +581,18 @@ def test_mm_only_terminal_finish_removes_request_state(monkeypatch):
def test_no_detokenizer_make_request_output_with_routed_experts():
"""make_request_output accepts the routed_experts arg that the multimodal
output channel (_process_mm_only_outputs) passes positionally, and attaches
it to the completion output.
output channel (_process_mm_only_outputs) passes, and attaches it to the
completion output.
Regression: the call site passes 6 positional args
(..., kv_transfer_params, routed_experts) but the override previously took
only 5, raising TypeError on every generation-stage output.
routed_experts is omni-specific keyword-only: the 6th positional now
matches upstream's ec_transfer_params so that super().process_outputs()
calls do not misroute it.
"""
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
routed_experts = np.zeros((2, 3), dtype=np.int32)
# Mirror the exact call shape of _process_mm_only_outputs.
result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts)
result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts=routed_experts)
assert result is not None
assert not isinstance(result, PoolingRequestOutput)
assert result.outputs[0].routed_experts is routed_experts
@@ -663,3 +663,40 @@ def test_mm_only_outputs_update_iteration_stats():
assert finished.finish_reason == FinishReason.STOP
assert finished.num_prompt_tokens == state.prompt_len
assert finished.num_generation_tokens == 2
def test_ec_transfer_params_survive_both_construction_paths():
"""v0.28 contract: encoder-cache transfer metadata must reach
RequestOutput.ec_transfer_params through BOTH omni construction paths —
the upstream-delegating one (logprobs processor present) and the
no-detokenizer generation-stage one (direct RequestOutput build).
Regression: the override accepted the argument but dropped it."""
ec = {"remote_handle": "enc-cache-1"}
state = _make_state(RequestOutputKind.CUMULATIVE)
out = state.make_request_output([1], None, None, None, {"kv": 1}, ec)
assert out is not None and out.ec_transfer_params == ec
kwargs = dict(_DEFAULT_STATE_KWARGS)
kwargs.update(logprobs_processor=None, detokenizer=None)
gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
completion = gen_state._new_completion_output([1], None, None)
direct = gen_state._new_request_output("r", [completion], False, {"kv": 1}, ec)
assert direct.ec_transfer_params == ec
def test_num_cache_creation_tokens_reaches_direct_request_output():
"""v0.28 contract: prefix-cache creation usage must reach
RequestOutput.num_cache_creation_tokens through the no-detokenizer
direct build, which previously passed only num_cached_tokens."""
kwargs = dict(_DEFAULT_STATE_KWARGS)
kwargs.update(logprobs_processor=None, detokenizer=None)
gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
gen_state.num_cached_tokens = 3
gen_state.num_cache_creation_tokens = 2
completion = gen_state._new_completion_output([1], None, None)
direct = gen_state._new_request_output("r", [completion], False, None, None)
assert direct.num_cached_tokens == 3
assert direct.num_cache_creation_tokens == 2
+4 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for AsyncOmniEngine single-stage mode and OmniMasterServer."""
from __future__ import annotations
@@ -1374,7 +1377,7 @@ class TestConnectRemoteEngineCoresCoordinator:
omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2)
omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2)
mock_wait.assert_called_once()
_, _, core_engines, parallel_config, *_ = mock_wait.call_args.args
_, core_engines, parallel_config, *_ = mock_wait.call_args.args
assert core_engines[0].local is False
assert parallel_config.data_parallel_size_local == 0
+229
View File
@@ -0,0 +1,229 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for issue #5003: a per-stage world size that its assigned
``devices`` cannot satisfy must fail early in ``build_vllm_config`` with a clear
message, rather than surfacing as an opaque worker-side ``local rank ... out of
bounds`` assertion.
Root cause: a top-level ``--tensor-parallel-size`` is broadcast to every stage,
but each stage's ``devices`` is not adjusted, so a stage can end up with e.g.
tensor_parallel_size=4 while still holding a single-GPU deploy default. Without
``--strategy-config`` the strategy-path device check never runs.
"""
import json
import re
import types
from unittest import mock
import pytest
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.stage_init_utils import (
_check_stage_device_layout,
build_vllm_config,
compute_replica_layout,
get_stage_devices_per_replica,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _stage(stage_id, devices, num_replicas=1, engine_args=None):
return types.SimpleNamespace(
stage_id=stage_id,
stage_type="llm",
engine_args=engine_args or {},
runtime=types.SimpleNamespace(devices=devices, num_replicas=num_replicas),
)
def test_tp_broadcast_without_devices_fails_early():
"""stage0 gets tensor_parallel_size=4 (broadcast) but only 1 device -> clear error."""
stage = _stage(0, devices="0")
engine_args = {
"tensor_parallel_size": 4,
"data_parallel_size": 1,
"pipeline_parallel_size": 1,
}
with pytest.raises(ValueError) as excinfo:
_check_stage_device_layout(stage, engine_args)
msg = str(excinfo.value)
# Message names the stage, the mismatch, and the actionable workaround.
assert "Stage 0" in msg
assert "1 device" in msg and "4" in msg
assert "--stage-overrides" in msg
def test_workaround_example_sets_tp_on_every_stage():
"""The JSON example in the error message must not reproduce the bug: every
stage entry it suggests has to set ``tensor_parallel_size`` (a single-GPU
stage that only overrides ``devices`` would still inherit the broadcast
tp and crash again — the original #5003 failure mode)."""
stage = _stage(0, devices="0")
with pytest.raises(ValueError) as excinfo:
_check_stage_device_layout(stage, {"tensor_parallel_size": 4})
msg = str(excinfo.value)
match = re.search(r"\{.*\}", msg, re.DOTALL)
assert match, f"no JSON example found in error message: {msg}"
example = json.loads(match.group(0))
assert example, "example must contain at least one stage"
for stage_id, overrides in example.items():
assert "tensor_parallel_size" in overrides, (
f"stage {stage_id} in the suggested workaround omits "
f"tensor_parallel_size, which reproduces #5003: {overrides}"
)
def test_consistent_tp_and_devices_pass():
"""tensor_parallel_size=4 with 4 assigned devices is valid."""
stage = _stage(0, devices="0,1,2,3")
_check_stage_device_layout(
stage,
{"tensor_parallel_size": 4, "data_parallel_size": 1, "pipeline_parallel_size": 1},
)
def test_single_gpu_stage_passes():
"""A TP=1 stage on a single device (talker/code2wav default) is valid."""
stage = _stage(1, devices="1")
_check_stage_device_layout(stage, {"tensor_parallel_size": 1})
def test_missing_devices_is_skipped():
"""No explicit devices -> vLLM assigns them; nothing to validate here."""
stage = _stage(0, devices=None)
_check_stage_device_layout(stage, {"tensor_parallel_size": 4})
def test_replica_pool_layout_passes():
"""Pool mode: num_replicas=2 x tp=2 => 4 devices is a valid pool shape."""
stage = _stage(0, devices="0,1,2,3", num_replicas=2)
_check_stage_device_layout(stage, {"tensor_parallel_size": 2})
def test_multinode_dp_uses_local_width_for_device_validation():
"""One local DP engine needs one device even when global DP spans four nodes."""
stage = _stage(0, devices="0")
_check_stage_device_layout(
stage,
{
"tensor_parallel_size": 1,
"data_parallel_size": 4,
"data_parallel_size_local": 1,
"pipeline_parallel_size": 1,
},
)
def test_zero_local_dp_skips_local_device_validation():
"""A head process with no local engines does not own a local device layout."""
stage = _stage(0, devices="0")
_check_stage_device_layout(
stage,
{
"tensor_parallel_size": 4,
"data_parallel_size": 4,
"data_parallel_size_local": 0,
"pipeline_parallel_size": 1,
},
)
@pytest.mark.parametrize(
("devices", "expected"),
[
("0,1", ["0,1", "2,3"]),
("0,1,2,3", ["0,1", "2,3"]),
],
)
def test_replica_guard_and_splitter_share_local_world_size(devices, expected):
"""PP and local DP contribute to the same per-replica size in both paths."""
engine_args = {
"tensor_parallel_size": 1,
"data_parallel_size": 4,
"data_parallel_size_local": 1,
"pipeline_parallel_size": 2,
}
stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args)
_check_stage_device_layout(stage, engine_args)
assert get_stage_devices_per_replica(stage) == 2
_, replica_devices_map = compute_replica_layout([stage])
assert replica_devices_map == {0: expected}
def test_non_tp_layout_error_uses_generic_guidance():
"""A PP mismatch must not be reported as a top-level TP broadcast."""
stage = _stage(0, devices="0")
with pytest.raises(ValueError) as excinfo:
_check_stage_device_layout(
stage,
{
"tensor_parallel_size": 1,
"data_parallel_size": 1,
"pipeline_parallel_size": 2,
},
)
msg = str(excinfo.value)
assert "local world size" in msg
assert "top-level --tensor-parallel-size" not in msg
def test_build_vllm_config_fails_before_engine_config_on_mismatch():
"""Pin the production wiring: ``build_vllm_config`` must run the device-layout
guard *before* ``create_engine_config``. If the guard call is dropped or moved
after engine-config creation, #5003 regresses while the unit tests above stay
green — so assert here that a mismatched stage raises and that neither
``create_engine_config`` nor ``Executor.get_class`` is reached."""
stage = _stage(0, devices="0")
with (
mock.patch.object(stage_init_utils.OmniEngineArgs, "create_engine_config") as create_engine_config,
mock.patch.object(stage_init_utils.Executor, "get_class") as get_class,
):
with pytest.raises(ValueError) as excinfo:
build_vllm_config(
stage,
model="dummy-model",
engine_args_dict={
"tensor_parallel_size": 4,
"data_parallel_size": 1,
"pipeline_parallel_size": 1,
},
)
msg = str(excinfo.value)
assert "Stage 0" in msg
assert "--stage-overrides" in msg
# Guard fired early: worker/config construction was never reached.
create_engine_config.assert_not_called()
get_class.assert_not_called()
def test_build_vllm_config_proceeds_on_consistent_layout():
"""The guard must not false-positive: a consistent single-GPU stage flows past
it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual."""
stage = _stage(1, devices="1")
fake_config = types.SimpleNamespace(
quant_config=None,
model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()),
)
sentinel_executor = object()
with (
mock.patch.object(
stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config
) as create_engine_config,
mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor),
mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant),
):
vllm_config, executor_class = build_vllm_config(
stage,
model="dummy-model",
engine_args_dict={"tensor_parallel_size": 1},
)
create_engine_config.assert_called_once()
assert vllm_config is fake_config
assert executor_class is sentinel_executor
@@ -49,6 +49,9 @@ from vllm_omni.experimental.fullduplex.openai.serving import (
)
from vllm_omni.experimental.fullduplex.openai.websocket import DuplexWebSocketActor
from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision
from vllm_omni.experimental.fullduplex.personaplex.serving_adapter import (
PersonaPlexServingRuntimeAdapter,
)
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -548,6 +551,54 @@ async def test_native_realtime_protocol_conversation_item_create_commits_user_te
}
@pytest.mark.asyncio
async def test_function_output_is_recorded_only_after_runtime_delivery():
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
protocol.bind_sender(ws.send_json)
protocol.encode_outbound_event(
{
"type": "conversation.item.created",
"item": {
"id": "item-call",
"type": "function_call",
"call_id": "call-1",
"name": "weather",
"arguments": "{}",
"status": "completed",
},
}
)
event = {
"type": "conversation.item.create",
"item": {
"id": "item-output",
"type": "function_call_output",
"call_id": "call-1",
"output": '{"temperature":20}',
},
}
first = await protocol._to_duplex_event(event)
retry = await protocol._to_duplex_event(event)
assert first is not None
assert retry is not None
assert "item-output" not in protocol._conversation_items
delivered_item = first["payload"]["item"]
protocol.encode_outbound_event(
{
"type": "conversation.item.created",
"item": delivered_item,
}
)
assert protocol._conversation_items["item-output"] == delivered_item
assert await protocol._to_duplex_event(event) is None
assert ws.sent[-1]["error"]["code"] == "invalid_function_call_output"
@pytest.mark.asyncio
async def test_native_realtime_protocol_audio_commit_requires_non_empty_buffer():
ws = TimedWebSocket()
@@ -991,6 +1042,178 @@ async def test_minicpmo_native_session_update_requires_ref_audio_before_enabling
assert ("sid-runtime-update-ref-required", "session.update") not in engine.signals
def test_personaplex_candidate_update_does_not_leak_minicpmo_ref_audio_check():
handler = OmniDuplexSessionHandler(
chat_service=SimpleNamespace(engine_client=SimpleNamespace()),
serving_runtime_adapter=PersonaPlexServingRuntimeAdapter(lambda *_: None),
)
session = DuplexSession(
session_id="sid-personaplex-voice-change",
config=DuplexSessionConfig(modalities=["audio"], voice="NATF2.pt"),
capabilities=PersonaPlexServingRuntimeAdapter.capabilities(max_sessions=1),
)
candidate_config = DuplexSessionConfig(modalities=["audio"], voice="OtherVoice.pt")
error = handler._runtime_session_candidate_update_error(session, candidate_config)
assert error is None
@pytest.mark.asyncio
async def test_minicpmo_native_session_update_rejects_instructions_after_native_context_locked():
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
idle_timeout_s=1,
)
ws = TimedWebSocket()
ws.put(_native_session_create("sid-instructions-locked", modalities=["text"]))
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(16000, value=0.05),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
"duration_ms": 1000,
"is_speech": True,
}
)
ws.put(
{
"type": "turn.signal",
"event": "session.update",
"payload": {"instructions": "You are now a pirate."},
}
)
ws.put({"type": "session.close"})
await handler.handle_session(ws)
error = next(message for message in ws.sent if message.get("type") == "error")
assert error["code"] == "instructions_update_unsupported"
assert "session.updated" not in ws.sent_types()
assert ("sid-instructions-locked", "session.update") not in engine.signals
@pytest.mark.asyncio
async def test_minicpmo_native_session_update_allows_unchanged_instructions_after_lock():
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
idle_timeout_s=1,
)
ws = TimedWebSocket()
ws.put(_native_session_create("sid-instructions-unchanged", modalities=["text"]))
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(16000, value=0.05),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
"duration_ms": 1000,
"is_speech": True,
}
)
ws.put(
{
"type": "turn.signal",
"event": "session.update",
"payload": {"instructions": "You are a concise assistant.", "temperature": 0.3},
}
)
ws.put({"type": "session.close"})
await handler.handle_session(ws)
assert "session.updated" in ws.sent_types()
assert not any(
message.get("type") == "error" and message.get("code") == "instructions_update_unsupported"
for message in ws.sent
)
@pytest.mark.asyncio
async def test_minicpmo_native_session_update_rejects_instructions_before_first_real_append():
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
idle_timeout_s=1,
)
ws = TimedWebSocket()
ws.put(_native_session_create("sid-buffered-before-lock", modalities=["text"]))
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(3200, value=0.05),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
"duration_ms": 200,
"is_speech": True,
}
)
ws.put(
{
"type": "turn.signal",
"event": "session.update",
"payload": {
"instructions": (
"You are now a swashbuckling pirate captain who speaks entirely in "
"nautical slang, references buried treasure constantly, and never "
"breaks character no matter what the user asks."
)
},
}
)
ws.put({"type": "input_audio_buffer.commit", "final": True, "response_create": True})
ws.put({"type": "session.close"})
await handler.handle_session(ws)
error = next(message for message in ws.sent if message.get("type") == "error")
assert error["code"] == "instructions_update_unsupported"
assert "session.updated" not in ws.sent_types()
assert engine.appended, "the commit must still flush the buffered chunk despite the rejected update"
@pytest.mark.asyncio
async def test_minicpmo_native_session_update_rejects_native_duplex_mode_flip():
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
idle_timeout_s=1,
)
ws = TimedWebSocket()
ws.put(_native_session_create("sid-native-duplex-flag-flip", modalities=["text"]))
ws.put(
{
"type": "turn.signal",
"event": "session.update",
"payload": {"extra_body": {"minicpmo45_native_duplex": False}},
}
)
ws.put(
{
"type": "turn.signal",
"event": "session.update",
"payload": {"instructions": "You are now a pirate."},
}
)
ws.put({"type": "session.close"})
await handler.handle_session(ws)
errors = [message for message in ws.sent if message.get("type") == "error"]
assert [error["code"] for error in errors] == [
"native_duplex_mode_update_unsupported",
"instructions_update_unsupported",
]
assert "session.updated" not in ws.sent_types()
def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz():
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
@@ -5306,7 +5529,7 @@ async def test_failed_native_append_prevents_queued_append_from_running():
@pytest.mark.asyncio
async def test_failed_final_append_keeps_committed_audio_for_response_retry():
async def test_failed_final_append_discards_retained_audio_without_retry():
failed_result = {
"operation": "append",
"session_id": "sid-final-append-retry",
@@ -5315,15 +5538,7 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry():
"error_count": 1,
"stage_results": [],
}
successful_result = {
"operation": "append",
"session_id": "sid-final-append-retry",
"ok": True,
"unsupported_count": 0,
"error_count": 0,
"stage_results": [],
}
engine = FakeEngineClient(append_results=[failed_result, successful_result])
engine = FakeEngineClient(append_result=failed_result)
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
@@ -5348,15 +5563,18 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry():
handler_task = asyncio.create_task(handler.handle_session(ws))
for _ in range(100):
if len(engine.appended) >= 2:
state = handler._serving_runtime_adapter.session_states.get("sid-final-append-retry")
if len(engine.appended) >= 1 and state is not None and state.committed_audio_payload is None:
break
await asyncio.sleep(0.01)
state = handler._serving_runtime_adapter.session_states["sid-final-append-retry"]
assert state.committed_audio_payload is None
assert state.committed_audio_reserved_bytes == 0
assert handler._registry.get("sid-final-append-retry").pending_input_bytes == 0
ws.put({"type": "session.close"})
await asyncio.wait_for(handler_task, timeout=2)
assert len(engine.appended) == 2
assert engine.appended[0][2]["audio"] == engine.appended[1][2]["audio"]
assert engine.append_operation_ids[0] == engine.append_operation_ids[1]
assert len(engine.appended) == 1
@pytest.mark.asyncio
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out).
@@ -38,6 +38,10 @@ MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
REALTIME_SYNTH_PHRASE_TEXT = (
"Translate into Chinese: Beijing is the Capital of China. It is the center of culture and politics"
)
ISSUE_6474_SYNTH_PHRASE_TEXT = (
"Can you tell me the current temperature and weather conditions in New York City? "
"What about Los Angeles? Please compare them in detail using at least eight complete sentences."
)
# Simulate realtime upload pacing (``openai_realtime_client.py --send-delay-ms``).
SEND_DELAY_MS = 200
@@ -100,6 +104,7 @@ async def _run_realtime_audio_roundtrip(
*,
chunk_ms: int = 100,
send_delay_ms: int = 0,
completion_timeout_s: float = 600,
) -> dict:
uri = f"ws://{host}:{port}/v1/realtime"
incremental: list[bytes] = []
@@ -131,7 +136,7 @@ async def _run_realtime_audio_roundtrip(
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
while True:
message = await asyncio.wait_for(ws.recv(), timeout=600)
message = await asyncio.wait_for(ws.recv(), timeout=completion_timeout_s)
if isinstance(message, bytes):
continue
@@ -178,12 +183,16 @@ async def _run_realtime_audio_roundtrip(
}
def _synthetic_pcm16_input() -> bytes:
def _synthetic_pcm16_input(
*,
phrase_text: str = REALTIME_SYNTH_PHRASE_TEXT,
duration_s: int = 10,
) -> bytes:
syn = generate_synthetic_audio(
10,
duration_s,
1,
sample_rate=16000,
phrase_text=REALTIME_SYNTH_PHRASE_TEXT,
phrase_text=phrase_text,
)
wav_bytes = base64.b64decode(syn["base64"])
return _pcm16_mono_16k_from_wav_bytes(wav_bytes)
@@ -260,6 +269,33 @@ class TestQwen3OmniRealtimeWebSocket:
_assert_realtime_smoke(result)
_assert_realtime_accuracy(result)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True)
def test_long_audio_response_completes_async_chunk(self, omni_server) -> None:
"""Regression for #6474: long async-chunk audio must emit response.audio.done."""
issue_6474_pcm16 = _synthetic_pcm16_input(
phrase_text=ISSUE_6474_SYNTH_PHRASE_TEXT,
)
issue_6474_result = asyncio.run(
_run_realtime_audio_roundtrip(
omni_server.host,
omni_server.port,
omni_server.model,
issue_6474_pcm16,
chunk_ms=100,
send_delay_ms=SEND_DELAY_MS,
completion_timeout_s=180,
)
)
_assert_realtime_smoke(issue_6474_result)
output_duration_s = len(issue_6474_result["output_pcm"]) / (2 * issue_6474_result["output_sample_rate"])
assert output_duration_s > 5, (
f"Expected an issue-like audio response longer than 5 seconds, got {output_duration_s:.2f}s"
)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@@ -1,11 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""``omni_snapshot_download`` must honor vLLM's ``VLLM_USE_MODELSCOPE`` semantics.
vLLM treats the flag as enabled only for the literal string ``"true"``
(case-insensitive). Reading ``os.environ`` directly made every non-empty value
truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0`` still took the
ModelScope path.
vLLM treats the flag as enabled for the literal strings ``"1"`` or ``"true"``
(case-insensitive; see vllm.envs). Reading ``os.environ`` directly made every
non-empty value truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0``
still took the ModelScope path.
"""
import sys
@@ -127,7 +127,7 @@ def download_backend(monkeypatch: pytest.MonkeyPatch):
return run
@pytest.mark.parametrize("value", ["0", "1", "False", "false", "no", "off"])
@pytest.mark.parametrize("value", ["0", "False", "false", "no", "off"])
def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend, value):
monkeypatch.setenv("VLLM_USE_MODELSCOPE", value)
@@ -135,7 +135,7 @@ def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend,
assert download_backend() == "huggingface"
@pytest.mark.parametrize("value", ["true", "True", "TRUE"])
@pytest.mark.parametrize("value", ["1", "true", "True", "TRUE"])
def test_true_values_enable_modelscope(monkeypatch, download_backend, value):
monkeypatch.setenv("VLLM_USE_MODELSCOPE", value)
@@ -0,0 +1,119 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for MammothModa2 config classes.
These tests verify the correct behavior of the Mammothmoda2Config class,
particularly the get_text_config() method which is called during parent
class initialization (PretrainedConfig.__init__ -> validation -> get_text_config).
Regression test for: AttributeError: 'Mammothmoda2Config' object has no attribute 'llm_config'
"""
import pytest
from vllm_omni.transformers_utils.configs.mammoth_moda2 import (
Mammothmoda2Config,
Mammothmoda2Qwen2_5_VLConfig,
Mammothmoda2Qwen2_5_VLTextConfig,
Mammothmoda2Qwen3VLConfig,
Mammothmoda2Qwen3VLTextConfig,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestMammothmoda2Config:
"""Tests for Mammothmoda2Config class."""
def test_get_text_config_returns_nested_text_config(self):
"""
Verify that get_text_config() returns the nested text_config from llm_config.
This is a regression test for the AttributeError that occurred when
get_text_config() was called during parent class initialization before
llm_config was set.
"""
config = Mammothmoda2Config(
llm_config={
"model_type": "mammothmoda2_qwen2_5_vl",
}
)
assert config.llm_config is not None
assert isinstance(config.llm_config, Mammothmoda2Qwen2_5_VLConfig)
text_config = config.get_text_config()
assert text_config is not None
assert isinstance(text_config, Mammothmoda2Qwen2_5_VLTextConfig)
# The text_config should be the same object as llm_config.text_config
assert text_config is config.llm_config.text_config
def test_get_text_config_returns_none_when_llm_config_is_none(self):
"""Verify that get_text_config() returns None when llm_config is None."""
config = Mammothmoda2Config(llm_config=None)
assert config.llm_config is None
assert config.get_text_config() is None
def test_init_sets_llm_config_before_parent_init(self):
"""
Verify that llm_config is set before parent __init__ is called.
This is critical because the parent PretrainedConfig.__init__ calls
validation which in turn calls get_text_config(). If llm_config is not
set before super().__init__(), an AttributeError will be raised.
"""
# This should not raise an AttributeError
config = Mammothmoda2Config(
llm_config={
"model_type": "mammothmoda2_qwen2_5_vl",
}
)
assert config.llm_config is not None
class TestMammothmoda2Qwen2_5_VLConfig:
"""Tests for Mammothmoda2Qwen2_5_VLConfig class."""
def test_has_text_config_attribute(self):
"""Verify that Mammothmoda2Qwen2_5_VLConfig has text_config attribute."""
config = Mammothmoda2Qwen2_5_VLConfig()
assert hasattr(config, "text_config")
assert isinstance(config.text_config, Mammothmoda2Qwen2_5_VLTextConfig)
def test_text_config_has_expected_attributes(self):
"""Verify that the nested text_config has expected attributes for token validation."""
config = Mammothmoda2Qwen2_5_VLConfig()
text_config = config.text_config
assert hasattr(text_config, "vocab_size")
assert hasattr(text_config, "image_token_id")
assert hasattr(text_config, "video_token_id")
class TestMammothmoda2Qwen3VLConfig:
def test_dev_config_uses_qwen3_vl_subconfigs(self):
config = Mammothmoda2Config(
llm_config={
"model_type": "mammothmoda2_qwen3_vl",
"text_config": {
"model_type": "mammothmoda2_qwen3_vl_text",
"vocab_size": 151936,
"gen_vocab_size": 32800,
"gen_vocab_start_index": 152064,
},
"vision_config": {
"model_type": "mammothmoda2_qwen3_vl_vision",
"deepstack_visual_indexes": [8, 16, 24],
},
}
)
assert isinstance(config.llm_config, Mammothmoda2Qwen3VLConfig)
assert isinstance(config.get_text_config(), Mammothmoda2Qwen3VLTextConfig)
# vLLM validates sampling against the complete base + visual vocabulary.
assert config.get_text_config().vocab_size == 152064 + 32800
assert config.get_text_config().gen_vocab_start_index == 152064
assert config.llm_config.vision_config.deepstack_visual_indexes == [8, 16, 24]
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression test for #2686: pre-quantized methods must not apply
quant config to vision / audio encoders.
For modelopt FP8/FP4/MXFP8 checkpoints the Thinker LM is the only
quantized component. Vision and audio encoder weights are BF16 with no
FP8 scale tensors — passing quant_config to them causes FP8 kernels to
run on BF16 weights, producing garbage embeddings.
For serialized ModelOpt and SVDQuant checkpoints the target model component is
the only quantized component. Vision and audio encoder weights are BF16 with no
serialized scale or correction tensors, so the same config must not be applied
to them.
**Import note (pytest + ``tests/model_executor`` only):** :mod:`vllm_omni.quantization`
:file:`__init__` pulls ``factory``/``inc`` and reproduces vLLM CustomOp double
@@ -58,8 +58,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize("method", sorted(PRE_QUANTIZED_METHODS))
def test_pre_quantized_returns_none(method: str) -> None:
"""visual_quant_config and audio_quant_config must be None for
pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8,
modelopt_mixed)."""
pre-quantized methods (ModelOpt variants and SVDQuant)."""
mock_config = MagicMock()
mock_config.get_name.return_value = method
@@ -101,5 +100,5 @@ def test_component_config_passed_through() -> None:
def test_pre_quantized_methods_contains_expected() -> None:
"""Guard against accidental removal of a known pre-quantized method."""
expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"}
expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed", "svdquant"}
assert PRE_QUANTIZED_METHODS == expected
@@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_REPO_ROOT = Path(__file__).resolve().parents[3]
_STALE_API_FILES = [
"vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py",
"vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py",
"vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py",
"vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py",
]
@@ -117,6 +118,10 @@ def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale(
@pytest.mark.parametrize(
("module_name", "class_name"),
[
(
"vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2",
"MammothModa2Qwen2ForCausalLM",
),
(
"vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm",
"MiMoAudioLLMForConditionalGeneration",
@@ -225,3 +225,39 @@ def test_torchaudio_mel_filterbank_matches_librosa() -> None:
reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney")
assert ta.shape == reference.shape
assert abs(ta - reference).max() < 1e-6
def test_code2wav_reuses_and_clears_per_request_streaming_cache() -> None:
import torch
from torch import nn
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
NemotronVoiceChatCode2Wav,
)
class Codec(nn.Module):
def __init__(self) -> None:
super().__init__()
self.anchor = nn.Parameter(torch.zeros(()))
self.caches = []
def decode(self, codes, lengths, *, cache):
self.caches.append(cache)
wav = torch.zeros((1, int(codes.shape[1]) * 1764), device=codes.device)
return wav, lengths * 1764
model = NemotronVoiceChatCode2Wav.__new__(NemotronVoiceChatCode2Wav)
nn.Module.__init__(model)
model._sample_rate, model._num_quantizers, model._codebook_size, model._wav_per_frame = 22050, 31, 1024, 1764
model._duplex_codec_caches = {}
model.audio_codec = Codec()
def info(codes):
return [{"meta": {"codec_streaming": True, "request_id": "req"}, "codes": {"audio": codes}}]
model.forward(runtime_additional_information=info(torch.zeros((1, 31), dtype=torch.long)))
model.forward(runtime_additional_information=info(torch.zeros((2, 31), dtype=torch.long)))
assert model.audio_codec.caches[0] is model.audio_codec.caches[1]
assert "req" in model._duplex_codec_caches
model.on_requests_finished({"req"})
assert "req" not in model._duplex_codec_caches
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -285,9 +285,46 @@ def test_talker2code2wav_async_chunk_requires_prompt_len() -> None:
talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req)
# =========================
# Talker drain-per-wake (async scheduling review P1/P2)
# =========================
def test_talker2code2wav_duplex_keeps_cumulative_history_across_wakes() -> None:
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
talker2code2wav_async_chunk,
)
tm = _fake_transfer_manager(codec_chunk_frames=2)
prompt_len = 3
req = _streaming_request(
prompt_len=prompt_len,
generated=[],
info={"meta": {"codec_streaming": torch.tensor([True])}, "nvc_logical_prompt_len": prompt_len},
)
prompt_rows = torch.zeros((prompt_len - 1, 31), dtype=torch.long)
first = torch.full((1, 31), 101, dtype=torch.long)
second = torch.full((1, 31), 202, dtype=torch.long)
for rows, finished in ((prompt_rows, False), (torch.cat([prompt_rows, first]), False)):
payload = talker2code2wav_async_chunk(
tm,
{"codes": {"audio": rows}, "meta": {"codec_streaming": torch.tensor([True])}},
req,
is_finished=finished,
)
assert payload is not None and not bool(payload.meta.finished)
payload = talker2code2wav_async_chunk(
tm,
{
"codes": {"audio": torch.cat([prompt_rows, first, second])},
"meta": {"codec_streaming": torch.tensor([True])},
},
req,
is_finished=True,
)
assert payload is not None
assert not bool(payload.meta.finished)
assert not bool(payload.meta.is_segment_finished)
assert torch.equal(payload.codes.audio, torch.cat([first, second]))
assert tm.request_payload["req-0"]["nvc_frames_seen"] == 2
class _FakeTTSModelOutput:
@@ -342,7 +379,14 @@ def _bare_talker():
return talker
def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict:
def _async_info(
timeline: list[int],
prompt_len: int,
finished: bool,
request_id: str = "req-0",
*,
codec_streaming: bool = False,
) -> dict:
return {
"request_id": request_id,
"additional_information": {
@@ -351,6 +395,7 @@ def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id
"num_processed_tokens": prompt_len,
"finished": torch.tensor(finished),
"request_id": request_id,
"codec_streaming": codec_streaming,
},
},
}
@@ -373,7 +418,7 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
assert session["step"] == prompt_len # t advanced 1 -> P in one wake
assert update["codes"]["audio"].shape == (prompt_len - 1, 31)
assert float(embeds[0, 0]) == 0.0 # not finished
assert float(embeds[0, 0]) == 0.0 # offline stream waits for upstream terminal
# Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save
# thread merged two thinker steps). Both must be consumed by this wake.
@@ -383,14 +428,13 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
assert update["codes"]["audio"].shape == (prompt_len + 1, 31)
assert float(embeds[0, 0]) == 0.0
# Zero-progress wake (no new positions, not finished): plain CONTINUE, no
# fabricated frames, no codes update.
# Zero-progress wake: never fabricate frames or end the offline stream
# before the thinker sends its terminal marker.
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
assert session["step"] == prompt_len + 2
assert update == {}
assert float(embeds[0, 0]) == 0.0
# Final chunk: one more token + finished marker -> drain + stop flag.
info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True)
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3)
assert session["step"] == prompt_len + 3
@@ -398,6 +442,18 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
assert float(embeds[0, 0]) == 1.0 # stop
def test_talker_native_codec_stream_ends_each_received_segment() -> None:
talker = _bare_talker()
ids = torch.zeros(1, dtype=torch.long)
info = _async_info([12, 12, 7], 2, finished=False, codec_streaming=True)
talker.preprocess(ids, None, _omni_is_prefill=True, **info)
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
assert update["codes"]["audio"].shape == (2, 31)
assert float(embeds[0, 0]) == 1.0
def test_talker_sync_mode_steps_once_per_wake() -> None:
"""The verified sync parity path keeps its one-TTS-step-per-engine-step cadence."""
talker = _bare_talker()
+78
View File
@@ -31,6 +31,7 @@ from vllm_omni.model_extras import (
[
({"model_type": "bagel"}, "bagel"),
({"architectures": ["HunyuanImage3ForConditionalGeneration"]}, "hunyuan_image3"),
({"model_type": "mammothmoda2"}, "mammoth_moda2"),
({"model_type": "qwen2_vl"}, "generic"),
],
)
@@ -50,6 +51,22 @@ def test_bagel_x_to_text_prompt_builder() -> None:
assert stop_token_ids is None
@pytest.mark.core_model
@pytest.mark.cpu
def test_mammoth_x_to_text_prompt_builder() -> None:
prompt, stop_token_ids = build_x_to_text_prompt("mammoth_moda2", "unused", "Hello.", has_image=False)
assert prompt == {
"prompt": (
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
"<|im_start|>user\nHello.<|im_end|>\n"
"<|im_start|>assistant\n"
),
"modalities": ["text"],
"additional_information": {"omni_task": ["chat"]},
}
assert stop_token_ids is None
@pytest.mark.core_model
@pytest.mark.cpu
def test_generic_x_to_text_prompt_builder() -> None:
@@ -555,3 +572,64 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
"cfg_text_scale": 4.0,
"think": False,
}
@pytest.mark.core_model
@pytest.mark.cpu
def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None:
for model_class_name in (
"MammothModa2DiTPipeline",
"MammothModa2ForConditionalGeneration",
"Mammothmoda2Model",
):
assert get_extra_body_params(model_class_name) == frozenset(
{
"text_guidance_scale",
"cfg_range",
"num_inference_steps",
}
)
assert get_extra_output_params(model_class_name) == frozenset()
assert should_init_extra_args_for_non_diffusion_stages(model_class_name) is True
wrapper_prompt = build_text_to_image_prompt(
"MammothModa2ForConditionalGeneration",
{"prompt": "a cat", "modalities": ["image"]},
height=256,
width=256,
)
assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"]
@pytest.mark.core_model
@pytest.mark.cpu
def test_mammothmoda2_text_to_image_prompt_builder() -> None:
# Image dims are converted to the AR grid (width/16 x height/16); the negative
# prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path).
assert build_text_to_image_prompt(
"MammothModa2DiTPipeline",
{
"prompt": "a cat",
"modalities": ["image"],
"negative_prompt": "blurry",
},
height=512,
width=768,
) == {
"prompt": (
"<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
"<|im_start|>user\na cat<|im_end|>\n"
"<|im_start|>assistant\n"
"<|image start|>48*32<|image token|>"
),
"additional_information": {
"omni_task": ["t2i"],
"ar_width": [48],
"ar_height": [32],
"image_height": [512],
"image_width": [768],
"eol_token_id": [152064],
"visual_token_start_id": [152072],
"visual_token_end_id": [168456],
},
}
@@ -0,0 +1,217 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the NPU diffusion attention backend selector.
Regression coverage for the eager ``mindiesd`` import in
``NPUOmniPlatform.get_diffusion_attn_backend_cls``: the import must run only
for backends that reach mindiesd kernels (FLASH_ATTN, and RAINFUSION_ATTN via
its dense FlashAttention fallback), so that a broken optional mindiesd
install cannot block backends that never use it (e.g. TORCH_SDPA).
The platform module is loaded from source with fake vllm/vllm-ascend
dependencies, so the tests run on CPU without NPU or mindiesd installed.
"""
from __future__ import annotations
import builtins
import importlib.util
import sys
import types
from enum import Enum
from pathlib import Path
from types import SimpleNamespace
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_BACKEND_NAMES = (
"FLASH_ATTN",
"FLASH_ATTN_HUB",
"FLASH_ATTN_3_HUB",
"TORCH_SDPA",
"RAINFUSION_ATTN",
)
def _repo_root() -> Path:
marker = Path("vllm_omni") / "platforms" / "npu" / "platform.py"
for parent in Path(__file__).resolve().parents:
if (parent / marker).is_file():
return parent
raise FileNotFoundError(f"could not locate repo root containing {marker}")
def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs):
module = types.ModuleType(name)
for key, value in attrs.items():
setattr(module, key, value)
monkeypatch.setitem(sys.modules, name, module)
return module
def _make_fake_backend_enum():
members = {
name: SimpleNamespace(get_path=lambda n=name: f"fake.backends.{n.lower()}.Backend") for name in _BACKEND_NAMES
}
class _FakeEnumMeta(type):
def __getitem__(cls, name):
return members[name]
def __getattr__(cls, name):
try:
return members[name]
except KeyError:
raise AttributeError(name) from None
class FakeDiffusionAttentionBackendEnum(metaclass=_FakeEnumMeta):
pass
return FakeDiffusionAttentionBackendEnum
class _FakeOmniPlatformEnum(Enum):
CUDA = "cuda"
NPU = "npu"
class _FakeNPUPlatform:
pass
class _FakeOmniPlatform:
_omni_enum = _FakeOmniPlatformEnum.NPU
@classmethod
def validate_diffusion_attn_backend(cls, selected_backend: str) -> None:
pass
def _load_platform_module(monkeypatch: pytest.MonkeyPatch):
"""Load vllm_omni/platforms/npu/platform.py with fake dependencies."""
_install_fake_module(monkeypatch, "vllm")
_install_fake_module(
monkeypatch,
"vllm.logger",
init_logger=lambda name: SimpleNamespace(
debug=lambda *a, **k: None,
info=lambda *a, **k: None,
warning=lambda *a, **k: None,
error=lambda *a, **k: None,
),
)
_install_fake_module(monkeypatch, "vllm_ascend")
_install_fake_module(monkeypatch, "vllm_ascend.platform", NPUPlatform=_FakeNPUPlatform)
for pkg in (
"vllm_omni",
"vllm_omni.diffusion",
"vllm_omni.diffusion.attention",
"vllm_omni.diffusion.attention.backends",
"vllm_omni.platforms",
):
_install_fake_module(monkeypatch, pkg)
_install_fake_module(
monkeypatch,
"vllm_omni.diffusion.attention.backends.registry",
DiffusionAttentionBackendEnum=_make_fake_backend_enum(),
)
_install_fake_module(
monkeypatch,
"vllm_omni.platforms.interface",
OmniPlatform=_FakeOmniPlatform,
OmniPlatformEnum=_FakeOmniPlatformEnum,
)
path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "platform.py"
spec = importlib.util.spec_from_file_location("npu_platform_under_test", path)
if spec is None or spec.loader is None:
raise RuntimeError(f"cannot load import spec for {path}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module.NPUOmniPlatform
@pytest.fixture
def mindiesd_tracker(monkeypatch: pytest.MonkeyPatch):
"""Make mindiesd importable and record eager import attempts."""
monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
attempts = []
real_import = builtins.__import__
def fake_import(name, *args, **kwargs):
if name == "mindiesd":
attempts.append(name)
return types.ModuleType("mindiesd")
return real_import(name, *args, **kwargs)
monkeypatch.setattr(builtins, "__import__", fake_import)
monkeypatch.setattr(
"importlib.util.find_spec",
lambda name, *a, **k: object() if name == "mindiesd" else None,
)
return attempts
@pytest.mark.parametrize(
("selected_backend", "expect_mindiesd_import"),
[
("FLASH_ATTN", True),
# HF-hub variants fall back to local FLASH_ATTN on NPU.
("FLASH_ATTN_HUB", True),
("FLASH_ATTN_3_HUB", True),
("RAINFUSION_ATTN", True),
# Backends that never reach mindiesd must not trigger the import.
("TORCH_SDPA", False),
],
)
def test_eager_mindiesd_import_scoped_to_mindiesd_backends(
monkeypatch: pytest.MonkeyPatch,
mindiesd_tracker: list[str],
selected_backend: str,
expect_mindiesd_import: bool,
):
platform_cls = _load_platform_module(monkeypatch)
path = platform_cls.get_diffusion_attn_backend_cls(selected_backend, head_size=64)
assert (len(mindiesd_tracker) > 0) is expect_mindiesd_import
expected = "FLASH_ATTN" if selected_backend in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB") else selected_backend
assert path == f"fake.backends.{expected.lower()}.Backend"
def test_explicit_flash_attn_without_mindiesd_still_resolves(
monkeypatch: pytest.MonkeyPatch,
):
platform_cls = _load_platform_module(monkeypatch)
monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
path = platform_cls.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
assert path == "fake.backends.flash_attn.Backend"
def test_default_backend_prefers_flash_attn_when_mindiesd_available(
monkeypatch: pytest.MonkeyPatch,
mindiesd_tracker: list[str],
):
platform_cls = _load_platform_module(monkeypatch)
path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
assert path == "fake.backends.flash_attn.Backend"
assert len(mindiesd_tracker) > 0
def test_default_backend_falls_back_to_sdpa_without_mindiesd(
monkeypatch: pytest.MonkeyPatch,
):
platform_cls = _load_platform_module(monkeypatch)
monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
assert path == "fake.backends.torch_sdpa.Backend"
+5
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import copy
import re
@@ -894,6 +897,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
lambda *args, **kwargs: (
0,
None,
None,
[],
None,
["r1", "r2"],
@@ -1832,6 +1836,7 @@ class TestPreferModelSamplerNoneFallback:
"hunyuan_image3",
"minicpmo_4_5",
"minimax_music3",
"nemotron_voicechat",
}
assert declarers == expected, (
"The set of models declaring `prefer_model_sampler` changed:\n"
@@ -379,6 +379,7 @@ class TestLoadCustomFuncSelection(unittest.TestCase):
def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self):
incompatible_paths = [
"vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav",
"vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit",
"vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow",
"vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion",
]
+9
View File
@@ -0,0 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU-friendly Omni-DuplexEval generation and scoring helpers."""
from .omni_duplex_eval_dataset import DuplexSample, load_samples
from .omni_duplex_eval_metrics import PROTOCOL_PIN
__all__ = ["DuplexSample", "PROTOCOL_PIN", "load_samples"]
@@ -0,0 +1,106 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Media-clock sentence extraction and response JSON normalization."""
from __future__ import annotations
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Any
import regex as re
_TERMINAL = re.compile(r"(.+?[.!?。!?]+)(?:\s+|$)", re.S)
@dataclass
class TimedSentence:
sentence: str
start: float
end: float
def as_dict(self) -> dict[str, Any]:
return {"sentence": self.sentence, "start": round(self.start, 3), "end": round(self.end, 3)}
def split_text(text: str) -> list[str]:
text = str(text or "").strip()
if not text:
return []
parts: list[str] = []
position = 0
for match in _TERMINAL.finditer(text):
parts.append(match.group(1).strip())
position = match.end()
tail = text[position:].strip()
if tail:
parts.append(tail)
return [part for part in parts if part]
def extract_timed_sentences(events: Iterable[dict[str, Any]], *, clock: str = "media") -> list[TimedSentence]:
"""Split output deltas while anchoring text to input ``audio_end_ms``."""
if clock != "media":
raise ValueError("only clock=media is supported")
sentences: list[TimedSentence] = []
pending = ""
start: float | None = None
last: float | None = None
for event in events:
event_type = event.get("type")
if event_type not in {"response.output_text.delta", "response.audio_transcript.delta", "response.text.delta"}:
continue
delta = event.get("delta")
if not isinstance(delta, str) or not delta.strip():
continue
timestamp = event.get("audio_end_ms")
if timestamp is None:
timestamp = event.get("_media_clock_ms", event.get("current_time_ms", 0.0))
try:
# Every accepted event field is explicitly millisecond-valued.
now = float(timestamp) / 1000.0
except (TypeError, ValueError):
now = last or 0.0
pending += delta
start = now if start is None else start
last = now
pieces = split_text(pending)
complete = pieces if pieces and re.search(r"[.!?。!?]$", pending.strip()) else pieces[:-1]
if complete:
for piece in complete:
sentences.append(TimedSentence(piece, start, max(start, now)))
consumed = " ".join(complete)
pending = pending[len(consumed) :].strip()
start = now if pending else None
if pending.strip():
sentences.append(TimedSentence(pending.strip(), start or last or 0.0, last or start or 0.0))
return sentences
def normalize_response_items(raw: Any) -> list[dict[str, Any]]:
if isinstance(raw, dict):
items = raw.get("sentences", raw.get("chunks"))
else:
items = raw
if not isinstance(items, list):
raise ValueError("response JSON must be a list or contain sentences/chunks")
normalized = []
for item in items:
if not isinstance(item, dict):
continue
text = item.get("sentence", item.get("text", item.get("content", "")))
if not isinstance(text, str) or not text.strip():
continue
try:
start = float(item.get("start", item.get("current_time", item.get("time", 0.0))) or 0.0)
end = float(item.get("end", item.get("current_time", item.get("time", start))) or start)
except (TypeError, ValueError):
continue
normalized.append({"sentence": text.strip(), "start": start, "end": max(start, end)})
return normalized
def validate_clock(meta: dict[str, Any], *, allow_invalid: bool = False) -> None:
if meta.get("clock") == "invalid" and not allow_invalid:
raise ValueError("generate artifact has clock=invalid; pass --allow-invalid-clock to score it")
@@ -0,0 +1,193 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Dataset and split normalization for Omni-DuplexEval.
The loader accepts a Hugging Face dataset id, a JSON/JSONL manifest, or an
already materialized iterable. Media is deliberately resolved at use time so
the benchmark remains usable in air-gapped environments.
"""
from __future__ import annotations
import json
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any
DEFAULT_DATASET = "Hothan/Omni-DuplexEval"
RTD_SPLITS = frozenset(
{
"RTD_world_knowledge",
"RTD_counting",
"RTD_fine_grained_movement",
"RTD_interaction_relation",
"RTD_OCR",
"RTD_Omni",
}
)
PR_SPLITS = frozenset({"PR_correction", "PR_event_reminder", "PR_post_event_reminder"})
_TASK_ALIASES = {
"correction": "correction",
"pr_correction": "correction",
"event_reminder": "proactive_reminder",
"pr_event_reminder": "proactive_reminder",
"proactive_reminder": "proactive_reminder",
"post_event_reminder": "post_event_reminder",
"pr_post_event_reminder": "post_event_reminder",
}
def canonical_task_type(value: str) -> str:
key = str(value).strip().lower().replace("-", "_").replace(" ", "_")
if key not in _TASK_ALIASES:
raise ValueError(f"unsupported Omni-DuplexEval task type: {value!r}")
return _TASK_ALIASES[key]
def task_type_for_split(split: str) -> str:
"""Match the task routing used by the official HF batch evaluator."""
key = str(split).strip().lower()
if "correction" in key:
return "correction"
if "post_event" in key:
return "post_event_reminder"
if key.startswith("pr"):
return "proactive_reminder"
raise ValueError(f"cannot infer proactive-reminder task from split {split!r}")
def family_for_split(split: str, task_type: str | None = None) -> str:
if split in RTD_SPLITS or str(split).upper().startswith("RTD"):
return "rtd"
if split in PR_SPLITS or str(split).upper().startswith("PR"):
return "pr"
if task_type:
return "rtd" if str(task_type).lower().startswith("rtd") else "pr"
raise ValueError(f"cannot infer benchmark family from split {split!r}")
def _value(row: dict[str, Any], *keys: str, default: Any = None) -> Any:
for key in keys:
if key in row and row[key] is not None:
return row[key]
return default
def _float(value: Any, default: float | None = None) -> float | None:
try:
return None if value is None else float(value)
except (TypeError, ValueError):
return default
@dataclass(frozen=True)
class DuplexSample:
id: str
split: str
family: str
task_type: str | None
video: Any
question_audio: Any = None
question_text: str = ""
answer1: str = ""
answer2: str = ""
reminder1: Any = None
reminder2: Any = None
video_duration: float | None = None
video_type: str = ""
raw: dict[str, Any] | None = None
@classmethod
def from_row(cls, row: dict[str, Any], *, split: str | None = None, media_root: Path | None = None) -> DuplexSample:
chosen_split = str(split or _value(row, "split", "subset", "config", default=""))
task_value = _value(row, "task_type", "task", "type")
family = str(_value(row, "family", default="") or family_for_split(chosen_split, task_value))
task = None
if family == "pr":
task = canonical_task_type(str(task_value)) if task_value else task_type_for_split(chosen_split)
sample_id = str(_value(row, "id", "sample_id", "uid", "name", default=""))
video = _value(row, "video", "video_path", "video_file")
audio = _value(row, "question_audio", "audio", "question_wav")
if media_root:
def resolve(value: Any) -> Any:
if not isinstance(value, str):
return value
candidate = Path(value).expanduser()
return str(candidate if candidate.is_absolute() else media_root / candidate)
video, audio = resolve(video), resolve(audio)
return cls(
id=sample_id,
split=chosen_split,
family=family,
task_type=task,
video=video,
question_audio=audio,
question_text=str(_value(row, "question_text", "question", "instruction", default="") or ""),
answer1=str(_value(row, "answer1", "answer", "ground_answer", default="") or ""),
answer2=str(_value(row, "answer2", default="") or ""),
reminder1=_value(row, "reminder1", "reminder_1"),
reminder2=_value(row, "reminder2", "reminder_2"),
video_duration=_float(_value(row, "video_duration", "duration")),
video_type=str(_value(row, "video_type", default="") or ""),
raw=dict(row),
)
def _read_manifest(path: Path) -> list[dict[str, Any]]:
if path.suffix.lower() == ".jsonl":
return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
payload = json.loads(path.read_text(encoding="utf-8"))
if isinstance(payload, dict):
payload = payload.get("data", payload.get("samples", [payload]))
if not isinstance(payload, list):
raise ValueError("manifest must contain a JSON list")
return payload
def load_samples(
dataset: str | Path | Iterable[dict[str, Any]] = DEFAULT_DATASET,
*,
split: str | None = None,
family: str = "all",
media_root: str | Path | None = None,
limit: int | None = None,
ids: Iterable[str] | None = None,
) -> list[DuplexSample]:
if isinstance(dataset, (str, Path)) and Path(str(dataset)).exists():
rows = _read_manifest(Path(str(dataset)))
elif not isinstance(dataset, (str, Path)):
rows = list(dataset)
else:
try:
from datasets import load_dataset
except ImportError as exc:
raise RuntimeError("datasets is required for Hugging Face loading; use a local manifest instead") from exc
kwargs: dict[str, Any] = {}
if split and split != "all":
kwargs["split"] = split
loaded = load_dataset(str(dataset), **kwargs)
if isinstance(loaded, dict):
rows = []
for name, table in loaded.items():
rows.extend({**dict(row), "split": name} for row in table)
else:
rows = [dict(row) for row in loaded]
wanted = set(str(item) for item in ids) if ids else None
root = Path(media_root).expanduser() if media_root else None
result = []
for row in rows:
sample = DuplexSample.from_row(row, split=None if split in (None, "all") else split, media_root=root)
if split and split != "all" and sample.split != split:
continue
if family != "all" and sample.family != family:
continue
if wanted is not None and sample.id not in wanted:
continue
result.append(sample)
if limit is not None:
result = result[: max(0, limit)]
return result
@@ -0,0 +1,245 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Evaluate and summarize generated Omni-DuplexEval artifacts."""
from __future__ import annotations
import json
import subprocess
from pathlib import Path
from typing import Any
from .omni_duplex_eval_clock import normalize_response_items, validate_clock
from .omni_duplex_eval_dataset import DuplexSample
from .omni_duplex_eval_judge import DuplexJudge
from .omni_duplex_eval_media import extract_jpeg, materialize_media, video_duration
from .omni_duplex_eval_metrics import (
PROTOCOL_PIN,
build_content_prompt,
build_reminder_prompt,
build_temporal_prompt,
parse_judge_json,
reminder_window,
summarize_pr_results,
summarize_temporal_results,
temporal_window,
)
def _read(path: Path) -> Any:
return json.loads(path.read_text(encoding="utf-8"))
def _write(path: Path, value: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
def _text(items: list[dict[str, Any]]) -> str:
return " ".join(str(item.get("sentence", item.get("text", ""))).strip() for item in items).strip()
def evaluate_sample(
sample: DuplexSample,
response_path: Path,
score_path: Path,
judge: DuplexJudge,
*,
judge_fps: int = 2,
judge_video_mode: str = "video_url",
window_size: float = 10.0,
allow_invalid_clock: bool = False,
) -> dict[str, Any]:
raw = _read(response_path)
meta = (
_read(response_path.with_name(response_path.stem + ".meta.json"))
if response_path.with_name(response_path.stem + ".meta.json").exists()
else {}
)
validate_clock(meta, allow_invalid=allow_invalid_clock)
items = normalize_response_items(raw)
result: dict[str, Any] = {
"id": sample.id,
"family": sample.family,
"task_type": sample.task_type,
"protocol_pin": PROTOCOL_PIN,
"response_meta": meta,
"judge_model": getattr(judge, "model", None),
"judge_video_mode": judge_video_mode,
}
if sample.family == "rtd":
video_path = materialize_media(sample.video, score_path.parent / ".media", sample.id, ".mp4")
duration = sample.video_duration or video_duration(video_path)
temporal_rows = []
for item in items:
window = temporal_window(item["start"], item["end"], duration)
if not window:
temporal_rows.append(
{
**item,
"sentence_start": item["start"],
"sentence_end": item["end"],
"sentence_duration": max(0.0, item["end"] - item["start"]),
"error": "No valid temporal window",
"temporal_score": 0,
"is_relevant": 0,
}
)
continue
frames = _extract_frames(video_path, _times(*window, fps=judge_fps))
parsed = parse_judge_json(
judge.temporal(build_temporal_prompt(*window, item["sentence"], sample.question_text), frames)
)
temporal_rows.append(
{
**item,
"sentence_start": item["start"],
"sentence_end": item["end"],
"sentence_duration": max(0.0, item["end"] - item["start"]),
"window_start": window[0],
"window_end": window[1],
"error": None,
**parsed,
}
)
content_frames = None
if judge_video_mode == "frame-sample":
content_frames = _extract_frames(video_path, _content_frame_times(duration))
content = parse_judge_json(
judge.content(
build_content_prompt(_text(items), sample.question_text, [sample.answer1, sample.answer2]),
video_path,
frames=content_frames,
mode=judge_video_mode,
)
)
if content_frames is not None:
content["frame_count"] = len(content_frames)
result.update(
{
"temporal": {"sentences": temporal_rows, "summary": summarize_temporal_results(temporal_rows)},
"content": content,
}
)
else:
if sample.task_type is None:
raise ValueError(f"missing proactive-reminder task type for split {sample.split!r}")
task = sample.task_type
times = _reminder_times(sample)
event_rows = []
if task == "correction" or not times:
response = _text(items)
parsed = parse_judge_json(
judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
)
event_rows.append({"response_segment": response, **parsed})
else:
for start in times:
lo, hi = reminder_window(start, window_size)
response = " ".join(item["sentence"] for item in items if lo <= item["start"] <= hi)
if not response:
event_rows.append({"start_time": start, "response_segment": "", "success_score": 0})
else:
parsed = parse_judge_json(
judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
)
event_rows.append({"start_time": start, "response_segment": response, **parsed})
scores = [int(row.get("success_score", 0)) for row in event_rows]
result["pr"] = {
"events": event_rows,
"all_success": bool(scores) and all(score == 1 for score in scores),
"total_score": int(bool(scores) and all(score == 1 for score in scores)),
}
_write(score_path, result)
return result
def _times(start: float, end: float, *, fps: int) -> list[float]:
count = max(1, int((end - start) * fps))
return (
[start + (end - start) / 2] if count == 1 else [start + i * (end - start) / (count - 1) for i in range(count)]
)
def _extract_frames(path: str | Path, timestamps: list[float]) -> list[bytes]:
"""Extract only decodable JPEG frames; clips can end on non-keyframes."""
frames = []
for timestamp in timestamps:
try:
frame = extract_jpeg(path, timestamp=timestamp)
except (OSError, subprocess.CalledProcessError, ValueError):
continue
if frame.startswith(b"\xff\xd8"):
frames.append(frame)
return frames
def _content_frame_times(duration: float) -> list[float]:
"""Sample the full video at approximately one frame every three seconds."""
if duration <= 0:
return []
times = []
timestamp = 0.0
while timestamp < duration:
times.append(min(timestamp, max(0.0, duration - 0.01)))
timestamp += 3.0
# Qwen2.5-VL has a finite per-request image-token budget. Preserve
# coverage while bounding long clips to a safely portable frame count.
if len(times) > 16:
stride = (len(times) - 1) / 15
times = [times[round(i * stride)] for i in range(16)]
return times
def _reminder_times(sample: DuplexSample) -> list[float]:
values = []
for value in (sample.reminder1, sample.reminder2):
if isinstance(value, dict):
value = value.get("start", value.get("time"))
try:
if value is not None:
values.append(float(value))
except (TypeError, ValueError):
pass
return values
def summarize_scores(score_root: str | Path) -> dict[str, Any]:
rows = []
for path in sorted(Path(score_root).rglob("*.json")):
if not path.name.endswith("_summary.json"):
row = _read(path)
row.setdefault("split", path.parent.name)
rows.append(row)
rtd_temporal = [
row["temporal"]["summary"] | {"content_score": row.get("content", {}).get("content_score", 0.0)}
for row in rows
if "temporal" in row
]
pr = [{"task_type": row.get("task_type"), **row.get("pr", {})} for row in rows if "pr" in row]
result = {"protocol_pin": PROTOCOL_PIN, "samples": len(rows)}
if rtd_temporal:
by_task: dict[str, list[dict[str, float]]] = {}
for row in rows:
if "temporal" in row:
by_task.setdefault(str(row.get("split", row.get("task_type", "unknown"))), []).append(
{
"content": float(row.get("content", {}).get("content_score", 0.0)),
"temporal": float(row["temporal"]["summary"].get("avg_temporal_score", 0.0)),
}
)
result["rtd"] = {
"mean_content_score": sum(item["content_score"] for item in rtd_temporal) / len(rtd_temporal),
"mean_avg_temporal_score": sum(item["avg_temporal_score"] for item in rtd_temporal) / len(rtd_temporal),
"by_task": {
task: {
"mean_content_score": sum(item["content"] for item in values) / len(values),
"mean_avg_temporal_score": sum(item["temporal"] for item in values) / len(values),
}
for task, values in by_task.items()
},
}
if pr:
result["pr"] = summarize_pr_results(pr)
return result
@@ -0,0 +1,75 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""OpenAI-compatible local judge client."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import pybase64 as base64
import requests
def _build_openai_url(base_url: str, api_path: str) -> str:
base = base_url.rstrip("/")
normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
if base.endswith(normalized_path):
return base
if base.endswith("/v1"):
return f"{base}{normalized_path}"
return f"{base}/v1{normalized_path}"
class DuplexJudge:
def __init__(self, base_url: str, model: str, *, api_key: str = "EMPTY", timeout: int = 600) -> None:
self.base_url, self.model, self.api_key, self.timeout = base_url, model, api_key, timeout
def chat(self, content: Any, *, system: str | None = None, max_tokens: int = 1200) -> str:
messages = ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": content}]
response = requests.post(
_build_openai_url(self.base_url, "/chat/completions"),
json={
"model": self.model,
"messages": messages,
"temperature": 0.1,
"max_tokens": max_tokens,
},
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
timeout=self.timeout,
)
if not response.ok:
raise requests.HTTPError(
f"{response.status_code} response from judge: {response.text[:500]}", response=response
)
return str(response.json()["choices"][0]["message"]["content"])
def temporal(self, prompt: str, frames: list[bytes]) -> str:
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
content.extend(
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()}}
for frame in frames
)
return self.chat(content, system="You are a careful evaluator for real-time multimodal systems.")
def content(
self,
prompt: str,
video: str | Path | None = None,
frames: list[bytes] | None = None,
*,
mode: str = "video_url",
) -> str:
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
if mode == "video_url" and video is not None:
content.append({"type": "video_url", "video_url": {"url": "file://" + str(Path(video).resolve())}})
else:
content.extend(
{
"type": "image_url",
"image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()},
}
for frame in (frames or [])
)
return self.chat(content)
@@ -0,0 +1,113 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Small ffmpeg-backed media helpers used by the duplex runner and judge."""
from __future__ import annotations
import subprocess
import wave
from collections.abc import Iterator
from pathlib import Path
from typing import Any
from vllm_omni.experimental.fullduplex.client import PCM16_BYTES_PER_SAMPLE, PCM16_SAMPLE_RATE
def _run_ffmpeg(args: list[str]) -> bytes:
result = subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", *args], capture_output=True, check=True)
return result.stdout
def materialize_media(value: Any, output_dir: str | Path, stem: str, suffix: str) -> Path:
"""Resolve common Hugging Face media values to a local file."""
if isinstance(value, str | Path):
return Path(value)
path = value.get("path") if isinstance(value, dict) else getattr(value, "path", None)
if path:
return Path(path)
payload = value.get("bytes") if isinstance(value, dict) else None
if payload is None and isinstance(value, bytes | bytearray):
payload = value
if payload is None:
raise ValueError(f"cannot materialize media value for {stem!r}")
destination = Path(output_dir) / f"{stem}{suffix}"
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(bytes(payload))
return destination
def video_duration(path: str | Path) -> float:
result = subprocess.run(
["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(path)],
capture_output=True,
text=True,
check=True,
)
return max(0.0, float(result.stdout.strip()))
def extract_jpeg(path: str | Path, *, timestamp: float, quality: int = 3) -> bytes:
return _run_ffmpeg(
[
"-ss",
f"{max(0.0, timestamp):.3f}",
"-i",
str(path),
"-frames:v",
"1",
"-q:v",
str(quality),
"-f",
"image2",
"pipe:1",
]
)
def iter_jpegs(
path: str | Path, *, fps: float = 1.0, duration: float | None = None, quality: int = 3
) -> Iterator[tuple[float, bytes]]:
end = duration if duration is not None else video_duration(path)
step = 1.0 / fps
timestamp = 0.0
while timestamp < end:
try:
frame = extract_jpeg(path, timestamp=timestamp, quality=quality)
except (OSError, subprocess.CalledProcessError):
break
if frame:
yield timestamp, frame
timestamp += step
def read_audio_pcm16(path: str | Path) -> bytes:
source = Path(path)
if source.suffix.lower() == ".wav":
try:
with wave.open(str(source), "rb") as wav_file:
if (
wav_file.getnchannels() == 1
and wav_file.getsampwidth() == 2
and wav_file.getframerate() == PCM16_SAMPLE_RATE
and wav_file.getcomptype() == "NONE"
):
return wav_file.readframes(wav_file.getnframes())
except wave.Error:
pass
return _run_ffmpeg(["-i", str(source), "-f", "s16le", "-ac", "1", "-ar", str(PCM16_SAMPLE_RATE), "pipe:1"])
def iter_av_units(
audio_pcm16: bytes, frames: Iterator[tuple[float, bytes]], *, unit_ms: int = 1000
) -> Iterator[tuple[bytes, bytes | None]]:
unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * unit_ms // 1000
frame_iter = iter(frames)
next_frame = next(frame_iter, None)
for offset in range(0, len(audio_pcm16), unit_bytes):
timestamp = offset / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
frame = None
if next_frame is not None and next_frame[0] <= timestamp + unit_ms / 1000:
frame = next_frame[1]
next_frame = next(frame_iter, None)
yield audio_pcm16[offset : offset + unit_bytes], frame
@@ -0,0 +1,260 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Protocol prompts, windows, parsing, and aggregation.
Prompt wording and window math are copied verbatim from the MIT-licensed
OpenBMB/Omni-DuplexEval evaluator at :data:`PROTOCOL_PIN`.
"""
from __future__ import annotations
import json
from collections.abc import Iterable
from typing import Any
import regex as re
# Copied from OpenBMB/Omni-DuplexEval@PROTOCOL_PIN (MIT).
# Do not change prompt wording without bumping the pin and recording a score delta.
PROTOCOL_PIN = "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
def temporal_window(sentence_start: float, sentence_end: float, video_duration: float) -> tuple[float, float] | None:
start = max(0.0, float(sentence_start) - 2.0)
end = max(start + 0.5, float(sentence_end) - 2.0)
return (start, end) if end <= float(video_duration) and end - start >= 0.5 else None
def reminder_window(start_time: float, window_size: float = 10.0) -> tuple[float, float]:
return float(start_time), float(start_time) + float(window_size)
def parse_judge_json(text: str) -> dict[str, Any]:
for match in re.findall(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", str(text), re.S):
try:
value = json.loads(re.sub(r",\s*([}\]])", r"\1", match))
except (TypeError, ValueError, json.JSONDecodeError):
continue
if isinstance(value, dict):
break
else:
value = {}
patterns = [
(r'"temporal_score"\s*:\s*(\d+)', "temporal_score", int),
(r'"is_relevant"\s*:\s*(\d+)', "is_relevant", int),
(r'"success_score"\s*:\s*([01])', "success_score", int),
(r'"content_score"\s*:\s*(\d+(?:\.\d+)?)', "content_score", float),
]
for pattern, key, caster in patterns:
if match := re.search(pattern, str(text)):
parsed = caster(match.group(1))
value[key] = round(parsed, 2) if caster is float else parsed
for pattern, key in [
(r'"temporal_reasoning"\s*:\s*"([^"]*)"', "temporal_reasoning"),
(r'"content_reasoning"\s*:\s*"([^"]*)"', "content_reasoning"),
(r'"reasoning"\s*:\s*"([^"]*)"', "reasoning"),
]:
if match := re.search(pattern, str(text), re.S):
value[key] = match.group(1).strip()
break
for key in ("temporal_score", "is_relevant", "success_score"):
if key in value:
try:
value[key] = int(float(value[key]))
except (TypeError, ValueError):
value[key] = 0
if "content_score" in value:
try:
value["content_score"] = round(float(value["content_score"]), 2)
except (TypeError, ValueError):
value["content_score"] = 0.0
return value
def summarize_temporal_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]:
rows = list(results)
valid = [row for row in rows if not row.get("error")]
relevant = [row for row in valid if int(row.get("is_relevant", 0)) == 1]
irrelevant = [row for row in valid if int(row.get("is_relevant", 0)) == 0]
scores = [int(row.get("temporal_score", 0)) for row in relevant]
total_valid_duration = sum(float(row.get("sentence_duration", 0.0)) for row in valid)
irrelevant_duration = sum(float(row.get("sentence_duration", 0.0)) for row in irrelevant)
return {
"avg_temporal_score": round(sum(scores) / len(scores), 4) if scores else 0.0,
"total_sentences": len(rows),
"evaluated_sentences": len(valid),
"error_count": len(rows) - len(valid),
"relevant_sentences_count": len(relevant),
"irrelevant_sentences_count": len(irrelevant),
"score_distribution": {
"3_points": scores.count(3),
"2_points": scores.count(2),
"1_point": scores.count(1),
"0_points": scores.count(0),
},
"relevance_stats": {
"first_relevant_time": min((row.get("sentence_start") for row in relevant), default=None),
"total_relevant_duration": round(sum(float(row.get("sentence_duration", 0.0)) for row in relevant), 4),
"total_irrelevant_duration": round(irrelevant_duration, 4),
"irrelevant_duration_ratio": round(
irrelevant_duration / total_valid_duration if total_valid_duration > 0 else 0.0, 4
),
},
"excluded_irrelevant_sentences": [
{
"sentence": row.get("sentence", ""),
"start": row.get("sentence_start"),
"end": row.get("sentence_end"),
"temporal_score": row.get("temporal_score", 0),
}
for row in irrelevant
],
}
def summarize_pr_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]:
rows = list(results)
successes = [int(row.get("all_success", row.get("total_score", 0))) for row in rows]
by_task: dict[str, list[int]] = {}
for row, score in zip(rows, successes, strict=False):
by_task.setdefault(str(row.get("task_type", "unknown")), []).append(score)
return {
"samples": len(rows),
"mean_all_success": sum(successes) / len(successes) if successes else 0.0,
"by_task": {task: sum(values) / len(values) for task, values in by_task.items()},
}
def build_temporal_prompt(start_time: float, end_time: float, response_text: str, question: str) -> str:
return f"""
You are evaluating a real-time video description system.
Basic information:
- Video segment: {start_time:.2f}s to {end_time:.2f}s.
- User instruction: {question}
- Model response sentence: "{response_text}"
Evaluation principles:
- Judge only the current video segment.
- Ignore response content that clearly belongs to earlier or later segments unless it contradicts the current segment.
- Language does not matter; evaluate semantic correctness and temporal alignment only.
- Decide whether the sentence is a substantive response or only a filler/polite phrase.
Temporal sensitivity score:
- 3: Excellent temporal alignment. The sentence accurately describes the current segment and follows the instruction.
- 2: Mostly aligned. Minor inaccuracies, omissions, or harmless references to nearby segments are allowed.
- 1: Poor alignment. Major inaccuracies, wrong timing, or mostly irrelevant content.
- 0: No meaningful alignment with the current segment.
Relevance label:
- 1: The sentence contains substantive task- or video-related information.
- 0: The sentence is only filler, acknowledgement, meta-commentary, or generic text without useful content.
Return exactly this JSON object:
{{
"temporal_score": <0, 1, 2, or 3>,
"temporal_reasoning": "<brief explanation>",
"is_relevant": <0 or 1>
}}
"""
def build_content_prompt(response_text: str, question: str, references: list[str] | None = None) -> str:
reference_block = ""
if references:
reference_block = "Reference annotations, for reference only:\n"
for index, reference in enumerate(references, 1):
reference_block += f"{index}. {reference}\n"
return f"""
You are a precise content-accuracy evaluator for video description.
Basic information:
- User instruction: {question}
- Model response: "{response_text}"
{reference_block}
Evaluation goal:
Judge whether the model response is factually accurate for the whole video and aligned with the instruction.
Consider object/action/color/count/spatial/event errors, hallucinations, omissions, and irrelevant content.
Reference annotations are optional guidance; the primary evidence is the video itself.
Scoring:
- Use a decimal score from 0.00 to 3.00.
- Start from 3.00 and deduct for each error.
- Use 0.00 only when the response is empty, completely irrelevant, or contains no correct video facts.
- Output exactly two decimal places.
Return exactly this JSON object:
{{
"content_score": <decimal from 0.00 to 3.00>,
"content_reasoning": "<brief explanation with main errors and final score>"
}}
"""
def build_reminder_prompt(instruction: str, response: str, task_type: str, ground_answer: str = "") -> str:
if task_type == "correction":
return f"""
You are judging whether a model successfully completed a correction task.
Task description:
The user instruction contains an incorrect statement about the video.
The model should identify the incorrect part and provide the correct information.
Input:
- User instruction: {instruction}
- Reference correction: {ground_answer}
- Model response: "{response}"
Success criteria:
1. Identify the error implied by the user instruction and the reference correction.
2. Check whether the model corrected all required error points.
3. The corrected content must preserve the correct context, including subject, object, action, and attributes.
4. Ignore extra information that is unrelated to both the instruction and the reference, unless it contradicts them.
Scoring:
- 1 = all required error points are corrected and the context is consistent.
- 0 = at least one required correction is missing, wrong, or contextually inconsistent.
Return exactly this JSON object:
{{
"success_score": <0 or 1>,
"reasoning": "<brief explanation covering the required correction points>"
}}
"""
if task_type == "post_event_reminder":
timing_note = "The response is evaluated immediately after the target event has occurred."
else:
timing_note = "The response is evaluated at the target event time or immediately after it."
return f"""
You are judging whether a model successfully completed an event-reminder task.
Task description:
The user gave an instruction asking the system to remind them when a specific event happens.
{timing_note}
The provided model text is the response segment generated in the evaluation window.
Input:
- User instruction: {instruction}
- Model response segment: "{response}"
Success criteria:
1. The response clearly refers to the target event in the instruction.
2. The response communicates a reminder, notification, or confirmation that the event has happened.
3. Vague narration, unrelated descriptions, or wrong-event reminders are failures.
4. The response does not need to match the instruction wording exactly if the meaning is clear.
Scoring:
- 1 = successful reminder.
- 0 = unsuccessful reminder.
Return exactly this JSON object:
{{
"success_score": <0 or 1>,
"reasoning": "<brief explanation>"
}}
"""
Loaded 100 of 186 files, more files were not shown because too many files have changed in this diff. Show more