Compare commits
19
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fc5bafb2af | ||
|
|
c6df39e9f7 | ||
|
|
94b1546ce3 | ||
|
|
19da238246 | ||
|
|
651b03225b | ||
|
|
da4a08b65b | ||
|
|
0e44fc7406 | ||
|
|
432b2c6ca3 | ||
|
|
c2723979c0 | ||
|
|
2d4724eb74 | ||
|
|
91dd7fcaec | ||
|
|
926e7fbfa2 | ||
|
|
8cfede4309 | ||
|
|
d3c990dca8 | ||
|
|
5ec576ab4c | ||
|
|
030fccb2bf | ||
|
|
816335cb46 | ||
|
|
59cf3add40 | ||
|
|
7439ed41c2 |
No files matched your search
@@ -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"
|
||||
'
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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}}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,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:
|
||||
|
||||
|
||||
@@ -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]
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
@@ -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+ |
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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] = {}
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
Reference in new issue
Block a user