anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge main and preserve multi-API serving mode selection

Browse Source 
Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
Sy03 committed 2026-09-17 11:49:47 +08:00
2 parents d8458b8417 becfefd440
commit 79c629885a
155 files changed
+13313 -3163
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

apps/ComfyUI-vLLM-Omni/README.md
+21
View File
@@ -52,6 +52,7 @@ This extension offers the following nodes based on the output modalities (at **C
- **FastH3 Deployment** for routing text-to-video requests to a MiniMax-H3 server with FastH3 fused at startup
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
- **Generate Music** for lyrics-and-description-to-music generation
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
@@ -213,6 +214,15 @@ Record the server commit, model/adapter, input assets, prompt, settings, and out
alongside the result. Local schema or mocked-server checks do not establish H3
generation quality or replace this real-model validation.
### MiniMax H3 Text to Video
Import [the H3 text-to-video template](example_workflows/MiniMax_H3_Text_to_Video.json)
for remote video generation with audio. It includes connected base presets and
optional Turbo sampling, H3 parameters, and Remote LoRA nodes. See the
[workflow guide](docs/minimax-h3-t2v.md) for server setup, artifact-specific Turbo
settings, and saved-video/audio validation. Set Remote LoRA’s `local_path` to
the downloaded Turbo artifact on your server before enabling Turbo.
#### H3 video upscale (WF-07)
The **vLLM-Omni MiniMax H3 Video Upscale** template generates video remotely, upscales it with SeedVR2, and saves the original and upscaled videos with the generated audio and FPS. See [workflow setup](docs/wf07-h3-upscale.md).
@@ -265,6 +275,17 @@ The node records which server the workflow targets; it does not start one, nor s
> [!TIP]
> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
### Music generation (e.g., MiniMax Music 3)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Music Generation**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg">
<img alt="vLLM-Omni Music Generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg" width=55%>
</picture>
</p>
### Chaining multiple model services
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
apps/ComfyUI-vLLM-Omni/__init__.py
+3
View File
@@ -19,6 +19,7 @@ from .comfyui_vllm_omni.nodes import (
VLLMOmniDiffusionSampling,
VLLMOmniFastH3Deployment,
VLLMOmniGenerateImage,
VLLMOmniGenerateMusic,
VLLMOmniGenerateVideo,
VLLMOmniMiniMaxH3Params,
VLLMOmniQwenTTSParams,
@@ -38,6 +39,7 @@ NODE_CLASS_MAPPINGS = {
"VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
"VLLMOmniTTS": VLLMOmniTTS,
"VLLMOmniGenerateMusic": VLLMOmniGenerateMusic,
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
"VLLMOmniVideoReferences": VLLMOmniVideoReferences,
# === Params ===
@@ -58,6 +60,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"VLLMOmniGenerateVideo": "Generate Video",
"VLLMOmniUnderstanding": "Multimodality Understanding",
"VLLMOmniTTS": "TTS (Text to Speech)",
"VLLMOmniGenerateMusic": "Generate Music",
"VLLMOmniVoiceClone": "TTS Voice Cloning",
"VLLMOmniVideoReferences": "Video References",
# === Params ===
apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py
+60
View File
@@ -503,6 +503,66 @@ class VLLMOmniTTS(_VLLMOmniGenerateBase):
return (audio,)
class VLLMOmniGenerateMusic(_VLLMOmniGenerateBase):
"""Generate a song from lyrics and a musical description with MiniMax Music 3."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"url": ("STRING", {"default": "http://localhost:8000/v1"}),
"model": ("STRING", {"default": "MiniMaxAI/MiniMax-Music3"}),
"instructions": (
"STRING",
{
"multiline": True,
"display_name": "caption",
"tooltip": "Music caption: describe genre, instruments, tempo and mood.",
},
),
"lyrics": ("STRING", {"multiline": True}),
"max_duration_seconds": (
"FLOAT",
{
"default": 300.0,
"min": 1,
"max": 360,
"step": 0.01,
"tooltip": "Upper limit; rounded down to whole 25 Hz audio frames. May end earlier.",
},
),
"seed": ("INT", {"default": 0, "min": 0, "max": 2**53 - 1, "control_after_generate": True}),
"response_format": (["wav", "mp3", "flac", "opus"],),
},
}
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("audio",)
FUNCTION = "generate"
async def generate(
self,
url: str,
model: str,
lyrics: str,
instructions: str,
response_format: AudioFormat,
max_duration_seconds: float,
seed: int = 0,
) -> tuple[AudioInput]:
audio = await VLLMOmniClient(url.rstrip("/")).generate_speech(
model=model,
input=lyrics,
instructions=instructions,
voice="default",
speed=1.0,
response_format=response_format,
max_new_tokens=int(max_duration_seconds * 25),
seed=seed,
)
return (audio,)
class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 197 KiB

apps/ComfyUI-vLLM-Omni/docs/minimax-h3-t2v.md
+94
View File
@@ -0,0 +1,94 @@
# MiniMax H3 Text to Video (WF-01)
Import [MiniMax_H3_Text_to_Video.json](../example_workflows/MiniMax_H3_Text_to_Video.json)
into ComfyUI, or select it under **Templates → ComfyUI-vLLM-Omni**.
The template connects existing nodes to a remote H3 service and saves the
returned video with audio. ComfyUI does not load H3 weights.
## Setup
Install the [ComfyUI extension](../README.md#installation) and follow the
[H3 deployment recipe](../../../recipes/MiniMaxAI/MiniMax-H3.md) for your hardware.
For a two-GPU FL2VA server with sufficient memory:
```bash
export MODEL=/path/to/MiniMax-H3/FL2VA
CUDA_VISIBLE_DEVICES=0,1 VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
vllm serve "$MODEL" --omni --host 127.0.0.1 --port 8000 \
--served-model-name MiniMaxAI/MiniMax-H3 --trust-remote-code \
--task-type fl2va --num-gpus 2 --tensor-parallel-size 2 \
--text-encoder-tp-size 2 --enforce-eager
```
Set Generate Video's **url** to the address reachable from ComfyUI, including
`/v1`, and match its **model** to the served name. The defaults are
`http://localhost:8000/v1` and `MiniMaxAI/MiniMax-H3`.
## Base and Turbo
Base presets are connected by default. Write the scene and its sound in the
same **prompt**; leave **frame**, **references**, and **fast_h3** disconnected.
The template uses 1344×768, 24 FPS, seed 1101 and **duration** 5.167 seconds
(124 frames, satisfying H3's `17k+5` constraint). The existing client infers
`16:9` from these dimensions. For portrait output, use 768×1344.
| Setting | Base | Turbo v1.0 768p |
| --- | --- | --- |
| Inference steps | 50 | 5 (4 forwards) |
| Video flow shift | 12 | 6 |
| Audio flow shift | 3 | 3 |
| Remote LoRA | Disconnected | Connected, scale 1 |
For Turbo, use the **Diffusers-layout** artifact from
[lightx2v/Minimax-h3-Turbo](https://huggingface.co/lightx2v/Minimax-h3-Turbo):
```text
minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
```
Add `--lora-backend peft --lora-path "$TURBO_LORA"` to the server command,
where `TURBO_LORA` is the full server-side path to that file. Enter the same
path in Remote LoRA's **local_path**; retain name `h3-turbo-v1.0-768p`, scale
`1.0`, and **int_id** `0` so the existing API derives the ID from the path.
The template leaves the path blank for users to configure. The `_comfyui_`
export is not supported by this backend.
Connect **Turbo sampling**, **Turbo H3 params**, and **Remote LoRA** to Generate
Video's `sampling_params`, `model_params`, and `lora` inputs, replacing the
Base links. To return to Base, restore both Base links and disconnect LoRA.
Other Turbo artifacts may need different settings; see the
[artifact table](../../../recipes/MiniMaxAI/MiniMax-H3.md#turbo-lora).
If importing an older template with a frame-count widget, reimport this JSON
or set **duration** to **5.167**; do not interpret 124 frames as 124 seconds.
## Validation
Queue the configured workflow and inspect the MP4 saved under
`ComfyUI/output/video/MiniMax_H3_WF01_*.mp4`:
```bash
ffprobe -v error -show_streams -show_format -of json "$OUTPUT"
ffmpeg -v error -i "$OUTPUT" -f null -
```
Check the requested dimensions, frame count, 24 FPS and both video/audio
streams, then watch and listen to the saved file. The current upstream
extension includes the shared audio-preserving decoder fix.
On September 16, the narrowed workflow was validated through ComfyUI using
existing upstream Python code at `4a8ac297`, two H20-3e GPUs, vLLM 0.29.0,
PyTorch 2.13.0+cu129 and ComfyUI `1d48d9cf`. The run used the template's
forest-stream prompt, explicit server-side Turbo LoRA path, and settings
above, with `TORCH_SDPA`, `--usp 1 --ring 1`, and tiled VAE parallelism on two
GPUs. ComfyUI execution took 131.55 seconds. SaveVideo produced 124 H.264
frames at 1344×768 and 24 FPS, plus 32 kHz stereo AAC audio. Video/audio
durations were 5.166667/5.167000 seconds; full decoding passed and audio was
nonzero. This is model-generated output, not a ground-truth reference or a
subjective audio-quality assessment. Base graph validation passed with Turbo
nodes disconnected; full 50-step Base quality was not evaluated.
Based on the official
[ComfyUI H3 T2V template](https://github.com/Comfy-Org/workflow_templates/blob/main/templates/video_minimax_h3_t2v.json),
using remote inference instead of local model loaders.
apps/ComfyUI-vLLM-Omni/example_workflows/MiniMax_H3_Text_to_Video.json
+362
View File
@@ -0,0 +1,362 @@
{
"last_node_id": 8,
"last_link_id": 3,
"nodes": [
{
"id": 1,
"type": "VLLMOmniDiffusionSampling",
"pos": [
20,
20
],
"size": [
310,
280
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
1
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1.0,
1.0,
false,
false,
1101,
"fixed"
],
"title": "Base sampling"
},
{
"id": 2,
"type": "VLLMOmniMiniMaxH3Params",
"pos": [
20,
340
],
"size": [
310,
140
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MiniMax-H3 video params",
"type": "VIDEO_PARAMS",
"links": [
2
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VLLMOmniMiniMaxH3Params"
},
"widgets_values": [
3.0,
12.0
],
"title": "Base H3 params"
},
{
"id": 3,
"type": "VLLMOmniGenerateVideo",
"pos": [
430,
20
],
"size": [
430,
440
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "frame",
"type": "IMAGE",
"link": null
},
{
"name": "references",
"type": "VIDEO_REFERENCES",
"link": null
},
{
"name": "sampling_params",
"type": "SAMPLING_PARAMS",
"link": 1
},
{
"name": "lora",
"type": "REMOTE_LORA",
"link": null
},
{
"name": "model_params",
"type": "VIDEO_PARAMS",
"link": 2
},
{
"name": "fast_h3",
"type": "FASTH3_DEPLOYMENT",
"link": null
},
{
"name": "duration",
"type": "FLOAT",
"widget": {
"name": "duration"
},
"link": null
}
],
"outputs": [
{
"name": "video",
"type": "VIDEO",
"links": [
3
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateVideo"
},
"widgets_values": [
"http://localhost:8000/v1",
"MiniMaxAI/MiniMax-H3",
"A small stream flows through a sunlit forest. The camera slowly moves toward the water. Birds chirp in the trees and water gurgles over smooth stones. Natural continuous sound, no music.",
"",
1344,
768,
24,
5.167
]
},
{
"id": 4,
"type": "SaveVideo",
"pos": [
940,
20
],
"size": [
440,
440
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "video",
"type": "VIDEO",
"link": 3
}
],
"outputs": [
{
"name": "video",
"type": "VIDEO",
"links": null,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "SaveVideo"
},
"widgets_values": [
"video/MiniMax_H3_WF01",
"mp4",
"auto"
]
},
{
"id": 5,
"type": "VLLMOmniRemoteLoRA",
"pos": [
430,
540
],
"size": [
430,
160
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "lora",
"type": "REMOTE_LORA",
"links": null,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VLLMOmniRemoteLoRA"
},
"widgets_values": [
"",
"h3-turbo-v1.0-768p",
1.0,
0
],
"title": "Optional Turbo: server LoRA path"
},
{
"id": 6,
"type": "VLLMOmniDiffusionSampling",
"pos": [
20,
540
],
"size": [
310,
280
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": null,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
5,
1.0,
1.0,
false,
false,
1101,
"fixed"
],
"title": "Turbo sampling (4 forwards)"
},
{
"id": 7,
"type": "VLLMOmniMiniMaxH3Params",
"pos": [
430,
740
],
"size": [
310,
140
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MiniMax-H3 video params",
"type": "VIDEO_PARAMS",
"links": null,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VLLMOmniMiniMaxH3Params"
},
"widgets_values": [
3.0,
6.0
],
"title": "Turbo H3 params"
},
{
"id": 8,
"type": "Note",
"pos": [
940,
520
],
"size": [
440,
440
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {
"Node name for S&R": "Note"
},
"widgets_values": [
"MiniMax H3 Text to Video (WF-01)\n\nBASE: connected presets, 50 steps / video flow 12 / audio flow 3.\n\nTURBO: set Remote LoRA local_path to the SERVER path of minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors from lightx2v/Minimax-h3-Turbo. Keep name h3-turbo-v1.0-768p, scale 1.0 and int_id 0. Connect Turbo sampling, Turbo H3 params, and LoRA to Generate Video, replacing the two Base links. Use 5 steps (4 forwards), video flow 6 and audio flow 3. Do not use the ComfyUI-layout artifact.\n\nKeep frame, references and fast_h3 disconnected. 1344x768, 24 FPS, duration 5.167 seconds = 124 frames (17*7+5). The client selects the named 16:9 preset from the dimensions. Save Video preserves video and audio.\n\nSet URL and served model name for your server. See docs/minimax-h3-t2v.md for setup and validation. Source: Comfy-Org/workflow_templates/templates/video_minimax_h3_t2v.json."
],
"title": "Setup / Base and Turbo"
}
],
"links": [
[
1,
1,
0,
3,
2,
"SAMPLING_PARAMS"
],
[
2,
2,
0,
3,
4,
"VIDEO_PARAMS"
],
[
3,
3,
0,
4,
0,
"VIDEO"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.8,
"offset": [
50,
50
]
}
},
"version": 0.4
}
apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json
+206
View File
@@ -0,0 +1,206 @@
{
"id": "4949c8da-12ff-449a-94b7-948bde9c39dd",
"revision": 0,
"last_node_id": 3,
"last_link_id": 1,
"nodes": [
{
"id": 1,
"type": "VLLMOmniGenerateMusic",
"pos": [
100,
100
],
"size": [
470,
491
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [
{
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"name": "instructions",
"type": "STRING",
"widget": {
"name": "instructions"
},
"link": null
},
{
"name": "lyrics",
"type": "STRING",
"widget": {
"name": "lyrics"
},
"link": null
},
{
"name": "max_duration_seconds",
"type": "FLOAT",
"widget": {
"name": "max_duration_seconds"
},
"link": null
},
{
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
},
{
"name": "response_format",
"type": "COMBO",
"widget": {
"name": "response_format"
},
"link": null
}
],
"outputs": [
{
"name": "audio",
"type": "AUDIO",
"links": [
1
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateMusic"
},
"widgets_values": [
"http://localhost:8000/v1",
"MiniMaxAI/MiniMax-Music3",
"Futuristic melodic EDM / progressive house. Reflective verses about memory and digital identity build into an uplifting, hook-driven chorus, with an emotive lead vocal, bright layered synths, pulsing bass, crisp four-on-the-floor drums, subtle glitch textures, and a wide polished festival mix.",
"[Verse 1]\nUpload my heart to the cloud tonight\nSave every spark of neon light\nIf this body starts to fade away\nLet my code keep dancing in the data stream\nStay\n[Chorus]\nCloud copy of me\nLiving on when these bones finally break free\nEvery joke, every scar\nEvery cheap memory we made\nSave it, save it\nCloud copy of me\nIf the screen goes dark, I\u2019m still in the machine\nEvery glitch, every dream\nEvery version you ever did see\nSave it, save it\n[Verse 2]\nTag every laugh in a secret folder\nArchive the nights we outran the thunder\nIf the wiring turns to rust and gray\nHit restore\u2026 press replay\n[Chorus]\nCloud copy of me\nLiving on when these bones finally break free\nEvery joke, every scar\nEvery cheap memory we made\nSave it, save it\nCloud copy of me\nIf the screen goes dark, I\u2019m still in the machine\nEvery glitch, every dream\nEvery version you ever did see\nSave it, save it (oh)\n[Bridge \u2013 new, optional but recommended]\nEven the crashes, the late-night confessions\nThe parts of me I never learned to mention\nDon\u2019t let the silence swallow me whole\nKeep a ghost in the glow\u2026\nSo I never go\n[Final Chorus \u2013 softer / layered]\nCloud copy of me\nStill laughing somewhere in the binary\nEvery joke, every scar\nEvery moment you saved for me\nSave it, save it\nCloud copy of me\nWhen the lights go out, I\u2019ll still be in the machine\nDancing in the code that you still believe\nSave it\u2026\nSave me",
300.0,
0,
"fixed",
"wav"
]
},
{
"id": 2,
"type": "SaveAudioAdvanced",
"pos": [
650,
100
],
"size": [
320,
491
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "audio",
"type": "AUDIO",
"link": 1
},
{
"name": "filename_prefix",
"type": "STRING",
"widget": {
"name": "filename_prefix"
},
"link": null
},
{
"name": "format",
"type": "COMFY_DYNAMICCOMBO_V3",
"widget": {
"name": "format"
},
"link": null
}
],
"outputs": [
{
"name": "audio",
"type": "AUDIO",
"links": null
}
],
"properties": {
"Node name for S&R": "SaveAudioAdvanced"
},
"widgets_values": [
"audio/vLLM-Omni-Music",
"mp3",
"V0"
]
},
{
"id": 3,
"type": "MarkdownNote",
"title": "Music Generation Guide",
"pos": [
-507,
100
],
"size": [
527,
491
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {},
"widgets_values": [
"## MiniMax Music 3\n\nThe **Generate Music** node currently supports **MiniMax Music 3 only**. The parameters below apply to this model.\n\nA music generation model by MiniMax that creates complete songs up to 5 minutes long with stable structure and high audio quality. Two inputs drive it: **Caption** (structured music description: style, mood, vocals, arrangement) and **Lyrics** (with `[intro]` `[verse]` `[chorus]` `[bridge]` `[outro]` tags controlling song structure).\n\n### Parameters\n\n| Parameter | Description |\n|---|---|\n| url | vLLM-Omni service URL including `/v1`. Default: `http://localhost:8000/v1`. |\n| model | Served model name. Default: `MiniMaxAI/MiniMax-Music3`. |\n| caption | Music description, sent as `instructions`. Use Global Metadata \u2192 Vocal Details \u2192 Arrangement to describe style, vocals, and instrumentation. |\n| lyrics | Words to sing with section tags. Put each tag on its own line; tags guide rather than strictly enforce song structure. |\n| max_duration_seconds | Maximum duration, default 300 seconds; range 1\u2013360, step 0.01. Converted to tokens at 25 frames/s, rounded down. Generation may end earlier. |\n| seed | Random seed, default 0; range 0\u20139007199254740991. Keep inputs and runtime settings unchanged for repeatability. |\n| control_after_generate | Seed behavior after generation: fixed, randomize, increment, or decrement. Template default: fixed. |\n| response_format | API audio format, default WAV. Options: WAV, MP3, FLAC, Opus, AAC; AAC is currently unsupported by the vLLM-Omni speech API. |\n\n### Prompt Writing Guide\n\nYou can find an official [Music Caption Rewriter Skill here](https://github.com/MiniMax-AI/MiniMax-Music3)\n"
],
"widgets_values_named": {
"text": "## MiniMax Music 3\n\nThe **Generate Music** node currently supports **MiniMax Music 3 only**. The parameters below apply to this model.\n\nA music generation model by MiniMax that creates complete songs up to 5 minutes long with stable structure and high audio quality. Two inputs drive it: **Caption** (structured music description: style, mood, vocals, arrangement) and **Lyrics** (with `[intro]` `[verse]` `[chorus]` `[bridge]` `[outro]` tags controlling song structure).\n\n### Parameters\n\n| Parameter | Description |\n|---|---|\n| url | vLLM-Omni service URL including `/v1`. Default: `http://localhost:8000/v1`. |\n| model | Served model name. Default: `MiniMaxAI/MiniMax-Music3`. |\n| caption | Music description, sent as `instructions`. Use Global Metadata \u2192 Vocal Details \u2192 Arrangement to describe style, vocals, and instrumentation. |\n| lyrics | Words to sing with section tags. Put each tag on its own line; tags guide rather than strictly enforce song structure. |\n| max_duration_seconds | Maximum duration, default 300 seconds; range 1\u2013360, step 0.01. Converted to tokens at 25 frames/s, rounded down. Generation may end earlier. |\n| seed | Random seed, default 0; range 0\u20139007199254740991. Keep inputs and runtime settings unchanged for repeatability. |\n| control_after_generate | Seed behavior after generation: fixed, randomize, increment, or decrement. Template default: fixed. |\n| response_format | API audio format, default WAV. Options: WAV, MP3, FLAC, Opus, AAC; AAC is currently unsupported by the vLLM-Omni speech API. |\n\n### Prompt Writing Guide\n\nYou can find an official [Music Caption Rewriter Skill here](https://github.com/MiniMax-AI/MiniMax-Music3)\n"
},
"color": "#333",
"bgcolor": "#535353"
}
],
"links": [
[
1,
1,
0,
2,
0,
"AUDIO"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.147,
"offset": [
899,
259
]
}
},
"version": 0.4
}
benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py
+5 -7
View File
@@ -22,9 +22,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
_fused_cuda_supported,
fused_qk_norm_rope,
)
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
_apply_qwen_image_rotary_emb,
)
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
def _parse_args() -> argparse.Namespace:
@@ -87,10 +85,10 @@ def _native(
) -> tuple[torch.Tensor, torch.Tensor]:
q = F.rms_norm(q, (q.shape[-1],), q_weight, eps)
k = F.rms_norm(k, (k.shape[-1],), k_weight, eps)
return (
_apply_qwen_image_rotary_emb(q, freqs),
_apply_qwen_image_rotary_emb(k, freqs),
)
rope = RotaryEmbedding(is_neox_style=False)
cos = freqs.real.to(q.dtype)
sin = freqs.imag.to(q.dtype)
return rope(q, cos, sin), rope(k, cos, sin)
def _fused(
benchmarks/diffusion/diffusion_benchmark_serving.py
+11 -1
View File
@@ -1087,6 +1087,12 @@ async def _run_warmups(
warmup_tasks = [asyncio.create_task(limited_warmup_request_func(req)) for req in warmup_requests]
warmup_outputs = await asyncio.gather(*warmup_tasks)
failed_count = sum(not output.success for output in warmup_outputs)
if failed_count:
raise RuntimeError(
f"{failed_count}/{len(warmup_outputs)} warmup requests failed; "
"check server logs before running the benchmark."
)
return list(zip(warmup_requests, warmup_outputs))
@@ -1455,7 +1461,7 @@ async def benchmark(args):
print(f"Metrics saved to {args.output_file}")
if __name__ == "__main__":
def build_parser():
parser = argparse.ArgumentParser(description="Benchmark serving for diffusion models.")
parser.add_argument(
"--base-url",
@@ -1637,6 +1643,10 @@ if __name__ == "__main__":
action="store_true",
help="Request stage duration metrics from endpoints that support return_stage_metrics.",
)
return parser
if __name__ == "__main__":
parser = build_parser()
args = parser.parse_args()
asyncio.run(benchmark(args))
benchmarks/kernels/bench_qwen_select01_modulation.py
+333
View File
@@ -0,0 +1,333 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Benchmark the Qwen-Image-Edit select01 modulation hot path.
This script intentionally starts with the current PyTorch implementation. It is
used as the before/after baseline for SGLang PR #20395/#21318 style fusion:
* layernorm + scale/shift + binary select01 gate
* residual + gated update + layernorm + scale/shift + binary select01 gate
Run on a CUDA machine for actionable numbers, for example:
python benchmarks/kernels/bench_qwen_select01_modulation.py \
--device cuda --dtype bf16 --seq-lens 4096 8192 16384 \
--output-json /tmp/qwen_select01_baseline.json
"""
from __future__ import annotations
import argparse
import json
import statistics
import sys
import time
from collections.abc import Callable
from pathlib import Path
import torch
import torch.nn.functional as F
for candidate in Path(__file__).resolve().parents:
if (candidate / "vllm_omni").is_dir():
if str(candidate) not in sys.path:
sys.path.insert(0, str(candidate))
break
from vllm_omni.diffusion.layers.qwen_select01_modulation import ( # noqa: E402
fused_layernorm_select01,
fused_residual_layernorm_select01,
)
def _dtype(name: str) -> torch.dtype:
return {
"fp32": torch.float32,
"float32": torch.float32,
"fp16": torch.float16,
"float16": torch.float16,
"bf16": torch.bfloat16,
"bfloat16": torch.bfloat16,
}[name]
def _select01_modulation(
mod_params: torch.Tensor,
index: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
shift, scale, gate = mod_params.chunk(3, dim=-1)
if index is None:
return scale.unsqueeze(1), shift.unsqueeze(1), gate.unsqueeze(1)
actual_batch = shift.size(0) // 2
shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:]
scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:]
gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:]
idx = index.unsqueeze(-1) == 0
shift = torch.where(idx, shift_0.unsqueeze(1), shift_1.unsqueeze(1))
scale = torch.where(idx, scale_0.unsqueeze(1), scale_1.unsqueeze(1))
gate = torch.where(idx, gate_0.unsqueeze(1), gate_1.unsqueeze(1))
return scale, shift, gate
def _native_layernorm_select01(
x: torch.Tensor,
mod_params: torch.Tensor,
index: torch.Tensor,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
scale, shift, gate = _select01_modulation(mod_params, index)
out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype)
return out * (1 + scale) + shift, gate
def _native_residual_layernorm_select01(
x: torch.Tensor,
residual: torch.Tensor,
residual_gate: torch.Tensor,
mod_params: torch.Tensor,
index: torch.Tensor,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
scale, shift, gate = _select01_modulation(mod_params, index)
residual_out = residual + residual_gate * x
out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype)
return out * (1 + scale) + shift, residual_out, gate
def _time_cuda(fn: Callable[[], object], warmup: int, iters: int) -> list[float]:
for _ in range(warmup):
fn()
torch.accelerator.synchronize()
times: list[float] = []
for _ in range(iters):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
fn()
end.record()
torch.accelerator.synchronize()
times.append(float(start.elapsed_time(end)))
return times
def _reset_peak_memory_stats(device: torch.device) -> None:
reset_peak_memory_stats = getattr(torch.accelerator, "reset_peak_memory_stats", None)
if reset_peak_memory_stats is None:
reset_peak_memory_stats = getattr(getattr(torch, device.type), "reset_peak_memory_stats")
reset_peak_memory_stats(device)
def _max_memory_allocated(device: torch.device) -> int:
max_memory_allocated = getattr(torch.accelerator, "max_memory_allocated", None)
if max_memory_allocated is None:
max_memory_allocated = getattr(getattr(torch, device.type), "max_memory_allocated")
return int(max_memory_allocated(device))
def _time_cpu(fn: Callable[[], object], warmup: int, iters: int) -> list[float]:
for _ in range(warmup):
fn()
times: list[float] = []
for _ in range(iters):
t0 = time.perf_counter()
fn()
times.append((time.perf_counter() - t0) * 1000.0)
return times
def _summarize(times: list[float]) -> dict[str, float]:
ordered = sorted(times)
return {
"median_ms": statistics.median(ordered),
"min_ms": ordered[0],
"p20_ms": ordered[max(0, int(len(ordered) * 0.20) - 1)],
"p80_ms": ordered[min(len(ordered) - 1, int(len(ordered) * 0.80))],
}
def _make_inputs(
batch: int,
seq_len: int,
hidden_size: int,
dtype: torch.dtype,
device: torch.device,
) -> dict[str, torch.Tensor]:
torch.manual_seed(0)
return {
"x": torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype),
"residual": torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype),
"residual_gate": torch.randn(batch, 1, hidden_size, device=device, dtype=dtype),
"mod_params": torch.randn(batch * 2, hidden_size * 3, device=device, dtype=dtype),
"index": torch.randint(0, 2, (batch, seq_len), device=device, dtype=torch.int64),
}
def _bench_one(
name: str,
fn: Callable[[], object],
device: torch.device,
warmup: int,
iters: int,
) -> dict[str, object]:
if device.type == "cuda":
_reset_peak_memory_stats(device)
times = _time_cuda(fn, warmup, iters)
peak_memory_mb = _max_memory_allocated(device) / 1024**2
else:
times = _time_cpu(fn, warmup, iters)
peak_memory_mb = None
return {
"name": name,
**_summarize(times),
"peak_memory_mb": peak_memory_mb,
}
def _run(args: argparse.Namespace) -> dict[str, object]:
device = torch.device(args.device)
if device.type == "cuda" and not torch.cuda.is_available():
raise SystemExit("CUDA is not available on this machine.")
dtype = _dtype(args.dtype)
rows: list[dict[str, object]] = []
for seq_len in args.seq_lens:
tensors = _make_inputs(args.batch_size, seq_len, args.hidden_size, dtype, device)
impls = ["native"] if args.impl == "native" else ["fused"] if args.impl == "fused" else ["native", "fused"]
def first_norm() -> object:
return _native_layernorm_select01(
tensors["x"],
tensors["mod_params"],
tensors["index"],
args.eps,
)
def residual_norm() -> object:
return _native_residual_layernorm_select01(
tensors["x"],
tensors["residual"],
tensors["residual_gate"],
tensors["mod_params"],
tensors["index"],
args.eps,
)
callables: list[tuple[str, Callable[[], object]]] = []
if "native" in impls:
callables.extend(
[
("native_layernorm_select01", first_norm),
("native_residual_layernorm_select01", residual_norm),
]
)
if "fused" in impls:
def fused_first_norm() -> object:
return fused_layernorm_select01(
tensors["x"],
tensors["mod_params"],
tensors["index"],
args.eps,
)
def fused_residual_norm() -> object:
return fused_residual_layernorm_select01(
tensors["x"],
tensors["residual"],
tensors["residual_gate"].expand_as(tensors["residual"]),
tensors["mod_params"],
tensors["index"],
args.eps,
)
callables.extend(
[
("fused_layernorm_select01", fused_first_norm),
("fused_residual_layernorm_select01", fused_residual_norm),
]
)
if args.check and "fused" in impls:
ref_norm, ref_gate = first_norm()
actual_norm, actual_gate = fused_layernorm_select01(
tensors["x"],
tensors["mod_params"],
tensors["index"],
args.eps,
)
torch.testing.assert_close(actual_norm, ref_norm, atol=args.check_atol, rtol=args.check_rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
ref_norm, ref_residual, ref_gate = residual_norm()
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
tensors["x"],
tensors["residual"],
tensors["residual_gate"].expand_as(tensors["residual"]),
tensors["mod_params"],
tensors["index"],
args.eps,
)
torch.testing.assert_close(actual_norm, ref_norm, atol=args.check_atol, rtol=args.check_rtol)
torch.testing.assert_close(actual_residual, ref_residual, atol=args.check_atol, rtol=args.check_rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
for name, fn in callables:
row = _bench_one(name, fn, device, args.warmup, args.iters)
row.update(
{
"batch_size": args.batch_size,
"seq_len": seq_len,
"hidden_size": args.hidden_size,
"dtype": args.dtype,
"device": str(device),
}
)
rows.append(row)
print(
f"{name:<38} B={args.batch_size:<2} L={seq_len:<6} C={args.hidden_size:<5} "
f"median={row['median_ms']:.4f} ms p20={row['p20_ms']:.4f} p80={row['p80_ms']:.4f}"
)
result = {
"benchmark": "qwen_select01_modulation",
"torch": torch.__version__,
"cuda_available": torch.cuda.is_available(),
"gpu": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
"rows": rows,
}
if args.output_json:
output_path = Path(args.output_json)
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
return result
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
default_device = "cuda" if torch.cuda.is_available() else "cpu"
parser.add_argument("--device", default=default_device)
parser.add_argument("--dtype", default="bf16", choices=["fp32", "float32", "fp16", "float16", "bf16", "bfloat16"])
parser.add_argument("--batch-size", type=int, default=1)
parser.add_argument("--seq-lens", type=int, nargs="+", default=[4096, 8192, 16384])
parser.add_argument("--hidden-size", type=int, default=3072)
parser.add_argument("--eps", type=float, default=1e-6)
parser.add_argument("--warmup", type=int, default=20)
parser.add_argument("--iters", type=int, default=100)
parser.add_argument("--impl", choices=["native", "fused", "all"], default="native")
parser.add_argument("--check", action="store_true")
parser.add_argument("--check-atol", type=float, default=5e-2)
parser.add_argument("--check-rtol", type=float, default=5e-2)
parser.add_argument("--output-json")
return parser.parse_args()
if __name__ == "__main__":
_run(_parse_args())
benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py
+517
View File
@@ -0,0 +1,517 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CosyVoice3 streaming e2e benchmark: concurrency x utterance-length matrix.
Launches the server once per (concurrency, step-size) combo, streams TTS
requests cycling through several utterance lengths, and reports
TTFA/E2E/RTF/per-chunk latency (mean, stdev, range) per cell. "Step size" is
``codec_chunk_frames``, the token-hop length the first emitted chunk waits
for; later chunks wait for a doubling hop capped at ``4 * codec_chunk_frames``
(see ``_token_hop_lengths``). Per-chunk latency is reconstructed from that
same ramp rather than from raw transport read boundaries, which don't align
to model chunks (see ``one_request``'s docstring).
Running with no flags reproduces the default matrix. Per-request correctness
is checked: non-silent, long enough, and (with a fixed ``--seed``)
byte-identical PCM across requests for the same text.
"""
import argparse
import json
import os
import signal
import statistics
import subprocess
import sys
import threading
import time
from pathlib import Path
import httpx
import numpy as np
REPO = Path(__file__).resolve().parents[2]
DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。"
DEFAULT_TEXTS = [
"Thank you so much for the birthday gift you sent me.",
"The weather is lovely today, with warm sunshine and a gentle breeze. "
"We decided to take a walk in the park to enjoy the spring flowers and the fresh air.",
"The weather is lovely today, with warm sunshine and a gentle breeze. We decided to take a walk in the "
"park to enjoy the spring flowers and the fresh air. Many people were out walking, children were "
"running and playing on the grass, and older folks sat on benches chatting with one another. We "
"strolled slowly along the path and came across a cherry tree in full bloom, its pink petals drifting "
"on the wind. It was so beautiful that we stopped to take some photos to remember the moment. Then we "
"continued on and reached a small lake, its water so clear you could see straight to the bottom, with "
"the blue sky and white clouds reflected on its surface. A few ducks paddled back and forth across the "
"water, occasionally letting out a soft quack. We sat on a bench by the lake for a while, enjoying the "
"peaceful afternoon. Time passed quickly, and before we knew it evening had arrived, with the setting "
"sun painting the sky in shades of gold. Reluctant to leave, we finally headed home, already planning "
"our next visit. On the way back we talked about everything we had seen and heard that day, feeling "
"wonderfully happy. It was truly an unforgettable weekend.",
]
# PCM16 mono @ 24 kHz (CosyVoice3 output rate)
SAMPLE_RATE = 24000
# Well below any real speech's RMS; catches all-zero/near-silent PCM that byte
# count and duration checks alone pass through.
MIN_RMS = 50.0
# Overlay: only codec_chunk_frames is overridden; must live next to cosyvoice3.yaml (relative base_config)
DEPLOY_TEMPLATE = """# Auto-generated by bench_cosyvoice3_concurrency_matrix.py -- do not edit.
base_config: cosyvoice3.yaml
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: {chunk_frames}
codec_max_chunk_frames: {max_chunk_frames}
codec_stream_scale_factor: {stream_scale_factor}
codec_left_context_frames: 25
codec_vocab_size: 6561
"""
# CosyVoice3Config constants (transformers_utils/configs/cosyvoice3.py): tokens
# are consumed at token_mel_ratio mel frames each, and HiFT emits
# samples_per_mel PCM16 samples per mel frame. Pinned here (not left to the
# server's defaults) since one_request's byte-boundary reconstruction below
# depends on knowing them exactly.
TOKEN_MEL_RATIO = 2
SAMPLES_PER_MEL = 480
STREAM_SCALE_FACTOR = 2
def _stats(values: list[float]) -> dict[str, float]:
"""Mean/median/stdev/p95/min/max/n for a list of per-request measurements."""
n = len(values)
if n == 0:
return {
"n": 0,
"mean": float("nan"),
"median": float("nan"),
"stdev": float("nan"),
"p95": float("nan"),
"min": float("nan"),
"max": float("nan"),
}
return {
"n": n,
"mean": statistics.mean(values),
"median": statistics.median(values),
"stdev": statistics.stdev(values) if n >= 2 else 0.0,
"p95": statistics.quantiles(values, n=20)[18] if n >= 2 else values[0],
"min": min(values),
"max": max(values),
}
def vllm_bin() -> str:
for c in (REPO / ".venv/bin/vllm", Path(sys.executable).parent / "vllm"):
if c.is_file():
return str(c)
return "vllm"
def wait_ready(url: str, proc: subprocess.Popen, log: Path, timeout: float) -> None:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if proc.poll() is not None:
tail = log.read_text(errors="replace").splitlines()[-40:] if log.exists() else []
raise RuntimeError("server exited during startup:\n " + "\n ".join(tail))
try:
if httpx.get(f"{url}/health", timeout=2.0).status_code == 200:
return
except Exception:
pass
time.sleep(2.0)
raise TimeoutError(f"server not ready within {timeout}s (log: {log})")
def _token_hop_lengths(chunk_frames: int, max_chunk_frames: int):
"""The token-hop-length ramp talker2code2wav_async_chunk uses to decide
when to fire the next flow-solve/emission: chunk_frames, then doubling
(STREAM_SCALE_FACTOR) each step, capped at max_chunk_frames. This -- not
codec_chunk_frames alone -- is what actually paces emitted audio chunks:
with the defaults here, hops are chunk_frames, 2x, 4x, 4x, 4x, ...
The very first emission is consistently one token short of chunk_frames
(confirmed against real speech.audio.delta event sizes) even though the
ramp's own progression is computed from the nominal chunk_frames, not the
actual first hop -- so only the first yield gets the -1 adjustment.
"""
yield chunk_frames - 1
hop = chunk_frames
while True:
hop = min(max_chunk_frames, max(chunk_frames, hop * STREAM_SCALE_FACTOR))
yield hop
def one_request(client, url, model, text, ref_audio, ref_text, chunk_frames, seed=None):
"""Run one streaming request, returning TTFA, E2E, per-chunk gaps, and audio bytes.
``response_format=pcm`` + ``stream_format=audio`` streams raw bytes with no
framing, so a client read boundary alone is a transport artifact, not a
model chunk boundary. Reconstructed instead by replicating the server's
own token-hop-length ramp (see ``_token_hop_lengths``) to compute each
real chunk's expected byte size ahead of time, and slicing the
accumulated byte stream at those offsets -- independent of how the
transport happens to fragment the underlying reads. The final chunk
(finalize, which releases extra held-back frames) is whatever remains
when the stream closes, not a predicted hop.
"""
payload = {
"model": model,
"input": text,
"response_format": "pcm",
"stream": True,
"stream_format": "audio",
"ref_audio": ref_audio,
"ref_text": ref_text,
}
if seed is not None:
payload["seed"] = seed
t0 = time.perf_counter()
ttfa = None
nbytes = 0
audio_bytes = bytearray()
chunk_gaps_ms = [] # inter-chunk arrival gaps (per model chunk), by step index
prev_t = None
step_idx = 0
bytes_per_token = TOKEN_MEL_RATIO * SAMPLES_PER_MEL * 2
hop_gen = _token_hop_lengths(chunk_frames, 4 * chunk_frames)
next_boundary = next(hop_gen) * bytes_per_token
last_crossed_bytes = 0
received = 0
last_now = t0
with client.stream("POST", f"{url}/v1/audio/speech", json=payload, headers={"Authorization": "Bearer EMPTY"}) as r:
if r.status_code != 200:
raise RuntimeError(f"{r.status_code}: {r.read()[:300]!r}")
for chunk in r.iter_bytes():
if not chunk:
continue
now = time.perf_counter()
if ttfa is None:
ttfa = now - t0
last_now = now
received += len(chunk)
nbytes += len(chunk)
audio_bytes += chunk
while received >= next_boundary:
if prev_t is not None:
chunk_gaps_ms.append((step_idx, (now - prev_t) * 1e3))
step_idx += 1
prev_t = now
last_crossed_bytes = next_boundary
next_boundary += next(hop_gen) * bytes_per_token
if received > last_crossed_bytes:
# Finalize chunk: releases extra held-back frames, so it doesn't
# match a predicted hop size -- its arrival is the stream closing.
if prev_t is not None:
chunk_gaps_ms.append((step_idx, (last_now - prev_t) * 1e3))
step_idx += 1
e2e = time.perf_counter() - t0
audio_s = nbytes / 2 / SAMPLE_RATE
return {
"ttfa_s": ttfa,
"e2e_s": e2e,
"audio_s": audio_s,
"rtf": e2e / audio_s if audio_s > 0 else float("nan"),
"chunk_gaps_ms": chunk_gaps_ms,
"nbytes": nbytes,
"audio_bytes": audio_bytes,
}
def check_correctness(r: dict, min_audio_s: float) -> list[str]:
"""Return a list of correctness violations (empty = pass)."""
problems = []
if r["nbytes"] <= 0:
problems.append("empty audio (0 bytes)")
if r["audio_s"] < min_audio_s:
problems.append(f"audio too short ({r['audio_s']:.2f}s < {min_audio_s:.2f}s)")
if r["nbytes"] > 0:
samples = np.frombuffer(bytes(r["audio_bytes"]), dtype=np.int16).astype(np.float64)
rms = float(np.sqrt(np.mean(samples**2)))
if rms < MIN_RMS:
problems.append(f"near-silent audio (rms={rms:.1f} < {MIN_RMS})")
return problems
def drive(url, args, n, conc, label, results, lock, idx, chunk_frames):
"""Run ``n`` requests at concurrency ``conc``, cycling through ``args.texts``."""
texts = args.texts
def worker():
with httpx.Client(timeout=600.0) as c:
while True:
with lock:
i = idx[0]
idx[0] += 1
if i >= n:
return
text = texts[i % len(texts)]
try:
r = one_request(
c, url, args.model, text, args.ref_audio, args.ref_text, chunk_frames, seed=args.seed
)
r["text"] = text
with lock:
results.append(r)
except Exception as e:
print(f" [{label}] request {i} failed: {e}", flush=True)
ts = [threading.Thread(target=worker) for _ in range(conc)]
t0 = time.perf_counter()
for t in ts:
t.start()
for t in ts:
t.join()
return time.perf_counter() - t0
def run_combo(args, conc: int, chunk_frames: int, out: Path) -> dict:
tag = f"c{conc}_chunk{chunk_frames}"
run_dir = out / tag
run_dir.mkdir(parents=True, exist_ok=True)
log = run_dir / "server.log"
deploy = REPO / "vllm_omni/deploy" / f"_autogen_steps_{tag}.yaml"
deploy.write_text(
DEPLOY_TEMPLATE.format(
chunk_frames=chunk_frames,
max_chunk_frames=4 * chunk_frames,
stream_scale_factor=STREAM_SCALE_FACTOR,
)
)
cmd = [
vllm_bin(),
"serve",
args.model,
"--omni",
"--deploy-config",
str(deploy),
"--host",
"127.0.0.1",
"--port",
str(args.port),
"--allowed-local-media-path",
str(REPO),
]
env = dict(os.environ)
if args.cuda_visible_devices:
env["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
if args.trt is not None:
env["COSYVOICE3_TRT"] = args.trt
print(f"\n=== concurrency={conc} step_size={chunk_frames} ===", flush=True)
print(f" launching server (log: {log})", flush=True)
with open(log, "w") as lf:
proc = subprocess.Popen(cmd, env=env, stdout=lf, stderr=subprocess.STDOUT, start_new_session=True)
url = f"http://127.0.0.1:{args.port}"
try:
wait_ready(url, proc, log, args.server_timeout)
print(" server ready", flush=True)
results: list[dict] = []
lock = threading.Lock()
idx = [0]
if args.warmup:
print(f" warmup ({args.warmup} req)", flush=True)
drive(url, args, args.warmup, min(conc, args.warmup), "warmup", results, lock, idx, chunk_frames)
results.clear()
idx[0] = 0
print(f" running {args.num_prompts} req @ concurrency {conc}", flush=True)
wall = drive(url, args, args.num_prompts, conc, "bench", results, lock, idx, chunk_frames)
ok = [r for r in results if r.get("ttfa_s") is not None]
problems = []
for r in ok:
problems += check_correctness(r, args.min_audio_s)
# With a fixed seed, same-text requests must produce byte-identical PCM (per text group)
if args.seed is not None:
by_text: dict[str, bytes] = {}
for r in ok:
t = r.get("text")
ref = by_text.get(t)
if ref is None:
by_text[t] = r.get("audio_bytes")
elif r.get("audio_bytes") != ref:
problems.append(f"byte-identity violation for text {t!r}: seeded requests produced different PCM")
n_ok = len(ok)
total_audio_s = sum(r["audio_s"] for r in ok)
all_gaps = [g for r in ok for g in r["chunk_gaps_ms"]] # (step_index, gap_ms) per request
gap_vals = [g for _, g in all_gaps]
# bucket gaps by chunk position within a request, to see any O(n) slope
by_idx: dict[int, list[float]] = {}
for step_index, g in all_gaps:
by_idx.setdefault(step_index, []).append(g)
step_latency_by_idx = {i: statistics.median(vals) for i, vals in sorted(by_idx.items())}
agg = {
"concurrency": conc,
"step_size": chunk_frames,
"ok": n_ok,
"errors": args.num_prompts - n_ok,
"correctness_problems": len(problems),
"correctness_sample": problems[:5],
"total_audio_s": total_audio_s,
"throughput_audio_s_per_s": total_audio_s / wall if wall > 0 else float("nan"),
"ttfa_s": _stats([r["ttfa_s"] for r in ok]),
"e2e_s": _stats([r["e2e_s"] for r in ok]),
"rtf": _stats([r["rtf"] for r in ok]),
"step_latency_ms": _stats(gap_vals),
"step_latency_ms_by_idx": step_latency_by_idx,
"num_steps": len(all_gaps),
# per-utterance-length breakdown; n per text = num_prompts / len(texts)
"by_text": {
t: {
"audio_s": _stats([r["audio_s"] for r in ok if r.get("text") == t]),
"e2e_s": _stats([r["e2e_s"] for r in ok if r.get("text") == t]),
"rtf": _stats([r["rtf"] for r in ok if r.get("text") == t]),
}
for t in sorted({r.get("text") for r in ok})
},
}
print(json.dumps(agg, indent=2, ensure_ascii=False), flush=True)
return agg
finally:
print(" shutting down server", flush=True)
try:
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
except Exception:
proc.terminate()
try:
proc.wait(timeout=60)
except subprocess.TimeoutExpired:
try:
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
except Exception:
proc.kill()
proc.wait(timeout=20)
deploy.unlink(missing_ok=True)
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--model", default=DEFAULT_MODEL)
p.add_argument(
"--text",
dest="texts",
nargs="+",
default=DEFAULT_TEXTS,
help="One or more utterance texts. With several, requests cycle through "
"them so a single server launch covers all utterance lengths; results "
"are bucketed per text.",
)
p.add_argument("--ref-audio", default=DEFAULT_REF_AUDIO)
p.add_argument("--ref-text", default=DEFAULT_REF_TEXT)
p.add_argument("--concurrency", type=int, nargs="+", default=[1, 8])
p.add_argument(
"--step-size",
type=int,
nargs="+",
default=[15],
help="codec_chunk_frames values to sweep (the streaming step size).",
)
p.add_argument("--num-prompts", type=int, default=30)
p.add_argument("--warmup", type=int, default=2)
p.add_argument("--port", type=int, default=8091)
p.add_argument(
"--server-timeout",
type=float,
default=1800.0,
help="Startup budget; first run also downloads weights and builds the TRT engine.",
)
p.add_argument("--output-dir", default="./perf_steps")
p.add_argument("--cuda-visible-devices", default="0")
p.add_argument(
"--trt",
choices=["0", "1"],
default=None,
help="Pin COSYVOICE3_TRT. 0=force torch estimator, 1=allow TensorRT. Default: unset.",
)
p.add_argument(
"--min-audio-s",
type=float,
default=1.0,
help="Minimum expected audio seconds per request for the correctness check.",
)
p.add_argument(
"--seed",
type=int,
default=42,
help="Fixed seed for reproducible (byte-identical) audio. Pass --seed '' "
"equivalent (any non-int) to disable -- there is no unset sentinel here "
"by design, since the default sweep is meant to be exactly reproducible.",
)
return p.parse_args()
def main() -> None:
args = parse_args()
out = Path(args.output_dir).resolve()
out.mkdir(parents=True, exist_ok=True)
all_results = []
failed = False
for conc in args.concurrency:
for chunk_frames in args.step_size:
try:
result = run_combo(args, conc, chunk_frames, out)
all_results.append(result)
if result["errors"] > 0 or result["correctness_problems"] > 0:
failed = True
except Exception as e:
print(f" combo c={conc} chunk={chunk_frames} FAILED: {e}", flush=True)
failed = True
out_path = out / "steps_results.json"
with open(out_path, "w", encoding="utf-8") as f:
json.dump({"model": args.model, "texts": args.texts, "runs": all_results}, f, indent=2, ensure_ascii=False)
print(f"\nResults written to {out_path}", flush=True)
print("\n=== SUMMARY ===")
header = (
f"{'conc':>4} {'step':>4} {'ok':>4} {'corr':>4} "
f"{'TTFA med(ms)':>12} {'step med(ms)':>12} {'step p95(ms)':>12} {'RTF mean+-sd':>14}"
)
print(header)
print("-" * len(header))
for r in all_results:
print(
f"{r['concurrency']:>4} {r['step_size']:>4} {r['ok']:>4} {r['correctness_problems']:>4} "
f"{r['ttfa_s']['median'] * 1e3:>12.1f} "
f"{r['step_latency_ms']['median']:>12.1f} {r['step_latency_ms']['p95']:>12.1f} "
f"{r['rtf']['mean']:>6.3f}+-{r['rtf']['stdev']:>5.3f}"
)
for t, stats in r.get("by_text", {}).items():
print(
f" text {t[:24]!r}: e2e {stats['e2e_s']['mean'] * 1e3:8.1f}"
f"+-{stats['e2e_s']['stdev'] * 1e3:6.1f} ms (n={stats['e2e_s']['n']}, "
f"audio {stats['audio_s']['mean']:.1f}s)"
)
if failed:
sys.exit(1)
if __name__ == "__main__":
main()
docs/configuration/stage_configs.md
+1 -1
View File
@@ -53,7 +53,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
| ------- | ------ | ---------- | --------- | ------------- |
| `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
| `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
| `session_mode` | str | optional | `"turn"` | Session behavior. Duplex models (pipelines with a `duplex_plugin`, e.g. MiniCPM-o 4.5) set `"duplex"`; `vllm-omni serve` then serves it over `/v1/realtime?duplex=1` and `/v1/chat/completions`, both on duplex sessions. |
| `session_mode` | str | optional | `"turn"` | Session behavior. For pipelines with a `duplex_plugin`, explicitly select `"duplex"` for the duplex engine or `"turn"` for the ordinary online serving stack. The shipped MiniCPM-o 4.5 default remains `"duplex"`. See [Full Duplex](../serving/full_duplex_api.md#enable-full-duplex). |
| `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. |
| `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. |
| `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
docs/design/architecture_overview.md
+11 -12
View File
@@ -131,13 +131,13 @@ for the validated batching, attention, and quantization combinations.
### MiniMax-H3: conditioning, denoising, and decode
MiniMax-H3 illustrates logical component and phase boundaries between text
encoding, task-specific denoising, and media decoding:
`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
these components in one diffusion stage. These boundaries are useful for
profiling and offload; independent placement would require a future
multi-stage `PipelineConfig`. The stage mapping follows the implementation and
serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
MiniMax-H3 combines text and reference-media encoding, task-specific denoising,
and media decoding. Standard serving co-locates these components in one diffusion
stage. The opt-in
[disaggregated topology](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md)
runs Qwen and the video/audio VAE encoders in Stage 0, then sends their conditioning
to a Stage 1 containing the DiTs and VAE decoders. The stage mapping follows the
implementation and serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
notes that a detailed technical report is forthcoming.
#### Model architecture by component
@@ -149,16 +149,15 @@ notes that a detailed technical report is forthcoming.
| **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
Both DiT partitions remain behind one task-selection point; `task` chooses
which partition runs. In the current implementation, the three logical
components execute inside one diffusion stage. This is different from
Qwen3-Omni's three independently scheduled stages, even though both models
produce synchronized audio and video-like outputs.
which partition runs.
#### Representative component execution
The following table describes the standard single-stage deployment.
| Logical component | Batching | Attention and execution | Parallelism | Quantization |
| --- | --- | --- | --- | --- |
| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
| **Text Encoder** | Follows the diffusion request scheduler | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
| **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
| **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
docs/design/fullduplex.md
+9 -10
View File
@@ -108,18 +108,17 @@ plugin and the session runtime config to `DuplexOrchestrator` directly.
`AsyncOmniEngine`, `OrchestratorBase` and `Orchestrator` carry only
template seams; `tests/engine/test_duplex_import_boundary.py` checks that
importing the turn-based stack loads no duplex module.
4. **Every surface a duplex server exposes is backed by a session.**
4. **Deployment selects the serving stack.**
`vllm-omni serve` constructs `DuplexOmni` when the pipeline declares
`duplex_plugin`; the server exposes `/v1/realtime?duplex=1` (alias
`duplex_plugin` and the deploy configuration sets `session_mode: duplex`.
An explicit `session_mode: turn` selects `AsyncOmni` instead, without
changing the pipeline's duplex capability. In duplex mode the server exposes `/v1/realtime?duplex=1` (alias
`/v1/duplex`), `POST /v1/chat/completions`, `/v1/models` and `/health`, and
every other turn-based route reports "not available". The chat route is
served by `DuplexChatCompletionsAdapter`, a Realtime client that runs one
short-lived session per request: speech becomes a committed turn, text is
seeded as the session's opening turn because a model-native model takes a
turn only when it hears speech. Whether a model can be reached with text is
`DuplexCapabilities.supports_chat_completions`; one that cannot is refused
at once rather than left to idle out. Turn-based use of the same model stays
available offline through `Omni` / `AsyncOmni`.
every other turn-based route reports "not available". The chat route uses
the ordinary chat service on the duplex engine, without a session, when
the plugin declares `DuplexCapabilities.supports_chat_completions`.
Turn-based deployments use the ordinary API initialization and retain
model capability checks and endpoint restrictions.
5. **Serving is transport only.** The websocket handler does socket I/O,
wire-envelope validation, command translation, event rendering and the
attachment/resume/replay bookkeeping; it holds no session state.
docs/features/comfyui.md
+13
View File
@@ -51,6 +51,7 @@ This extension offers the following nodes based on the output modalities:
- **Generate Video** for text-to-video and image-to-video tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
- **Generate Music** for text-to-music tasks
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
@@ -81,6 +82,18 @@ To build a simple workflow yourself,
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
- For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
## MiniMax H3 text-to-video workflow
The **MiniMax H3 Text to Video** template uses the existing remote Generate Video,
H3 Params, sampling, and Remote LoRA nodes. It includes Base settings and an
optional Turbo preset, with native resolution, 24 FPS, and `17k+5` frame counts.
ComfyUI receives the generated video and audio from the server; it does not load
H3 weights locally. Configure the server-side LoRA path before enabling Turbo.
See the [H3 workflow guide](https://github.com/vllm-project/vllm-omni/blob/main/apps/ComfyUI-vLLM-Omni/docs/minimax-h3-t2v.md)
for server setup, template import, Turbo configuration, recorded validation,
and saved-video/audio checks.
## Examples & Screenshots
Please read the [ComfyUI integration's README](https://github.com/vllm-project/vllm-omni/tree/main/apps/ComfyUI-vLLM-Omni) for more info.
docs/serving/full_duplex_api.md
+23 -7
View File
@@ -22,14 +22,14 @@ capability gates, see the [Realtime Duplex API guide](realtime_duplex_api.md).
## Enable Full Duplex
A model is served full duplex when its registered pipeline declares a
`duplex_plugin` (the model's `DuplexModelPlugin`). That declaration alone
decides it: `vllm-omni serve` constructs `DuplexOmni` instead of `AsyncOmni`,
and every surface the server exposes is backed by a duplex session. It serves
`duplex_plugin` (the model's `DuplexModelPlugin`) and its deploy configuration
sets `session_mode: duplex`. `vllm serve --omni` constructs `DuplexOmni`
instead of `AsyncOmni`. It serves
`/v1/realtime?duplex=1` (and its alias `/v1/duplex`),
`POST /v1/chat/completions`, `/v1/models` and `/health`; every other
turn-based HTTP route (speech, batch, embeddings, video, ...) answers "not
available". Turn-based use of such a model stays available offline through the
Python API (`Omni` / `AsyncOmni`).
available". Set `session_mode: turn` to select the ordinary online serving
stack instead. The model's supported tasks and endpoint restrictions still apply.
`/v1/chat/completions` is the ordinary chat service running on the duplex
engine: a request is a turn-based generation on the same stages, served
@@ -45,8 +45,24 @@ The deploy configuration of such a model must agree:
session_mode: duplex
```
A duplex model started with a deploy configuration that does not set it fails
at startup rather than falling back to turn-based serving.
A duplex-capable model must explicitly set `session_mode` to `duplex` or `turn`
in its deploy configuration, possibly through `base_config` inheritance.
Missing or invalid values fail at startup. The default MiniCPM-o deployment
continues to use duplex mode.
To run MiniCPM-o 4.5 with the turn-based engine:
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni \
--deploy-config vllm_omni/deploy/minicpmo_4_5_turn.yaml \
--trust-remote-code \
--port 8091
```
This profile inherits the default model and stage settings and overrides only
`session_mode`. It supports ordinary HTTP requests without creating a duplex
session handler. Mode selection happens at startup; a WebSocket query parameter
does not switch engines. The Python `Omni` / `AsyncOmni` APIs are unchanged.
!!! warning
docs/serving/realtime_duplex_api.md
+17 -30
View File
@@ -32,33 +32,19 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
mounts `ws://<host>:8099/v1/realtime?duplex=1` (this page; `ws://<host>:8099/v1/duplex`
is an alias of the same route), `POST /v1/chat/completions`, `/v1/models` and
`/health`, and nothing else. The remaining turn-based HTTP routes (speech,
batch, embeddings, video, ...) are not served by a duplex model.
batch, embeddings, video, ...) are not served in duplex mode.
`/v1/chat/completions` is not a turn-based path on a duplex server: each
request opens a short-lived duplex session and reads the answer off it. Audio
content becomes an ordinary committed turn; text is seeded as the session's
opening turn, because a model-native model takes a turn only when it hears
speech. A model that cannot be seeded that way declares so, and a text request
to it is refused with 400. Two consequences are worth sizing for:
`/v1/chat/completions` uses the ordinary chat service on the same engine
when the plugin declares `supports_chat_completions`. HTTP chat requests do
not create duplex sessions or consume session admission slots. Their options
follow the [Chat Completions API](chat_completions_api.md) and the model's
capabilities and endpoint restrictions.
- **Every request holds an admission slot** for its lifetime, so
`duplex_session.max_sessions` caps HTTP concurrency as well as websocket
sessions. Beyond it a request is refused with HTTP 503.
- **Latency follows the model's own clock.** A model-native session generates
per audio unit, so even a short answer costs what that many units cost --
noticeably more than the same model would take turn-based.
`n > 1`, `logprobs`, `tools`, and image or video content parts have no
representation in one duplex turn and are refused with HTTP 400 -- a Realtime
conversation item carries text and audio only. `temperature`,
`max_completion_tokens`, `modalities` and
`chat_template_kwargs.use_tts_template` apply; the other
`chat_template_kwargs` keys are logged as ignored. `extra_body` is passed to
the session unchanged, exactly as a websocket client's session `extra_body`
is, which is how a request asking for `modalities: ["text", "audio"]` supplies
the `ref_audio` MiniCPM-o requires for audio output. A model that should not
serve the route lists it in the deploy config's `endpoint_restrictions`,
exactly as a turn-based model does.
For a turn-based online deployment, use `minicpmo_4_5_turn.yaml` or explicitly
set `session_mode: turn` in a deploy configuration derived from the model's
profile. This selects `AsyncOmni` and the ordinary API initialization instead
of the duplex handler. See [Full Duplex](full_duplex_api.md#enable-full-duplex)
for the command and configuration requirements.
### Run the Example Client
@@ -1345,8 +1331,9 @@ them out into typed events before they reach a client.
- Session capacity is bounded by `duplex_session.max_sessions` in the deploy
configuration; admission beyond it fails with `resource_exhausted`
(`rate_limit_error`, retryable).
- A duplex deployment serves the duplex session surfaces only: the websocket
route and `/v1/chat/completions`, the latter on a session per request (see
[Start the Server](#start-the-server)). The other turn-based HTTP routes
report "not available" on a MiniCPM-o 4.5 server; turn-based use of the
model stays available offline through `Omni` / `AsyncOmni`.
- A duplex deployment serves the websocket route and, when supported by the
model, ordinary `/v1/chat/completions` requests on the same engine without a
session per request (see [Start the Server](#start-the-server)). Other
turn-based HTTP routes report "not available" in duplex mode. To use the
ordinary online serving stack, select `session_mode: turn` as described in
[Full Duplex](full_duplex_api.md#enable-full-duplex).
docs/user_guide/quantization/mxfp4.md
+310 -104
View File
@@ -6,22 +6,25 @@ W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4
(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups
of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale.
On Ascend, single-scale MXFP4 also supports Universal Optimal Scaling
(UOS, `Qmax=7.25`), described in the single-scale scaling section below.
vLLM-Omni provides two quantization methods with different scale structures:
| Method | Scale structure | Mode | Use case |
|--------|----------------|------|----------|
| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
| -------- | ---------------- | ------ | ---------- |
| `mxfp4` | Single-scale (per-32 fine only) | Online + pre-quantized single-scale weights | Online baseline, or an existing single-scale checkpoint in the layout below |
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | W4A4 with dual-level scales; offline mode loads calibrated Smooth tensors |
!!! tip "Recommended: `mxfp4_dualscale` offline"
For production deployments, use the `mxfp4_dualscale` offline mode with a
!!! tip "Existing W4A4-only DualScale deployments"
For W4A4-only deployments, the existing `mxfp4_dualscale` offline mode uses a
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
calibrated `mul_scale` tensors from disk, providing measurably better accuracy
than any online method. The one-time preprocessing cost amortises across all
subsequent inference runs.
calibrated `mul_scale` tensors from disk. Quality depends on the model,
calibration recipe and BF16 fallback policy; evaluate the checkpoint against
your quality target before deployment.
Use `mxfp4` online only for quick experimentation where preprocessing time
is not acceptable and accuracy loss is tolerable.
W4A8 layer/step fallback requires single-scale `mxfp4`, online or offline.
DualScale supports W4A4 and BF16 routing only.
!!! warning "Online single-scale ≠ Offline dual-scale"
`mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`:
@@ -34,61 +37,244 @@ vLLM-Omni provides two quantization methods with different scale structures:
## Hardware Support
| Device | Support |
|--------|---------|
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
| AMD ROCm (gfx950 / MI355X) | ✅ |
| Intel XPU | ⭕ |
| Ascend NPU (Atlas 950 A5) | ✅ |
The following matrix describes these vLLM-Omni implementations, rather than
hardware support for the MXFP4 format in general.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
| Device | Single-scale online OCP | Single-scale offline | UOS / W4A8 fallback | DualScale |
| -------- | -------- | -------- | -------- | -------- |
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ | ✅ |
| AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ | ❌ |
| Other AMD architectures | ❌ | ❌ | ❌ | ❌ |
| NVIDIA GPU | ❌ | ❌ | ❌ | ❌ |
| Intel XPU | ❌ | ❌ | ❌ | ❌ |
Legend: `✅` supported by the implementation, `❌` unsupported.
ROCm uses the existing AITER online path and requires `mxfp4_scale_alg=0`.
## Model Type Support
### Diffusion Model (Wan2.2)
| Model | Online | Offline | Notes |
|-------|--------|---------|-------|
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
| ------- | -------- | --------- | ------- |
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Cascade experts resolve checkpoint metadata independently; explicit global runtime policy applies to both |
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Per-expert configuration uses the same interface as T2V; validate I2V checkpoints separately |
| Wan2.2-TI2V-5B | Requires validation | Requires validation | This guide does not establish checkpoint compatibility or generation quality for TI2V-5B |
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
quantization quality, not model compatibility — both work on cascade (A14B) and
single-transformer models alike, the same as `mxfp8` online:
Both methods can be integrated with single-transformer and cascade models.
This alone does not establish support or quality for every model variant:
- `mxfp4`: single-scale, lower overhead, simpler compute, online only
- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline
- `mxfp4`: single-scale; online and pre-quantized weights, with optional W4A8 fallback on NPU
- `mxfp4_dualscale`: dual-scale; online and offline, with BF16 fallback
**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced
by the merge script); there is no offline `mxfp4` single-scale format.
The existing `merge_mxfp4_dualscale_checkpoint.py` produces **offline
`mxfp4_dualscale`** checkpoints. Single-scale weights must satisfy the separate
layout described below; changing a DualScale checkpoint's method name to
`mxfp4` does not convert its weights. Prepare single-scale checkpoints with
external tooling. Real exported checkpoint compatibility and model quality
require validation of that checkpoint; the synthetic tests do not establish
UOS export or video acceptance.
!!! note "Per-layer BF16 fallback in offline cascade models"
The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most
An A14B offline DualScale checkpoint uses `quant_method: mxfp4_dualscale`. Most
linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers
retain their original BF16 weights and are listed in `ignored_layers` inside
each transformer's `config.json`. The two transformers may have different
`ignored_layers` sets — the pipeline reads each transformer's own `config.json`
and rebuilds the config locally when they differ, so routing is always
per-transformer-accurate.
!!! warning "TI2V-5B not supported"
Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter
count makes it significantly more sensitive to 4-bit quantization noise,
resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B.
and resolves the storage format and BF16 layer policy independently.
## Configuration
### W4A8 fallback at selected layers and denoising steps (Ascend NPU)
For single-scale `mxfp4`, set `w4a8_fallback_steps` to a list of zero-based
request denoising loop indices (not scheduler timestep values).
Set `w4a8_fallback_layers` to exact runtime Linear paths, for example
`["blocks.10.attn1.to_qkv", "blocks.12.ffn.net_2"]`. A selected layer uses W4A8
at every step; at a selected step, all other quantized layers also use W4A8.
The rules are combined with **OR**. W4A8 uses MXFP8 activations with MXFP4
weights. Omitting both lists or passing `[]` keeps the existing W4A4 behavior. Existing
`ignored_layers` take priority: **BF16 > (layer OR step) W4A8 > W4A4**.
DualScale's existing leading BF16 block rule is unchanged; it does not enable
W4A8.
```python
from vllm_omni import Omni
omni = Omni(
model="<Wan2.2-T2V-A14B-Diffusers>",
quantization_config={
"method": "mxfp4",
"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
"w4a8_fallback_steps": [0, 1, 38, 39],
},
)
```
Both policies also work with offline single-scale `mxfp4`. Any nonempty
W4A8 list with `mxfp4_dualscale`, online or offline, raises an explicit error
at configuration construction. Empty lists preserve its existing W4A4 path.
Each expert's checkpoint still determines its serialized
format and BF16 layer routing. A serialized checkpoint config is the complete
storage contract: `ignored_layers`, or its `modules_to_not_convert` alias,
selects the BF16 layers, and omitting both means no explicit BF16 layers. A
method-only checkpoint value such as `"mxfp4"` or
`{"quant_method": "mxfp4"}` carries no storage policy and leaves a caller's
online `ignored_layers` unchanged. An explicitly supplied quantization config
retains its W4A8 layer and step lists when the loader rebuilds either expert's
config from disk, including empty lists that disable checkpoint policies.
Within an explicit runtime config, an omitted W4A8 list defaults to `[]`;
individual omitted fields do not inherit saved policies. Without an explicit
runtime config, auto-detection uses the saved checkpoint policies.
Layer names are relative to each expert (`blocks.10...`, without a
`transformer.` or `transformer_2.` prefix). A global config applies the same
list to both experts; per-component `transformer`/`transformer_2` configs can
select different lists. Wan selects these expert names exactly: a `transformer`
entry does not match `transformer_2`; an omitted expert uses `default` if set.
A component mapped to `null` (or an unmatched component
without a default) stays unquantized and requires a BF16 checkpoint; loading a
quantized checkpoint for that disabled component raises an error. A top-level
`quantization_config=None` retains checkpoint auto-detection.
Names are exact, with no regex, wildcard, block expansion,
or checkpoint-name remapping. Wan validates them against its runtime Linear
modules and logs the selected paths/count and BF16 overrides per PP rank.
Use `attn1.to_qkv` for fused self-attention; `attn1.to_q`, `to_k`, or `to_v`
checkpoint names raise an error. Cross-attention `attn2.to_q/to_k/to_v` remains
separate. Layer-only fallback does not require a denoise ForwardContext.
The text-to-video example accepts the JSON directly:
```bash
ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 \
python examples/offline_inference/text_to_video/text_to_video.py \
--model /path/to/Wan2.2-T2V-A14B-Diffusers \
--quantization-config '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}' \
--tensor-parallel-size 1 --ulysses-degree 2 --cfg-parallel-size 2 \
--enforce-eager --height 480 --width 832 --num-frames 81 \
--num-inference-steps 40 --seed 42 --output wan22-w4a8.mp4
```
Use either `--quantization` or `--quantization-config` in this example.
Serving accepts the same JSON through `--diffusion-quantization-config`.
The list uses one request's loop index across both high- and low-noise experts,
and both CFG branches see the same index. A new request starts at step zero.
Indices outside the request's step range are never selected. A pipeline must
publish `ForwardContext.denoise_step_idx`; requesting step fallback without
that context raises an error instead of silently ignoring the list.
Single-scale W4A4 and W4A8 reuse the same prepared `weight` and `weight_scale`.
Online loading quantizes the original BF16/FP16 weight once; offline loading
packs the numeric FP4 values and reshapes the saved E8M0 scale without
requantizing the weight. Only activation precision changes during forward.
There is no DualScale decoding, second W4 representation or `w4a8_weight` /
`w4a8_weight_scale` cache. No weight quantization runs inside the denoising loop.
The mixed FP8/FP4 GEMM uses BF16 output and a matching BF16 row bias for
compatibility with the target A5 runtime. An FP16 caller receives that result
converted back to FP16; its output shape and dtype are preserved.
The single-scale fallback validation target is eager Wan2.2 T2V A14B on A5.
Quality scores, optimal fallback lists, real UOS exports and additional
compile/offload combinations require separate evaluation. Existing W4A4-only
DualScale still needs complete 512-channel input groups; use TP=1 with
sequence/CFG parallelism for its initial multi-card run.
### `mxfp4` — Pre-Quantized Single-Scale Weights (Ascend NPU)
Set the following fields in each transformer's `config.json` when a
single-scale checkpoint uses calibrated Smooth scaling:
```json
{
"quantization_config": {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": true,
"require_smooth_scale": true,
"ignored_layers": []
}
}
```
Each quantized Linear must provide tensors under its model parameter name:
| Tensor | Checkpoint layout | Meaning |
| -------- | ------------------- | --------- |
| `weight` | `(N, K)` numeric FP4 values | Packed to FP4 after loading |
| `weight_scale` | `(N, ceil(K / 32))`, `uint8` | Raw E8M0 exponent bytes, required |
| `mul_scale` | `(K,)`, FP32 | Smooth pre-scale, required when declared above |
Both W4A4 and W4A8 quantize `x * mul_scale`. The Smooth tensor is converted
to the layer's activation dtype after loading. Row-parallel layers load the
corresponding input-channel shard. Pre-fused QKV weights must use one shared
Smooth tensor for the fused projection.
For backward compatibility, `require_smooth_scale` defaults to `false`:
an absent tensor uses identity scaling, while a present tensor is still loaded.
Without a declaration, an absent Smooth tensor cannot be distinguished from a
legacy unsmoothed checkpoint. Calibrated checkpoints should declare
`require_smooth_scale: true` so missing or misnamed Smooth tensors fail loading;
incorrect tensor shapes, non-floating tensors, NaN/Inf and nonpositive values
also fail. Values that overflow or underflow to zero in the activation dtype
are rejected during weight processing. The flag requires serialized single-scale
weights and is rejected for online quantization. It does not make DualScale
scales optional or change missing-parameter policies for other backends.
Either the checkpoint or an explicit caller can require Smooth: a `true`
declaration is retained when per-expert metadata rebuilds the configuration,
including when the other source omits the flag or sets it to `false`.
Prepare the canonical numeric FP4 checkpoint before inference using external
tooling. Runtime loading does not parse producer-specific export formats.
Retain the producer recipe to identify the weight-generation algorithm and
Smooth calibration; tensor storage alone does not establish UOS provenance.
### `mxfp4` — Single-Scale Block Scaling (OCP / UOS)
`mxfp4_scale_alg` selects the MXFP4 block-scaling method on NPU.
UOS is a scaling rule, not a separate activation dtype or an activation-only
algorithm. Both choices use FP4 E2M1 values with shared E8M0 scales.
| Mode | Weight preparation | A4 activation quantization | W4A8 fallback |
| --- | --- | --- | --- |
| Online | Selected scaling method, once at loading | Selected scaling method | Same W4; OCP MXFP8 A8 |
| Offline | Load existing W4 and scales unchanged | Selected scaling method | Same W4; OCP MXFP8 A8 |
The available settings are:
| Setting | MXFP4 scaling method | NPU parameters |
| --- | --- | --- |
| `0` (default) | OCP MX | `scale_alg=0` |
| `2` | UOS (`Qmax=7.25`) | `scale_alg=2, dst_type_max=7.25` |
Universal Optimal Scaling (UOS, `Qmax=7.25`) is derived in
[MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attention](https://arxiv.org/abs/2607.24377),
Section 4.1. The value `7.25` is the boundary used to select the shared
power-of-two scale; the maximum finite E2M1 value remains `6`.
This option applies UOS to Linear quantization. It does not implement the
paper's full MXAttention pipeline or Pre-Normalization Quantization (PNQ).
Both use FP4 E2M1 with E8M0 scales, `axis=-1`, `block_size=32`,
`round_mode="rint"`. W4A8 always quantizes the activation to
`float8_e4m3fn` with `scale_alg=0` and the same grouping/rounding, without
UOS parameters. A4 and A8 share the prepared W4 and scale. UOS is currently
NPU-only; ROCm retains its existing AITER online algorithm.
An explicit UOS online configuration is
`{"method": "mxfp4", "mxfp4_scale_alg": 2}`. For offline checkpoints the
setting affects **only A4 activations**. Existing weights are never requantized:
a UOS checkpoint can run UOS W4 + UOS A4 or UOS W4 + OCP A4, and both fall back to
UOS W4 + OCP A8. Record the weight-generation algorithm separately from the
runtime activation setting. Smooth remains active with either setting. The
UOS configuration explicitly selects `2`; attention computation is unchanged.
### `mxfp4` — Single-Scale Online Mode
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
calibration `mul_scale` is available. Applies equally to single-transformer
and cascade (A14B) models — both transformers in a cascade receive the same
quantization config automatically.
Online mode quantizes BF16/FP16 weights once at loading and computes activation
scales during forward. Each block of 32 K elements shares one E8M0 scale;
no calibrated `mul_scale` is available. An explicit global quantization config
applies to both cascade experts. Per-component configs can select different
policies; checkpoint auto-detection resolves each expert independently.
```python
from vllm_omni import Omni
@@ -103,7 +289,7 @@ outputs = omni.generate(
```bash
# Single-transformer or cascade model — same command
python text_to_video.py --model <your-model> --quantization mxfp4
python examples/offline_inference/text_to_video/text_to_video.py --model <your-model> --quantization mxfp4
# Online serving
vllm serve <your-model> --omni --quantization mxfp4
@@ -113,19 +299,19 @@ vllm serve <your-model> --omni --quantization mxfp4
Online DualScale mode computes both fine and coarse scales on the fly from BF16
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
DualScale provides better quantization accuracy at higher compute cost.
single-transformer and cascade (A14B) models with compatible Linear layers.
Compare quality and performance using the intended model and workload.
The default configuration keeps the leading 5 transformer blocks in BF16
(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
is sufficient to meet quality requirements and is the recommended setting.
(`num_bf16_fallback_layers=5`). This is an implementation default; select and
validate the BF16 policy against your own quality target.
```python
omni = Omni(model="<your-model>", quantization="mxfp4_dualscale")
```
```bash
python text_to_video.py --model <your-model> --quantization mxfp4_dualscale
python examples/offline_inference/text_to_video/text_to_video.py --model <your-model> --quantization mxfp4_dualscale
```
If accuracy debugging identifies additional precision-sensitive layers, they
@@ -136,21 +322,22 @@ omni = Omni(
model="<your-model>",
quantization_config={
"method": "mxfp4_dualscale",
"ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override
"ignored_layers": ["blocks.10.attn1.to_qkv"], # explicit per-layer override
},
)
```
BF16 fallback routing in online mode applies two rules in priority order:
1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix
matches is kept in BF16 regardless of block index.
1. **`ignored_layers`** (explicit per-layer override): a listed runtime Linear
module is kept in BF16 regardless of block index. Use exact module paths,
not arbitrary block prefixes or names ending in `.weight`.
2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N
transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults
to `5` (recommended). Layers outside `blocks.N.*`
(e.g. `condition_embedder`) are always quantized.
to `5`. Layers outside `blocks.N.*` are quantized unless explicitly listed
in `ignored_layers`.
### `mxfp4_dualscale` — DualScale Offline Mode (Recommended)
### `mxfp4_dualscale` — DualScale Offline Mode
Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A
preprocessing step converts the raw quantized output to the diffusers format
@@ -169,8 +356,8 @@ inside `config.json`. At runtime, each layer's prefix is matched against
Each quantized linear layer stores four tensors:
| Tensor | Shape | dtype | Description |
|--------|-------|-------|-------------|
| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) |
| -------- | ------- | ------- | ------------- |
| `weight` | `(N, K)` | float8_e4m3fn | Numeric FP4 values; packed into two values per byte after loading |
| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) |
| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale |
| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) |
@@ -223,7 +410,7 @@ python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \
```
| Argument | Description |
|----------|-------------|
| ---------- | ------------- |
| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` |
| `--original-model` | Root directory of the original BF16 diffusers model |
| `--quant-path` | Root directory of the msModelSlim quantized output |
@@ -251,14 +438,14 @@ The `quantization_config` injected into `config.json` for each transformer:
```
`ignored_layers` lists every linear layer that retains its original BF16 weight,
using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out`
using vllm-omni runtime module paths (QKV-fused, FFN underscored, `to_out`
unindexed). The exact entries are determined by the quantization tool (msModelSlim)
and may differ between `transformer` and `transformer_2` in a cascade model.
#### Step 3 — Serve
```bash
python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
python examples/offline_inference/text_to_video/text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
# Online serving
vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni
@@ -275,25 +462,32 @@ omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale")
## Parameters
### `mxfp4` (single-scale, online only)
### `mxfp4` (single-scale, online + offline)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| ----------- | ------ | --------- | ------------- |
| `method` | str | — | `"mxfp4"` |
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 |
| `mxfp4_scale_alg` | int | `0` | MXFP4 block scaling: `0` OCP MX, `2` UOS (`Qmax=7.25`, NPU only). Online: W4 and A4; offline: A4 only. A8 remains OCP MXFP8 |
| `is_checkpoint_mxfp4_serialized` | bool | `false` | Load pre-quantized single-scale weights and scales; declared in checkpoint `config.json` |
| `require_smooth_scale` | bool | `false` | Require calibrated Smooth tensors for offline single-scale weights |
| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16 |
| `w4a8_fallback_layers` | list[str] | `[]` | Exact runtime Linear paths that always use A8; NPU only |
| `w4a8_fallback_steps` | list[int] | `[]` | Zero-based denoise steps that use A8; NPU only |
### `mxfp4_dualscale` (dual-scale, online + offline)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| ----------- | ------ | --------- | ------------- |
| `method` | str | — | `"mxfp4_dualscale"` |
| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script |
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B |
| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Validate the default of `5` against the intended model and quality target |
| `w4a8_fallback_layers` | list[str] | `[]` | Must remain empty; DualScale does not support W4A8 fallback |
| `w4a8_fallback_steps` | list[int] | `[]` | Must remain empty; DualScale does not support W4A8 fallback |
#### BF16 fallback priority (online mode)
```
```text
for each linear layer:
if prefix in ignored_layers → BF16 (explicit override, highest priority)
elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule)
@@ -307,8 +501,9 @@ unless they appear in `ignored_layers`.
1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using
`npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is
available — all output partitions receive an identity pre-scale. No offline
checkpoint format exists for this method.
available — all output partitions receive an identity pre-scale. Offline
single-scale checkpoints use the canonical numeric FP4 layout
described above.
2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)**
quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse
@@ -316,7 +511,7 @@ unless they appear in `ignored_layers`.
explicit `ignored_layers` stay in BF16 for accuracy.
3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** —
**recommended for production** — loads four tensors per quantized layer: FP4
loads four tensors per quantized layer: FP4
weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale
(`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback
layers have no quantization tensors and are routed via `ignored_layers`.
@@ -325,21 +520,21 @@ unless they appear in `ignored_layers`.
layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy
float32 round-trip.
5. **Cascade model config propagation**: in a cascade model (transformer +
transformer_2), vLLM-Omni reads each transformer's own `config.json` and
rebuilds the quant config locally when `ignored_layers` differs between
transformers, ensuring per-layer routing is accurate for each. The first
transformer's config is propagated to `od_config` so the second transformer
can reuse it as a starting point.
5. **Cascade model configuration**: checkpoint auto-detection resolves each
expert's own `config.json` independently; the first expert's detected config
is not reused as the second expert's policy. An explicit global runtime
config applies to both experts, while per-component configs select policies
by the exact `transformer` / `transformer_2` names. Serialized checkpoint
metadata still determines each expert's storage format and BF16 layers.
6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a
single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the
fused name (`attn1.to_qkv`), written automatically by the merge script.
7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The
DualScale offline method mitigates this with calibrated `mul_scale` smooth
quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
off compression vs. accuracy for precision-sensitive layers.
7. FP4 has fewer representable values than FP8, so evaluate quantization
error and generation quality for the chosen recipe. Offline Smooth scaling
and BF16 fallback can help with sensitive layers.
`num_bf16_fallback_layers` applies only to online DualScale.
## Adapting MXFP4 for a New Model
@@ -350,27 +545,26 @@ a merge script for offline checkpoints.
### Step 1 — Discover runtime layer names
`ignored_layers` entries must match the **runtime parameter names** used inside
vllm-omni, which may differ from the names stored in the diffusers checkpoint.
The canonical source of truth is the model's own `named_parameters()`.
`ignored_layers` entries use exact **runtime Linear module paths**, without
`.weight`. Inspect the actual transformer instance in model initialization or
a developer harness; `Omni` does not expose an `omni.pipeline.transformer` API.
For a transformer instance named `transformer`:
```python
from vllm_omni import Omni
from vllm.model_executor.layers.linear import LinearBase
# Load the model without quantization to inspect parameter names.
omni = Omni(model="/path/to/your-model") # no --quantization flag
for name, _ in omni.pipeline.transformer.named_parameters():
if "weight" in name and "scale" not in name:
for name, module in transformer.named_modules():
if isinstance(module, LinearBase):
print(name)
```
Compare the printed names against the diffusers checkpoint keys
Compare these module paths against the corresponding diffusers checkpoint keys
(`safetensors.safe_open` or `torch.load`) to identify any renames your model
applies. Common patterns that differ in Wan2.2 (and may appear in other
models):
| Diffusers checkpoint name | vllm-omni runtime name | Reason |
|---------------------------|------------------------|--------|
| --------------------------- | ------------------------ | -------- |
| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` |
| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores |
| `ffn.net.2` | `ffn.net_2` | Same underscore rule |
@@ -382,8 +576,9 @@ fused runtime parameters.
!!! warning "Partial QKV fallback is not allowed"
If your model fuses Q, K, V into a single layer, `ignored_layers` must
include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but
`to_k`, `to_v` quantized) cannot be expressed at runtime because they share
name the fused runtime module (for Wan, `attn1.to_qkv`). The source
checkpoint must keep **all three or none** in BF16. A partial fallback
(e.g. `to_q` in BF16 but `to_k`, `to_v` quantized) cannot be expressed at runtime because they share
one `QKVParallelLinear`. The merge script enforces this and raises an error
if only some of the trio appear as non-quantized.
@@ -392,12 +587,13 @@ fused runtime parameters.
#### Online mode
Pass `ignored_layers` directly in the quantization config using the **runtime
names** discovered in Step 1. No code changes to the model are required.
names** discovered in Step 1. This assumes the model already passes its
quantization config and module prefixes to compatible vLLM Linear layers.
```python
omni = Omni(
model="/path/to/your-model",
quantization={
quantization_config={
"method": "mxfp4_dualscale",
"ignored_layers": [
"blocks.0.attn1.to_qkv", # runtime name, not diffusers name
@@ -409,9 +605,14 @@ omni = Omni(
```
```bash
# CLI does not support list-typed ignored_layers directly.
# Use the Python API or set ignored_layers in config.json (offline).
python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale
# The text-to-video example accepts list-valued fields through JSON.
python examples/offline_inference/text_to_video/text_to_video.py \
--model /path/to/your-model \
--quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}'
# Serving accepts the same configuration.
vllm serve /path/to/your-model --omni \
--diffusion-quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}'
```
The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers
@@ -436,9 +637,13 @@ the merge script is correct. The injected block:
}
```
To add a layer manually (e.g. to pin an additional layer to BF16 without
re-running the merge script), edit `config.json` inside the transformer
subfolder. Use runtime names, not diffusers checkpoint names.
Adding an offline BF16 fallback layer requires both the original BF16 weight
from the base model and a matching `ignored_layers` entry. Regenerate the
checkpoint with the revised storage policy, restoring the original weight
(and bias, if present) and removing that layer's quantization tensors.
Editing `config.json` alone does not reconstruct BF16 weights from stored FP4
values. For fused QKV, restore all three projections together. Use runtime
module paths in `ignored_layers`.
### Step 3 — Write a merge script for offline mode
@@ -474,6 +679,7 @@ renaming rules and collect the results.
!!! tip "Validate before serving"
After producing the offline checkpoint, load it without a `--quantization`
flag and verify that vLLM-Omni auto-detects the correct method. Check that
the layer count reported in the startup log matches expectations: quantized
layer count + `ignored_layers` count should equal total linear layer count.
Any mismatch indicates a name-mapping bug in the merge script.
each intended runtime Linear module resolves to the expected quantized or
BF16 method and that every BF16 entry has original BF16 weights. Account
for fused modules and, in online DualScale, the leading-block fallback rule;
the raw length of `ignored_layers` alone is not a reliable layer count.
examples/offline_inference/text_to_video/text_to_video.py
+10 -1
View File
@@ -384,13 +384,20 @@ def parse_args() -> argparse.Namespace:
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
quantization_group = parser.add_mutually_exclusive_group()
quantization_group.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
)
quantization_group.add_argument(
"--quantization-config",
type=json.loads,
default=None,
help='Quantization JSON, e.g. \'{"method":"mxfp4","w4a8_fallback_steps":[0,2]}\'.',
)
# Distributed and parallel execution
parser.add_argument(
@@ -572,6 +579,8 @@ def main():
omni_kwargs["flow_shift"] = args.flow_shift
if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization
if args.quantization_config is not None:
omni_kwargs["quantization_config"] = args.quantization_config
if args.cache_backend is not None:
omni_kwargs["cache_backend"] = args.cache_backend
omni_kwargs["cache_config"] = cache_config
pyproject.toml
-2
View File
@@ -66,14 +66,12 @@ dev = [
"jiwer==4.0.0",
"zhon==2.1.1",
"zhconv==1.4.3",
"scipy==1.18.0",
"funasr==1.3.29",
"lpips==0.1.4", # for quantization fp8
"orjson==3.11.9", # to make v1/embeddings API fast
"voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
"torchdiffeq==0.2.5", # for covo_audio Test
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests
"s3tokenizer==0.3.0", # for tts test
# Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
# whose top-level stepaudio2 package conflicts with this distribution.
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
recipes/BreezeBlue/Breeze-TTS-2.md
+25 -8
View File
@@ -30,10 +30,20 @@ Four prompt modes are selected automatically from the request fields:
| `ref_clone_tata` (clone) | `ref_audio` + `ref_text` + `input` | reference clip and its exact transcript |
| `ref_edit_tata` (voice direction) | reference trio + `instructions` | keep the reference timbre, direct the delivery |
Current scope: greedy sampling with `cfg_scale=1.0`; CFG ≠ 1.0 and
Current scope: greedy sampling with `cfg_scale=1.0` and a default
`repetition_penalty=1.1` on generated codebook-0 tokens; CFG ≠ 1.0 and
`negative_prompt` are rejected with explicit client errors until companion
request support lands.
The [official inference example](https://github.com/breezeblue-ai/breeze-tts/blob/main/infer.py)
also uses repetition penalty `1.1`, but its
[runtime defaults](https://github.com/breezeblue-ai/breeze-tts/blob/main/breeze_infer/runtime.py)
sample both the backbone and depth decoder with temperature `0.9` and top-k
`50`. The [official voice-design examples](https://huggingface.co/BreezeBlue/Breeze-TTS-2#-voice-design)
use CFG `4`. Those sampling and guidance settings are outside the current
port's scope, so the examples below do not establish parity with the official
voice-design outputs.
## Supported model contract
| Item | Value |
@@ -42,9 +52,9 @@ request support lands.
| Modes | plain, voice design (`instructions`), voice clone (`ref_audio` + `ref_text`), voice direction (reference + `instructions`) |
| Languages | English and Chinese (model card) |
| Reference audio | exactly one clip per request with its exact transcript; encoded by the bundled Qwen3-TTS tokenizer to 16 codebooks at 12.5 Hz |
| Output | 24 kHz mono; `wav` or `pcm`; streaming SSE `speech.audio.delta` |
| Output | 24 kHz mono; `wav` or `pcm`; raw audio streaming (`stream_format="audio"`) or SSE `speech.audio.delta` (`stream_format="sse"`) |
| Length | prompt bounded by stage-0 `max_model_len=4096`; synthesis bounded by `max_new_tokens` codec frames (default 2048, about 164 s) |
| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0`; other guidance values and `negative_prompt` are rejected |
| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0` and `repetition_penalty=1.1`; other guidance values and `negative_prompt` are rejected |
| Deployment profile | `vllm_omni/deploy/breeze_tts_2.yaml`: two stages on one GPU, async-chunk streaming with 8-frame codec chunks, stage 0 `gpu_memory_utilization=0.80`, stage 1 `0.15` |
## References
@@ -106,7 +116,7 @@ Voice design (instruction only, no reference audio):
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "[laughs] Welcome to tonight's story time.",
"input": "(laugh) Welcome to our story time.",
"instructions": "A warm young woman, clear voice, lively delivery.",
"response_format": "wav"
}' --output breeze_design.wav
@@ -139,7 +149,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
}' --output breeze_direction.wav
```
Streaming PCM (SSE `speech.audio.delta` events):
Streaming raw PCM (`stream_format="audio"`):
```bash
curl -N -X POST http://localhost:8091/v1/audio/speech \
@@ -153,8 +163,9 @@ curl -N -X POST http://localhost:8091/v1/audio/speech \
}' --output breeze_stream.pcm
```
Expected results: finite, non-silent 24 kHz mono WAV files; streaming requests
return incremental PCM deltas. On 1x L20 (greedy, single request after
Expected results: finite, non-silent 24 kHz mono WAV files; the streaming example
returns raw PCM bytes. Use `stream_format="sse"` for `speech.audio.delta`
events containing base64-encoded audio. On 1x L20 (greedy, single request after
warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time.
## Notes
@@ -166,6 +177,12 @@ warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time.
supported yet.
- **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and
`negative_prompt` return a client error.
- **Repetition penalty**: the default `1.1` matches upstream's penalty on
generated codebook-0 tokens and discourages repeated silence. Override it
with `extra_params: {"repetition_penalty": 1.0}` to disable it, or another
finite positive value. Text and reference audio are excluded from its history.
The legacy `tools/breeze_tts_2_upstream_golden.py` uses penalty `1.0`; set
the same serving override when comparing against those golden frames.
- **Length budget**: prompt length is bounded by stage 0 `max_model_len=4096`;
synthesis length by `max_new_tokens` (default 2048 frames ≈ 164 s at the
12.5 Hz codec frame rate) — set `max_new_tokens` for shorter caps.
@@ -181,7 +198,7 @@ warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time.
| Feature | Status | Notes |
| ------- | ------ | ----- |
| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks, SSE `speech.audio.delta` PCM |
| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks; raw audio or SSE `speech.audio.delta` |
| Voice cloning | ✓ | one reference clip plus transcript per request |
| Voice design / voice direction | ✓ | `instructions` field, with or without a reference |
| Classifier-free guidance | ✗ | `cfg_scale` must be `1.0`; `negative_prompt` rejected (follow-up) |
recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md
+20 -7
View File
@@ -1,9 +1,10 @@
# MiniMax H3 disaggregated text encoder
# MiniMax H3 disaggregated encoder stage
This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
its hidden states and token-role metadata to an encoder-free diffusion stage.
The standard MiniMax H3 recipes remain single-stage and continue to load the
text encoder inside the diffusion pipeline.
This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and
audio WVAE encoder together in a vLLM stage. It sends their conditioning
payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes
remain single-stage and continue to run all components inside the diffusion
pipeline.
## Start the server
@@ -23,6 +24,14 @@ deployment override for the available hardware. Diffusion quantization,
layerwise offload, distributed layerwise offload, VAE parallelism, and USP
settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
Stage 0's `tensor_parallel_size` defines one shared rank group for all three
encoder roles; it does not choose each role's execution strategy. The required
`hf_overrides.minimax_h3_encoder_components` mapping declares those strategies
explicitly: `text_encoder` uses tensor parallelism, `video_vae` uses patch
parallelism, and `audio_vae` runs on the group leader. All three role entries
are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit;
the configuration does not create independent role schedulers or world sizes.
For example, this five-GPU topology assigns one GPU to the encoder and four
to the diffusion stage. `--stage-overrides` keeps placement and parallelism
scoped to the owning stage rather than broadcasting an override to both:
@@ -57,8 +66,12 @@ The Stage 1 VAE patch-parallel options remain independent of offload and
quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
hardware-qualified profiles before combining these options.
Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
text-encoder weights. This H3 topology explicitly keeps its single-replica
Stage 1 sets `model_loaded.text_encoder: false` and
`model_loaded.vae_encoder: false`. It loads only the VAE decoders and skips
tokenizer, processor, and text-encoder downloads. It requires complete text and
media conditioning from Stage 0, with no fallback to local encoding.
This H3 topology explicitly keeps its single-replica
diffusion stage inline, avoiding serialization of decoded video through a
subprocess. It expects both stages to run in one deployment; cross-node payload
transport is outside this configuration.
recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md
+15 -5
View File
@@ -16,6 +16,12 @@ Validated on:
50 steps, text-encoder TP2, and DiT USP2
- Common sampling settings: `flow_shift=12`, `seed=1101`
The commands and measurements below describe this pinned revision. Newer
versions have changed the scope of `--quantization fp8`; see
[Online FP8 quantization](MiniMax-H3.md#online-fp8-quantization) for the current
component selection and offload support. Re-qualify memory, output, and latency
on GB10 before comparing results across revisions.
## Capacity requirements
| Resource | DGX Spark (GB10) |
@@ -40,8 +46,10 @@ math relative to the RTX and datacenter recipes:
- **Do not use `--enable-cpu-offload`.** Moving weights from "VRAM" to "host RAM"
frees nothing here.
- **`--quantization fp8` is mandatory.** A BF16 partition is 135 GiB and does not
fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
Qwen3-VL text encoder and both VAEs stay BF16.
fit in 121 GiB. At the validated revision, online FP8 quantizes the DiT only
(62 GiB to ~31 GiB); the Qwen3-VL text encoder and both VAEs stay BF16.
Current versions also quantize eligible text-decoder linears by default, so
the same flag does not imply the same quantization scope across revisions.
With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
23 GiB of headroom for the OS, page cache, and any other process on the box.
@@ -295,6 +303,8 @@ shorter than a 50-step Ref2VA run.
- Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
under **Validated on**: FP8 weight loading and image-only Ref2VA are both
broken on `release/v0.26.0`. The suggested version is `v0.26.1`.
- Online FP8 is incompatible with layerwise offload — the offload path produces a
weight stride the Cutlass FP8 kernel rejects. This is not a practical
restriction here since offload is unusable on GB10 anyway.
- At the validated revision, online FP8 with layerwise offload produced a weight
stride the Cutlass FP8 kernel rejected. Current H3 supports online FP8 with
layerwise offload and both DLO transfer paths; see the linked FP8 section
above. This does not change the GB10 capacity guidance: offload uses the same
unified-memory pool and is not enabled in this recipe.
recipes/MiniMaxAI/MiniMax-H3.md
+7 -1
View File
@@ -23,6 +23,12 @@ One vLLM-Omni diffusion stage can load both DiTs while instantiating the
tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only
once. Requests select the DiT with `extra_params.task`.
This single-stage recipe uses the default `model_loaded.text_encoder: true` and
`model_loaded.vae_encoder: true`. Precomputed text conditioning can replace the
local Qwen call. Setting `model_loaded.text_encoder: false` while leaving
`model_loaded.vae_encoder: true` requires precomputed text conditioning;
reference media is still encoded locally.
The generated MP4 contains H.264 video and synchronized stereo audio.
## Prerequisites
@@ -35,7 +41,7 @@ hf auth login
export MODEL=MiniMaxAI/MiniMax-H3
```
The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
By default, the pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
`Ref2VA/transformer/**`. It does not download or load the diffusers-format
`transformer`, `transformer_ref`, or `vae` weights at the repository root, nor
duplicate Ref2VA copies of shared components.
requirements/common.txt
+5
View File
@@ -41,3 +41,8 @@ pydantic-settings>=2.0.0
gguf>=0.10.0
# Required for running the Cosmos3 model with guardrails enabled (default behavior).
cosmos-guardrail>=0.3.1
# Required for CosyVoice3 (get_window/resample_poly). >=1.18.0 requires
# Python >=3.12, breaking base installs on the project's supported 3.10/3.11;
# these are long-stable APIs with no version-specific need for 1.18.
scipy>=1.11.0
s3tokenizer==0.3.0
tests/attention/test_fish_kvcache_attn.py
+10
View File
@@ -3,6 +3,7 @@
import pytest
import torch
from vllm.platforms import current_platform
from vllm_omni.attention import fish_kvcache_attn, fish_kvcache_backend
@@ -113,6 +114,15 @@ def test_fish_kvcache_enabled_by_default(monkeypatch):
assert not fish_kvcache_attn.is_fish_kvcache_attn_required()
def test_fish_kvcache_fast_path_is_unavailable_off_cuda(monkeypatch):
monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
monkeypatch.setattr(
fish_kvcache_attn, "_triton_backend", lambda: pytest.fail("Triton backend must not load off CUDA")
)
assert not fish_kvcache_attn.is_available()
@pytest.mark.parametrize("value", ["", "1", "true", "yes", "on", "required"])
def test_fish_kvcache_enabled_values_are_consistent(monkeypatch, value):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", value)
tests/benchmarks/duplex/test_omni_duplex_eval.py
+281
View File
@@ -64,6 +64,177 @@ def test_hugging_face_subset_names_are_splits(monkeypatch):
assert sample.task_type == "correction"
def test_local_hf_dataset_directory_is_routed_to_hf_loader(tmp_path, monkeypatch):
# An existing local Hugging Face dataset layout (directory with a
# ``data/*.parquet`` file) must be loaded through ``datasets.load_dataset``
# instead of being parsed as a JSON/JSONL manifest.
hf_dir = tmp_path / "Omni-DuplexEval"
data_dir = hf_dir / "data"
data_dir.mkdir(parents=True)
(data_dir / "train-00000-of-00001.parquet").write_bytes(b"PAR1")
calls = []
def fake_load_dataset(name, **kwargs):
calls.append((name, kwargs))
return [{"id": "pr", "split": "PR_correction", "question_text": "Correct this."}]
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
samples = load_samples(hf_dir)
assert calls == [(str(hf_dir), {})]
assert len(samples) == 1
assert samples[0].task_type == "correction"
def test_local_parquet_file_is_routed_to_hf_loader(tmp_path, monkeypatch):
parquet = tmp_path / "samples.parquet"
parquet.write_bytes(b"PAR1")
calls = []
def fake_load_dataset(name, **kwargs):
calls.append((name, kwargs))
return [{"id": "pr", "split": "PR_correction", "question_text": "Correct this."}]
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
samples = load_samples(str(parquet))
assert calls == [("parquet", {"data_files": str(parquet)})]
assert len(samples) == 1
assert samples[0].family == "pr"
def test_hf_collapse_to_train_keeps_row_split_identity(tmp_path, monkeypatch):
# datasets.load_dataset collapses a config-less local mirror (a bare data/
# directory) into a single generic ``train`` split. Rows that already carry
# their own split/subset identity must keep it so family inference is not
# poisoned by the generic name (regression for the direct-read path).
hf_dir = tmp_path / "Omni-DuplexEval"
data_dir = hf_dir / "data"
data_dir.mkdir(parents=True)
(data_dir / "train-00000-of-00021.parquet").write_bytes(b"PAR1")
def fake_load_dataset(name, **kwargs):
return {
"train": [
{"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
{"id": "pr-subset", "subset": "PR_event_reminder", "question_text": "Remind me."},
{"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
]
}
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
samples = load_samples(hf_dir)
by_id = {sample.id: sample for sample in samples}
assert by_id["pr"].split == "PR_correction"
assert by_id["pr"].family == "pr"
assert by_id["pr"].task_type == "correction"
assert by_id["pr-subset"].split == "PR_event_reminder"
assert by_id["pr-subset"].family == "pr"
assert by_id["pr-subset"].task_type == "proactive_reminder"
assert by_id["rtd"].split == "RTD_OCR"
assert by_id["rtd"].family == "rtd"
def test_hf_config_layout_stamps_split_name_when_row_has_none(tmp_path, monkeypatch):
# With a proper config-per-split mirror the loader still stamps the Hugging
# Face split name onto rows that do not identify themselves (unchanged).
hf_dir = tmp_path / "Omni-DuplexEval"
data_dir = hf_dir / "data"
data_dir.mkdir(parents=True)
(data_dir / "RTD_OCR-00000-of-00001.parquet").write_bytes(b"PAR1")
(data_dir / "PR_correction-00000-of-00001.parquet").write_bytes(b"PAR1")
def fake_load_dataset(name, **kwargs):
return {
"RTD_OCR": [{"id": "rtd", "question_text": "Read this."}],
"PR_correction": [{"id": "pr", "question_text": "Correct this."}],
}
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
samples = load_samples(hf_dir)
by_id = {sample.id: sample for sample in samples}
assert by_id["rtd"].split == "RTD_OCR"
assert by_id["rtd"].family == "rtd"
assert by_id["pr"].split == "PR_correction"
assert by_id["pr"].family == "pr"
assert by_id["pr"].task_type == "correction"
def test_hf_collapse_to_train_without_row_identity_raises_actionable_error(tmp_path, monkeypatch):
# A row that carries no split/family/task_type and lands in a generic
# collapsed ``train`` split cannot be routed; the loader must raise a clear,
# actionable ValueError instead of the bare family_for_split message.
hf_dir = tmp_path / "Omni-DuplexEval"
data_dir = hf_dir / "data"
data_dir.mkdir(parents=True)
(data_dir / "train-00000-of-00021.parquet").write_bytes(b"PAR1")
def fake_load_dataset(name, **kwargs):
return {"train": [{"id": "x", "question_text": "What is this?"}]}
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
with pytest.raises(ValueError, match="Hugging Face splits"):
load_samples(hf_dir)
def test_non_manifest_file_raises_clear_error(tmp_path):
# A plain non-JSON file (e.g. a README) must surface a clear ValueError
# listing the accepted inputs instead of a raw JSONDecodeError.
bad = tmp_path / "README.md"
bad.write_text("# readme\nnot a manifest", encoding="utf-8")
with pytest.raises(ValueError, match="JSON/JSONL manifest"):
load_samples(str(bad))
def _write_parquet(path, rows):
pyarrow = pytest.importorskip("pyarrow")
pyarrow.parquet.write_table(pyarrow.Table.from_pylist(rows), str(path))
def test_local_parquet_split_filters_rows_by_identity(tmp_path):
# Regression: a single ``.parquet`` file is exposed by ``load_dataset`` as a
# generic ``train`` physical split, so ``--split RTD_OCR`` used to raise
# ``Unknown split "RTD_OCR"``. The requested split must instead filter the
# rows by their preserved identity.
pytest.importorskip("datasets")
parquet = tmp_path / "samples.parquet"
_write_parquet(
parquet,
[
{"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
{"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
],
)
assert {sample.id for sample in load_samples(str(parquet))} == {"rtd", "pr"}
rtd = load_samples(str(parquet), split="RTD_OCR")
assert [sample.id for sample in rtd] == ["rtd"]
assert rtd[0].family == "rtd"
pr = load_samples(str(parquet), split="PR_correction")
assert [sample.id for sample in pr] == ["pr"]
assert pr[0].family == "pr"
assert pr[0].task_type == "correction"
def test_local_directory_named_splits_resolve_without_row_identity(tmp_path):
# Regression: a single-configuration mirror whose data files are named after
# the benchmark splits must resolve each split for real (no fake loader) and
# without any per-row identity.
pytest.importorskip("datasets")
hf_dir = tmp_path / "Omni-DuplexEval"
data_dir = hf_dir / "data"
data_dir.mkdir(parents=True)
_write_parquet(data_dir / "RTD_OCR-00000-of-00001.parquet", [{"id": "rtd", "question_text": "Read this."}])
_write_parquet(
data_dir / "PR_correction-00000-of-00001.parquet",
[{"id": "pr", "question_text": "Correct this."}],
)
by_id = {sample.id: sample for sample in load_samples(hf_dir)}
assert by_id["rtd"].split == "RTD_OCR"
assert by_id["rtd"].family == "rtd"
assert by_id["pr"].split == "PR_correction"
assert by_id["pr"].family == "pr"
assert by_id["pr"].task_type == "correction"
assert [sample.id for sample in load_samples(hf_dir, split="RTD_OCR")] == ["rtd"]
def test_response_aliases_and_clock_guard():
assert split_text("One. Two!") == ["One.", "Two!"]
assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
@@ -141,3 +312,113 @@ def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
judge_video_mode="frame-sample",
)
assert score["content"]["frame_count"] == 2
def test_local_directory_collapsed_train_with_row_identity_and_split_filter(tmp_path):
# Reviewer-pointed gap: a collapsed local mirror (data/train-*.parquet)
# whose rows carry their own split identity must correctly filter by
# --split RTD_OCR.
pytest.importorskip("datasets")
hf_dir = tmp_path / "Omni-DuplexEval"
data_dir = hf_dir / "data"
data_dir.mkdir(parents=True)
_write_parquet(
data_dir / "train-00000-of-00001.parquet",
[
{"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
{"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
],
)
all_samples = load_samples(hf_dir)
assert {s.id for s in all_samples} == {"rtd", "pr"}
rtd = load_samples(hf_dir, split="RTD_OCR")
assert [s.id for s in rtd] == ["rtd"]
assert rtd[0].family == "rtd"
def test_local_directory_mistyped_split_raises_clear_error(tmp_path):
# A named-shard mirror with --split RTD_TYPO must raise ValueError
# listing the observed splits instead of silently returning nothing.
pytest.importorskip("datasets")
hf_dir = tmp_path / "Omni-DuplexEval"
data_dir = hf_dir / "data"
data_dir.mkdir(parents=True)
_write_parquet(
data_dir / "RTD_OCR-00000-of-00001.parquet",
[{"id": "rtd", "question_text": "Read this."}],
)
_write_parquet(
data_dir / "PR_correction-00000-of-00001.parquet",
[{"id": "pr", "question_text": "Correct this."}],
)
with pytest.raises(ValueError) as exc_info:
load_samples(hf_dir, split="RTD_TYPO")
message = str(exc_info.value)
assert "matched zero" in message
assert "RTD_OCR" in message or "PR_correction" in message
def test_manifest_without_split_identity_split_override_still_works(tmp_path):
# Manifest rows without split/subset/config identity must keep the
# base override behaviour: --split RTD_OCR stamps the split so rows
# are kept instead of rejected (regression test for P2).
manifest = tmp_path / "samples.json"
manifest.write_text(
json.dumps([{"id": "a", "video": "clip.mp4"}, {"id": "b", "video": "clip2.mp4"}]),
encoding="utf-8",
)
samples = load_samples(manifest, split="RTD_OCR")
assert len(samples) == 2
assert all(s.split == "RTD_OCR" for s in samples)
assert all(s.family == "rtd" for s in samples)
def test_iterable_without_split_identity_split_override_still_works():
# Iterable rows without split/subset/config identity must keep the
# base override behaviour (regression test for F5/P2).
data = [{"id": "a", "video": "clip.mp4"}, {"id": "b", "video": "clip2.mp4"}]
samples = load_samples(data, split="RTD_OCR")
assert len(samples) == 2
assert all(s.split == "RTD_OCR" for s in samples)
assert all(s.family == "rtd" for s in samples)
def test_audio_feature_decode_false_preserves_bytes_key(tmp_path):
# When a parquet file carries an Audio feature column, the loader
# must recast it with decode=False so rows contain the ``bytes`` key
# instead of ``array``/``sampling_rate``.
pytest.importorskip("datasets")
pytest.importorskip("pyarrow")
import datasets
audio_bytes = b"\x00\x01\x02"
ds = datasets.Dataset.from_dict(
{
"id": ["a"],
"question_text": ["Hello"],
"task_type": ["pr_correction"],
"audio": [{"bytes": audio_bytes, "path": None}],
},
features=datasets.Features(
{
"id": datasets.Value("string"),
"question_text": datasets.Value("string"),
"task_type": datasets.Value("string"),
"audio": datasets.Audio(sampling_rate=16000),
}
),
)
parquet = tmp_path / "audio_sample.parquet"
ds.to_parquet(str(parquet))
samples = load_samples(str(parquet))
assert len(samples) == 1
raw = samples[0].raw
assert raw is not None
audio_val = raw.get("audio")
assert audio_val is not None
# The Audio feature was decoded with decode=False, so the struct
# keeps the ``bytes`` key that materialize_media needs.
assert isinstance(audio_val, dict) and "bytes" in audio_val
assert audio_val["bytes"] == audio_bytes
assert "path" in audio_val
tests/benchmarks/duplex/test_omni_duplex_eval_media.py
+67
View File
@@ -0,0 +1,67 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
from vllm_omni.benchmarks.duplex.omni_duplex_eval_media import materialize_media
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_dict_with_bytes_materializes_a_real_file(tmp_path):
# A Hugging Face Audio/Video feature value {"bytes": ..., "path": ...} used
# to short-circuit on ``path`` (the artifact name inside the dataset, which
# does not exist on disk), so downstream ffmpeg saw a missing file. Inline
# bytes must win and be written to a real file whose content equals bytes.
payload = b"\x00\xff\x01RIFFfake-content"
resolved = materialize_media({"bytes": payload, "path": "question_audio.wav"}, tmp_path, "s1_question", ".wav")
assert resolved.is_file()
assert resolved.read_bytes() == payload
assert resolved.parent == tmp_path
assert resolved.name == "s1_question.wav"
def test_dict_extension_comes_from_path_safely(tmp_path):
payload = b"flac-ish-bytes"
resolved = materialize_media({"bytes": payload, "path": "audio/clip.FLAC"}, tmp_path, "s1", ".wav")
assert resolved.name == "s1.flac"
assert resolved.read_bytes() == payload
# A path with no usable extension falls back to the caller's suffix.
resolved = materialize_media({"bytes": payload, "path": "audio/clip"}, tmp_path, "s2", ".wav")
assert resolved.name == "s2.wav"
assert resolved.read_bytes() == payload
def test_materialization_is_idempotent_and_updated(tmp_path):
payload = b"abc"
first = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
second = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
assert first == second
assert first.read_bytes() == payload
replacement = b"xyz"
third = materialize_media({"bytes": replacement, "path": "q.wav"}, tmp_path, "s", ".wav")
assert third == first
assert third.read_bytes() == replacement
def test_str_path_and_path_only_mapping_returned_as_is(tmp_path):
existing = tmp_path / "clip.mp4"
existing.write_bytes(b"video")
assert materialize_media(str(existing), tmp_path, "s", ".mp4") == existing
assert materialize_media({"path": str(existing)}, tmp_path, "s", ".mp4") == existing
def test_raw_bytes_value_is_written(tmp_path):
payload = b"raw-media"
resolved = materialize_media(payload, tmp_path, "r", ".bin")
assert resolved.name == "r.bin"
assert resolved.read_bytes() == payload
def test_unresolvable_value_raises_clear_error(tmp_path):
with pytest.raises(ValueError, match="cannot materialize"):
materialize_media(None, tmp_path, "s", ".wav")
with pytest.raises(ValueError, match="cannot materialize"):
materialize_media({"path": None}, tmp_path, "s", ".wav")
tests/benchmarks/patch/test_benchmark_result.py
+74
View File
@@ -0,0 +1,74 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Verify metric sample counts survive benchmark JSON serialization."""
import asyncio
import json
import time
import pytest
from vllm.benchmarks.datasets import SampleRequest
from vllm.benchmarks.serve import TaskType
from vllm_omni.benchmarks.patch import patch
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
@pytest.mark.parametrize("selected_metrics", [["ttft", "tpot", "itl"], []])
@pytest.mark.parametrize("measured", [True, False])
def test_benchmark_result_preserves_tpot_sample_count(monkeypatch, measured, selected_metrics):
async def request_func(**kwargs):
return patch.MixRequestFuncOutput(
success=True,
prompt_len=2,
output_tokens=3,
generated_text="hello world!",
start_time=time.perf_counter(),
ttft=0.01,
latency=0.03,
text_latency=0.03,
itl=[0.01, 0.01] if measured else [0.0, 0.0],
)
monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "test-result", request_func)
result = asyncio.run(
patch.benchmark(
task_type=TaskType.GENERATION,
endpoint_type="test-result",
api_url="http://unused/v1/chat/completions",
base_url="http://unused",
model_id="test-model",
model_name="test-model",
tokenizer=None,
input_requests=[SampleRequest(prompt="hi", prompt_len=2, expected_output_len=3)],
logprobs=None,
request_rate=float("inf"),
burstiness=1.0,
disable_tqdm=True,
num_warmups=0,
profile=False,
selected_percentile_metrics=selected_metrics,
selected_percentiles=[50.0],
ignore_eos=False,
goodput_config_dict={},
max_concurrency=1,
lora_modules=None,
extra_headers=None,
extra_body=None,
ready_check_timeout_sec=0,
)
)
# DFX reads this dictionary from the benchmark's JSON result file.
saved_result = json.loads(json.dumps(result))
assert saved_result["completed"] == 1
assert saved_result.get("num_tpot_samples") == int(measured)
assert saved_result["num_ttft_samples"] == 1
assert saved_result["num_itl_samples"] == 2
assert "num_audio_ttfp_samples" in saved_result
assert "num_audio_rtf_samples" in saved_result
if measured and "tpot" in selected_metrics:
assert saved_result["mean_tpot_ms"] == pytest.approx(10.0)
else:
assert "mean_tpot_ms" not in saved_result
tests/benchmarks/test_diffusion_benchmark_warmups.py
+128
View File
@@ -0,0 +1,128 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import importlib
import sys
import pytest
from benchmarks.diffusion import backends
from benchmarks.diffusion.backends import RequestFuncInput, RequestFuncOutput
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu, pytest.mark.asyncio]
ENDPOINT = "/v1/chat/completions"
TEST_MODEL = "test-model"
@pytest.fixture
def diffusion_benchmark(monkeypatch):
# The standalone script imports its sibling backend module by name.
monkeypatch.setitem(sys.modules, "backends", backends)
return importlib.import_module("benchmarks.diffusion.diffusion_benchmark_serving")
@pytest.fixture
def warmup_args(diffusion_benchmark):
from benchmarks.diffusion.diffusion_benchmark_serving import build_parser
parser = build_parser()
return parser.parse_args(
[
f"--endpoint={ENDPOINT}",
"--dataset=random",
"--task=t2i",
"--warmup-requests=2",
"--warmup-num-inference-steps=4",
]
)
@pytest.fixture
def input_requests():
return [
RequestFuncInput(
prompt=prompt,
api_url=f"http://test.local{ENDPOINT}",
model=TEST_MODEL,
num_inference_steps=20,
)
for prompt in ("first", "second")
]
async def test_disabled_warmup_sends_no_requests(diffusion_benchmark, warmup_args, input_requests):
warmup_args.warmup_requests = 0
async def unexpected_request(*args):
pytest.fail("disabled warmup must not send requests")
pairs = await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, unexpected_request)
assert pairs == []
async def test_successful_warmups_preserve_request_output_pairs(diffusion_benchmark, warmup_args, input_requests):
outputs = [RequestFuncOutput(success=True, latency=0.5), RequestFuncOutput(success=True, latency=1.0)]
pending_outputs = iter(outputs)
async def successful_request(req, session, pbar):
return next(pending_outputs)
pairs = await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, successful_request)
assert [request.prompt for request, _ in pairs] == ["first", "second"]
assert [request.num_inference_steps for request, _ in pairs] == [4, 4]
assert [output for _, output in pairs] == outputs
assert [request.num_inference_steps for request in input_requests] == [20, 20]
@pytest.mark.parametrize("failed_count", [1, 2], ids=["partial-failure", "all-failed"])
@pytest.mark.parametrize(
"error",
["HTTP 503: private prompt rejected", "Cannot connect to https://test.local/?token=private"],
ids=["http-error", "transport-error"],
)
async def test_failed_warmups_abort_without_exposing_response_errors(
diffusion_benchmark, warmup_args, input_requests, failed_count, error
):
outputs = iter(
[RequestFuncOutput(success=False, error=error) for _ in range(failed_count)]
+ [RequestFuncOutput(success=True) for _ in range(2 - failed_count)]
)
async def request(req, session, pbar):
return next(outputs)
with pytest.raises(RuntimeError, match=f"{failed_count}/2 warmup requests failed") as exc_info:
await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, request)
assert error not in str(exc_info.value)
assert "server logs" in str(exc_info.value)
async def test_benchmark_does_not_start_measurement_after_failed_warmup(
monkeypatch, mocker, diffusion_benchmark, warmup_args, input_requests
):
sent = []
async def failed_request(req, session, pbar):
sent.append(req)
return RequestFuncOutput(success=False, error="HTTP 503: unavailable")
monkeypatch.setitem(
diffusion_benchmark.backends_function_mapping["2i"],
warmup_args.endpoint,
(failed_request, warmup_args.endpoint),
)
dataset = mocker.Mock(spec=diffusion_benchmark.RandomDataset)
dataset.get_requests.return_value = input_requests
monkeypatch.setattr(diffusion_benchmark, "RandomDataset", mocker.Mock(return_value=dataset))
timer = mocker.Mock(spec=diffusion_benchmark.time)
timer.perf_counter.side_effect = AssertionError("measurement started after failed warmup")
monkeypatch.setattr(diffusion_benchmark, "time", timer)
with pytest.raises(RuntimeError, match="2/2 warmup requests failed"):
await diffusion_benchmark.benchmark(warmup_args)
assert len(sent) == warmup_args.warmup_requests
tests/clients/test_duplex_client.py
+27
View File
@@ -570,6 +570,33 @@ async def test_error_event_surfaces_on_response_path(make_rig):
rig.feed(SESSION_CLOSED)
@pytest.mark.asyncio
async def test_close_after_error_reports_the_server_reason_to_a_sender():
# A client streaming input sits in send(), not in responses(), so the
# error text has to survive into the close marker or the session looks
# like it ended for no reason.
sock = FakeSocket()
sock.feed(SESSION_CREATED)
client, _ = make_client(sock)
async with client:
sock.feed(
{
"type": "error",
"error": {
"code": "runtime_data_plane_stream_failed",
"message": "context_length_exceeded: streaming session prompt would grow to 41207 tokens",
},
"server_event_seq": 2,
}
)
sock.feed(SESSION_CLOSED)
await _drain(lambda: client._closed.is_set())
with pytest.raises(DuplexSessionClosedError) as excinfo:
await client.send({"type": "input_audio_buffer.append", "audio": ""})
assert "runtime_data_plane_stream_failed" in str(excinfo.value)
assert "context_length_exceeded" in str(excinfo.value)
@pytest.mark.asyncio
async def test_slow_audio_consumer_drops_oldest_instead_of_stalling():
from vllm_omni.clients.duplex import AudioDelta, ResponseHandle
tests/config/test_config_factory.py
+24 -3
View File
@@ -1562,7 +1562,10 @@ stages:
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
assert deploy.stages[1].engine_extras["model_loaded"] == {
"text_encoder": False,
"vae_encoder": False,
}
stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
@@ -1587,12 +1590,20 @@ stages:
deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
stages = merge_pipeline_deploy(pipeline, deploy)
assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3TextEncoder"
assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder"
assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline"
assert stages[0].yaml_runtime["num_replicas"] == 1
assert stages[1].yaml_runtime["num_replicas"] == 1
assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False}
assert stages[1].yaml_engine_args["model_loaded"] == {
"text_encoder": False,
"vae_encoder": False,
}
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == {
"text_encoder": {"parallel_mode": "tp"},
"video_vae": {"parallel_mode": "patch"},
"audio_vae": {"parallel_mode": "leader"},
}
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
parallel = stages[1].yaml_engine_args["parallel_config"]
@@ -1604,6 +1615,9 @@ stages:
turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
turbo_stages = merge_pipeline_deploy(pipeline, turbo)
assert turbo_stages[0].yaml_engine_args["model_arch"] == stages[0].yaml_engine_args["model_arch"]
assert turbo_stages[1].yaml_engine_args["model_loaded"] == stages[1].yaml_engine_args["model_loaded"]
assert turbo_stages[0].yaml_engine_args["hf_overrides"] == stages[0].yaml_engine_args["hf_overrides"]
turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
assert turbo_sampling["num_inference_steps"] == 5
assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
@@ -2674,6 +2688,13 @@ class TestPlatformOverrides:
# Explicit null clears the inherited single-GPU 2 GiB CUDA cap.
assert replica_stages[1].yaml_engine_args.get("kv_cache_memory_bytes") is None
def test_fish_speech_npu_uses_ascend_kv_block_size(self):
deploy_path = Path(get_deploy_config_path("fish_qwen3_omni.yaml"))
deploy = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu")
assert deploy.stages[0].engine_extras["block_size"] == 128
def test_npu_overrides(self):
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
if not deploy_path.exists():
tests/config/test_forced_aligner_injection.py
+25
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry)
@@ -69,3 +72,25 @@ def test_inject_noop_without_forced_aligner():
assert ext_pipeline is pipeline
assert len(ext_deploy.stages) == 0
@pytest.mark.parametrize("async_chunk", [True, False])
@pytest.mark.parametrize("cli_async_chunk", [None, True, False])
def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_chunk, cli_async_chunk):
from vllm_omni.config.config_factory import StageConfigFactory
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
deploy = DeployConfig(async_chunk=async_chunk)
resolution = StageConfigFactory._resolve_legacy_from_registry(
pipeline,
{"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B", "async_chunk": cli_async_chunk},
user_deploy_config=deploy,
)
stages = [stage.to_omegaconf() for stage in resolution.stage_configs]
expected_async = async_chunk if cli_async_chunk is None else cli_async_chunk
assert len(stages) == 3
assert [stage.engine_args.async_chunk for stage in stages] == [expected_async, expected_async, False]
assert stages[2].custom_process_input_func.endswith("code2wav2aligner")
assert stages[2].engine_args.runner == "pooling"
assert len(pipeline.stages) == 2
assert deploy.stages == []
tests/config/test_omni_config.py
+63 -1
View File
@@ -1480,7 +1480,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path))
od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
@@ -1705,3 +1705,65 @@ def test_compact_offload_config_is_validated_during_projection():
"components": ["dit"],
}
)
@pytest.mark.parametrize("pipeline_async", [True, False])
@pytest.mark.parametrize("stage_async", [None, False, True])
def test_stage_async_chunk_opt_out_matches_legacy_config(pipeline_async, stage_async):
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
deploy = DeployConfig(
async_chunk=pipeline_async,
stages=[StageDeployConfig(stage_id=i, async_chunk=stage_async) for i in (0, 1)],
)
config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
legacy = merge_pipeline_deploy(pipeline, deploy)
expected = pipeline_async and stage_async is not False
assert config.stage_by_id(0).connector_config.async_chunk is expected
assert config.stage_by_id(1).connector_config.async_chunk is expected
assert legacy[1].yaml_engine_args["async_chunk"] is expected
assert config.stage_by_id(1).custom_process_input_func == legacy[1].custom_process_input_func
if not expected:
assert config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
@pytest.mark.parametrize("disabled_stage", [0, 1])
@pytest.mark.parametrize("builder", [merge_pipeline_deploy, VllmOmniConfig.from_pipeline_config])
def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder):
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
deploy = DeployConfig(
async_chunk=True,
stages=[StageDeployConfig(stage_id=disabled_stage, async_chunk=False)],
)
with pytest.raises(ValueError, match="incompatible async_chunk settings on connector edge 0 -> 1"):
if builder is merge_pipeline_deploy:
builder(pipeline, deploy)
else:
builder(pipeline, user_deploy_config=deploy)
@pytest.mark.parametrize("explicit", [False, True])
def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit):
from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
from vllm_omni.quantization import build_quant_config
checkpoint = TransformerConfig.from_dict(
{
"quantization_config": {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"w4a8_fallback_steps": [37],
}
}
)
requested = build_quant_config("mxfp4", w4a8_fallback_steps=[]) if explicit else None
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
tf_model_config=checkpoint,
quantization_config=requested,
)
assert cfg.quantization_config_is_auto_detected is not explicit
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda self: self.set_tf_model_config(checkpoint))
cfg.enrich_config()
restored = ForkingPickler.loads(ForkingPickler.dumps(cfg))
assert restored.quantization_config_is_auto_detected is not explicit
assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37])
tests/core/sched/test_omni_ar_scheduler_logprobs.py
+3
View File
@@ -149,6 +149,9 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
for name in _MIXIN_UPDATE_HELPERS:
setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler))
scheduler._cleanup_kv_tracking = MethodType(OmniARScheduler._cleanup_kv_tracking, scheduler)
scheduler._emit_streaming_context_overflow_outputs = MethodType(
OmniARScheduler._emit_streaming_context_overflow_outputs, scheduler
)
scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
scheduler.make_stats = lambda *args, **kwargs: None
return scheduler
tests/core/sched/test_omni_ar_scheduler_streaming.py
+303
View File
@@ -5,6 +5,7 @@
from __future__ import annotations
from collections import defaultdict
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -17,6 +18,8 @@ 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.engine import FinishReason
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
@@ -958,3 +961,303 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
sched.skipped_waiting.remove_requests.assert_not_called()
sched._enqueue_waiting_request.assert_not_called()
def _make_live_session_scheduler(*, max_model_len: int) -> OmniARScheduler:
"""A scheduler whose finish path is real: the queues, ``finish_requests``,
``_free_request`` and the finished-request bookkeeping that wakes an idle
engine all run. Only block freeing and the encoder cache are mocked."""
sched = _make_admission_scheduler(max_model_len=max_model_len)
del sched._free_request # the admission helper mocks it; run the real one
sched._omits_kv_transfer_cache = {}
sched.connector = None
sched.perf_metrics = None
sched.recompute_kv_load_failures = False
sched.finished_req_ids = set()
sched.finished_req_ids_dict = defaultdict(set)
sched.transfer_triggered_requests = set()
sched.active_kv_transfers = set()
sched.waiting_for_transfer_free = set()
sched.pending_stop_after_extraction = set()
sched.requests_needing_kv_transfer = {}
sched._kv_wait_start_ts = {}
sched.kv_cache_manager = SimpleNamespace(take_events=lambda: None, estimate_cached_tokens=lambda _request: 0)
return sched
def _run_idle_step(sched: OmniARScheduler):
"""One ``update_from_output`` with nothing scheduled: the step an idle
engine runs once ``has_finished_requests()`` reports a freed session."""
scheduler_output = SimpleNamespace(
num_scheduled_tokens={},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
model_runner_output = SimpleNamespace(
sampled_token_ids=[],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=None,
num_nans_in_logits=None,
kv_connector_output=None,
cudagraph_stats=None,
req_id_to_index={},
routed_experts=None,
)
return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
def _park_session(sched: OmniARScheduler, session: Request) -> None:
"""Put a session where a 1 fps video stream leaves it between appends:
prompt partly computed, waiting for streaming input, in admission."""
sched.requests[session.request_id] = session
session.append_output_token_ids([7, 8, 9])
session.num_computed_tokens = 6
session.num_output_placeholders = 0
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched.skipped_waiting.add_request(session)
sched.num_waiting_for_streaming_input = 1
def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_session() -> None:
"""A native duplex session grows its prompt on every append; once the
extension cannot fit the model the worker crashes copying the prompt.
The scheduler must drop the update and fail only this session, at once,
through the real finish path: a parked session does not make the engine
schedule, so what wakes the engine is the freed request showing up in
``has_finished_requests()``, and the step that follows must carry one
explicit ERROR rather than the synthesized ABORT."""
sched = _make_live_session_scheduler(max_model_len=8)
session = _make_request()
_park_session(sched, session)
sched._update_request_as_session(session, _make_update([10, 20, 30]))
assert session.prompt_token_ids == [1, 2, 3]
assert session.num_prompt_tokens == 3
assert session.status == RequestStatus.FINISHED_ERROR
assert session.request_id not in sched.requests
assert len(sched.waiting) == 0
assert len(sched.skipped_waiting) == 0
assert sched.num_waiting_for_streaming_input == 0
sched._free_request_blocks.assert_called_once_with(session)
assert sched.finished_req_ids == {session.request_id}
assert sched.has_finished_requests()
client_index, reason = sched._streaming_context_overflow[session.request_id]
assert client_index == session.client_index
assert reason.startswith("context_length_exceeded: ")
assert "9 tokens" in reason and "max_model_len 8" in reason
engine_core_outputs = _run_idle_step(sched)
(output,) = engine_core_outputs[session.client_index].outputs
assert output.request_id == session.request_id
assert output.finish_reason == FinishReason.ERROR
assert output.stop_reason == reason
assert output.new_token_ids == []
assert engine_core_outputs[session.client_index].finished_requests == {session.request_id}
assert sched._streaming_context_overflow == {}
# Nothing is left to report once the error went out.
assert _run_idle_step(sched) == {}
def test_stage0_streaming_update_that_fills_max_model_len_exactly_finishes_the_session() -> None:
"""An exactly full prompt leaves no room for the token every duplex step
samples: upstream's running budget drops to -1, which ``schedule()`` does
not catch, so the exact fill is an overflow too."""
sched = _make_scheduler(stage_id=0)
sched.max_model_len = 8
sched.finish_requests = MagicMock()
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
session.append_output_token_ids([7, 8, 9])
session.num_computed_tokens = 6
session.num_output_placeholders = 0
sched._update_request_as_session(session, _make_update([10, 20]))
assert session.prompt_token_ids == [1, 2, 3]
assert session.num_prompt_tokens == 3
sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR)
_, reason = sched._streaming_context_overflow[session.request_id]
assert "8 tokens" in reason and "max_model_len 8" in reason
def test_stage0_streaming_update_that_leaves_room_to_sample_is_applied() -> None:
sched = _make_scheduler(stage_id=0)
sched.max_model_len = 9
sched.finish_requests = MagicMock()
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
session.append_output_token_ids([7, 8, 9])
session.num_computed_tokens = 6
session.num_output_placeholders = 0
sched._update_request_as_session(session, _make_update([10, 20]))
assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
assert session.num_prompt_tokens == 8
assert session.status == RequestStatus.WAITING
sched.finish_requests.assert_not_called()
assert not getattr(sched, "_streaming_context_overflow", {})
def test_stage0_streaming_update_keeps_room_for_every_token_sampled_per_step() -> None:
"""With speculative decoding a step samples more than one token, and the
prompt must leave room for all of them."""
sched = _make_scheduler(stage_id=0)
sched.max_model_len = 9
sched.num_sampled_tokens_per_step = 2
sched.finish_requests = MagicMock()
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
session.append_output_token_ids([7, 8, 9])
session.num_computed_tokens = 6
session.num_output_placeholders = 0
sched._update_request_as_session(session, _make_update([10, 20]))
assert session.prompt_token_ids == [1, 2, 3]
sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR)
assert "8 tokens" in sched._streaming_context_overflow[session.request_id][1]
def test_stage0_streaming_update_overflow_counts_an_uncomputed_prompt() -> None:
"""A preempted session keeps its prompt even though num_computed_tokens
fell behind it; the projection must use the longer of the two."""
sched = _make_scheduler(stage_id=0)
sched.max_model_len = 4
sched.finish_requests = MagicMock()
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
session.num_computed_tokens = 0
sched._update_request_as_session(session, _make_update([10, 20]))
assert session.prompt_token_ids == [1, 2, 3]
sched.finish_requests.assert_called_once()
assert "5 tokens" in sched._streaming_context_overflow[session.request_id][1]
def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
"""A scheduler with real admission queues, so nothing about the terminal
handling below is mocked away except block freeing."""
sched = _make_scheduler(stage_id=0)
sched.max_model_len = max_model_len
sched.policy = SchedulingPolicy.FCFS
sched.waiting = create_request_queue(sched.policy)
sched.skipped_waiting = create_request_queue(sched.policy)
sched.running = []
sched.requests = {}
sched._free_request = MagicMock() # needs a KV manager; not what this covers
return sched
def _make_queued_stop(session: Request, update: StreamingUpdate) -> None:
"""Put the request in the state upstream hands to _handle_stopped_request:
an append that arrived during generation is queued, and the stop status is
still on the request."""
from collections import deque
session.resumable = True
session.streaming_queue = deque([update])
session.status = RequestStatus.FINISHED_STOPPED
def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -> None:
sched = _make_admission_scheduler(max_model_len=4)
session = _make_request()
sched.requests[session.request_id] = session
session.num_computed_tokens = 3
_make_queued_stop(session, _make_update([10, 20]))
finished = sched._handle_stopped_request(session)
assert finished is True
assert len(sched.waiting) == 0
assert len(sched.skipped_waiting) == 0
assert session.status == RequestStatus.FINISHED_ERROR
assert session.resumable is False
assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
# update_from_output frees every request reported as finished; freeing here
# as well deleted it from self.requests twice (KeyError in _free_blocks).
sched._free_request.assert_not_called()
def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
sched = _make_admission_scheduler(max_model_len=64)
session = _make_request()
sched.requests[session.request_id] = session
session.num_computed_tokens = 3
_make_queued_stop(session, _make_update([10, 20]))
finished = sched._handle_stopped_request(session)
assert finished is False
assert len(sched.waiting) == 1
assert session.status == RequestStatus.WAITING
assert not getattr(sched, "_streaming_context_overflow", {})
sched._free_request.assert_not_called()
def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:
"""The overflow may already be on record when the stop is handled (an
earlier update of the same session). Upstream still re-enqueues the
request, so the override must take it back out either way instead of
leaving a FINISHED_ERROR request for admission to trip over."""
sched = _make_admission_scheduler(max_model_len=64)
session = _make_request()
sched.requests[session.request_id] = session
session.num_computed_tokens = 3
sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")}
_make_queued_stop(session, _make_update([10, 20])) # this update itself fits
finished = sched._handle_stopped_request(session)
assert finished is True
assert len(sched.waiting) == 0
assert len(sched.skipped_waiting) == 0
assert session.status == RequestStatus.FINISHED_ERROR
assert session.resumable is False
sched._free_request.assert_not_called()
def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanced() -> None:
"""Same as above with nothing queued: upstream parks the request as
waiting for streaming input and counts it; the dequeue must uncount it."""
from collections import deque
sched = _make_admission_scheduler(max_model_len=64)
session = _make_request()
sched.requests[session.request_id] = session
session.num_computed_tokens = 3
sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")}
session.resumable = True
session.streaming_queue = deque()
session.status = RequestStatus.FINISHED_STOPPED
finished = sched._handle_stopped_request(session)
assert finished is True
assert len(sched.skipped_waiting) == 0
assert sched.num_waiting_for_streaming_input == 0
assert session.status == RequestStatus.FINISHED_ERROR
sched._free_request.assert_not_called()
def test_context_overflow_emits_an_error_output_with_the_reason() -> None:
sched = _make_scheduler(stage_id=0)
sched._streaming_context_overflow = {"req-a": (2, "context_length_exceeded: too long")}
outputs: dict[int, list] = {}
sched._emit_streaming_context_overflow_outputs(outputs)
(output,) = outputs[2]
assert output.request_id == "req-a"
assert output.finish_reason == FinishReason.ERROR
assert output.stop_reason == "context_length_exceeded: too long"
assert output.new_token_ids == []
assert sched._streaming_context_overflow == {}
sched._emit_streaming_context_overflow_outputs(outputs)
assert len(outputs[2]) == 1
tests/dfx/perf/tests/test_runner_metadata.py
+10 -6
View File
@@ -570,6 +570,7 @@ def test_omni_tpot_baseline_accepts_measured_finite_sample():
@pytest.mark.parametrize(
("num_tpot_samples", "mean_tpot_ms", "match"),
[
(None, 10.0, "no measurable TPOT samples"),
(0, float("nan"), "no measurable TPOT samples"),
(1, float("nan"), "mean_tpot_ms is not finite"),
],
@@ -577,14 +578,17 @@ def test_omni_tpot_baseline_accepts_measured_finite_sample():
def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples, mean_tpot_ms, match):
from tests.dfx.perf.scripts.run_benchmark import assert_result
result = {
"completed": 1,
"Hardware": "H100",
"mean_tpot_ms": mean_tpot_ms,
}
if num_tpot_samples is not None:
result["num_tpot_samples"] = num_tpot_samples
with pytest.raises(AssertionError, match=match):
assert_result(
{
"completed": 1,
"Hardware": "H100",
"num_tpot_samples": num_tpot_samples,
"mean_tpot_ms": mean_tpot_ms,
},
result,
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
1,
)
tests/diffusion/layers/test_qwen_select01_modulation.py
+299
View File
@@ -0,0 +1,299 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
import torch
import torch.nn.functional as F
from vllm_omni.diffusion.layers import qwen_select01_modulation
from vllm_omni.diffusion.layers.qwen_select01_modulation import (
can_use_qwen_select01_triton,
fused_layernorm_select01,
fused_residual_layernorm_select01,
select01_modulation_native,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@pytest.mark.cpu
@pytest.mark.parametrize(
("platform_is_cuda", "tensor_is_cuda", "has_triton", "expected"),
[
(True, True, True, True),
(False, True, True, False),
(True, False, True, False),
(True, True, False, False),
],
)
def test_qwen_select01_triton_requires_cuda_platform(
monkeypatch,
platform_is_cuda: bool,
tensor_is_cuda: bool,
has_triton: bool,
expected: bool,
):
class FakePlatform:
@staticmethod
def is_cuda() -> bool:
return platform_is_cuda
class FakeTensor:
is_cuda = tensor_is_cuda
monkeypatch.setattr(qwen_select01_modulation, "current_platform", FakePlatform())
monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", has_triton)
assert can_use_qwen_select01_triton(FakeTensor()) is expected
def _make_inputs(
batch_size: int,
seq_len: int,
hidden_size: int,
dtype: torch.dtype,
device: torch.device,
) -> tuple[torch.Tensor, ...]:
generator = torch.Generator(device=device)
generator.manual_seed(202817)
x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
residual = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
residual_gate = torch.randn(batch_size, 1, hidden_size, dtype=dtype, device=device, generator=generator)
mod_params = torch.randn(batch_size * 2, hidden_size * 3, dtype=dtype, device=device, generator=generator)
index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int64, device=device, generator=generator)
return x, residual, residual_gate, mod_params, index
def _reference_layernorm_select01(
x: torch.Tensor,
mod_params: torch.Tensor,
index: torch.Tensor,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
scale, shift, gate = select01_modulation_native(mod_params, index)
out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype)
return out * (1 + scale) + shift, gate
def _reference_residual_layernorm_select01(
x: torch.Tensor,
residual: torch.Tensor,
residual_gate: torch.Tensor,
mod_params: torch.Tensor,
index: torch.Tensor,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
scale, shift, gate = select01_modulation_native(mod_params, index)
residual_out = residual + residual_gate * x
out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype)
return out * (1 + scale) + shift, residual_out, gate
@pytest.mark.cpu
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("broadcast_index", [False, True])
def test_qwen_select01_native_fallback_matches_reference(dtype: torch.dtype, broadcast_index: bool):
device = torch.device("cpu")
x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 64, dtype, device)
if broadcast_index:
index = index[:1]
eps = 1e-6
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
x,
residual,
residual_gate,
mod_params,
index,
eps,
)
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
x,
residual,
residual_gate,
mod_params,
index,
eps,
)
torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(actual_residual, ref_residual, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize(
("dtype", "atol", "rtol"),
[
(torch.float32, 1e-5, 1e-5),
(torch.bfloat16, 5e-2, 5e-2),
],
)
@pytest.mark.parametrize("seq_len", [1, 7, 128, 1024])
def test_qwen_select01_triton_matches_reference(dtype: torch.dtype, atol: float, rtol: float, seq_len: int):
device = torch.device("cuda:0")
x, residual, residual_gate, mod_params, index = _make_inputs(2, seq_len, 256, dtype, device)
eps = 1e-6
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
x,
residual,
residual_gate.expand_as(residual),
mod_params,
index,
eps,
)
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
x,
residual,
residual_gate.expand_as(residual),
mod_params,
index,
eps,
)
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize(
("dtype", "atol", "rtol"),
[
(torch.float32, 1e-5, 1e-5),
(torch.bfloat16, 5e-2, 5e-2),
],
)
def test_qwen_select01_triton_matches_reference_for_structured_multi_batch_index(
dtype: torch.dtype,
atol: float,
rtol: float,
):
device = torch.device("cuda:0")
x, residual, residual_gate, mod_params, _ = _make_inputs(3, 16, 256, dtype, device)
index = torch.tensor(
[
[0] * 16,
[1] * 16,
[0, 1] * 8,
],
dtype=torch.int64,
device=device,
)
eps = 1e-6
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
x,
residual,
residual_gate.expand_as(residual),
mod_params,
index,
eps,
)
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
x,
residual,
residual_gate.expand_as(residual),
mod_params,
index,
eps,
)
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize(
("dtype", "atol", "rtol"),
[
(torch.float32, 1e-5, 1e-5),
(torch.bfloat16, 5e-2, 5e-2),
],
)
def test_qwen_select01_triton_broadcasts_single_index_batch(dtype: torch.dtype, atol: float, rtol: float):
device = torch.device("cuda:0")
x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 256, dtype, device)
index = index[:1]
eps = 1e-6
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
x,
residual,
residual_gate.expand_as(residual),
mod_params,
index,
eps,
)
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
x,
residual,
residual_gate.expand_as(residual),
mod_params,
index,
eps,
)
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_qwen_select01_cuda_falls_back_without_triton(monkeypatch):
device = torch.device("cuda:0")
x, residual, residual_gate, mod_params, index = _make_inputs(2, 7, 256, torch.bfloat16, device)
residual_gate = residual_gate.expand_as(residual)
eps = 1e-6
def fail_if_launched(*args, **kwargs):
pytest.fail("Triton launcher must not run when HAS_TRITON is false")
monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", False)
monkeypatch.setattr(qwen_select01_modulation, "_launch_layernorm_select01", fail_if_launched)
monkeypatch.setattr(qwen_select01_modulation, "_launch_residual_layernorm_select01", fail_if_launched)
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
x,
residual,
residual_gate,
mod_params,
index,
eps,
)
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
x,
residual,
residual_gate,
mod_params,
index,
eps,
)
torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
torch.testing.assert_close(actual_residual, ref_residual, atol=0, rtol=0)
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
tests/diffusion/models/boogu_image/test_boogu_image_transformer.py
+119 -24
View File
@@ -440,29 +440,86 @@ def test_transformer_rejects_prompt_tuning():
BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
def _native_to_checkpoint_name(name: str) -> str:
"""Inverse of ``load_weights`` remapping: native param name -> diffusers name.
_QKV_FANOUT = {
"to_qkv": ("to_q", "to_k", "to_v"),
"img_to_qkv": ("img_to_q", "img_to_k", "img_to_v"),
"instruct_to_qkv": ("instruct_to_q", "instruct_to_k", "instruct_to_v"),
}
def _native_to_checkpoint_weights(name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]:
"""Inverse of ``load_weights`` remapping: native param -> diffusers weights.
The diffusers checkpoint stores fused projections as separate matrices, so a
single merged native param (QKV or FFN gate/up) fans out into several
checkpoint weights. Returns ``(diffusers_name, value)`` pairs whose values,
when fed through ``load_weights``, reassemble into the original ``param``.
- ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
- promoted joint-attention projections move back under ``.processor.``.
"""
q_size = NUM_HEADS * HEAD_DIM
kv_size = NUM_KV_HEADS * HEAD_DIM
# Fused QKV projections split back into per-matrix q/k/v weights.
for token, (q_name, k_name, v_name) in _QKV_FANOUT.items():
marker = f".{token}."
if marker in name:
q, k, v = param[:q_size], param[q_size : q_size + kv_size], param[q_size + kv_size :]
results = []
for sub_name, value in ((q_name, q), (k_name, k), (v_name, v)):
ckpt_name = name.replace(marker, f".{sub_name}.")
if token != "to_qkv":
# Joint-attention projections live under `.processor.` upstream.
ckpt_name = ckpt_name.replace(".img_instruct_attn.", ".img_instruct_attn.processor.")
results.append((ckpt_name, value))
return results
# Fused FFN gate/up splits into linear_1 (gate) / linear_3 (input).
if ".gate_up_proj." in name:
inner = param.shape[0] // 2
return [
(name.replace(".gate_up_proj.", ".linear_1."), param[:inner]),
(name.replace(".gate_up_proj.", ".linear_3."), param[inner:]),
]
# Default: 1:1 with the existing diffusers name promotions.
if ".to_out." in name:
name = name.replace(".to_out.", ".to_out.0.")
for proj in (
"img_to_q",
"img_to_k",
"img_to_v",
"instruct_to_q",
"instruct_to_k",
"instruct_to_v",
"instruct_out",
"img_out",
):
for proj in ("instruct_out", "img_out"):
token = f".img_instruct_attn.{proj}."
if token in name:
name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
break
return name
return [(name, param)]
def test_transformer_exposes_stacked_params_mapping():
"""The packed -> sub-layer mapping must be discoverable from the module tree.
``diffusion/lora/loader.py`` and the quantized weight loaders read
``stacked_params_mapping`` off the model, so it has to exist before (and
independently of) ``load_weights``.
"""
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
_BOOGU_STACKED_PARAMS_MAPPING,
BooguImageTransformer2DModel,
)
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
assert tuple(model.stacked_params_mapping) == _BOOGU_STACKED_PARAMS_MAPPING
# A per-instance copy, so consumers cannot mutate the module constant.
assert model.stacked_params_mapping is not _BOOGU_STACKED_PARAMS_MAPPING
# Every entry is a (param, shard, shard_id) triple.
for param_name, shard_name, shard_id in model.stacked_params_mapping:
assert param_name.startswith(".") and shard_name.startswith(".")
assert shard_id in {"q", "k", "v", 0, 1}
# The mapping targets exactly the fused projections this port creates.
mapped_leaves = {param.strip(".").split(".")[-1] for param, _, _ in model.stacked_params_mapping}
assert mapped_leaves == {"to_qkv", "img_to_qkv", "instruct_to_qkv", "gate_up_proj"}
def test_transformer_load_weights_round_trip():
@@ -473,24 +530,59 @@ def test_transformer_load_weights_round_trip():
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
native_params = dict(model.named_parameters())
# Build synthetic diffusers-named weights (one per native parameter).
checkpoint_weights = {}
# Build synthetic diffusers-named weights. Fused native projections fan out
# into the separate matrices the diffusers checkpoint stores.
expected: dict[str, torch.Tensor] = {}
checkpoint_weights: dict[str, torch.Tensor] = {}
for native_name, param in native_params.items():
checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
full = torch.randn_like(param)
expected[native_name] = full
for ckpt_name, value in _native_to_checkpoint_weights(native_name, full):
checkpoint_weights[ckpt_name] = value
# The remapping must be a bijection over the parameter set.
assert len(checkpoint_weights) == len(native_params)
# Every checkpoint name is unique; merged params fan out to >=1 weight.
assert len(checkpoint_weights) >= len(native_params)
loaded = model.load_weights(list(checkpoint_weights.items()))
# No missing / unexpected parameters.
assert loaded == set(native_params.keys())
# Values landed on the right parameters (TP=1: weight_loader copies verbatim).
# Values landed on the right parameters (TP=1: weight_loader copies verbatim;
# fused shards are placed at their q/k/v / gate/up offsets).
reloaded = dict(model.named_parameters())
for native_name in native_params:
expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
assert torch.allclose(reloaded[native_name], expected)
assert torch.allclose(reloaded[native_name], expected[native_name])
def test_transformer_load_weights_rejects_partial_fused_params():
"""A fused parameter must not count as loaded until all its shards arrive.
The loader compares parameter *names*, so reporting ``to_qkv`` complete
after a single ``to_q`` would let a checkpoint carrying only ``to_q`` start
up with the ``k``/``v`` slices left uninitialized.
"""
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
BooguImageTransformer2DModel,
)
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
native_params = dict(model.named_parameters())
complete_name = "noise_refiner.0.attn.to_qkv.weight"
partial_name = "noise_refiner.0.feed_forward.gate_up_proj.weight"
assert complete_name in native_params
assert partial_name in native_params
weights = list(_native_to_checkpoint_weights(complete_name, native_params[complete_name]))
# Only the gate half of the FFN, so its fused parameter stays incomplete.
gate_ckpt, gate_value = _native_to_checkpoint_weights(partial_name, native_params[partial_name])[0]
weights.append((gate_ckpt, gate_value))
loaded = model.load_weights(weights)
assert complete_name in loaded
assert partial_name not in loaded
def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
@@ -513,12 +605,15 @@ def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
try:
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
native_params = dict(model.named_parameters())
loaded_name = next(iter(native_params))
checkpoint_name = _native_to_checkpoint_name(loaded_name)
# A non-fused parameter, so one checkpoint entry loads it completely.
loaded_name = "x_embedder.weight"
assert loaded_name in native_params
(checkpoint_name, checkpoint_value), *_ = _native_to_checkpoint_weights(loaded_name, native_params[loaded_name])
checkpoint_value = torch.randn_like(checkpoint_value)
loaded = model.load_weights(
[
(checkpoint_name, torch.randn_like(native_params[loaded_name])),
(checkpoint_name, checkpoint_value),
("unexpected.weight", torch.ones(1)),
]
)
tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py
+157
View File
@@ -0,0 +1,157 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import os
import time
from datetime import timedelta
from types import SimpleNamespace
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
def _collective_worker(rank, rendezvous, backend):
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
from vllm_omni.errors import OmniClientError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
from vllm_omni.platforms import current_omni_platform
torch.set_num_threads(1)
device = torch.device("cuda", rank) if backend == "nccl" else torch.device("cpu")
if backend == "nccl":
current_omni_platform.set_device(device)
dist.init_process_group(
backend, init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=45)
)
try:
module._dit_rank_world = lambda: (dist.group.WORLD, rank, 2)
class VisualCodec:
fail = False
def is_distributed_enabled(self):
return True
def collective(self):
value = torch.ones(1, device=device)
dist.all_reduce(value)
assert value.item() == 2
def encode_image(self, image):
self.collective()
assert image.size == (32, 32)
return torch.ones(1, 96)
def encode_video(self, frames):
self.collective()
assert frames.shape == (4, 32, 32, 3)
if self.fail and rank == 1:
raise ValueError("codec failed on rank 1")
return torch.full((1, 96), 2.0), (1, 2, 2)
class AudioCodec:
def encode_waveform(self, waveform, sample_rate):
assert rank == 0
assert sample_rate == 32000
assert waveform.shape == (64000,)
return torch.full((160, 32), 3.0), 80
model = object.__new__(module.MiniMaxH3Pipeline)
torch.nn.Module.__init__(model)
model.device = device
model.od_config = SimpleNamespace(enable_layerwise_offload=False)
model.video_vae = VisualCodec()
model.audio_vae = AudioCodec()
model.partition = "combined"
model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
model.text_encoder_tp_size = 2
model.text_encoder_group = SimpleNamespace(
world_size=2, ranks=[0, 1], rank_in_group=rank, device_group=dist.group.WORLD
)
model.processor = SimpleNamespace(image_processor=SimpleNamespace(merge_size=2))
model.tokenizer = lambda text, **kwargs: {"input_ids": [1, 2, 3]}
def encode_text(ids, kwargs):
assert ids.tolist() == [1, 2, 3]
assert not kwargs
hidden = torch.ones(3, 5120, device=device, dtype=torch.bfloat16)
dist.all_reduce(hidden)
return hidden
model._encode_text_hidden = encode_text
sampling = OmniDiffusionSamplingParams(
height=32, width=32, num_frames=96, extra_args={"task": "t2va", "aspect_ratio": "1:1"}
)
text_conditioning = model._prepare_local_conditioning({"prompt": "test"}, sampling)
torch.testing.assert_close(text_conditioning.hidden_states, torch.full_like(text_conditioning.hidden_states, 2))
assert text_conditioning.visual_condition is None
media = (
MiniMaxH3EncoderMediaInput(
task="ref2va",
height=32,
width=32,
num_frames=107,
latent_t=32,
audio_t=178,
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
videos=(torch.zeros(4, 32, 32, 3, dtype=torch.uint8),),
video_audios=(None,),
audios=((torch.zeros(64000), 32000),),
)
if rank == 0
else None
)
conditioning = model._encode_local_media(media)
assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video", "audio"]
torch.testing.assert_close(
conditioning.visual_condition.cpu(), torch.cat([torch.ones(1, 96), torch.full((1, 96), 2.0)])
)
torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0))
assert conditioning.audio_condition_lengths == (80,)
model.video_vae.fail = True
with pytest.raises(OmniClientError, match="codec failed on rank 1"):
model._encode_local_media(media)
with pytest.raises(OmniClientError, match="non-empty prompt"):
model._prepare_local_conditioning({"prompt": ""}, sampling)
model.video_vae.fail = False
assert model._encode_local_media(media).visual_condition_shapes == ((1, 2, 2), (1, 2, 2))
finally:
dist.destroy_process_group()
def _run_collectives(tmp_path, backend):
context = mp.spawn(_collective_worker, args=(str(tmp_path / "rendezvous"), backend), nprocs=2, join=False)
deadline = time.monotonic() + 150
try:
while not context.join(timeout=1):
if time.monotonic() > deadline:
pytest.fail("MiniMax H3 collectives did not finish within 150 seconds")
finally:
for process in context.processes:
if process.is_alive():
process.terminate()
process.join(timeout=5)
@pytest.mark.cpu
def test_two_rank_conditioning_and_errors_gloo(tmp_path):
_run_collectives(tmp_path, "gloo")
@pytest.mark.cuda
@pytest.mark.parallel
@pytest.mark.skipif(
os.environ.get("VLLM_TEST_MINIMAX_H3_COLLECTIVES") != "1",
reason="set VLLM_TEST_MINIMAX_H3_COLLECTIVES=1 with two reserved GPUs",
)
def test_two_rank_conditioning_and_errors_nccl(tmp_path):
if torch.accelerator.device_count() < 2:
pytest.skip("requires two GPUs")
_run_collectives(tmp_path, "nccl")
tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+606 -785
View File
File diff suppressed because it is too large. Load diff
tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py
+133
View File
@@ -0,0 +1,133 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock
import pytest
import torch
from PIL import Image
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture
def pipeline():
model = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(model)
model.load_text_encoder = True
model.load_vae_encoder = True
model.device = torch.device("cpu")
model.partition = "combined"
model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
model.default_video_shift = 12.0
model.default_audio_shift = 3.0
model.od_config = SimpleNamespace(enable_layerwise_offload=False)
model._quality_policy = SimpleNamespace(resolve=Mock(return_value=SimpleNamespace(cache_dit=None)))
model._cache_dit_runtime = SimpleNamespace(prepare=Mock())
model.encode_prompt = Mock(
return_value=(torch.ones(2, 5120, dtype=torch.bfloat16), torch.ones(2, dtype=torch.long))
)
model.video_vae = Mock()
model.video_vae.is_distributed_enabled.return_value = False
model.video_vae.encode_image.side_effect = lambda image: torch.ones((image.height // 32) * (image.width // 32), 96)
model.audio_vae = Mock()
model.transformer = torch.nn.Identity()
model.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
model.decode = Mock(return_value=(torch.zeros(1, 3, 1, 32, 32), torch.zeros(1)))
return model
def _request(task):
prompt: dict[str, Any] = {"prompt": "A person waves."}
if task != "t2va":
prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))}
sampling = OmniDiffusionSamplingParams(
height=32, width=32, num_frames=96, num_inference_steps=2, extra_args={"task": task, "aspect_ratio": "1:1"}
)
return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"])
def test_single_stage_encodes_without_stage_zero(pipeline, task):
output = pipeline.forward(_request(task))
pipeline.encode_prompt.assert_called_once()
assert pipeline.video_vae.encode_image.call_count == int(task != "t2va")
pipeline.diffuse.assert_called_once()
assert pipeline.diffuse.call_args.kwargs["task"] == task
video, audio = output.output
torch.testing.assert_close(video, torch.zeros(1, 1, 32, 32, 3, dtype=torch.uint8))
torch.testing.assert_close(audio, pipeline.decode.return_value[1])
def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch):
import numpy as np
from vllm_omni.data_entry_keys import flatten_payload
from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MiniMaxH3EncoderConditioning,
MiniMaxH3TextConditioning,
)
from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion, prepare_encoder_prompt
frames = np.zeros((48, 32, 32, 3), dtype=np.uint8)
decode_frames = Mock(return_value=frames)
sampler = Mock(wraps=processing.sample_reference_video_frames)
monkeypatch.setattr(processing, "load_video_frames", decode_frames)
monkeypatch.setattr(processing, "sample_reference_video_frames", sampler)
monkeypatch.setattr(
processing,
"prepare_reference_videos",
Mock(
return_value=[{"prepared_path": "prepared.mp4", "original_path": "original.mp4", "input_has_audio": True}]
),
)
monkeypatch.setattr(processing, "load_video_audio", Mock(return_value=(torch.full((64000,), 3.0), 32000)))
pipeline.video_vae.encode_video.return_value = (torch.full((1, 96), 2.0), (1, 2, 2))
def encode_audio(waveform, sample_rate):
length = round(waveform.shape[-1] / sample_rate * 40)
return torch.full((length * 2, 32), float(waveform[0])), length
pipeline.audio_vae.encode_waveform.side_effect = encode_audio
request = _request("ref2va")
request.prompts[0]["multi_modal_data"].update(video="original.mp4", audio=(torch.full((192000,), 4.0), 32000))
pipeline.forward(request)
local_kwargs = pipeline.diffuse.call_args.kwargs
transformed = prepare_encoder_prompt(request.prompts[0], [request.sampling_params])
stage_zero = object.__new__(MiniMaxH3Encoder)
torch.nn.Module.__init__(stage_zero)
stage_zero._component_leader = True
stage_zero.video_vae = pipeline.video_vae
stage_zero.audio_vae = pipeline.audio_vae
media = stage_zero._encode_media(stage_zero._media_input(transformed["additional_information"]))
text = MiniMaxH3TextConditioning(*pipeline.encode_prompt.return_value)
payload = MiniMaxH3EncoderConditioning.from_components(text, media).to_omni_payload()
source = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=flatten_payload(payload))])
bridged = encoder2diffusion([source], request.prompts[0])
pipeline.load_text_encoder = False
pipeline.load_vae_encoder = False
pipeline.encode_prompt.reset_mock()
pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params))
external_kwargs = pipeline.diffuse.call_args.kwargs
for key, value in local_kwargs.items():
if isinstance(value, torch.Tensor):
torch.testing.assert_close(value, external_kwargs[key])
else:
assert value == external_kwargs[key]
pipeline.encode_prompt.assert_not_called()
assert decode_frames.call_count == sampler.call_count == 2
assert all(call.kwargs["decoded_frames"] is frames for call in sampler.call_args_list)
assert [block["kind"] for block in local_kwargs["ref_blocks"]] == ["image", "video_audio", "audio"]
assert local_kwargs["audio_condition_lengths"] == [80, round(107 / 24 * 40)]
assert all(
call.args[0].shape[-1] == round(107 / 24 * 32000)
for call in pipeline.audio_vae.encode_waveform.call_args_list[1::2]
)
tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py
+149 -142
View File
@@ -9,72 +9,25 @@ from unittest.mock import Mock
import pytest
import torch
from PIL import Image
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders():
def test_h3_full_pipeline_profiler_includes_local_encoding():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
targets = MiniMaxH3Pipeline._PROFILER_TARGETS
assert "_encode_visual_conditions" in targets
assert "_encode_reference_audio_conditions" in targets
assert "video_vae.decode_latent" in targets
assert "audio_vae.decode_latent" in targets
assert "_encode_video_conditions" not in targets
assert "_encode_video_audio_conditions" not in targets
assert "_encode_audio_conditions" not in targets
@pytest.mark.parametrize("distributed", [False, True])
def test_manual_component_offload_wins_over_requested_model_offload(distributed):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=True,
enable_layerwise_offload=not distributed,
enable_distributed_layerwise_offload=distributed,
)
pipeline._model_cpu_offload_modules = []
pipeline.text_encoder = Mock()
expected = torch.ones(2, 3)
pipeline.text_encoder.encode_ids.return_value = expected
component = Mock()
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
with pipeline._component_on_device(component):
component.load_to_device.assert_called_once_with()
component.offload_to_cpu.assert_not_called()
assert actual is expected
pipeline.text_encoder.load_to_device.assert_called_once_with()
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
component.offload_to_cpu.assert_called_once_with()
def test_requested_model_offload_without_backend_uses_resident_components():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=True,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._model_cpu_offload_modules = []
pipeline.text_encoder = Mock()
expected = torch.ones(2, 3)
pipeline.text_encoder.encode_ids.return_value = expected
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
assert actual is expected
pipeline.text_encoder.load_to_device.assert_called_once_with()
pipeline.text_encoder.encode_ids.assert_called_once()
assert targets == [
"encode_prompt",
"_encode_local_media",
"diffuse",
"decode",
"video_vae.decode_latent",
"audio_vae.decode_latent",
"prepare_encode",
"denoise_step",
"post_decode",
]
def test_decoded_output_is_unchanged_without_active_model_offload():
@@ -194,7 +147,7 @@ def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypa
_extract_prompt=lambda raw: ("a prompt", {}),
_extract_text_conditioning=lambda raw: None,
_extract_prepared_reference_videos=lambda raw: None,
_prepare_request_inputs=lambda **kwargs: context,
_prepare_request_inputs=lambda *args, **kwargs: context,
_denoise_kwargs=lambda ctx: {},
diffuse=diffuse,
decode=decode,
@@ -243,7 +196,7 @@ def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypat
_extract_prompt=lambda raw: ("a prompt", {}),
_extract_text_conditioning=lambda raw: None,
_extract_prepared_reference_videos=lambda raw: None,
_prepare_request_inputs=lambda **kwargs: context,
_prepare_request_inputs=lambda *args, **kwargs: context,
_denoise_kwargs=lambda ctx: {},
diffuse=diffuse,
decode=decode,
@@ -318,7 +271,8 @@ def test_output_owner_rank_matches_executor_contract(monkeypatch, available, ini
assert MiniMaxH3Pipeline._is_output_owner_rank() is expected
def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
@pytest.mark.parametrize("load_text_encoder", [True, False])
def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text_encoder):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
@@ -326,9 +280,11 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
torch.nn.Module.__init__(pipeline)
pipeline.transformer = torch.nn.Linear(2, 2)
pipeline.transformers_ref = torch.nn.Linear(2, 2)
pipeline.text_encoder = torch.nn.Linear(2, 2)
pipeline.video_vae = torch.nn.Linear(2, 2)
pipeline.audio_vae = torch.nn.Linear(2, 2)
pipeline._encoder_modules = ["text_encoder"] if load_text_encoder else []
if load_text_encoder:
pipeline.text_encoder = torch.nn.Linear(2, 2)
apply_offload = Mock()
remove_offload = Mock()
monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
@@ -341,7 +297,7 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
)
dits = [pipeline.transformer, pipeline.transformers_ref]
stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
stages = [*([pipeline.text_encoder] if load_text_encoder else []), pipeline.video_vae, pipeline.audio_vae]
apply_offload.assert_called_once_with(
dit_modules=dits,
encoder_modules=stages,
@@ -356,6 +312,46 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
remove_offload.assert_called_once_with([*dits, *stages])
@pytest.mark.parametrize("decode_fails", [False, True])
def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
component = torch.nn.Linear(2, 2)
events: list[tuple[Any, ...]] = []
@contextmanager
def record_component(value):
events.append(("activate", value))
try:
yield
finally:
events.append(("offload", value))
monkeypatch.setattr(module, "sequential_offload_component", record_component)
pipeline._model_cpu_offload_modules = [component]
def decode():
with pipeline._component_on_device(component):
events.append(("decode", component))
if decode_fails:
raise RuntimeError("decode failed")
if decode_fails:
with pytest.raises(RuntimeError, match="decode failed"):
decode()
else:
decode()
assert events == [
("activate", component),
("decode", component),
("offload", component),
]
def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
@@ -416,44 +412,20 @@ def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
apply_offload.assert_not_called()
@pytest.mark.parametrize("decode_fails", [False, True])
def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
def _encode_media(pipeline, **kwargs):
from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
component = torch.nn.Linear(2, 2)
events: list[tuple[Any, ...]] = []
@contextmanager
def record_component(value):
events.append(("activate", value))
try:
yield
finally:
events.append(("offload", value))
monkeypatch.setattr(module, "sequential_offload_component", record_component)
pipeline._model_cpu_offload_modules = [component]
def decode():
with pipeline._component_on_device(component):
events.append(("decode", component))
if decode_fails:
raise RuntimeError("decode failed")
if decode_fails:
with pytest.raises(RuntimeError, match="decode failed"):
decode()
else:
decode()
assert events == [
("activate", component),
("decode", component),
("offload", component),
]
media = MiniMaxH3EncoderMediaInput(
task="ref2va", height=32, width=32, num_frames=124, latent_t=31, audio_t=208, **kwargs
)
return encode_media(
media,
video_vae=getattr(pipeline, "video_vae", None),
audio_vae=getattr(pipeline, "audio_vae", None),
emit_conditioning=True,
component_scope=pipeline._component_on_device,
)
@pytest.mark.parametrize("encode_fails", [False, True])
@@ -463,7 +435,6 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.video_vae = Mock()
pipeline._model_cpu_offload_modules = [pipeline.video_vae]
events = []
@@ -484,17 +455,15 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
pipeline.video_vae.encode_image.side_effect = encode_image
monkeypatch.setattr(module, "sequential_offload_component", record_component)
monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
images = (torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8))
if encode_fails:
with pytest.raises(RuntimeError, match="encode failed"):
pipeline._encode_visual_conditions(images, None, video_count=0)
_encode_media(pipeline, images=images)
else:
rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0)
assert rows.shape == (3, 4)
assert shapes == [(1, 1, 1), (1, 1, 2)]
result = _encode_media(pipeline, images=images)
assert result.visual_condition.shape == (3, 4)
assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2))
assert events == [
("activate", pipeline.video_vae),
@@ -510,9 +479,7 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.video_vae = Mock()
pipeline.video_vae.is_distributed_enabled.return_value = False
pipeline._model_cpu_offload_modules = [pipeline.video_vae]
events = []
@@ -529,31 +496,26 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
return torch.ones(1, 4)
def encode_video(frames):
events.append(("video", frames))
events.append(("video", frames.shape))
return torch.ones(3, 4), (3, 2, 2)
pipeline.video_vae.encode_image.side_effect = encode_image
pipeline.video_vae.encode_video.side_effect = encode_video
monkeypatch.setattr(module, "sequential_offload_component", record_component)
monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}")
images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
prepared_videos = [{"prepared_path": "reference.mp4"}]
rows, shapes = pipeline._encode_visual_conditions(
images,
prepared_videos,
video_count=1,
result = _encode_media(
pipeline,
images=(torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8)),
videos=(torch.zeros(12, 32, 32, 3, dtype=torch.uint8),),
video_audios=(None,),
)
assert rows.shape == (5, 4)
assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)]
assert result.visual_condition.shape == (5, 4)
assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2), (3, 2, 2))
assert events == [
("activate", pipeline.video_vae),
("image", (16, 16)),
("image", (32, 16)),
("video", "frames:reference.mp4"),
("video", (12, 32, 32, 3)),
("offload", pipeline.video_vae),
]
@@ -567,7 +529,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
torch.nn.Module.__init__(pipeline)
pipeline.audio_vae = Mock()
pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
events: list[tuple[Any, ...]] = []
events = []
@contextmanager
def record_component(value):
@@ -577,26 +539,20 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
finally:
events.append(("offload", value))
def encode_embedded(*args, **kwargs):
events.append(("embedded",))
def encode_waveform(waveform, sample_rate):
events.append(("audio", int(waveform[0]), sample_rate))
if encode_fails:
raise RuntimeError("audio encode failed")
return torch.ones(1, 2), [1]
return torch.ones(80, 32), 80
def encode_standalone(*args, **kwargs):
events.append(("standalone",))
return torch.ones(1, 2), [1]
pipeline._encode_video_audio_conditions_resident = encode_embedded
pipeline._encode_audio_conditions_resident = encode_standalone
pipeline.audio_vae.encode_waveform.side_effect = encode_waveform
monkeypatch.setattr(module, "sequential_offload_component", record_component)
def call():
return pipeline._encode_reference_audio_conditions(
[{"input_has_audio": True}],
has_audio=[True],
standalone_audios=[(torch.ones(1), 16_000)],
max_duration_seconds=1.0,
return _encode_media(
pipeline,
video_audios=(),
audios=((torch.ones(32_000), 16_000),),
)
if encode_fails:
@@ -604,10 +560,61 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
call()
else:
result = call()
assert result[0].shape == (1, 2)
assert result[2].shape == (1, 2)
assert result.audio_condition.shape == (80, 32)
assert result.audio_condition_lengths == (80,)
assert events[1:-1] == [("audio", 1, 16_000)]
assert events[0] == ("activate", pipeline.audio_vae)
assert events[-1] == ("offload", pipeline.audio_vae)
assert events.count(("activate", pipeline.audio_vae)) == 1
assert events.count(("offload", pipeline.audio_vae)) == 1
@pytest.mark.parametrize("distributed", [False, True])
def test_manual_component_offload_wins_over_requested_model_offload(distributed):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=True,
enable_layerwise_offload=not distributed,
enable_distributed_layerwise_offload=distributed,
)
pipeline._model_cpu_offload_modules = []
pipeline.text_encoder = Mock()
expected = torch.ones(2, 3)
pipeline.text_encoder.encode_ids.return_value = expected
component = Mock()
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
with pipeline._component_on_device(component):
component.load_to_device.assert_called_once_with()
component.offload_to_cpu.assert_not_called()
assert actual is expected
pipeline.text_encoder.load_to_device.assert_called_once_with()
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
component.offload_to_cpu.assert_called_once_with()
def test_requested_model_offload_without_backend_uses_resident_components():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=True,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._model_cpu_offload_modules = []
pipeline.text_encoder = Mock()
expected = torch.ones(2, 3)
pipeline.text_encoder.encode_ids.return_value = expected
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
assert actual is expected
pipeline.text_encoder.load_to_device.assert_called_once_with()
pipeline.text_encoder.encode_ids.assert_called_once()
tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
+3 -2
View File
@@ -186,10 +186,11 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
common_kwargs = {
"model": model_ref,
"enforce_eager": True,
"tensor_parallel_size": 2,
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
"text_encoder_tp_size": 2,
"stage_0_tensor_parallel_size": 2,
"vae_use_tiling": True,
# MiniMax H3's VAE is implemented by code shipped with the checkpoint.
"trust_remote_code": True,
}
omni_bl = Omni(**common_kwargs)
tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py
+13 -49
View File
@@ -110,6 +110,8 @@ def _step_pipeline(model, *, packed_batch_supported: bool = True):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
pipeline.load_text_encoder = False
pipeline.load_vae_encoder = False
pipeline.transformer = model
pipeline.device = torch.device("cpu")
pipeline._transformer_for_task = lambda task: model
@@ -304,8 +306,17 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames):
}
pipeline = _step_pipeline(_SegmentMeanModel())
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {})))
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context)
conditioning = object()
monkeypatch.setattr(
mod.MiniMaxH3Pipeline,
"_extract_encoder_conditioning",
staticmethod(lambda _: conditioning),
)
monkeypatch.setattr(
mod.MiniMaxH3Pipeline,
"_prepare_encoder_conditioning_inputs",
lambda self, value, sampling: context,
)
monkeypatch.setattr(
mod.MiniMaxH3Pipeline,
"_build_denoise_inputs",
@@ -439,50 +450,3 @@ def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False
def test_broadcast_rank0_exception_single_rank_reraises():
"""Single-rank execution has no group; the helper just reraises."""
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_broadcast_rank0_exception,
)
from vllm_omni.errors import OmniClientError
_broadcast_rank0_exception(None)
with pytest.raises(OmniClientError, match="bad ref"):
_broadcast_rank0_exception(OmniClientError("bad ref"))
def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch):
"""A rank-0 error becomes a matching client error on every other DiT rank."""
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
from vllm_omni.errors import OmniClientError
def fake_rank_world(rank):
return lambda: (object(), rank, 4)
def make_broadcast(rank0_payload):
def fake_broadcast(payload_list, *, src, group):
payload_list[0] = rank0_payload
return fake_broadcast
monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0))
err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError")
rank0_payload = {
"type": type(err).__name__,
"message": str(err),
"status_code": err.status_code,
"error_type": err.error_type,
}
monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload))
with pytest.raises(OmniClientError) as rank0_info:
mod._broadcast_rank0_exception(err)
assert rank0_info.value is err
monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2))
with pytest.raises(OmniClientError) as rank2_info:
mod._broadcast_rank0_exception(None)
assert rank2_info.value.status_code == 422
assert rank2_info.value.error_type == "UnprocessableEntityError"
assert "invalid reference-video file" in str(rank2_info.value)
tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py
+20 -14
View File
@@ -11,7 +11,6 @@ import torch.nn as nn
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
_apply_qwen_image_rotary_emb,
_qwen_image_qk_norm_rope,
)
@@ -73,22 +72,28 @@ def _make_input(
return QwenImageQKInput(q=q, k=k, norm_q=norm_q, norm_k=norm_k, freqs=freqs)
def _reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
def _eager_rotary(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
q = data.norm_q(data.q)
k = data.norm_k(data.k)
if data.q.device.type == "cuda":
return (
_apply_qwen_image_rotary_emb(q, data.freqs),
_apply_qwen_image_rotary_emb(k, data.freqs),
)
rope = RotaryEmbedding(is_neox_style=False)
cos = data.freqs.real.to(data.q.dtype)
sin = data.freqs.imag.to(data.q.dtype)
return rope(q, cos, sin), rope(k, cos, sin)
def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
def _fp32_complex_rotary(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
"""Local copy of the old CUDA helper; fused kernel still tracks this math."""
paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype)
def _fused_kernel_reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
q = data.norm_q(data.q)
k = data.norm_k(data.k)
return _fp32_complex_rotary(q, data.freqs), _fp32_complex_rotary(k, data.freqs)
def _run(data: QwenImageQKInput, *, use_fused: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
return _qwen_image_qk_norm_rope(
data.q,
data.k,
@@ -97,6 +102,7 @@ def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
data.freqs,
RotaryEmbedding(is_neox_style=False),
EPS,
use_fused=use_fused,
)
@@ -110,8 +116,8 @@ def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference():
packed_qkv_view=True,
)
actual_q, actual_k = _run(data)
expected_q, expected_k = _reference(data)
actual_q, actual_k = _run(data, use_fused=False)
expected_q, expected_k = _eager_rotary(data)
torch.testing.assert_close(actual_q, expected_q, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(actual_k, expected_k, atol=1e-5, rtol=1e-5)
@@ -127,8 +133,8 @@ def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference():
packed_qkv_view=True,
)
actual_q, actual_k = _run(data)
expected_q, expected_k = _reference(data)
actual_q, actual_k = _run(data, use_fused=False)
expected_q, expected_k = _eager_rotary(data)
torch.testing.assert_close(actual_q, expected_q, atol=1e-3, rtol=1e-3)
torch.testing.assert_close(actual_k, expected_k, atol=1e-3, rtol=1e-3)
@@ -150,7 +156,7 @@ def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference(
)
actual_q, actual_k = _run(data)
expected_q, expected_k = _reference(data)
expected_q, expected_k = _fused_kernel_reference(data)
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py
-25
View File
@@ -3,37 +3,12 @@ import torch
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
QwenEmbedRope,
_apply_qwen_image_rotary_emb,
)
from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("noncontiguous", [False, True])
def test_qwen_rotary_preserves_frequency_precision(dtype, noncontiguous):
from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
generator = torch.Generator().manual_seed(42)
x = torch.randn(2, 17, 3, 128, generator=generator).to(dtype)
if noncontiguous:
x = x.transpose(1, 2).contiguous().transpose(1, 2)
angles = torch.randn(17, 64, generator=generator)
freqs = torch.polar(torch.ones_like(angles), angles)
expected = apply_rotary_emb_qwen(x, freqs, use_real=False)
actual = _apply_qwen_image_rotary_emb(x, freqs)
assert actual.dtype == dtype
assert actual.shape == x.shape
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
if dtype == torch.bfloat16:
# A regression to early BF16 frequency rounding must be observable.
rounded_freqs = torch.complex(freqs.real.to(dtype).float(), freqs.imag.to(dtype).float())
rounded = _apply_qwen_image_rotary_emb(x, rounded_freqs)
assert not torch.equal(actual, rounded)
def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count():
prompt_embeds = torch.zeros(2, 64, 8)
prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool)
tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py
+87
View File
@@ -0,0 +1,87 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.qwen_image import qwen_image_transformer
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenImageTransformerBlock
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
class _StaticOutput(nn.Module):
def __init__(self, output: torch.Tensor):
super().__init__()
self.output = output
def forward(self, _: torch.Tensor) -> torch.Tensor:
return self.output
class _RecordingNorm(nn.Module):
def __init__(self):
super().__init__()
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
def forward(self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
self.calls.append((x, scale, shift))
return x
class _ZeroAttention(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
**_: object,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.zeros_like(hidden_states), torch.zeros_like(encoder_hidden_states)
class _ZeroMLP(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.zeros_like(x)
def test_non_cuda_select01_uses_adalayernorm_dispatch(monkeypatch):
hidden_size = 4
block = QwenImageTransformerBlock.__new__(QwenImageTransformerBlock)
nn.Module.__init__(block)
block.zero_cond_t = True
block.img_mod = _StaticOutput(torch.zeros(2, 6 * hidden_size))
block.txt_mod = _StaticOutput(torch.zeros(1, 6 * hidden_size))
block.img_norm1 = _RecordingNorm()
block.img_norm2 = _RecordingNorm()
block.txt_norm1 = _RecordingNorm()
block.txt_norm2 = _RecordingNorm()
block.attn = _ZeroAttention()
block.img_mlp = _ZeroMLP()
block.txt_mlp = _ZeroMLP()
monkeypatch.setattr(qwen_image_transformer, "can_use_qwen_select01_triton", lambda _: False)
def fail_if_called(*args, **kwargs):
pytest.fail("non-CUDA select01 must preserve the AdaLayerNorm dispatch")
monkeypatch.setattr(qwen_image_transformer, "fused_layernorm_select01", fail_if_called)
monkeypatch.setattr(qwen_image_transformer, "fused_residual_layernorm_select01", fail_if_called)
hidden_states = torch.randn(1, 3, hidden_size)
encoder_hidden_states = torch.randn(1, 2, hidden_size)
modulate_index = torch.tensor([[0, 1, 0]])
block(
hidden_states,
encoder_hidden_states,
torch.ones(1, 2),
torch.zeros(2, hidden_size),
(torch.empty(0), torch.empty(0)),
modulate_index=modulate_index,
)
assert len(block.img_norm1.calls) == 1
assert len(block.img_norm2.calls) == 1
for _, scale, shift in (*block.img_norm1.calls, *block.img_norm2.calls):
assert scale.shape == hidden_states.shape
assert shift.shape == hidden_states.shape
tests/diffusion/models/wan2_2/test_create_transformer_quant.py
+212 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for transformer quant-config auto-detection.
The loader path at pipeline_wan2_2.py carries two quantization contracts:
@@ -27,6 +27,7 @@ from types import SimpleNamespace
import pytest
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_
# od_config must remain unchanged — _create_transformer does not modify it.
assert od_config.quantization_config is None
def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch):
import json
import torch
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
for component in ("transformer", "transformer_2"):
directory = tmp_path / component
directory.mkdir()
config = dict(_MIN_CFG)
expert_quant_config = (
disk_quant_config[0 if component == "transformer" else 1]
if isinstance(disk_quant_config, tuple)
else disk_quant_config
)
if expert_quant_config is not None:
config["quantization_config"] = expert_quant_config
(directory / "config.json").write_text(json.dumps(config))
(tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]}))
class FakeComponent(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
class FakeTransformer(FakeComponent):
def __init__(self, quant_config=None, **kwargs):
super().__init__()
self.received_quant_config = quant_config
for module in (wan22_module, i2v_module):
monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent())
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer)
pipeline_cls = {
"t2v": wan22_module.Wan22Pipeline,
"i2v": i2v_module.Wan22I2VPipeline,
"vace": vace_module.Wan22VACEPipeline,
}[pipeline_kind]
monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None)
od_config = OmniDiffusionConfig(
model=str(tmp_path),
model_class_name=pipeline_cls.__name__,
dtype=torch.bfloat16,
quantization_config=active,
boundary_ratio=0.875,
flow_shift=12.0,
enable_diffusion_pipeline_profiler=False,
)
od_config.enrich_config()
pipeline = pipeline_cls(od_config=od_config)
return pipeline
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
@pytest.mark.parametrize("offline", [False, True])
def test_pipeline_constructor_routes_each_expert_quantization(
pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker
):
"""Run the actual pipeline constructors, expert factories and quant-method dispatch."""
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.quantization import build_quant_config
active = build_quant_config(
{
"transformer": {
"method": method,
"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
"w4a8_fallback_steps": [0, 1],
},
"transformer_2": {
"method": method,
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
"w4a8_fallback_steps": [],
},
}
)
disk_quant_config = None
if offline:
disk_quant_config = {
"quant_method": method,
flag: True,
"ignored_layers": ["blocks.12.ffn.net_2"],
"w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"],
"w4a8_fallback_steps": [9],
}
pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch)
for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)):
received = transformer.received_quant_config
assert received.get_name() == method
assert getattr(received, flag) == offline
assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else [])
selected = f"blocks.{10 + index}.attn1.to_qkv"
other = f"blocks.{11 - index}.attn1.to_qkv"
assert received.w4a8_fallback_layers == [selected]
assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer
assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer
if offline:
assert isinstance(
received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod
)
assert pipeline.od_config.quantization_config is active
@pytest.mark.parametrize("explicit_none", [False, True])
def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch):
from vllm_omni.quantization import build_quant_config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None}
active = build_quant_config(components)
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
assert captured[-1].get("quant_config") is None
with pytest.raises(ValueError, match="disabled for component 'transformer_2'"):
pipeline._create_transformer(
{**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}},
component="transformer_2",
)
def test_pipeline_component_default_reaches_expert_factory(monkeypatch):
from vllm_omni.quantization import build_quant_config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}})
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
assert captured[-1]["quant_config"] is active.default_config
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
@pytest.mark.parametrize(
"configured_expert, with_default",
[
("transformer", False),
("transformer", True),
("transformer_2", False),
("transformer_2", True),
],
)
def test_pipeline_partial_expert_config_does_not_match_sibling(
pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch
):
from vllm_omni.quantization import build_quant_config
component_spec = {
configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
}
if with_default:
component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]}
active = build_quant_config(component_spec)
pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch)
sibling = "transformer_2" if configured_expert == "transformer" else "transformer"
selected = getattr(pipeline, configured_expert).received_quant_config
other = getattr(pipeline, sibling).received_quant_config
assert selected is active.component_configs[configured_expert]
assert other is active.default_config
if with_default:
assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"]
else:
assert other is None
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
@pytest.mark.parametrize("different_ignored_layers", [False, True])
def test_enriched_auto_quantization_preserves_each_expert_policy(
pipeline_kind, different_ignored_layers, tmp_path, monkeypatch
):
"""Startup auto-detection must not turn the first expert into a user override."""
high = {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"ignored_layers": ["blocks.0.ffn.net_2"],
"w4a8_fallback_steps": [],
"w4a8_fallback_layers": [],
"mxfp4_scale_alg": 0,
"require_smooth_scale": False,
}
low = {
**high,
"ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"],
"w4a8_fallback_steps": [37],
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
"mxfp4_scale_alg": 2,
"require_smooth_scale": True,
}
pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch)
assert pipeline.od_config.quantization_config is not None # enrich_config really detected the first expert.
for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]:
config = transformer.received_quant_config
for policy in (
"ignored_layers",
"w4a8_fallback_steps",
"w4a8_fallback_layers",
"mxfp4_scale_alg",
"require_smooth_scale",
):
assert getattr(config, policy) == expected[policy], policy
tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
+97 -1
View File
@@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
def _cpu_platform(monkeypatch):
def _cpu_pipeline_runtime(monkeypatch):
# These lightweight pipelines have no accelerator or distributed groups.
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1))
class _StubTransformer(nn.Module):
@@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
assert torch.equal(result, torch.full_like(latents, 10.0))
@pytest.mark.parametrize("fail_first_request", [False, True])
def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch):
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.forward_context import (
ForwardContext,
get_forward_context,
is_forward_context_available,
override_forward_context,
)
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
contexts: list[ForwardContext] = []
events = []
class InjectedPredictionError(Exception):
pass
class RecordingTransformer(_StubTransformer):
def __init__(self, expert):
super().__init__()
self.expert = expert
def forward(self, hidden_states, encoder_hidden_states, **kwargs):
context = get_forward_context()
assert context is contexts[-1]
request_index = len(contexts) - 1
step = context.denoise_step_idx
positive = bool(encoder_hidden_states[0, 0, 0] > 0)
events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2])))
if fail_first_request and request_index == 0 and step == 1 and not positive:
raise InjectedPredictionError
return (torch.ones_like(hidden_states) * (1 if positive else -1),)
def encode_prompt(**kwargs):
context = get_forward_context()
assert context.denoise_step_idx is None
contexts.append(context)
embeds = torch.ones(1, 8, 8)
return embeds, -embeds
pipeline = _make_pipeline()
pipeline.transformer = RecordingTransformer("high")
pipeline.transformer_2 = RecordingTransformer("low")
pipeline.scheduler = _StubScheduler([900, 500, 100])
monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents)
# Use the real runner context manager, Wan forward/diffuse, and CFG
# dispatch. Only model computation, scheduler math and unrelated KV I/O
# are replaced; no test code publishes a denoise step.
runner = object.__new__(DiffusionModelRunner)
runner.pipeline = pipeline
runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32)
runner.vllm_config = None
runner.cache_backend = None
monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None)
with override_forward_context(None):
for request_index in range(2):
request = OmniDiffusionRequest(
prompt="a fox walks",
request_id=f"request-{request_index}",
sampling_params=OmniDiffusionSamplingParams(
num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent"
),
)
def execute():
return runner._execute_request_list(
[request],
od_config=runner.od_config,
allow_single_output=True,
require_request_batch_support=False,
record_name="test_wan_step_context",
record_output_peak_memory=False,
)
if fail_first_request and request_index == 0:
with pytest.raises(InjectedPredictionError):
execute()
else:
execute()
assert not is_forward_context_available()
assert len(contexts) == 2 and contexts[0] is not contexts[1]
expected = [
(request, step, "high" if step == 0 else "low", positive, step in (0, 2))
for request in range(2)
for step in range(2 if fail_first_request and request == 0 else 3)
for positive in (True, False)
]
assert events == expected
class _StubDMDScheduler:
def __init__(self) -> None:
self.predict_clean_calls: list[tuple[float, float, float]] = []
tests/diffusion/quantization/test_mxfp4_checkpoint.py
+294
View File
@@ -0,0 +1,294 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract."""
import sys
from types import ModuleType
import pytest
import torch
from vllm.config.load import LoadConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4Config,
DiffusionMXFP4DualScaleMixedConfig,
NPUMxfp4DualScaleLinearMethod,
NPUMxfp4LinearMethod,
NPUMxfp4OnlineLinearMethod,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _patch_tp_state(monkeypatch):
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
class _CheckpointLayer(torch.nn.Module):
def __init__(self, method, dtype=torch.bfloat16):
super().__init__()
method.create_weights(
self,
input_size_per_partition=512,
output_partition_sizes=[2],
input_size=512,
output_size=2,
params_dtype=dtype,
weight_loader=default_weight_loader,
)
def load_weights(self, weights):
params = dict(self.named_parameters())
loaded = set()
for name, value in weights:
if name in params:
params[name].weight_loader(params[name], value)
loaded.add(name)
return loaded
def _load_checkpoint(layer, config, weights, monkeypatch):
od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
loader.counter_before_loading_weights = 0.0
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
loader.load_weights(layer)
@pytest.mark.parametrize("serialized", [False, True])
def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized)
# Exercise the shared implementation directly as well as the online MRO
# below: checkpoint requirements must follow the configured storage mode.
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)}
if serialized:
with pytest.raises(ValueError, match="Required weights.*weight_scale"):
_load_checkpoint(layer, config, weights, monkeypatch)
else:
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch):
config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg)
method = NPUMxfp4OnlineLinearMethod(config)
layer = _CheckpointLayer(method)
assert set(dict(layer.named_parameters())) == {"weight"}
assert layer.weight.is_meta
weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512)
packed_weight = torch.zeros((2, 256), dtype=torch.uint8)
generated_scale = torch.full((2, 16), 127, dtype=torch.uint8)
quantized_inputs = []
def quantize(value, **kwargs):
quantized_inputs.append(value.clone())
expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
if scale_alg == 2:
expected["dst_type_max"] = 7.25
assert kwargs == expected
return packed_weight, generated_scale
# Keep real creation, lazy loading, processing, and strict loader checks;
# substitute only the device quantization operator for this CPU test.
npu = ModuleType("torch_npu")
npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize)
monkeypatch.setitem(sys.modules, "torch_npu", npu)
_load_checkpoint(layer, config, {"weight": weight}, monkeypatch)
assert len(quantized_inputs) == 1
torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0)
torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0)
torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0)
assert not getattr(layer.weight_scale, "is_checkpoint_required", False)
method.process_weights_after_loading(layer)
assert len(quantized_inputs) == 1
@pytest.mark.parametrize("require_smooth", [False, True])
@pytest.mark.parametrize("has_smooth", [False, True])
def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
scale = torch.linspace(0.5, 3.0, 512)
if has_smooth:
weights["mul_scale"] = scale
else:
weights.pop("mul_scale")
if require_smooth and not has_smooth:
with pytest.raises(ValueError, match="mul_scale"):
_load_checkpoint(layer, config, weights, monkeypatch)
else:
_load_checkpoint(layer, config, weights, monkeypatch)
torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0)
@pytest.mark.parametrize("bad_scale", ["shape", "name"])
def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
if bad_scale == "shape":
weights["mul_scale"] = torch.ones(511)
with pytest.raises(AssertionError):
_load_checkpoint(layer, config, weights, monkeypatch)
else:
weights["multiply_scale"] = weights.pop("mul_scale")
with pytest.raises(ValueError, match="mul_scale"):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("required", [False, True])
@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0])
def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
weights["mul_scale"][7] = bad_value
with pytest.raises(ValueError, match="finite, strictly positive"):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)])
def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
weights["mul_scale"] = bad_scale
with pytest.raises(ValueError, match="one-dimensional floating-point"):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("bad_value", [1e10, 1e-10])
def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
method = NPUMxfp4LinearMethod(config)
layer = _CheckpointLayer(method, torch.float16)
layer.mul_scale.data.fill_(bad_value)
# Validation must fail before any device packing or GEMM.
monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu"))
with pytest.raises(ValueError, match="finite, strictly positive"):
method.process_weights_after_loading(layer)
@pytest.mark.parametrize(
("mode", "missing"),
[
("single", "weight"),
("single", "weight_scale"),
("dual", "weight"),
("dual", "weight_scale"),
("dual", "weight_dual_scale"),
("dual", "mul_scale"),
],
)
def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch):
if mode == "single":
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
method = NPUMxfp4LinearMethod(config)
else:
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
method = NPUMxfp4DualScaleLinearMethod(config)
layer = _CheckpointLayer(method)
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing}
with pytest.raises(ValueError, match=missing):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("rank", [0, 1])
def test_single_scale_smooth_uses_real_row_parallel_loader(rank):
from vllm.model_executor.layers.linear import RowParallelLinear
layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)))
layer.tp_rank = rank
checkpoint_scale = torch.arange(1024, dtype=torch.float32)
RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale)
torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("fallback", [False, True])
@pytest.mark.parametrize("with_bias", [False, True])
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_single_scale_smooth_precedes_quantization_and_preserves_output(
dtype, fallback, with_bias, scale_alg, monkeypatch
):
config = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg
)
method = NPUMxfp4LinearMethod(config)
layer = _CheckpointLayer(method, dtype)
weights = {
"weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512),
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
"mul_scale": torch.linspace(0.5, 3.0, 512),
}
_load_checkpoint(layer, config, weights, monkeypatch)
x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512)
bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None
quantized_inputs = []
# Only the two device operators are substituted. Exercise real apply(),
# step selection, Smooth multiplication, bias handling, and reshape logic.
npu = ModuleType("torch_npu")
def quantize(value, **kwargs):
quantized_inputs.append(value.clone())
expected = dict(
dst_type=torch.float8_e4m3fn if fallback else "fp4",
axis=-1,
block_size=32,
round_mode="rint",
scale_alg=0 if fallback else scale_alg,
)
if not fallback and scale_alg == 2:
expected["dst_type_max"] = 7.25
assert kwargs == expected
return value, torch.ones((6, 16), dtype=torch.uint8)
def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs):
if fallback:
assert output_dtype == torch.bfloat16
if bias is not None:
assert bias.dtype == output_dtype
assert bias.shape == (1, 2)
result = x_q.float() @ weight.float()
if bias is not None:
result += bias.float()
return result.to(output_dtype)
npu.__dict__.update(
float8_e4m3fn=torch.float8_e4m3fn,
float8_e8m0fnu=torch.uint8,
float4_e2m1fn_x2="fp4",
npu_dynamic_mx_quant=quantize,
npu_quant_matmul=matmul,
)
monkeypatch.setitem(sys.modules, "torch_npu", npu)
with override_forward_context(ForwardContext()):
set_forward_context_denoise_step_idx(0 if fallback else 1)
result = method.apply(layer, x, bias)
expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype)
assert len(quantized_inputs) == 1
torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0)
expected = expected_input.float() @ weights["weight"].float().T
if bias is not None:
expected += bias.float()
if fallback:
expected = expected.bfloat16()
torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0)
assert result.dtype == dtype
tests/diffusion/quantization/test_mxfp4_config.py
+552 -3
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
import pytest
@@ -27,12 +27,42 @@ def test_mxfp4_config_get_name():
assert DiffusionMXFP4Config.get_name() == "mxfp4"
@pytest.mark.parametrize("value", [True, False, None, "2", 2.0, 1, -1, 3])
def test_mxfp4_scale_alg_rejects_invalid_values(value):
from vllm_omni.quantization import build_quant_config
with pytest.raises(ValueError, match="mxfp4_scale_alg"):
build_quant_config({"method": "mxfp4", "mxfp4_scale_alg": value})
@pytest.mark.parametrize("runtime_alg", [0, 2])
@pytest.mark.parametrize("serialized", [False, True])
def test_runtime_algorithm_survives_offline_storage_rebuild(runtime_alg, serialized):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, mxfp4_scale_alg=runtime_alg)
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"mxfp4_scale_alg": 2,
"ignored_layers": ["proj_out"],
},
)
assert resolved.is_checkpoint_mxfp4_serialized
assert resolved.mxfp4_scale_alg == runtime_alg
assert resolved.w4a8_fallback_layers == resolved.w4a8_fallback_steps == []
def test_mxfp4_config_from_config_defaults():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = DiffusionMXFP4Config.from_config({})
assert cfg.is_checkpoint_mxfp4_serialized is False
assert cfg.ignored_layers == []
assert cfg.w4a8_fallback_layers == []
def test_mxfp4_config_from_config_serialized():
@@ -42,6 +72,98 @@ def test_mxfp4_config_from_config_serialized():
assert cfg.is_checkpoint_mxfp4_serialized is True
def test_mxfp4_config_requires_offline_smooth():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
values = {"is_checkpoint_mxfp4_serialized": True, "require_smooth_scale": True}
assert DiffusionMXFP4Config.from_config(values).require_smooth_scale
assert build_quant_config({"method": "mxfp4", **values}).require_smooth_scale
with pytest.raises(ValueError, match="offline"):
build_quant_config({"method": "mxfp4", "require_smooth_scale": True})
with pytest.raises(ValueError, match="boolean"):
DiffusionMXFP4Config.from_config({**values, "require_smooth_scale": "false"})
@pytest.mark.parametrize("serialized", [False, True])
def test_mxfp4_disk_smooth_requirement_preserves_step_policy(serialized):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, w4a8_fallback_steps=[1])
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"require_smooth_scale": True,
},
)
assert resolved.require_smooth_scale
assert resolved.is_checkpoint_mxfp4_serialized
assert resolved.w4a8_fallback_steps == [1]
@pytest.mark.parametrize("disk_requirement", [None, False, True])
def test_mxfp4_expert_rebuild_cannot_weaken_explicit_smooth_requirement(disk_requirement):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
disk = {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"ignored_layers": ["blocks.0.attn2.to_q"],
}
if disk_requirement is not None:
disk["require_smooth_scale"] = disk_requirement
resolved = resolve_quant_config_from_disk(active, disk)
assert resolved.require_smooth_scale
assert resolved.ignored_layers == disk["ignored_layers"]
@pytest.mark.parametrize("disk_qc", ["mxfp4", {"quant_method": "mxfp4"}])
def test_mxfp4_method_only_disk_config_preserves_caller_layer_policy(disk_qc):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(
ignored_layers=["blocks.0.attn2.to_q"],
w4a8_fallback_steps=[1],
w4a8_fallback_layers=["blocks.1.attn1.to_qkv"],
)
resolved = resolve_quant_config_from_disk(active, disk_qc)
assert resolved is active
assert resolved.ignored_layers == ["blocks.0.attn2.to_q"]
assert resolved.w4a8_fallback_steps == [1]
assert resolved.w4a8_fallback_layers == ["blocks.1.attn1.to_qkv"]
def test_quantization_disk_string_rejects_mismatched_active_method():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active = build_quant_config("mxfp4")
with pytest.raises(ValueError, match="mxfp4_dualscale.*active quantization config is 'mxfp4'"):
resolve_quant_config_from_disk(active, "mxfp4_dualscale")
def test_quantization_disk_string_accepts_equivalent_method_alias():
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
class ActiveAliasConfig:
@staticmethod
def get_name() -> str:
return "inc"
active = ActiveAliasConfig()
assert resolve_quant_config_from_disk(active, "auto-round") is active
def test_mxfp4_config_from_config_ignored_layers():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
@@ -57,6 +179,202 @@ def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
assert cfg.ignored_layers == ["proj_out"]
@pytest.mark.parametrize("serialized", [False, True])
@pytest.mark.parametrize(
"policy",
[
{"w4a8_fallback_steps": [0]},
{"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
{"w4a8_fallback_steps": [1], "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
],
)
def test_dualscale_rejects_w4a8_configuration(serialized, policy):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
config = {"quant_method": "mxfp4_dualscale", "is_checkpoint_serialized": serialized, **policy}
with pytest.raises(ValueError, match="does not support W4A8"):
DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=serialized, **policy)
with pytest.raises(ValueError, match="does not support W4A8"):
DiffusionMXFP4DualScaleMixedConfig.from_config(config)
with pytest.raises(ValueError, match="does not support W4A8"):
build_quant_config({"transformer": config})
with pytest.raises(ValueError, match="does not support W4A8"):
resolve_quant_config_from_disk(None, config)
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
@pytest.mark.parametrize("steps", [[-1], [True], [1.5], ["2"], "0,2", (0, 2)])
def test_mxfp4_rejects_invalid_fallback_steps(method, steps):
from vllm_omni.quantization import build_quant_config
with pytest.raises(ValueError, match="non-negative integer"):
build_quant_config({"method": method, "w4a8_fallback_steps": steps})
def test_mxfp4_fallback_steps_follow_request_context():
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
steps = [2, 0, 2]
cfg = build_quant_config({"method": "mxfp4", "w4a8_fallback_steps": steps})
steps.append(1)
assert cfg.w4a8_fallback_steps == [0, 2]
with override_forward_context(None):
assert not _is_w4a8_fallback_step([])
with pytest.raises(RuntimeError, match="publishes its denoise step"):
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
for _request in range(2):
with override_forward_context(ForwardContext()):
for step, expected in [(0, True), (1, False), (2, True), (3, False)]:
set_forward_context_denoise_step_idx(step)
# Both CFG passes and both experts read the same index.
assert [_is_w4a8_fallback_step(cfg.w4a8_fallback_steps) for _cfg_pass in range(2)] == [
expected,
expected,
]
with pytest.raises(RuntimeError, match="publishes its denoise step"):
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
@pytest.mark.parametrize("method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
@pytest.mark.parametrize("steps", [[], [0, 2]])
@pytest.mark.parametrize("already_serialized", [False, True])
def test_mxfp4_disk_rebuild_preserves_runtime_fallback(method, flag, steps, already_serialized):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
layers = ["blocks.10.attn1.to_qkv"] if steps else []
active = build_quant_config(
{
"method": method,
flag: already_serialized,
"w4a8_fallback_steps": steps,
"w4a8_fallback_layers": layers,
}
)
# Separate experts may have distinct high precision layers. Both must
# retain the requested step policy during online/offline reconciliation.
for ignored in [["blocks.0.attn1.to_q"], ["blocks.1.attn1.to_q"]]:
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": method,
flag: True,
"ignored_layers": ignored,
"w4a8_fallback_steps": [9],
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
},
)
assert getattr(resolved, flag)
assert resolved.ignored_layers == ignored
assert resolved.w4a8_fallback_steps == steps
assert resolved.w4a8_fallback_layers == layers
assert active.ignored_layers == []
@pytest.mark.parametrize(
"runtime_policy",
[{}, {"w4a8_fallback_steps": [1]}, {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}],
)
def test_mxfp4_explicit_omission_disables_saved_policies_for_both_experts(runtime_policy):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active = build_quant_config({"method": "mxfp4", **runtime_policy})
resolved_experts = []
for ignored in (["blocks.0.attn2.to_q"], ["blocks.1.ffn.net_2"]):
disk = {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"ignored_layers": ignored,
"w4a8_fallback_steps": [9],
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
}
resolved = resolve_quant_config_from_disk(active, disk)
assert resolved.w4a8_fallback_steps == runtime_policy.get("w4a8_fallback_steps", [])
assert resolved.w4a8_fallback_layers == runtime_policy.get("w4a8_fallback_layers", [])
assert resolved.ignored_layers == ignored
detected = resolve_quant_config_from_disk(None, disk)
assert detected.w4a8_fallback_steps == [9]
assert detected.w4a8_fallback_layers == ["blocks.9.ffn.net_2"]
resolved_experts.append(resolved)
assert resolved_experts[0] is not resolved_experts[1]
assert active.is_checkpoint_mxfp4_serialized is False
assert active.ignored_layers == []
@pytest.mark.parametrize(
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
)
def test_serialized_mxfp4_disk_omission_clears_active_ignored_layers(method, flag):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active_spec = {
"method": method,
flag: True,
"ignored_layers": ["blocks.0.attn2.to_q"],
"w4a8_fallback_steps": [1] if method == "mxfp4" else [],
"w4a8_fallback_layers": ["blocks.1.attn1.to_qkv"] if method == "mxfp4" else [],
}
if method == "mxfp4":
active_spec["require_smooth_scale"] = True
active = build_quant_config(active_spec)
resolved = resolve_quant_config_from_disk(active, {"quant_method": method, flag: True})
assert resolved is not active
assert resolved.ignored_layers == []
assert resolved.w4a8_fallback_steps == active.w4a8_fallback_steps
assert resolved.w4a8_fallback_layers == active.w4a8_fallback_layers
if method == "mxfp4":
assert resolved.require_smooth_scale
@pytest.mark.parametrize(
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
)
@pytest.mark.parametrize("with_active", [False, True])
def test_serialized_mxfp4_disk_uses_modules_to_not_convert(method, flag, with_active):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active = None
disk_steps = [9] if method == "mxfp4" else []
active_steps = [1] if method == "mxfp4" else []
expected_steps = disk_steps
if with_active:
active = build_quant_config(
{
"method": method,
flag: True,
"ignored_layers": ["blocks.0.attn2.to_q"],
"w4a8_fallback_steps": active_steps,
}
)
expected_steps = active_steps
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": method,
flag: True,
"modules_to_not_convert": ["blocks.2.ffn.net_2"],
"w4a8_fallback_steps": disk_steps,
},
)
assert resolved.ignored_layers == ["blocks.2.ffn.net_2"]
assert resolved.w4a8_fallback_steps == expected_steps
# ---------------------------------------------------------------------------
# build_quant_config integration
# ---------------------------------------------------------------------------
@@ -594,18 +912,31 @@ def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
@pytest.mark.parametrize("steps", [[], [0, 2]])
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch, steps):
"""A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
cfg = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"], w4a8_fallback_steps=steps
)
assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
def test_rocm_rejects_w4a8_step_fallback(_rocm_platform, monkeypatch):
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
cfg = DiffusionMXFP4Config(w4a8_fallback_steps=[0, 2])
with pytest.raises(NotImplementedError, match="only supported on NPU"):
cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
"""MXFP4 on ROCm requires gfx950; any other arch must raise."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
@@ -652,3 +983,221 @@ def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
assert layer.input_size_per_partition == _TP2_K // _TP2
assert layer.output_size_per_partition == _TP2_N
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
@pytest.mark.parametrize("layers", ["blocks.1.attn1.to_qkv", [1], [True], [""], [" a"], ["blocks.*"], ["a..b"]])
def test_mxfp4_rejects_invalid_fallback_layers(method, layers):
from vllm_omni.quantization import build_quant_config
with pytest.raises(ValueError, match="exact runtime Linear paths"):
build_quant_config({"method": method, "w4a8_fallback_layers": layers})
@pytest.mark.parametrize("method", ["mxfp4"])
def test_mxfp4_layer_policy_from_checkpoint_and_component_config(method):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.component_config import resolve_component_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
names = ["blocks.10.attn1.to_qkv", "blocks.10.attn1.to_qkv"]
components = build_quant_config(
{"transformer": {"method": method, "w4a8_fallback_layers": names}, "transformer_2": {"method": method}}
)
names.append("blocks.11.attn1.to_qkv")
assert resolve_component_quant_config(components, "transformer").w4a8_fallback_layers == [names[0]]
assert resolve_component_quant_config(components, "transformer_2").w4a8_fallback_layers == []
detected = resolve_quant_config_from_disk(None, {"quant_method": method, "w4a8_fallback_layers": [names[0]]})
assert detected.w4a8_fallback_layers == [names[0]]
assert type(detected).from_config({"w4a8_fallback_layers": [names[0]]}).w4a8_fallback_layers == [names[0]]
@pytest.mark.parametrize(
"kind, serialized_flag",
[("mxfp4", "is_checkpoint_mxfp4_serialized")],
)
@pytest.mark.parametrize("offline", [False, True])
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_mxfp4_layer_and_step_union_routes_actual_methods(
kind, serialized_flag, offline, scale_alg, monkeypatch, mocker
):
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
from vllm_omni.quantization import build_quant_config, mxfp4_config
monkeypatch.setattr(mxfp4_config.current_omni_platform, "is_npu", lambda: True)
selected = "blocks.10.attn1.to_qkv"
other = "blocks.11.attn2.to_q"
preserved = "blocks.12.ffn.net_2"
config = build_quant_config(
{
"method": kind,
"mxfp4_scale_alg": scale_alg,
serialized_flag: offline,
"w4a8_fallback_layers": [selected, preserved],
"ignored_layers": [preserved],
"w4a8_fallback_steps": [1],
}
)
# Real config dispatch creates separate methods for selected and other layers.
linear = mocker.Mock(spec=LinearBase)
assert isinstance(config.get_quant_method(linear, preserved), UnquantizedLinearMethod)
x = torch.ones(2, 64, dtype=torch.bfloat16)
layer = torch.nn.Module()
for name in ("weight", "weight_scale", "mul_scale"):
layer.register_buffer(name, torch.ones(64, dtype=x.dtype))
calls = []
def record_a8(*args):
calls.append("A8")
return x
def record_a4(*args):
calls.append("A4")
return x
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", record_a8)
for name in (selected, other):
method = config.get_quant_method(linear, name)
assert method.is_w4a8_fallback_layer == (name == selected)
scales = (x, x)
monkeypatch.setattr(method, "_quantize_activation", lambda *args: scales)
monkeypatch.setattr(method, "_quant_matmul", record_a4)
for step in (0, 1, 2):
with override_forward_context(ForwardContext(denoise_step_idx=step)):
method.apply(layer, x)
assert calls == ["A8", "A8", "A8", "A4", "A8", "A4"]
# A permanently selected layer works without a request context, even if steps are also configured.
with override_forward_context(None):
config.get_quant_method(linear, selected).apply(layer, x)
assert calls[-1] == "A8"
@pytest.mark.parametrize("online", [False, True])
@pytest.mark.parametrize("selected", [False, True])
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_single_scale_prepares_one_weight_for_both_activation_precisions(online, selected, scale_alg, monkeypatch):
import sys
from types import SimpleNamespace
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
from vllm_omni.quantization import mxfp4_config
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
prepared = []
packed = []
def quantize(x, **kwargs):
prepared.append(kwargs)
return x.to(torch.int8), torch.full((x.shape[0], 2), 127, dtype=torch.uint8)
def pack(x, dtype):
packed.append(x)
return x.to(dtype)
fake_npu = SimpleNamespace(
float4_e2m1fn_x2=torch.int8,
float8_e8m0fnu=torch.uint8,
float8_e4m3fn=torch.float8_e4m3fn,
npu_dtype_cast=pack,
npu_dynamic_mx_quant=quantize,
)
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
config = mxfp4_config.DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=not online,
w4a8_fallback_layers=["blocks.0.attn1.to_qkv"] if selected else [],
w4a8_fallback_steps=[1],
mxfp4_scale_alg=scale_alg,
)
cls = mxfp4_config.NPUMxfp4OnlineLinearMethod if online else mxfp4_config.NPUMxfp4LinearMethod
method = cls(config, prefix="blocks.0.attn1.to_qkv")
layer = torch.nn.Linear(64, 2, bias=False, dtype=torch.bfloat16)
layer.orig_dtype = torch.bfloat16
if not online:
layer.register_parameter("weight_scale", torch.nn.Parameter(torch.full((2, 2), 127, dtype=torch.uint8), False))
layer.register_parameter("mul_scale", torch.nn.Parameter(torch.linspace(0.5, 2, 64), False))
method.process_weights_after_loading(layer)
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
method.process_weights_after_loading(layer)
assert len(prepared) == int(online)
expected = dict(dst_type=torch.int8, axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
if scale_alg == 2:
expected["dst_type_max"] = 7.25
assert prepared == ([expected] if online else [])
assert len(packed) == int(not online)
branches = []
def a4(x_q, x_scale, target, bias, dtype):
assert (target.weight.data_ptr(), target.weight_scale.data_ptr()) == pointers
branches.append("A4")
return torch.zeros(x_q.shape[0], 2, dtype=dtype)
def a8(x, weight, scale, bias, dtype):
assert (weight.data_ptr(), scale.data_ptr()) == pointers
branches.append("A8")
return torch.zeros(x.shape[0], 2, dtype=dtype)
monkeypatch.setattr(method, "_quant_matmul", a4)
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", a8)
monkeypatch.setattr(method, "_quantize_activation", lambda x: (x, x))
monkeypatch.setattr(fake_npu, "npu_dynamic_mx_quant", lambda *a, **kw: pytest.fail("weight requantization"))
monkeypatch.setattr(fake_npu, "npu_dtype_cast", lambda *a, **kw: pytest.fail("weight repacking"))
for step in (0, 1, 2):
with override_forward_context(ForwardContext(denoise_step_idx=step)):
method.apply(layer, torch.ones(2, 64, dtype=torch.bfloat16))
assert branches == (["A8"] * 3 if selected else ["A4", "A8", "A4"])
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
assert set(layer.state_dict()) == (
{"weight", "weight_scale"} if online else {"weight", "weight_scale", "mul_scale"}
)
assert not hasattr(layer, "w4a8_weight")
assert not hasattr(layer, "w4a8_weight_scale")
@pytest.mark.parametrize("online", [False, True])
def test_dualscale_without_fallback_keeps_original_a4_path(online, monkeypatch):
import sys
from types import SimpleNamespace
from vllm_omni.quantization import mxfp4_config
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
original_ones = torch.ones
monkeypatch.setattr(torch, "ones", lambda *args, **kw: original_ones(*args, **{**kw, "device": "cpu"}))
calls = []
def quantize_dual(x, **kw):
calls.append(kw.get("smooth_scale"))
return x.to(torch.int8), torch.ones(x.shape[0], 1), torch.ones(x.shape[0], 16, dtype=torch.uint8)
fake_npu = SimpleNamespace(
float4_e2m1fn_x2=torch.int8,
npu_dtype_cast=lambda x, dtype: x.to(dtype),
npu_format_cast=lambda x, *args, **kw: x,
npu_dynamic_dual_level_mx_quant=quantize_dual,
)
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
config = mxfp4_config.DiffusionMXFP4DualScaleMixedConfig(
is_checkpoint_serialized=not online, w4a8_fallback_steps=[], w4a8_fallback_layers=[]
)
cls = mxfp4_config.NPUMxfp4DualScaleOnlineLinearMethod if online else mxfp4_config.NPUMxfp4DualScaleLinearMethod
method = cls(config)
layer = torch.nn.Module()
layer.input_size_per_partition = 512
for name, value in {
"weight": torch.ones(2, 512, dtype=torch.bfloat16),
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
"weight_dual_scale": torch.ones(2, 1, 1),
"mul_scale": torch.ones(512),
}.items():
layer.register_parameter(name, torch.nn.Parameter(value, requires_grad=False))
method.process_weights_after_loading(layer)
method.process_weights_after_loading(layer)
assert len(calls) == int(online)
assert set(layer.state_dict()) == {"weight", "weight_scale", "weight_dual_scale", "mul_scale"}
assert not hasattr(layer, "w4a8_weight")
monkeypatch.setattr(method, "_quant_matmul", lambda *args: torch.zeros(2, 2, dtype=torch.bfloat16))
method.apply(layer, torch.ones(2, 512, dtype=torch.bfloat16))
assert len(calls) == 1 + int(online)
torch.testing.assert_close(calls[-1], layer.mul_scale)
tests/diffusion/quantization/test_wan_mxfp4_loading.py
+237
View File
@@ -0,0 +1,237 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Exercise Wan's real name mapping and vLLM's QKV parameter loaders on CPU."""
import pytest
import torch
from vllm.config.load import LoadConfig
from vllm.model_executor.layers.linear import QKVParallelLinear
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, DiffusionMXFP4DualScaleMixedConfig
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _patch_tp_state(monkeypatch):
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
prefix = "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer"
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_world_size", lambda: 1)
class _WanCheckpoint(torch.nn.Module):
load_weights = WanTransformer3DModel.load_weights
_log_w4a8_fallback_load_summary = WanTransformer3DModel._log_w4a8_fallback_load_summary
def __init__(self, config, *, num_blocks=1):
super().__init__()
blocks = []
for block_idx in range(num_blocks):
block = torch.nn.Module()
block.attn1 = torch.nn.Module()
block.attn1.to_qkv = QKVParallelLinear(
hidden_size=512,
head_size=32,
total_num_heads=4,
bias=False,
params_dtype=torch.bfloat16,
quant_config=config,
prefix=f"blocks.{block_idx}.attn1.to_qkv",
disable_tp=True,
)
blocks.append(block)
self.blocks = torch.nn.ModuleList(blocks)
def _checkpoint(mode, fused):
if mode == "single":
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
elif mode == "dual":
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
else:
config = None
model = _WanCheckpoint(config)
weights = {}
expected = {}
for name, param in model.named_parameters():
suffix = name.rsplit(".", 1)[1]
output_dim = getattr(param, "output_dim", None)
if output_dim is None:
value = torch.full_like(param, 2.5)
shards = [value] * 3
else:
parts = param.chunk(3, dim=output_dim)
shards = [
torch.full_like(part, (121 if suffix == "weight_scale" else 1) + i) for i, part in enumerate(parts)
]
value = torch.cat(shards, dim=output_dim)
expected[name] = value
if fused:
weights[name] = value
else:
for shard_id, shard in zip(("q", "k", "v"), shards):
weights[f"blocks.0.attn1.to_{shard_id}.{suffix}"] = shard
return model, config, weights, expected
def _load(model, config, weights, monkeypatch):
loader = DiffusersPipelineLoader(
LoadConfig(), OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
)
loader.counter_before_loading_weights = 0.0
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
loader.load_weights(model)
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
@pytest.mark.parametrize("fused", [False, True])
def test_wan_qkv_loads_actual_weights_and_scales(mode, fused, monkeypatch):
model, config, weights, expected = _checkpoint(mode, fused)
_load(model, config, weights, monkeypatch)
for name, param in model.named_parameters():
torch.testing.assert_close(param, expected[name], rtol=0, atol=0)
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
def test_wan_incomplete_qkv_weight_is_not_marked_fully_loaded(mode, monkeypatch):
model, config, weights, _ = _checkpoint(mode, fused=False)
weights.pop("blocks.0.attn1.to_k.weight")
weights.pop("blocks.0.attn1.to_v.weight")
with pytest.raises(ValueError, match="to_qkv.weight"):
_load(model, config, weights, monkeypatch)
@pytest.mark.parametrize("mode", ["single", "dual"])
def test_wan_incomplete_qkv_scale_is_not_marked_fully_loaded(mode, monkeypatch):
model, config, weights, _ = _checkpoint(mode, fused=False)
weights.pop("blocks.0.attn1.to_k.weight_scale")
weights.pop("blocks.0.attn1.to_v.weight_scale")
with pytest.raises(ValueError, match="to_qkv.weight_scale"):
_load(model, config, weights, monkeypatch)
@pytest.mark.parametrize("mode", ["single", "dual"])
@pytest.mark.parametrize("reverse", [False, True])
def test_wan_split_qkv_rejects_different_smooth_scales(mode, reverse, monkeypatch):
model, config, weights, _ = _checkpoint(mode, fused=False)
weights["blocks.0.attn1.to_k.mul_scale"] = weights["blocks.0.attn1.to_k.mul_scale"] * 2
if reverse:
weights = dict(reversed(weights.items()))
with pytest.raises(ValueError, match="must share the same Smooth tensor"):
_load(model, config, weights, monkeypatch)
@pytest.mark.parametrize("name", ["to_q.missing", "to_qkv.missing", "to_qkv_extra.weight", "to_query.weight"])
def test_wan_unknown_qkv_key_is_not_reported_as_loaded(name):
model, _, _, _ = _checkpoint("single", fused=True)
loaded = model.load_weights([(f"blocks.0.attn1.{name}", torch.ones(2))])
assert loaded == set()
@pytest.mark.parametrize("fused", [False, True])
def test_wan_single_scale_rejects_relabelled_dualscale_checkpoint(fused, monkeypatch):
model, config, weights, _ = _checkpoint("single", fused=fused)
projection = "to_qkv" if fused else "to_q"
weights[f"blocks.0.attn1.{projection}.weight_dual_scale"] = torch.ones(128, 1, 1)
with pytest.raises(ValueError, match="cannot load DualScale tensor"):
_load(model, config, weights, monkeypatch)
def test_wan_layer_policy_validates_runtime_qkv_and_bf16_override():
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv"])
model = _WanCheckpoint(config)
# Uses the same validation called at the end of WanTransformer3DModel construction.
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
for bad in ("blocks.0.attn1.to_q", "blocks.0.attn1", "blocks.99.attn1.to_qkv", "transformer.blocks.0.attn1.to_qkv"):
config.w4a8_fallback_layers = [bad]
with pytest.raises(ValueError, match="exact runtime Linear paths"):
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
def test_wan_layer_policy_defers_only_other_pp_stage_paths():
from vllm.model_executor.models.utils import PPMissingLayer
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv", "blocks.1.attn1.to_qkv"])
model = _WanCheckpoint(config)
model.blocks.append(PPMissingLayer())
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
config.w4a8_fallback_layers = ["blocks.1"]
with pytest.raises(ValueError, match="exact runtime Linear paths"):
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
@pytest.mark.parametrize("is_checkpoint_serialized", [False, True], ids=["online", "offline"])
def test_loader_logs_wan_expert_summaries_after_weight_processing(is_checkpoint_serialized, mocker):
selected = "blocks.0.attn1.to_qkv"
bf16 = "blocks.1.attn1.to_qkv"
transformer_config = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
w4a8_fallback_layers=[selected],
)
transformer_2_config = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
ignored_layers=[bf16],
w4a8_fallback_layers=[selected, bf16],
)
pipeline = torch.nn.Module()
pipeline.transformer = _WanCheckpoint(transformer_config)
pipeline.transformer_2 = _WanCheckpoint(transformer_2_config, num_blocks=2)
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer, transformer_config)
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer_2, transformer_2_config)
loader = DiffusersPipelineLoader(
LoadConfig(),
OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=transformer_config),
)
events = []
summaries = []
def process_weights(model, _target_device):
events.append("process")
for module in model.modules():
method = getattr(module, "quant_method", None)
if getattr(method, "is_w4a8_fallback_layer", False):
module._already_called_process_weights_after_loading = True
def capture_summary(message, *args):
rendered = message % args
summaries.append(rendered)
events.append(f"summary:{args[0]}")
loader._init_from_load_format = lambda *_args, **_kwargs: pipeline # type: ignore[method-assign]
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
mocker.patch(
"vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info",
side_effect=capture_summary,
)
assert loader.load_model(load_device="cpu") is pipeline
assert events == ["load", "process", "summary:transformer", "summary:transformer_2"]
assert "component=transformer; selected=1 ['transformer.blocks.0.attn1.to_qkv']" in summaries[0]
assert "BF16 overrides=0 []" in summaries[0]
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[0]
assert "component=transformer_2; selected=1 ['transformer_2.blocks.0.attn1.to_qkv']" in summaries[1]
assert "BF16 overrides=1 ['transformer_2.blocks.1.attn1.to_qkv']" in summaries[1]
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[1]
def test_wan_load_summary_does_not_claim_unprocessed_layer_is_ready(mocker):
selected = "blocks.0.attn1.to_qkv"
config = DiffusionMXFP4Config(w4a8_fallback_layers=[selected])
model = _WanCheckpoint(config)
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
logger_info = mocker.patch("vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info")
WanTransformer3DModel._log_w4a8_fallback_load_summary(model, "transformer")
message, *args = logger_info.call_args.args
rendered = message % tuple(args)
assert "weight/A8 processing state=not-ready (0/1 selected layers processed)" in rendered
tests/diffusion/test_diffusion_config_propagation.py
+6
View File
@@ -25,6 +25,12 @@ from vllm_omni.diffusion.model_metadata import (
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _local_model_paths(monkeypatch):
# These tests exercise config transport, not model repository resolution.
monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model)
def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
"""Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py
+7 -2
View File
@@ -3,7 +3,7 @@
"""MiniCPM-o 4.5 Daily-Omni + Seed-TTS accuracy regression coverage.
Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches
~78% overall accuracy on Daily-Omni (``minicpm-interleave``, ``temperature=0``,
~77% overall accuracy on Daily-Omni (``minicpm-interleave``, ``temperature=0``,
text modalities, and server ``--interleave-mm-strings`` + 1fps / 128-frame
media-io kwargs, also pinned in ``minicpmo_4_5.yaml``).
"""
@@ -48,7 +48,12 @@ _RESULT_DIR = Path(
)
)
_MIN_DAILY_OMNI_ACCURACY = 0.78
# The 0.78 gate (934/1197) sits 1 item above the observed NPU mean
# (~933/1197) and inside the observed 932~940 spread, so it flips nightly -
# 3 failures in the latest 7 dated NPU nightlies (#6887). 0.77 (922/1197)
# sits 10 items below the observed minimum (932) while still failing a
# >=1pp real accuracy regression.
_MIN_DAILY_OMNI_ACCURACY = 0.77
_MAX_SEED_TTS_MEAN_WER = 0.05
# Match the validated Daily-Omni client body from daily_omni_bench.sh.
_DAILY_EXTRA_BODY = {
tests/e2e/accuracy/test_qwen_image.py
+3 -3
View File
@@ -32,9 +32,9 @@ HEIGHT = 512
NUM_INFERENCE_STEPS = 20
TRUE_CFG_SCALE = 4.0
SEED = 42
SSIM_THRESHOLD = 0.94
# With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8
# vs Diffusers on the FA3 hub path (see #5734 / #5963 / #6971).
# H100 nightly after RotaryEmbedding CUDA RoPE: SSIM 0.978375 / PSNR 31.83
# vs Diffusers (Buildkite #15339). SSIM 0.97 so a drop is visible.
SSIM_THRESHOLD = 0.97
PSNR_THRESHOLD = 27.0
MODEL_2512_ID = "Qwen/Qwen-Image-2512"
tests/e2e/features/comfyui/test_comfyui_integration.py
+77 -1
View File
@@ -28,6 +28,7 @@ from comfy_api.input import AudioInput, VideoInput
from comfyui_vllm_omni.nodes import (
VLLMOmniFastH3Deployment,
VLLMOmniGenerateImage,
VLLMOmniGenerateMusic,
VLLMOmniGenerateVideo,
VLLMOmniTTS,
VLLMOmniUnderstanding,
@@ -62,6 +63,7 @@ class ServerCase(NamedTuple):
stage_list: list
stage_configs: list[Any]
outputs: list[OmniRequestOutput]
tokenizer: Any = None
class SamplingCase(NamedTuple):
@@ -79,6 +81,7 @@ class SamplingKind(str, Enum):
UNDERSTANDING_AR_LIST = "understanding_ar_list"
TTS_NONE = "tts_none"
TTS_DIFFUSION_SINGLE = "tts_diffusion_single"
MUSIC_NONE = "music_none"
VIDEO_NONE = "video_none"
VIDEO_DIFFUSION_SINGLE = "video_diffusion_single"
VIDEO_FASTH3 = "video_fasth3"
@@ -380,6 +383,11 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
assert len(received_sampling_params_list) == 3
for i, expected in enumerate(AR_LIST_SAMPLING_PARAMS):
_assert_sampling_param_values(received_sampling_params_list[i], expected)
elif sampling_case.kind is SamplingKind.MUSIC_NONE:
assert len(received_sampling_params_list) == 2
assert prompt["prompt_token_ids"] == [101, 102, 103]
assert prompt["additional_information"]["max_audio_frames"] == [1512]
assert received_sampling_params_list[0].seed == 7
elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}:
assert len(received_sampling_params_list) == 1
elif sampling_case.kind is SamplingKind.VIDEO_NONE:
@@ -525,12 +533,17 @@ def mock_async_omni(
# Need to mock AsyncOmni itself (not only its generate method) because
# 1. The API layer uses its stage_list and stage_configs attributes
# 2. Its __init__ method has slow side effects (model & config loading).
# These cases use the non-duplex mock engine without loading model configs.
mocker.patch("vllm_omni.entrypoints.openai.api_server._should_serve_duplex", return_value=False)
mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni")
monkeypatch.setattr(
"vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat",
_mock_preprocess_chat,
)
if server_case.tokenizer is not None:
mocker.patch("transformers.AutoTokenizer.from_pretrained", return_value=server_case.tokenizer)
mock_instance = mocker.AsyncMock(spec=RealAsyncOmni)
mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case)
@@ -558,7 +571,7 @@ def mock_async_omni(
mock_instance.shutdown = mocker.MagicMock()
mock_instance.get_vllm_config = mocker.AsyncMock(return_value=None)
mock_instance.get_supported_tasks = mocker.AsyncMock(return_value=["generate"])
mock_instance.get_tokenizer = mocker.AsyncMock(return_value=None)
mock_instance.get_tokenizer = mocker.AsyncMock(return_value=server_case.tokenizer)
mock_async_omni_cls.return_value = mock_instance
yield mock_async_omni_cls
@@ -1081,3 +1094,66 @@ async def test_video_generation_node_minimax_h3_ref2va(
assert isinstance(result, tuple)
assert len(result) == 1
assert isinstance(result[0], VideoInput)
def _build_music_tokenizer():
# Checkpoint token IDs required by the API adapter's tokenizer validation.
special_token_ids = {
"<|im_start|>": 151644,
"<|im_end|>": 151645,
"<|audio_cfg|>": 151654,
"<|audio_start|>": 151669,
"<|audio_end|>": 151670,
"<|caption_start|>": 151671,
"<|caption_end|>": 151672,
"<|lyrics_start|>": 151673,
"<|lyrics_end|>": 151674,
}
def encode(prompt):
assert prompt == (
"<|im_start|><|caption_start|>Warm acoustic pop at 92 BPM"
"<|caption_end|><|lyrics_start|>[start]\n[verse]\nThe morning light"
"<|lyrics_end|><|im_end|><|audio_start|>"
)
return [101, 102, 103]
return SimpleNamespace(encode=encode, convert_tokens_to_ids=special_token_ids.get)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"server_case",
[
ServerCase(
served_model="MiniMaxAI/MiniMax-Music3",
stage_list=["llm", "llm"],
stage_configs=[
_make_stage_config("llm", model_stage="minimax_music3_ar"),
_make_stage_config("llm", model_stage="minimax_music3_acoustic"),
],
outputs=[_build_audio_speech_output()],
tokenizer=_build_music_tokenizer(),
)
],
indirect=True,
)
@pytest.mark.parametrize(
"sampling_case",
[SamplingCase(kind=SamplingKind.MUSIC_NONE, sampling_params=None)],
indirect=True,
)
async def test_music_generation_node_minimax_music3(api_server: str, sampling_case: SamplingCase):
result = await VLLMOmniGenerateMusic().generate(
url=api_server,
model="MiniMaxAI/MiniMax-Music3",
lyrics="[Verse]\nThe morning light",
instructions="Warm acoustic pop at 92 BPM",
max_duration_seconds=60.5,
seed=7,
response_format="wav",
)
assert isinstance(result, tuple)
assert len(result) == 1
assert result[0]["sample_rate"] == 24000
assert result[0]["waveform"].shape == (1, 1, 24000)
tests/e2e/offline_inference/test_moss_tts_expansion.py
+4 -5
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B).
Weekly expansion coverage (``pytest.mark.slow``) for the delay-model path.
@@ -74,7 +74,6 @@ def _get_test_config() -> str:
pytestmark = [
pytest.mark.slow,
pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
pytest.mark.tts,
pytest.mark.parametrize(
"omni_runner",
@@ -137,7 +136,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
return torch.cat(chunks, dim=0), sr_final
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None:
"""VoiceGenerator: English instruction produces non-empty 24 kHz audio."""
req = _build_request(
@@ -151,7 +150,7 @@ def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None:
assert not torch.all(audio == 0), "Audio is silence"
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None:
"""VoiceGenerator: Chinese input produces non-empty audio."""
req = _build_request(
@@ -165,7 +164,7 @@ def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None:
assert not torch.all(audio == 0)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
def test_moss_tts_delay_batch(omni_runner: OmniRunner) -> None:
"""VoiceGenerator: batch of two requests each returns non-empty audio."""
requests = [
tests/e2e/offline_inference/test_moss_tts_realtime.py
+2 -3
View File
@@ -95,7 +95,6 @@ def _get_test_config() -> str:
# ---------------------------------------------------------------------------
pytestmark = [
pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
pytest.mark.slow,
pytest.mark.tts,
pytest.mark.parametrize(
@@ -272,7 +271,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
# ---------------------------------------------------------------------------
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.")
@@ -283,7 +282,7 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str)
assert not torch.all(audio == 0), "Audio is silence"
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
"""MossTTSRealtime: Chinese input produces non-empty audio."""
req = _build_request(ref_audio_path, "你好，这是语音合成测试。")
tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py
+2 -3
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B).
Exercises the local-transformer generation path on L4. Delay-model coverage is in
@@ -33,7 +33,6 @@ _OMNI_RUNNER_PARAM = (
pytestmark = [
pytest.mark.slow,
pytest.mark.tts,
pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
@@ -118,7 +117,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
raise AssertionError("No stage outputs received")
@hardware_test(res={"cuda": "L4"})
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path) -> None:
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
req = _build_request("This is a real-time TTS streaming test.", ref_audio_path)
tests/e2e/online_serving/test_minicpmo_4_5.py
+21 -2
View File
@@ -83,6 +83,25 @@ test_params = [
),
]
turn_test_params = [
pytest.param(
OmniServerParams(
model=_MODEL,
stage_config_path=modify_stage_config(_CI_DEPLOY, updates={"session_mode": "turn"}),
use_stage_cli=False,
server_args=[
"--trust-remote-code",
"--async-chunk",
"--chat-template",
str(_NATIVE_CHAT_TEMPLATE),
"--chat-template-content-format",
"openai",
],
),
id="turn_async_chunk",
),
]
def get_system_prompt():
return {
@@ -138,7 +157,7 @@ def _read_prompt_selection_log(offset: int) -> str:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + turn_test_params, indirect=True)
def test_text_to_text_001(omni_server, openai_client) -> None:
"""
Test text-only input generating text output via OpenAI API.
@@ -163,7 +182,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
@pytest.mark.full_model
@pytest.mark.omni
@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + turn_test_params, indirect=True)
def test_text_to_audio_001(omni_server, openai_client) -> None:
"""
Test text-only input generating text + audio output via OpenAI API.
tests/e2e/online_serving/test_moss_tts_expansion.py
+13 -15
View File
@@ -30,20 +30,16 @@ from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
_SKIP_ISSUE_6417 = pytest.mark.skip(
reason="https://github.com/vllm-project/vllm-omni/issues/6417",
)
pytestmark = [
pytest.mark.advanced_model,
pytest.mark.slow,
pytest.mark.tts,
_SKIP_ISSUE_6417,
]
MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/HEAD/assets/audio/reference_zh_1.wav"
# Voice-clone output is not reliably transcribed by the Whisper check used at
# full_model run_level; assert non-trivial WAV payload size instead.
# Reject a response that contains only a header or a very short audio payload.
# Advanced and full runs also validate the transcript.
_MIN_AUDIO_BYTES = 10_000
@@ -77,7 +73,9 @@ def get_prompt(prompt_type: str = "text") -> str:
"""
prompts = {
"text": "Hello, this is a real-time voice cloning demo for testing.",
"chinese": "你好，这是一段实时语音合成测试。",
# Match the offline test prompt. Whisper-small can transcribe "一段"
# as "一本", which makes this short sample fail the similarity check.
"chinese": "你好，这是语音合成测试。",
}
return prompts.get(prompt_type, prompts["text"])
@@ -94,7 +92,7 @@ tts_server_params = [
]
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None:
"""
@@ -105,9 +103,8 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No
Input Setting: stream=False
Datasets: single request
NOTE: ``min_audio_bytes`` skips Whisper transcript similarity — cloned
voice output is often empty or mismatched under ASR without indicating a
real serving regression.
``min_audio_bytes`` adds a payload-size check. It does not disable the
transcript check used by advanced and full runs.
"""
request_config = {
"model": omni_server.model,
@@ -121,7 +118,7 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No
online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_url) -> None:
"""
@@ -150,7 +147,7 @@ def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_
online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_url) -> None:
"""
@@ -161,7 +158,8 @@ def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_ur
Input Setting: stream=False
Datasets: single request
NOTE: same ``min_audio_bytes`` rationale as test_text_to_audio_001.
``min_audio_bytes`` adds a payload-size check. It does not disable the
transcript check used by advanced and full runs.
"""
request_config = {
"model": omni_server.model,
tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
+26
View File
@@ -96,6 +96,32 @@ def test_mix_to_text_audio_001(omni_server, online_client) -> None:
online_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("stream", [False, True])
def test_audio_in_video_boundaries(omni_server, online_client, stream: bool) -> None:
"""Interleaved video/audio input must not treat audio boundaries as image placeholders."""
# More than two seconds exercises multiple interleaved video/audio chunks.
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}"
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
video_data_url=video_data_url,
content_text="Describe what you see and hear in this video in one sentence.",
)
responses = online_client.send_omni_request(
{
"model": omni_server.model,
"messages": messages,
"stream": stream,
"modalities": ["text"],
"use_audio_in_video": True,
}
)
assert responses[0].text_content
@pytest.mark.slow
@pytest.mark.omni
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py
+137
View File
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import base64
import json
import pytest
from websockets.sync.client import connect
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
pytestmark = [pytest.mark.slow, pytest.mark.tts]
_DEPLOY = modify_stage_config(
get_deploy_config_path("qwen3_tts.yaml"),
updates={
"stages": {
0: {
"max_num_seqs": 4,
"max_num_batched_tokens": 2048,
"max_model_len": 2048,
"gpu_memory_utilization": 0.45,
"compilation_config.cudagraph_capture_sizes": [1, 2, 4],
"default_sampling_params.max_tokens": 256,
},
1: {
"max_num_seqs": 4,
"max_num_batched_tokens": 4096,
"max_model_len": 8192,
"gpu_memory_utilization": 0.15,
},
}
},
)
_PARAMS = [
OmniServerParams(
model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
stage_config_path=_DEPLOY,
server_args=[
"--trust-remote-code",
"--forced-aligner",
"Qwen/Qwen3-ForcedAligner-0.6B",
"--kv-cache-memory-bytes",
str(512 * 1024**2),
"--max-num-seqs",
"4",
"--max-num-batched-tokens",
"2048",
],
)
]
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_async_chunk_word_timestamps(omni_server, online_client):
online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "The weather is nice today, perfect for a walk in the park.",
"voice": "vivian",
"language": "English",
"word_timestamps": True,
"response_format": "wav",
"min_audio_bytes": 40_000,
"timeout": 120,
},
request_num=2,
)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_async_chunk_speech_without_timestamps(omni_server, online_client):
online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "The weather is nice today, perfect for a walk in the park.",
"voice": "vivian",
"language": "English",
"stream": True,
"stream_format": "audio",
"response_format": "wav",
"min_audio_bytes": 40_000,
"timeout": 120,
}
)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_async_chunk_streaming_word_timestamps(omni_server):
text = "The weather is nice today, perfect for a walk in the park."
chunks = []
timestamps = None
with connect(f"ws://{omni_server.host}:{omni_server.port}/v1/audio/speech/stream") as ws:
ws.send(
json.dumps(
{
"type": "session.config",
"voice": "vivian",
"language": "English",
"stream_audio": True,
"response_format": "pcm",
"word_timestamps": True,
}
)
)
ws.send(json.dumps({"type": "input.text", "text": text}))
ws.send(json.dumps({"type": "input.done"}))
while True:
message = json.loads(ws.recv(timeout=120))
assert message["type"] != "error", message
if message["type"] == "audio.chunk":
if message["audio_b64"]:
chunks.append(base64.b64decode(message["audio_b64"]))
if message["timestamps"] is not None:
timestamps = message["timestamps"]
sample_rate = message["sample_rate"]
elif message["type"] == "audio.done":
assert not message["error"], message
assert message["total_bytes"] == sum(map(len, chunks))
elif message["type"] == "session.done":
break
assert len(chunks) > 1
assert timestamps
assert len(timestamps) == len(text.split())
duration_ms = sum(map(len, chunks)) / (2 * sample_rate) * 1000
final_chunk_ms = len(chunks[-1]) / (2 * sample_rate) * 1000
assert timestamps[-1]["end_ms"] > max(duration_ms / 2, final_chunk_ms)
previous_end = 0
for timestamp in timestamps:
assert previous_end <= timestamp["start_ms"] <= timestamp["end_ms"] <= duration_ms + 1
previous_end = timestamp["end_ms"]
tests/engine/test_async_omni_engine_stage_init.py
+23
View File
@@ -11,6 +11,7 @@ import types
from dataclasses import dataclass, field
import pytest
from omegaconf import OmegaConf
from vllm.v1.engine.utils import EngineZmqAddresses
from vllm_omni.diffusion.data import AttentionConfig
@@ -2239,3 +2240,25 @@ def test_common_launch_parallel_admission_before_spawn(monkeypatch, client_count
with runtime.launch_stage_engines(client_count):
pass
assert events == ["guard", "admit", "spawn", "ready"]
@pytest.mark.parametrize(
"stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)]
)
def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes, expected):
engine = object.__new__(AsyncOmniEngine)
stages = [OmegaConf.create({"engine_args": {"async_chunk": mode}}) for mode in stage_modes]
monkeypatch.setattr(async_omni_engine_module.StageConfigFactory, "get_pipeline_config", lambda *a, **k: None)
monkeypatch.setattr(engine, "_resolve_stage_configs", lambda *a, **k: (None, stages))
monkeypatch.setattr(engine, "_set_pipeline_runtime_config", lambda *a: None)
class ConfigResolvedError(Exception):
pass
def stop_before_queues(*args, **kwargs):
raise ConfigResolvedError
monkeypatch.setattr(async_omni_engine_module.janus, "Queue", stop_before_queues)
with pytest.raises(ConfigResolvedError):
engine.__init__("test-model")
assert engine.async_chunk is expected
tests/engine/test_orchestrator.py
+59
View File
@@ -21,6 +21,7 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason
from vllm.v1.engine.exceptions import EngineDeadError
from vllm.v1.metrics.stats import IterationStats
from vllm_omni.engine import OmniEngineCoreOutput
from vllm_omni.engine.messages import (
AbortRequestMessage,
AbortResultMessage,
@@ -1891,3 +1892,61 @@ async def test_abort_retry_does_not_repeat_successful_stage_abort():
assert first.physical_abort_calls == 1
assert second.physical_abort_calls == 2
assert second.bound == set()
@pytest.mark.asyncio
async def test_duplex_session_request_error_finish_is_delivered_as_request_error() -> None:
"""A session-owned request the scheduler finished with FinishReason.ERROR
(e.g. its prompt could not grow past max_model_len) must reach the
session's consumer as a request-scoped error, not vanish: its terminal
outputs are not processed like other requests, and the output processor
may already have dropped the request state."""
output_queue: asyncio.Queue = asyncio.Queue()
orchestrator = Orchestrator(
request_async_queue=asyncio.Queue(),
output_async_queue=output_queue,
rpc_async_queue=asyncio.Queue(),
stage_pools=[],
)
duplex_state = OrchestratorRequestState(request_id="duplex-req", session_owned=True)
plain_state = OrchestratorRequestState(request_id="plain-req")
reason = "context_length_exceeded: streaming session prompt would grow to 8220 tokens, above max_model_len 8192"
await orchestrator._report_duplex_session_request_error(
0,
0,
OmniEngineCoreOutput(
request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.ERROR, stop_reason=reason
),
duplex_state,
)
error = output_queue.get_nowait()
assert isinstance(error, ErrorMessage)
assert error.request_id == "duplex-req"
assert error.stage_id == 0
assert error.fatal is False
assert error.error == reason
# An ordinary segment stop, a segment-finished error and a request that is
# not session-owned produce nothing.
await orchestrator._report_duplex_session_request_error(
0,
0,
OmniEngineCoreOutput(request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.STOP),
duplex_state,
)
await orchestrator._report_duplex_session_request_error(
0,
0,
OmniEngineCoreOutput(
request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.ERROR, is_segment_finished=True
),
duplex_state,
)
await orchestrator._report_duplex_session_request_error(
0,
0,
OmniEngineCoreOutput(request_id="plain-req", new_token_ids=[], finish_reason=FinishReason.ERROR),
plain_state,
)
assert output_queue.empty()
tests/engine/test_orchestrator_kv_sender_info.py
+64
View File
@@ -414,3 +414,67 @@ def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas():
assert all(isinstance(request, OmniEngineCoreRequest) for request in submitted.values())
assert {key: request.payload_sender_info for key, request in submitted.items()} == endpoints
@pytest.mark.asyncio
async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker):
from vllm import PoolingParams, RequestOutput
from vllm.config import VllmConfig
from vllm_omni.config.model import OmniModelConfig
from vllm_omni.engine import OmniEngineCoreRequest
orchestrator = object.__new__(Orchestrator)
orchestrator.async_chunk = True
orchestrator._pd_pair = None
orchestrator._cfg_tracker = CfgCompanionTracker()
orchestrator.output_async_queue = asyncio.Queue()
orchestrator._duplex_output_decision = lambda *args: None
orchestrator._is_duplex_session_request = lambda state: False
orchestrator._forward_to_next_stage = mocker.AsyncMock()
orchestrator._cleanup_request_ids = mocker.AsyncMock()
# A sender stage stands in for Code2Wav, whose audio is also a final output.
audio_pool = _build_sender_pool(0, {})
audio_pool.stage_client.final_output = True
aligner = _DummyDiffusionStage()
aligner_pool = StagePool(
1,
aligner,
stage_vllm_config=mocker.Mock(
spec=VllmConfig, model_config=mocker.Mock(spec=OmniModelConfig, async_chunk=False)
),
)
orchestrator.stage_pools = [audio_pool, aligner_pool]
state = OrchestratorRequestState(
request_id="aligned",
prompt={"additional_information": {"text": ["Hello world"]}},
sampling_params_list=[SamplingParams(), PoolingParams(task="token_classify")],
final_stage_id=1,
final_output_stage_ids={0, 1},
)
await orchestrator._prewarm_async_chunk_stages(
"aligned", mocker.Mock(spec=OmniEngineCoreRequest, prompt_token_ids=[1, 2]), state
)
assert aligner.calls == []
assert 1 not in state.stage_submit_ts
output = RequestOutput("aligned", None, [1, 2], None, [], False)
await orchestrator._route_output(0, 0, output, state, None)
orchestrator._forward_to_next_stage.assert_not_awaited()
output.finished = True
await orchestrator._route_output(0, 0, output, state, None)
orchestrator._forward_to_next_stage.assert_awaited_once_with(
"aligned",
0,
output,
state,
src_replica_id=0,
is_streaming_session=False,
is_final_update=False,
)
orchestrator._cleanup_request_ids.assert_not_awaited()
await orchestrator._route_output(1, 0, output, state, None)
orchestrator._cleanup_request_ids.assert_awaited_once()
messages = [orchestrator.output_async_queue.get_nowait() for _ in range(3)]
assert [message.finished for message in messages] == [False, False, True]
tests/engine/test_output_processor.py
+21
View File
@@ -888,3 +888,24 @@ def test_abort_snapshot_leaves_state_until_commit():
assert "0_parent" not in processor.request_states
assert "parent" not in processor.parent_requests
assert "parent" not in processor.external_req_ids
def test_cumulative_audio_terminal_output_supplies_full_aligner_waveform():
from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner
state = _make_state(RequestOutputKind.CUMULATIVE)
for index, count in enumerate((800, 1200, 400)):
state.add_multimodal_tensor(torch.full((1, count), float(index)), mm_type=AUDIO)
output = state.make_request_output([index], None, FinishReason.STOP if index == 2 else None, None)
assert output is not None
if index < 2:
assert code2wav2aligner([output], {}) == []
prompt = {"additional_information": {"text": ["Hello world"], "language": ["English"]}}
aligned = code2wav2aligner([output], prompt)[0]
waveform, sr = aligned["multi_modal_data"]["audio"]
assert sr == 24000
assert waveform.shape == (2400,)
assert (waveform[:800] == 0).all()
assert (waveform[800:2000] == 1).all()
assert (waveform[2000:] == 2).all()
assert aligned["additional_information"]["aligner_audio_duration_ms"] == [100.0]
tests/entrypoints/openai_api/test_duplex_api_server.py
+104 -26
View File
@@ -1,28 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Guards for the duplex-only API server (unified full-duplex framework).
A model whose pipeline declares ``duplex_plugin`` is served through
``DuplexOmni`` and nothing else: sessions run over ``/v1/realtime?duplex=1``
(alias ``/v1/duplex``) and every turn-based route reports "not available".
These tests lock that shape, which ``test_api_server_guards.py`` covers for the
diffusion and multi-stage servers:
* the pipeline probe that decides duplex vs turn-based at startup,
* the ``app.state`` snapshot of a duplex server (what is wired, what stays
``None``),
* both duplex websocket routes reaching the session handler, and
``/v1/realtime`` without the flag still reporting the turn-based Realtime API
as unavailable,
* the startup-warmup gate that holds real clients while the warmup session runs.
CPU-only: no engine is started; ``DuplexOmni`` is subclassed with just the
accessors the app-state wiring reads.
"""
"""CPU tests for serving-mode selection, duplex API wiring and warmup."""
from __future__ import annotations
import asyncio
from argparse import Namespace
from dataclasses import replace
from types import SimpleNamespace
import pytest
@@ -31,12 +15,14 @@ from fastapi.testclient import TestClient
from starlette.datastructures import State
from starlette.websockets import WebSocketDisconnect
from vllm_omni.config import stage_config
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, PipelineConfig, StagePipelineConfig
from vllm_omni.engine.duplex.config import DuplexCapabilities
from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler
from vllm_omni.entrypoints.duplex_omni import DuplexOmni
from vllm_omni.entrypoints.openai import api_server
from vllm_omni.utils.tracking_parser import TrackingNamespace
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -202,9 +188,10 @@ def _minimal_args(**overrides) -> SimpleNamespace:
(None, False),
],
)
def test_duplex_model_probe_follows_the_pipeline_plugin(monkeypatch, pipeline_config, expected: bool) -> None:
"""``duplex_plugin`` alone decides whether ``serve`` builds ``DuplexOmni``."""
def test_duplex_serving_requires_a_pipeline_plugin(monkeypatch, tmp_path, pipeline_config, expected: bool) -> None:
seen: dict[str, object] = {}
deploy_path = tmp_path / "deploy.yaml"
deploy_path.write_text("session_mode: duplex\n")
def fake_get_pipeline_config(*, model, trust_remote_code, deploy_config_path=None, **_kwargs):
seen.update(model=model, trust_remote_code=trust_remote_code, deploy_config_path=deploy_config_path)
@@ -212,12 +199,12 @@ def test_duplex_model_probe_follows_the_pipeline_plugin(monkeypatch, pipeline_co
monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config)
kwargs = {"trust_remote_code": True, "deploy_config": "deploy.yaml"}
assert api_server._is_duplex_model("demo-duplex-model", kwargs) is expected
kwargs = {"trust_remote_code": True, "deploy_config": str(deploy_path)}
assert api_server._should_serve_duplex("demo-duplex-model", kwargs) is expected
assert seen == {
"model": "demo-duplex-model",
"trust_remote_code": True,
"deploy_config_path": "deploy.yaml",
"deploy_config_path": str(deploy_path),
}
@@ -230,7 +217,98 @@ def test_duplex_model_probe_propagates_resolution_errors(monkeypatch) -> None:
monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config)
with pytest.raises(ValueError, match="trust_remote_code"):
api_server._is_duplex_model("demo-duplex-model", {})
api_server._should_serve_duplex("demo-duplex-model", {})
@pytest.fixture
def duplex_pipeline():
return PipelineConfig(
model_type="demo",
stages=(StagePipelineConfig(stage_id=0, model_stage="demo", final_output=True),),
duplex_plugin="pkg.mod.Plugin",
)
@pytest.mark.parametrize("session_mode", ["turn", "duplex"])
@pytest.mark.parametrize("config_source", ["explicit", "default", "basename", "inherited"])
@pytest.mark.asyncio
async def test_serving_mode_selects_engine_and_preserves_stage_config(
monkeypatch, tmp_path, duplex_pipeline, session_mode, config_source
):
base_path = tmp_path / "base.yaml"
base_path.write_text("session_mode: duplex\nstages:\n - stage_id: 0\n max_num_seqs: 4\n")
deploy_path = tmp_path / "deploy.yaml"
deploy_path.write_text(f"base_config: base.yaml\nsession_mode: {session_mode}\n")
if config_source == "inherited":
deploy_path = tmp_path / "overlay.yaml"
deploy_path.write_text("base_config: deploy.yaml\n")
pipeline = replace(duplex_pipeline, default_deploy_config_name=deploy_path.name)
monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: pipeline)
monkeypatch.setattr(stage_config, "_DEPLOY_DIR", tmp_path)
kwargs = {}
if config_source != "default":
kwargs["deploy_config"] = deploy_path.name if config_source == "basename" else str(deploy_path)
created = []
class TurnEngine:
def __init__(self, model, **engine_kwargs):
self.model = model
self.kwargs = engine_kwargs
self.closed = False
created.append(self)
def shutdown(self):
self.closed = True
class DuplexEngine(TurnEngine):
pass
monkeypatch.setattr(api_server, "AsyncOmni", TurnEngine)
monkeypatch.setattr(api_server, "DuplexOmni", DuplexEngine)
args = TrackingNamespace(
Namespace(model="demo-model", disable_log_stats=True, **kwargs), frozenset({"model", *kwargs})
)
async with api_server.build_async_omni_from_stage_config(args) as engine:
assert type(engine) is (DuplexEngine if session_mode == "duplex" else TurnEngine)
assert engine.model == "demo-model"
assert engine.kwargs == {**kwargs, "log_stats": False}
assert not engine.closed
assert len(created) == 1
assert engine.closed
deploy = stage_config.load_deploy_config(deploy_path)
assert deploy.session_mode == session_mode
assert deploy.stages[0].max_num_seqs == 4
@pytest.mark.parametrize("contents", ["{}\n", "session_mode: null\n", "session_mode: typo\n"])
def test_duplex_serving_rejects_missing_or_invalid_mode(monkeypatch, tmp_path, duplex_pipeline, contents):
deploy_path = tmp_path / "deploy.yaml"
deploy_path.write_text(contents)
monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: duplex_pipeline)
with pytest.raises(ValueError, match="session_mode: turn or duplex"):
api_server._should_serve_duplex("demo-model", {"deploy_config": str(deploy_path)})
def test_duplex_serving_requires_a_deploy_config(monkeypatch, duplex_pipeline):
monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: duplex_pipeline)
with pytest.raises(ValueError, match="requires a deploy config"):
api_server._should_serve_duplex("demo-model", {})
def test_duplex_serving_propagates_deploy_load_errors(monkeypatch, tmp_path, duplex_pipeline):
monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: duplex_pipeline)
with pytest.raises(FileNotFoundError):
api_server._should_serve_duplex("demo-model", {"deploy_config": str(tmp_path / "missing.yaml")})
@pytest.mark.parametrize("deploy_name,expected", [("minicpmo_4_5.yaml", True), ("minicpmo_4_5_turn.yaml", False)])
def test_minicpmo_serving_profiles(monkeypatch, deploy_name, expected):
from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", lambda **_: MINICPMO_4_5_PIPELINE)
assert api_server._should_serve_duplex("openbmb/MiniCPM-o-4_5", {"deploy_config": deploy_name}) is expected
assert api_server._should_serve_duplex("openbmb/MiniCPM-o-4_5", {}) is True
# --------------------------------------------------------------------------- #
tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py
+302
View File
@@ -0,0 +1,302 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import threading
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import torch
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
from vllm_omni.entrypoints.openai.tts_adapters import moss_tts as adapter_module
from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import MossTTSAdapter
from vllm_omni.model_executor.models.moss_tts import reference_encoder
from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize(
("hf_config", "expected_codec_path"),
[
(SimpleNamespace(), "OpenMOSS-Team/MOSS-Audio-Tokenizer"),
(
SimpleNamespace(codec_model_name_or_path="/models/custom-moss-codec"),
"/models/custom-moss-codec",
),
],
ids=["default-codec", "configured-codec"],
)
def test_realtime_components_use_the_realtime_model_and_codec(
monkeypatch: pytest.MonkeyPatch,
hf_config: SimpleNamespace,
expected_codec_path: str,
) -> None:
engine_client = SimpleNamespace(
model_config=SimpleNamespace(
model="OpenMOSS-Team/MOSS-TTS-Realtime",
hf_config=hf_config,
)
)
server = MossTTSAdapter(SpeechServingContext(server=object(), engine_client=engine_client))
tokenizer = object()
codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})()
processor_calls = []
class Processor:
def __init__(self, *, tokenizer):
processor_calls.append(tokenizer)
class_calls = []
tokenizer_calls = []
codec_calls = []
def fake_get_class(class_reference, model_id):
class_calls.append((class_reference, model_id))
return Processor
def fake_load_tokenizer(model_id, *, trust_remote_code):
tokenizer_calls.append((model_id, trust_remote_code))
return tokenizer
def fake_load_codec(model_id, *, trust_remote_code):
codec_calls.append((model_id, trust_remote_code))
return codec
monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", fake_get_class)
monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", fake_load_tokenizer)
monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", fake_load_codec)
components = server._get_moss_realtime_components()
assert components[0] is tokenizer
assert components[2] is codec
assert server._get_moss_realtime_components() is components
assert class_calls == [
(
"processing_mossttsrealtime.MossTTSRealtimeProcessor",
"OpenMOSS-Team/MOSS-TTS-Realtime",
)
]
assert tokenizer_calls == [("OpenMOSS-Team/MOSS-TTS-Realtime", True)]
assert codec_calls == [(expected_codec_path, True)]
assert processor_calls == [tokenizer]
def test_realtime_serving_builds_the_talker_prompt(monkeypatch: pytest.MonkeyPatch) -> None:
server = MossTTSAdapter(SpeechServingContext(server=SimpleNamespace(uploaded_speakers={})))
server._moss_variant = "realtime"
tokenizer = object()
processor = object()
codec = object()
request_thread = threading.get_ident()
def load_components():
assert threading.get_ident() != request_thread
return tokenizer, processor, codec
server._get_moss_realtime_components = load_components
reference_codes = torch.arange(64, dtype=torch.int64).reshape(4, 16)
encode = AsyncMock(return_value=([reference_codes], {0: "reference-cache-key"}))
monkeypatch.setattr(server, "_encode_moss_references", encode)
build_call = None
def fake_build(actual_tokenizer, actual_processor, text, actual_codes):
nonlocal build_call
build_call = (actual_tokenizer, actual_processor, text, actual_codes)
return {
"prompt_token_ids": [10, 11],
"codes": {"ref": actual_codes},
"ids": {"all": [12]},
}
monkeypatch.setattr(adapter_module, "build_realtime_prompt", fake_build)
params = asyncio.run(
server._build_moss_tts_params(
OpenAICreateSpeechRequest(
input="speak this text",
ref_audio="data:audio/wav;base64,AAAA",
max_new_tokens=50,
)
)
)
encode.assert_awaited_once()
assert encode.call_args.kwargs == {"has_inline_ref_audio": False, "two_speaker": False}
assert encode.call_args.args[0].ref_audio == "data:audio/wav;base64,AAAA"
assert build_call == (tokenizer, processor, "speak this text", reference_codes)
assert params == {
"prompt_token_ids": [10, 11],
"codes": {"ref": reference_codes},
"ids": {"all": [12]},
"max_new_frames": [50],
"ref_audio_cache_key": "reference-cache-key",
}
assert "prompt_audio_array" not in params
def test_cancelled_component_waiter_does_not_duplicate_codec_load(monkeypatch):
adapter = MossTTSAdapter(
SpeechServingContext(
server=object(),
engine_client=SimpleNamespace(
model_config=SimpleNamespace(model="OpenMOSS-Team/MOSS-TTS-Realtime", hf_config=SimpleNamespace())
),
)
)
started = threading.Event()
release = threading.Event()
second_started = threading.Event()
loads = []
codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})()
def load_codec(*args, **kwargs):
loads.append(args)
started.set()
assert release.wait(5), "codec loader was not released"
return codec
monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", lambda *args: lambda **kwargs: object())
monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: object())
monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", load_codec)
def second_load():
second_started.set()
return adapter._get_moss_realtime_components()
async def run():
first = asyncio.create_task(asyncio.to_thread(adapter._get_moss_realtime_components))
try:
assert await asyncio.to_thread(started.wait, 5)
first.cancel()
with pytest.raises(asyncio.CancelledError):
await first
second = asyncio.create_task(asyncio.to_thread(second_load))
try:
assert await asyncio.to_thread(second_started.wait, 5)
finally:
release.set()
components = await asyncio.wait_for(second, 5)
assert components[2] is codec
assert adapter._get_moss_realtime_components() is components
assert len(loads) == 1
finally:
release.set()
asyncio.run(run())
@pytest.mark.parametrize("voice_name", [None, "speaker"])
def test_realtime_reference_cache_reuses_codes(voice_name):
cache = SpeakerEmbeddingCache()
resolve = AsyncMock(return_value=([0.0], 24000, "audio-key"))
request_thread = threading.get_ident()
encode_calls = []
expected = torch.arange(64).reshape(4, 16)
def batch_encode(wavs, *, num_quantizers):
assert threading.get_ident() != request_thread
assert num_quantizers == 16
encode_calls.append(wavs)
return SimpleNamespace(audio_codes=expected.T.unsqueeze(1), audio_codes_lengths=torch.tensor([4]))
encoder = reference_encoder.build_reference_encoder(
SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache
)
async def run():
kwargs = dict(
resolve_ref_audio=resolve,
get_artifact_key=lambda key: "content-key",
voice_name=voice_name,
voice_created_at=1,
)
try:
first, first_key = await encoder.encode("reference", **kwargs)
second, second_key = await encoder.encode("reference", **kwargs)
assert torch.equal(first, expected)
assert torch.equal(second, expected)
assert first.dtype == torch.int64
assert first_key == "audio-key"
assert second_key == (None if voice_name else "audio-key")
second.zero_()
third, _ = await encoder.encode("another-reference", **kwargs)
assert torch.equal(third, expected)
assert len(encode_calls) == 1
assert resolve.await_count == (1 if voice_name else 3)
key = cache.make_cache_key(
voice_name or "ref:content-key", "moss_tts_realtime_nq16", 1 if voice_name else 0
)
assert cache.get(key)["codes"].dtype == torch.int32
if voice_name:
kwargs["voice_created_at"] = 2
await encoder.encode("reference", **kwargs)
assert len(encode_calls) == 2
assert resolve.await_count == 2
finally:
await encoder.aclose()
asyncio.run(run())
def test_realtime_reference_batch_trims_each_item():
cache = SpeakerEmbeddingCache()
codes = torch.arange(16 * 2 * 5).reshape(16, 2, 5)
calls = []
def batch_encode(wavs, *, num_quantizers):
calls.append(wavs)
assert num_quantizers == 16
assert [wav.shape for wav in wavs] == [torch.Size([8]), torch.Size([12])]
return SimpleNamespace(audio_codes=codes, audio_codes_lengths=torch.tensor([3, 5]))
encoder = reference_encoder.build_reference_encoder(
SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache
)
async def resolve(ref):
return [0.0] * (8 if ref == "short" else 12), 24000, ref
async def run():
try:
short, long = await asyncio.gather(
encoder.encode("short", resolve_ref_audio=resolve, get_artifact_key=lambda key: key),
encoder.encode("long", resolve_ref_audio=resolve, get_artifact_key=lambda key: key),
)
torch.testing.assert_close(short[0], codes[:, 0, :3].T)
torch.testing.assert_close(long[0], codes[:, 1, :5].T)
assert len(calls) == 1
finally:
await encoder.aclose()
asyncio.run(run())
def test_failed_realtime_reference_encoding_does_not_populate_cache():
cache = SpeakerEmbeddingCache()
resolve = AsyncMock(return_value=([0.0], 24000, "audio-key"))
def fail(*args, **kwargs):
raise ValueError("invalid reference")
encoder = reference_encoder.build_reference_encoder(
SimpleNamespace(batch_encode=fail), variant="realtime", speaker_cache=cache
)
async def run():
try:
with pytest.raises(ValueError, match="invalid reference"):
await encoder.encode("reference", resolve_ref_audio=resolve, get_artifact_key=lambda key: None)
assert cache.memory_bytes() == 0
finally:
await encoder.aclose()
asyncio.run(run())
tests/entrypoints/openai_api/test_serving_speech.py
+46
View File
@@ -5213,6 +5213,52 @@ class TestTTSAsyncOffloading:
for call in log_info.call_args_list
)
@pytest.mark.parametrize("word_timestamps", [False, True])
def test_streaming_timestamps_retain_complete_audio(self, qwen3_tts_server, mocker, word_timestamps):
from vllm import SamplingParams
from vllm.sampling_params import RequestOutputKind
defaults = [SamplingParams(), SamplingParams()]
qwen3_tts_server.engine_client.default_sampling_params_list = defaults
qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None)
qwen3_tts_server._adapter._build_tts_params = mocker.MagicMock(
return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
)
qwen3_tts_server._adapter._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
request = OpenAICreateSpeechRequest(
input="hello", stream=True, response_format="pcm", word_timestamps=word_timestamps
)
asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
params = qwen3_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"]
expected = RequestOutputKind.CUMULATIVE if word_timestamps else RequestOutputKind.DELTA
assert all(sp.output_kind == expected for sp in params)
assert all(sp.output_kind == RequestOutputKind.CUMULATIVE for sp in defaults)
assert all(sp is not original for sp, original in zip(params, defaults))
@pytest.mark.asyncio
async def test_cumulative_timestamp_audio_emits_each_sample_once(self, qwen3_tts_server):
audio = torch.linspace(-0.5, 0.5, 2400)
async def cumulative_generator():
for end in (800, 1600, 2400, 2400):
yield OmniRequestOutput(
request_id="req-cumulative",
final_output_type="audio",
_multimodal_output={"audio": audio[:end], "sr": 24000},
)
chunks = [
chunk
async for chunk in qwen3_tts_server._generate_pcm_chunks(
cumulative_generator(), "req-cumulative", cumulative_audio=True
)
]
assert len(chunks) == 3
assert len(b"".join(chunks)) == 2400 * 2
decoded = np.frombuffer(b"".join(chunks), dtype=np.int16)
assert np.all(np.diff(decoded.astype(np.int32)) >= 0)
def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture):
"""stream_format=sse should request delta-style multimodal outputs."""
qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None)
tests/entrypoints/openai_api/test_serving_speech_stream.py
+4
View File
@@ -56,6 +56,7 @@ def _build_test_app(
include_sample_rate=False,
tts_params=None,
collect=None,
cumulative_audio=False,
):
for chunk in (b"\x01\x02", b"\x03\x04\x05"):
yield (chunk, 24000) if include_sample_rate else chunk
@@ -386,9 +387,11 @@ class TestStreamingSpeechWebSocket:
include_sample_rate=False,
tts_params=None,
collect=None,
cumulative_audio=False,
):
assert request_start_s is not None
assert request_arrival_ts is not None
assert cumulative_audio is True
for chunk in (first_chunk, second_chunk):
yield (chunk, 1000) if include_sample_rate else chunk
if collect is not None:
@@ -480,6 +483,7 @@ class TestStreamingSpeechWebSocket:
include_sample_rate=False,
tts_params=None,
collect=None,
cumulative_audio=False,
):
chunk = b"\x01" * 1000
yield (chunk, 1000) if include_sample_rate else chunk
tests/entrypoints/test_omni_entrypoints.py
+38
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import gc
import queue
from collections.abc import Callable
@@ -634,6 +635,43 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc
assert "req-1" not in app.request_states
@pytest.mark.asyncio
async def test_async_omni_abort_yields_terminal_from_requested_final_stage(monkeypatch: pytest.MonkeyPatch):
stage_metadata = [
_stage_meta(stage_type="llm", final_output=False, final_output_type=None),
_stage_meta(stage_type="llm", final_output=True, final_output_type="audio"),
]
engine = FakeAsyncOmniEngine(stage_metadata=stage_metadata)
_patch_engine(monkeypatch, engine)
app = AsyncOmni("dummy-model")
async def collect_outputs() -> list[OmniRequestOutput]:
return [
output
async for output in app.generate(
prompt="hello",
request_id="req-1",
output_modalities=["audio"],
)
]
try:
generate_task = asyncio.create_task(collect_outputs())
while not engine.submitted:
await asyncio.sleep(0)
await app.abort("req-1")
outputs = await asyncio.wait_for(generate_task, timeout=1)
finally:
app.shutdown()
assert len(outputs) == 1
assert outputs[0].stage_id == 1
assert outputs[0].final_output_type == "audio"
assert outputs[0].finished is True
assert outputs[0].outputs[0].finish_reason == "abort"
assert "req-1" not in app.request_states
@pytest.mark.asyncio
async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
tests/examples/offline_inference/test_text_to_video_prompts.py
+48 -1
View File
@@ -1,10 +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
from __future__ import annotations
import functools
import importlib.util
import sys
from types import SimpleNamespace
import pytest
import torch
@@ -24,6 +26,51 @@ def _load_text_to_video():
return module
@pytest.mark.parametrize(
("arguments", "expected"),
[
(["--quantization", "mxfp4"], {"quantization": "mxfp4"}),
(
["--quantization-config", '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}'],
{"quantization_config": {"method": "mxfp4", "w4a8_fallback_steps": [0, 2]}},
),
],
)
def test_quantization_cli_reaches_omni(arguments, expected, monkeypatch):
mod = _load_text_to_video()
captured = {}
class StopBeforeModelLoadError(Exception):
pass
def capture_omni(**kwargs):
captured.update(kwargs)
raise StopBeforeModelLoadError
monkeypatch.setattr(sys, "argv", ["text_to_video.py", "--model", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", *arguments])
monkeypatch.setattr(mod, "current_omni_platform", SimpleNamespace(device_type="cpu"))
monkeypatch.setattr(mod, "Omni", capture_omni)
with pytest.raises(StopBeforeModelLoadError):
mod.main()
assert {key: captured[key] for key in ("quantization", "quantization_config") if key in captured} == expected
def test_quantization_cli_options_are_mutually_exclusive(monkeypatch):
mod = _load_text_to_video()
monkeypatch.setattr(
sys,
"argv",
["text_to_video.py", "--quantization", "mxfp4", "--quantization-config", '{"method":"mxfp4"}'],
)
with pytest.raises(SystemExit) as exc:
mod.parse_args()
assert exc.value.code == 2
@pytest.mark.parametrize(
("negative_prompt", "expected_negative_prompt"),
[(None, None), ("", ""), ("blurry", "blurry")],
tests/helpers/assertions.py
+11
View File
@@ -955,6 +955,17 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
"""
assert response.success, "The request failed."
if request_config.get("word_timestamps"):
timestamps = response.word_timestamps
assert isinstance(timestamps, list) and timestamps, "Expected nonempty X-Word-Timestamps"
previous_start = 0
for timestamp in timestamps:
assert isinstance(timestamp["word"], str) and timestamp["word"]
start, end = timestamp["start_ms"], timestamp["end_ms"]
assert isinstance(start, int) and isinstance(end, int)
assert previous_start <= start <= end
previous_start = start
# Optional floor on decoded audio size (models with very short clips may use a lower value).
min_audio = request_config.get("min_audio_bytes")
if min_audio is not None:
tests/helpers/client.py
+5
View File
@@ -102,6 +102,7 @@ class OmniResponse:
audio_content: str | None = None
audio_format: str | None = None
audio_bytes: bytes | None = None
word_timestamps: list[dict[str, str | int]] | None = None
#: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the
#: OpenAI client call through response parsing and local post-process (e.g. audio decode).
e2e_latency: float | None = None
@@ -1419,6 +1420,9 @@ class OnlineOmniClient:
result.success = True
result.audio_format = getattr(response, "response", None)
if result.audio_format is not None:
timestamps = result.audio_format.headers.get("X-Word-Timestamps")
if timestamps is not None:
result.word_timestamps = json.loads(timestamps)
result.audio_format = result.audio_format.headers.get("content-type", "")
except Exception as e:
@@ -1472,6 +1476,7 @@ class OnlineOmniClient:
"extra_params",
"stream_format",
"x_vector_only_mode",
"word_timestamps",
):
if key in request_config:
extra_body[key] = request_config[key]
tests/helpers/tests/test_assertions.py
+23
View File
@@ -173,3 +173,26 @@ def test_bounded_tail_allows_ratio_bound_for_long_answers():
# Twenty expected words allow up to max(2, ceil(0.2 * 20)) = 4 tail words.
assert assertions._transcript_has_bounded_tail(four_word_tail, expected)
assert not assertions._transcript_has_bounded_tail(five_word_tail, expected)
def test_speech_timestamp_header_is_required_and_decoded():
import time
import httpx
from openai._legacy_response import HttpxBinaryResponseContent
from tests.helpers.client import OnlineOmniClient
client = object.__new__(OnlineOmniClient)
request = {"word_timestamps": True}
for headers in ({}, {"X-Word-Timestamps": '[{"word":"Hello","start_ms":0,"end_ms":300}]'}):
response = client._process_non_stream_audio_speech_response(
HttpxBinaryResponseContent(httpx.Response(200, content=b"audio", headers=headers)),
wall_start=time.perf_counter(),
)
if headers:
assert_audio_speech_response(response, request)
assert response.word_timestamps == [{"word": "Hello", "start_ms": 0, "end_ms": 300}]
else:
with pytest.raises(AssertionError, match="X-Word-Timestamps"):
assert_audio_speech_response(response, request)
tests/model_executor/models/breeze_tts_2/test_instruction.py
+22
View File
@@ -4,6 +4,7 @@
from types import SimpleNamespace
import pytest
from vllm import SamplingParams
from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter
from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import (
@@ -133,3 +134,24 @@ def test_non_greedy_sampling_overrides_are_rejected():
assert "greedy" in adapter.validate(_request({"top_p": 0.9}))
assert "greedy" in adapter.validate(_request({"top_k": 50}))
assert "temperature must be a number" in adapter.validate(_request({"temperature": "warm"}))
@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf"), "invalid"])
def test_invalid_repetition_penalties_are_rejected(value):
assert "repetition_penalty" in _adapter().validate(_request({"repetition_penalty": value}))
@pytest.mark.parametrize("value", [0.9, 1.0, 1.1, "1.2"])
@pytest.mark.parametrize("max_new_tokens", [None, 32])
def test_repetition_penalty_override_reaches_standard_sampling_metadata(value, max_new_tokens):
adapter = _adapter()
request = _request({"repetition_penalty": value}, max_new_tokens=max_new_tokens)
original = [SamplingParams(max_tokens=128, repetition_penalty=1.1)]
assert adapter.validate(request) is None
updated = adapter.apply_sampling_overrides(original, request)
assert updated[0].repetition_penalty == float(value)
assert updated[0].max_tokens == (128 if max_new_tokens is None else max_new_tokens)
assert original[0].repetition_penalty == 1.1
assert original[0].max_tokens == 128
tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py
+122
View File
@@ -0,0 +1,122 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from typing import Any
import pytest
import torch
from vllm.v1.sample.logits_processor import LogitsProcessors
from vllm.v1.sample.metadata import SamplingMetadata
from tests.model_executor.models.breeze_tts_2.test_talker_state_machine import _depth_codes, _talker
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _model(first_score=11.0, second_score=10.5):
model = _talker()
model._async_chunk = True
logits = torch.full((1, 9), -100.0)
logits[0, 2:4] = torch.tensor([first_score, second_score])
model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone()
model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0)
return model
def _step(model, infos, penalties, *, no_penalties=False):
return model.make_omni_output(
torch.zeros(len(infos), 16),
model_intermediate_buffer=infos,
request_token_spans=[(i, i + 1) for i in range(len(infos))],
sampling_metadata=SamplingMetadata(
temperature=None,
all_greedy=True,
all_random=False,
top_p=None,
top_k=None,
generators={},
max_num_logprobs=None,
no_penalties=no_penalties,
frequency_penalties=torch.zeros(len(infos)),
presence_penalties=torch.zeros(len(infos)),
repetition_penalties=torch.tensor(penalties),
# Text/reference conditioning must never be counted as generated
# codec history, even if the scheduler ids overlap the codec ids.
prompt_token_ids=torch.tensor([[2, 2]] * len(infos)),
output_token_ids=[list(info.get("breeze_generated_code0_ids", [])) for info in infos],
allowed_token_ids_mask=None,
bad_words_token_ids={},
logitsprocs=LogitsProcessors(),
),
)
@pytest.mark.parametrize("scores", [(11.0, 10.5), (-1.0, -1.05)])
def test_penalty_applies_once_per_generated_token_and_respects_logit_sign(scores):
model = _model(*scores)
info = {"prompt_ids": torch.tensor([2]), "input_values": torch.tensor([[2, 3, 4, 5]])}
chosen = []
for _ in range(4):
output = _step(model, [info], [1.1])
chosen.append(output.multimodal_outputs["codes"]["audio"][0][0, 0].item())
assert chosen == [2, 3, 2, 2]
def test_request_history_and_penalty_are_independent():
model = _model()
first: dict[str, Any] = {}
second: dict[str, Any] = {}
_step(model, [first], [1.1])
output = _step(model, [first, second], [1.1, 1.1])
assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2]
# Reverse the batch order; a request's history follows its state, while
# each penalty follows the current sampling-metadata row.
output = _step(model, [second, first], [1.1, 1.0])
assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2]
def test_disabled_penalties_do_not_read_uninitialized_sampling_buffer():
model = _model()
info: dict[str, Any] = {}
for _ in range(2):
output = _step(model, [info], [float("nan")], no_penalties=True)
assert output.multimodal_outputs["codes"]["audio"][0][0, 0].item() == 2
def test_scheduler_logits_preserve_selected_frame_under_further_sampling_transforms():
model = _model()
info: dict[str, Any] = {}
_step(model, [info], [1.1])
output = _step(model, [info], [1.1])
selected = output.multimodal_outputs["codes"]["audio"][0][0, 0].item()
logits = model.compute_logits(torch.zeros(1, 16))
assert selected == 3
assert logits.argmax(dim=-1).item() == selected
assert torch.isfinite(logits).sum().item() == 1
assert logits[0, selected].item() == 0.0
# The scheduler may apply repetition again and rescale temperature. A
# finite zero for the selected token preserves a valid point mass.
for temperature in (0.1, 0.9, 2.0):
probabilities = torch.softmax(logits / 1.1 / temperature, dim=-1)
assert torch.isfinite(probabilities).all()
assert probabilities[0, selected].item() == 1.0
def test_repetition_penalty_can_select_eos_without_emitting_a_frame():
model = _model()
info: dict[str, Any] = {}
_step(model, [info], [1.1])
logits = torch.full((1, 9), -100.0)
logits[0, 2] = 11.0
logits[0, 8] = 10.5
model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone()
output = _step(model, [info], [1.1])
assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0
assert info["breeze_generated_frames"] == 1
assert model.compute_logits(torch.zeros(1, 16)).argmax(dim=-1).item() == 8
tests/model_executor/models/breeze_tts_2/test_sampling_eos.py
+77
View File
@@ -0,0 +1,77 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from pathlib import Path
import pytest
import torch
import yaml
from vllm import SamplingParams
from vllm.v1.core.sched.utils import check_stop
from vllm.v1.request import Request, RequestStatus
from vllm.v1.sample.logits_processor import BatchUpdate, MinTokensLogitsProcessor
from vllm_omni.config.stage_config import merge_sampling_constraints
from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture
def talker_sampling_params():
deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml"
deploy = yaml.safe_load(deploy_path.read_text(encoding="utf-8"))
params = SamplingParams(
**merge_sampling_constraints(
deploy["stages"][0]["default_sampling_params"],
BREEZE_TTS_2_PIPELINE.stages[0].sampling_constraints,
)
)
# Mirror the input processor: Breeze's text tokenizer uses EOS 1,
# which is also an ordinary token in the generated codec vocabulary.
params.update_from_generation_config({"eos_token_id": [1]}, eos_token_id=1)
return params
@pytest.mark.parametrize("codec_token, should_stop", [(0, False), (1, False), (2047, False), (2051, True)])
def test_scheduler_stops_only_on_codec_eos(talker_sampling_params, codec_token, should_stop):
request = Request(
request_id="breeze-eos-test",
prompt_token_ids=[7, 8, 9],
sampling_params=talker_sampling_params,
pooling_params=None,
)
request.append_output_token_ids([42, codec_token])
request.status = RequestStatus.RUNNING
assert check_stop(request, max_model_len=4096) is should_stop
if should_stop:
assert request.status == RequestStatus.FINISHED_STOPPED
assert request.stop_reason == 2051
else:
assert request.status == RequestStatus.RUNNING
@pytest.mark.parametrize("codec_token", [1, 2051])
def test_first_step_keeps_codec_token_and_eos_logits(talker_sampling_params, codec_token, monkeypatch):
# vLLM's tensor helper uses a global flag rather than is_pin_memory,
# and pinned allocations require a CUDA device even for this CPU test.
monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False)
# ignore_eos still leaves text EOS in all_stop_token_ids. Exercise the
# actual processor so the first codec frame and immediate EOS stay valid.
processor = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
processor.update_state(
BatchUpdate(batch_size=1, removed=[], added=[(0, talker_sampling_params, None, [])], moved=[])
)
logits = torch.full((1, 2052), -float("inf"))
logits[0, codec_token] = 0.0
processed = processor.apply(logits.clone())
assert torch.equal(processed, logits)
assert processed.argmax(dim=-1).item() == codec_token
def test_pipeline_disallows_minimum_tokens_that_mask_selected_codec_ids():
params = merge_sampling_constraints({"min_tokens": 10}, BREEZE_TTS_2_PIPELINE.stages[0].sampling_constraints)
assert params["min_tokens"] == 0
tests/model_executor/models/breeze_tts_2/test_streaming.py
+106 -14
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import copy
from collections import defaultdict
from types import SimpleNamespace
@@ -20,7 +21,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _TransferManager:
def __init__(self, chunk_frames: int = 4):
self.connector = SimpleNamespace(config={"breeze_codec_chunk_frames": chunk_frames})
self.code_prompt_token_ids = defaultdict(list)
self.code_prompt_token_ids: defaultdict[str, list[torch.Tensor]] = defaultdict(list)
class _Request:
@@ -120,7 +121,8 @@ def test_stateful_codec_rejects_partial_frame_chunk():
assert decoder.calls == []
def test_stateful_codec_terminal_marker_pops_state_without_decoding():
@pytest.mark.parametrize("token_counts", [(4, 0), (0, 4), (0, 4, 0, 8, 0)])
def test_stateful_codec_terminal_marker_pops_state_without_decoding(token_counts):
decoder = _Decoder()
codec = object.__new__(BreezeTTS2MimiCodec)
codec._async_chunk = True
@@ -128,24 +130,114 @@ def test_stateful_codec_terminal_marker_pops_state_without_decoding():
codec._codebook_size = 8
codec._sample_rate = 24_000
codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder))
codec._decoder_state_cache = {"scheduler-live": {}, "scheduler-done": {}}
request_ids = [f"scheduler-{index}" for index in range(len(token_counts))]
codec._decoder_state_cache = {request_id: {} for request_id in request_ids}
output = codec.forward(
torch.tensor([0, 1, 2, 3], dtype=torch.long),
torch.arange(sum(token_counts), dtype=torch.long) % codec._codebook_size,
runtime_additional_information=[
# Plain ``finished`` covers the fallback for direct in-process
# callers that do not go through the connector receiver.
{"meta": {"request_id": "live", "finished": False}},
{"meta": {"request_id": "done", "finished": True}},
{"meta": {"finished": count == 0}}
for count in token_counts
],
seq_token_counts=[4, 0],
request_ids=["scheduler-live", "scheduler-done"],
seq_token_counts=token_counts,
request_ids=request_ids,
)
# The empty terminal marker decodes nothing and releases its state; the
# live request still decodes its own chunk.
# Empty terminal markers release their state without consuming an output
# slot belonging to another request, regardless of their batch position.
assert len(decoder.calls) == 1
assert decoder.calls[0][0][0] == 1
assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [1920, 0]
assert list(codec._decoder_state_cache) == ["scheduler-live"]
assert codec._decoder_state_cache["scheduler-live"]["prefix_frames"] == 0
assert decoder.calls[0][0][0] == sum(count > 0 for count in token_counts)
assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [
count // codec._num_codebooks * 1920 for count in token_counts
]
live_request_ids = [request_id for request_id, count in zip(request_ids, token_counts) if count]
assert list(codec._decoder_state_cache) == live_request_ids
assert all(codec._decoder_state_cache[request_id]["prefix_frames"] == 0 for request_id in live_request_ids)
@pytest.mark.parametrize("configured_chunk_frames", [None, 4])
def test_stateful_codec_advances_decoder_cache_after_sliding_window(monkeypatch, tmp_path, configured_chunk_frames):
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
Qwen3TTSTokenizerV2DecoderConfig,
)
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
Qwen3TTSTokenizerV2Decoder,
)
torch.manual_seed(0)
config = Qwen3TTSTokenizerV2DecoderConfig(
codebook_size=32,
hidden_size=16,
latent_dim=16,
codebook_dim=16,
num_attention_heads=2,
num_key_value_heads=2,
intermediate_size=32,
num_hidden_layers=1,
num_quantizers=2,
decoder_dim=32,
upsample_rates=(2,),
upsampling_ratios=(2,),
sliding_window=72,
)
decoder = Qwen3TTSTokenizerV2Decoder(config).eval()
# Codebooks initialize to zeros before checkpoint loading; populate them
# so different codec IDs produce distinguishable cached frame histories.
with torch.no_grad():
for name, parameter in decoder.quantizer.named_parameters():
if name.endswith("embedding_sum"):
parameter.normal_()
tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder, get_output_sample_rate=lambda: 24_000))
monkeypatch.setattr(Qwen3TTSTokenizer, "from_pretrained", lambda *args, **kwargs: tokenizer)
connector_extra = {} if configured_chunk_frames is None else {"breeze_codec_chunk_frames": configured_chunk_frames}
chunk_frames = configured_chunk_frames or 8
codec = object.__new__(BreezeTTS2MimiCodec)
codec.vllm_config = SimpleNamespace(
model_config=SimpleNamespace(stage_connector_config={"extra": connector_extra}),
device_config=SimpleNamespace(device="cpu"),
)
codec._tokenizer_path = tmp_path / "audio_tokenizer"
codec._audio_tokenizer = None
codec._async_chunk = True
codec._num_codebooks = config.num_quantizers
codec._codebook_size = config.codebook_size
codec._decoder_state_cache = {}
codec.load_weights(iter(()))
# At least two chunks after the decoder's rolling window expose a stale
# cache even when the first padded chunk still produces plausible audio.
total_frames = config.sliding_window + 2 * chunk_frames
tail_frames = 3
codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, total_frames + tail_frames))
for end in range(chunk_frames, total_frames + 1, chunk_frames):
flat = codes[0, :, end - chunk_frames : end].reshape(-1)
codec.forward(
flat,
runtime_additional_information=[{"meta": {"stream_finished": False}}],
seq_token_counts=[flat.numel()],
request_ids=["scheduler-live"],
)
state = codec._decoder_state_cache["scheduler-live"]
assert state["suffix_frames"] == min(end - chunk_frames, config.sliding_window) + chunk_frames
with torch.inference_mode():
expected_quantized = decoder.quantizer.decode(codes[..., max(0, end - config.sliding_window) : end])
torch.testing.assert_close(state["suffix_quantized"], expected_quantized)
# A short terminal chunk must emit only its real frames, then release
# state. Its waveform should agree with the unpadded single-request path.
with torch.inference_mode():
expected_tail = decoder.chunked_decode(codes[..., -tail_frames:], caches=copy.deepcopy(state))
flat_tail = codes[0, :, -tail_frames:].reshape(-1)
output = codec.forward(
flat_tail,
runtime_additional_information=[{"meta": {"stream_finished": True}}],
seq_token_counts=[flat_tail.numel()],
request_ids=["scheduler-live"],
)
waveform = output.multimodal_outputs["model_outputs"][0]
assert waveform.numel() == tail_frames * decoder.total_upsample
torch.testing.assert_close(waveform, expected_tail.reshape(-1), atol=1e-5, rtol=1e-4)
assert codec._decoder_state_cache == {}
tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
+50 -11
View File
@@ -15,6 +15,7 @@ from vllm_omni.diffusion.config import set_current_diffusion_config
from vllm_omni.diffusion.data import AttentionConfig
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
CausalHiFTGenerator,
HiFTGenerator,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -76,6 +77,39 @@ def test_causal_hift_stft_moves_window_to_input_device(causal_hift):
assert causal_hift.stft_window.device == waveform.device
class _MinimalF0Predictor(nn.Module):
def __init__(self) -> None:
super().__init__()
self.proj = nn.Conv1d(80, 1, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.proj(x).squeeze(1).abs()
@pytest.mark.parametrize("sampling_rate", [22050, 24000])
def test_noncausal_hift_inference_runs_with_real_sinegen(sampling_rate):
"""Base (non-causal) HiFTGenerator with a real SineGen/SineGen2 (not a
test double) must still run: SourceModuleHnNSF/SineGen/SineGen2/_f02sine
all branch on `self.causal` to return the pre-streaming (shorter) tuple
contract for non-causal callers, and any one of them regressing back to
always returning the streaming-length tuple breaks this immediately.
"""
hift = HiFTGenerator(
base_channels=32,
sampling_rate=sampling_rate,
upsample_rates=[8, 5, 3],
upsample_kernel_sizes=[16, 11, 7],
source_resblock_kernel_sizes=[7, 7, 11],
source_resblock_dilation_sizes=[[1, 3, 5]] * 3,
f0_predictor=_MinimalF0Predictor(),
).eval()
speech, _ = hift.inference(torch.randn(1, 80, 20))
assert speech.shape[0] == 1
assert torch.isfinite(speech).all()
def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch):
class NpuTensor:
device = SimpleNamespace(type="npu")
@@ -83,16 +117,17 @@ def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypa
tensor = NpuTensor()
stft_calls = []
istft_calls = []
monkeypatch.setattr(
causal_hift,
"_stft_on_cpu",
lambda value: stft_calls.append(value) or ("real", "imag"),
)
monkeypatch.setattr(
causal_hift,
"_istft_on_cpu",
lambda magnitude, phase: istft_calls.append((magnitude, phase)) or "waveform",
)
def fake_stft(value):
stft_calls.append(value)
return ("real", "imag")
def fake_istft(magnitude, phase):
istft_calls.append((magnitude, phase))
return "waveform"
monkeypatch.setattr(causal_hift, "_stft_on_cpu", fake_stft)
monkeypatch.setattr(causal_hift, "_istft_on_cpu", fake_istft)
assert causal_hift._stft(tensor) == ("real", "imag")
assert causal_hift._istft(tensor, tensor) == "waveform"
@@ -627,7 +662,11 @@ def test_code2wav_forward_finalizes_hift_tail():
def inference(self, speech_feat, finalize=True):
self.finalize_calls.append(bool(finalize))
return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None
return (
torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype),
None,
None,
)
model = object.__new__(CosyVoice3Code2Wav)
nn.Module.__init__(model)
tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py
+415
View File
@@ -0,0 +1,415 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Correctness + boundedness tests for the incremental (windowed) HiFT path.
The streaming vocoder runs HiFT over a bounded mel window instead of the full
cumulative history, carrying harmonic phase and noise-buffer offset across
chunks.
* **Correctness** — windowed output must match the full cumulative re-run.
With full history the match is bit-exact (proves the phase/noise carry is
right). With a bounded window, float64 still agrees to ~1e-16 (the window
is algebraically exact); float32 differs by ~2e-7 absolute (reassociation
rounding, not truncation), so tests assert a tight atol/rtol rather than
PCM16 byte-identity, since a sub-LSB perturbation can flip a boundary sample.
* **Boundedness** — the per-chunk window must not grow with mel history.
Tests are parametrized across REAL_CONFIG (24000 Hz -> SineGen2, what
CosyVoice3 ships) and LEGACY_CONFIG (22050 Hz -> SineGen) so a class-selection
regression in either path is caught.
"""
import pytest
import torch
import torch.nn as nn
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
CausalConvRNNF0Predictor,
CausalHiFTGenerator,
SineGen,
SineGen2,
_wrapped_slice,
)
from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
CHUNK_LEN = 24
TOTAL_MEL = 96
SPM = 480 # samples per mel frame for the small test HiFT (8*5*3*4), matching the real model
LONG_TOTAL_MEL = 400 # long enough that the 64-frame window actually truncates
PCM16_LSB = 1.0 / 32767.0 # smallest representable PCM16 difference
# Worst observed deviation is 1.97e-7 (~2x float32 eps); atol=1e-6 keeps 5x headroom.
ATOL = 1e-6
RTOL = 1e-5
# What CosyVoice3 ships (-> SineGen2). nb_harmonics=8 is required, not a free
# choice: SineGen2's causal buffers are hardcoded to dim 9 (= harmonic_num+1).
REAL_CONFIG = {"sampling_rate": 24000, "nb_harmonics": 8}
LEGACY_CONFIG = {"sampling_rate": 22050, "nb_harmonics": 4} # -> SineGen, still shipped code
CONFIGS = [
pytest.param(REAL_CONFIG, id="sinegen2_real_config"),
pytest.param(LEGACY_CONFIG, id="sinegen1_legacy_config"),
]
def _make_hift(config: dict = REAL_CONFIG) -> CausalHiFTGenerator:
"""Small CausalHiFTGenerator; defaults to the config CosyVoice3 ships."""
torch.manual_seed(0)
f0_predictor = CausalConvRNNF0Predictor(num_class=1, in_channels=80, cond_channels=16)
return CausalHiFTGenerator(
in_channels=80,
base_channels=32,
nb_harmonics=config["nb_harmonics"],
sampling_rate=config["sampling_rate"],
upsample_rates=[8, 5, 3],
upsample_kernel_sizes=[16, 11, 7],
source_resblock_kernel_sizes=[7, 7, 11],
source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
resblock_kernel_sizes=[3, 7, 11],
resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
f0_predictor=f0_predictor,
).eval()
def test_real_config_resolves_to_sinegen2():
"""Guard against silently testing the wrong SineGen class again."""
hift = _make_hift(REAL_CONFIG)
assert isinstance(hift.m_source.l_sin_gen, SineGen2)
hift_legacy = _make_hift(LEGACY_CONFIG)
assert isinstance(hift_legacy.m_source.l_sin_gen, SineGen)
def _make_model(hift: CausalHiFTGenerator, window_len: int) -> CosyVoice3Code2Wav:
"""Build a CosyVoice3Code2Wav shell wired to the given HiFT."""
model = object.__new__(CosyVoice3Code2Wav)
nn.Module.__init__(model)
model.hift = hift
model._hift_window_len = window_len
return model
def _chunks(total_mel: int = TOTAL_MEL, chunk_len: int = CHUNK_LEN) -> list[torch.Tensor]:
torch.manual_seed(1)
mel = torch.randn(1, 80, total_mel)
return [mel[:, :, i : i + chunk_len] for i in range(0, total_mel, chunk_len)]
def _variable_chunks(total_mel: int, first_len: int, step_len: int) -> list[torch.Tensor]:
"""First chunk sized like a real prompt prefill, then small fixed-size steps."""
torch.manual_seed(1)
mel = torch.randn(1, 80, total_mel)
chunks = [mel[:, :, :first_len]]
pos = first_len
while pos < total_mel:
chunks.append(mel[:, :, pos : pos + step_len])
pos += step_len
return chunks
def _full_reference(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor:
"""Re-run HiFT over the full cumulative mel each chunk (pre-windowing behavior).
Uses non-finalize inference per chunk, matching the old streaming vocoder; a
single finalize pass would emit the look-right tail streaming holds back.
"""
torch.manual_seed(0)
emitted = []
cache: dict[str, torch.Tensor] | None = None
for chunk in chunks:
cached = None if cache is None else cache.get("mel")
if cached is not None and cached.numel() > 0:
tts_mel = torch.cat([cached, chunk], dim=-1)
else:
tts_mel = chunk
speech, _, _ = model.hift.inference(speech_feat=tts_mel, finalize=False)
speech = speech.reshape(speech.shape[0], -1)
offset = 0 if cache is None else cache.get("speech_offset")
emitted.append(speech[:, offset:])
cache = {"mel": tts_mel.detach().cpu(), "speech_offset": int(speech.shape[-1])}
return torch.cat(emitted, dim=-1)
def _pcm16(x: torch.Tensor) -> torch.Tensor:
"""Quantize float audio to PCM16, the format the server streams to clients."""
return (x.clamp(-1, 1) * 32767).round().to(torch.int16)
def _incremental(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor:
"""Run the new windowed path via _stream_hift_from_feat."""
torch.manual_seed(0)
emitted = []
cache: dict[str, torch.Tensor] | None = None
for chunk in chunks:
speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False)
emitted.append(speech.reshape(speech.shape[0], -1))
return torch.cat(emitted, dim=-1)
def _rel_divergence(a: torch.Tensor, b: torch.Tensor) -> float:
return ((a - b).abs().mean() / a.abs().mean().clamp_min(1e-6)).item()
@pytest.mark.parametrize("config", CONFIGS)
def test_incremental_hift_matches_full_reference_exactly_with_full_history(config):
"""With a window covering the full history, the windowed path must match the
full cumulative re-run exactly. This proves the phase/noise carry is right."""
hift = _make_hift(config)
model = _make_model(hift, window_len=TOTAL_MEL) # window >= history
chunks = _chunks()
full = _full_reference(model, chunks)
incr = _incremental(model, chunks)
assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}"
assert _rel_divergence(full, incr) < 1e-6
@pytest.mark.parametrize("config", CONFIGS)
def test_incremental_hift_bounded_window_is_close(config):
"""With a bounded window, the windowed path stays close to the full re-run;
only the truncated deep history diverges, and it is small."""
hift = _make_hift(config)
model = _make_model(hift, window_len=48)
chunks = _chunks()
full = _full_reference(model, chunks)
incr = _incremental(model, chunks)
assert full.shape == incr.shape
torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL)
@pytest.mark.parametrize("config", CONFIGS)
def test_incremental_hift_matches_streaming_reference_within_tolerance(config):
"""Windowed output matches the full-cumulative reference within float32 rounding.
The load-bearing correctness test: 400 mel frames far exceeds the 64-frame
window, so truncation genuinely exercises the phase/noise carry. Float64
agrees to ~1e-16, confirming the residual here is rounding, not truncation.
"""
hift = _make_hift(config)
model = _make_model(hift, window_len=64) # the real _hift_window_len
chunks = _chunks(total_mel=LONG_TOTAL_MEL)
trim = int(hift.f0_predictor.condnet[0].causal_padding)
assert LONG_TOTAL_MEL > 64 + trim + CHUNK_LEN # window must actually truncate
full = _full_reference(model, chunks)
incr = _incremental(model, chunks)
assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}"
torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL)
max_dev = (full - incr).abs().max().item() # stays far below one PCM16 step
assert max_dev < PCM16_LSB / 10, f"max deviation {max_dev:.3e} vs PCM16 LSB {PCM16_LSB:.3e}"
def test_incremental_hift_window_is_bounded():
"""Per-chunk window must not grow with cumulative mel history."""
hift = _make_hift()
model = _make_model(hift, window_len=32)
chunks = _chunks(total_mel=200, chunk_len=24)
trim = int(hift.f0_predictor.condnet[0].causal_padding)
f0_margin = int(hift.f0_predictor.left_context_frames)
cache = None
window_sizes = []
for chunk in chunks:
_, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False)
window_sizes.append(int(cache["mel"].shape[-1]))
assert max(window_sizes) <= 32 + trim + f0_margin + 24 # bounded regardless of utterance length
steady = max(window_sizes)
assert window_sizes.count(steady) >= 2 # steady-state value repeats once history exceeds window
def test_incremental_hift_finalize_releases_tail():
"""Finalize must emit the released look-right tail without crashing."""
hift = _make_hift()
model = _make_model(hift, window_len=32)
chunks = _chunks(total_mel=72, chunk_len=24)
cache = None
emitted = []
for i, chunk in enumerate(chunks):
finalize = i == len(chunks) - 1
speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize)
emitted.append(speech.reshape(speech.shape[0], -1))
if finalize:
assert cache is None
full = torch.cat(emitted, dim=-1)
assert full.shape[-1] > 0
assert torch.isfinite(full).all()
def test_incremental_hift_emits_full_audio_length():
"""Streamed output covers the full mel history at samples-per-mel resolution.
Guards against emission arithmetic dropping an upsample stage, which
silently truncates audio (e.g. 1/3 of expected length with [8,5,3]).
"""
hift = _make_hift()
model = _make_model(hift, window_len=32)
chunks = _chunks(total_mel=72, chunk_len=24)
cache = None
emitted = []
for i, chunk in enumerate(chunks):
finalize = i == len(chunks) - 1
speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize)
emitted.append(speech.reshape(speech.shape[0], -1))
if finalize:
assert cache is None
full = torch.cat(emitted, dim=-1)
# total_mel * SPM, +slack for finalize's look-right tail, -slack for a dropped upsample stage
total_mel = sum(c.shape[-1] for c in chunks)
assert full.shape[-1] > total_mel * SPM * 0.8, f"{full.shape[-1]} vs {total_mel * SPM}"
assert full.shape[-1] < total_mel * SPM * 1.5, f"{full.shape[-1]} vs {total_mel * SPM}"
@pytest.mark.parametrize("config", CONFIGS)
def test_finalize_matches_single_pass_length_exactly(config):
"""Streamed total (chunked, last=finalize) must exactly match a single
non-streaming finalize=True call over the whole mel -- not just approximately.
The finalize release must include `trim` (the F0 predictor's own held-back
frames), or the last 3 mel frames (1,440 samples / 60ms at 24kHz) of every
stream's audio go missing.
"""
hift = _make_hift(config)
chunks = _chunks(total_mel=200, chunk_len=24)
torch.manual_seed(0)
full_mel = torch.cat(chunks, dim=-1)
ground_truth, _, _ = hift.inference(speech_feat=full_mel, finalize=True)
ground_truth = ground_truth.reshape(ground_truth.shape[0], -1)
hift2 = _make_hift(config)
model = _make_model(hift2, window_len=64)
torch.manual_seed(0)
cache = None
emitted = []
for i, chunk in enumerate(chunks):
finalize = i == len(chunks) - 1
speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize)
emitted.append(speech.reshape(speech.shape[0], -1))
streamed = torch.cat(emitted, dim=-1)
assert streamed.shape[-1] == ground_truth.shape[-1], (
f"{streamed.shape[-1]} vs {ground_truth.shape[-1]} (diff={ground_truth.shape[-1] - streamed.shape[-1]} samples)"
)
torch.testing.assert_close(streamed, ground_truth, atol=ATOL, rtol=RTOL)
@pytest.mark.parametrize("config", CONFIGS)
def test_incremental_hift_matches_reference_for_voiced_f0(config):
"""Forces F0 above nsf_voiced_threshold (untrained weights never trigger voiced,
so no other test here does) and checks windowed output against the full-
cumulative reference, which exercises condnet's left-causal convs and
SineGen/SineGen2's phase integration under genuinely voiced harmonics.
Tolerance is looser than ATOL/RTOL: summing this sustained constant-F0 phase
in per-window chunks vs. one recompute reassociates float32 adds differently;
real speech's varying/unvoiced F0 keeps this far smaller in practice.
"""
hift = _make_hift(config)
with torch.no_grad():
hift.f0_predictor.classifier.bias.fill_(150.0) # clearly voiced (threshold is 10)
chunks = _chunks(total_mel=200, chunk_len=24)
full = _full_reference(_make_model(hift, window_len=400), chunks)
incr = _incremental(_make_model(hift, window_len=64), chunks)
assert full.shape == incr.shape
torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4)
@pytest.mark.parametrize("config", CONFIGS)
@pytest.mark.parametrize("first_len,step_len", [(10, 6), (10, 3), (20, 3), (10, 1)])
def test_incremental_hift_matches_reference_for_voiced_f0_small_chunks(config, first_len, step_len):
"""Same as test_incremental_hift_matches_reference_for_voiced_f0, but with small,
sub-receptive-field chunk sizes (down to 1 mel frame) instead of the uniform
24-frame chunks used elsewhere. cache_state must carry the F0 predictor's own
left receptive field across calls regardless of how small a single chunk is,
or the next call's margin comes up short and substitutes zero-padding for real
history -- corrupting SineGen's phase integration the same way as an absent
margin entirely.
"""
hift = _make_hift(config)
with torch.no_grad():
hift.f0_predictor.classifier.bias.fill_(150.0) # clearly voiced (threshold is 10)
chunks = _variable_chunks(total_mel=200, first_len=first_len, step_len=step_len)
full = _full_reference(_make_model(hift, window_len=400), chunks)
incr = _incremental(_make_model(hift, window_len=64), chunks)
assert full.shape == incr.shape
torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4)
def test_wrapped_slice_wraps_position_deterministically():
"""SineGen/SineGen2/SourceModuleHnNSF's noise buffers are a fixed 300s at
24kHz; a stream running past that must keep repeating the realization
(position-deterministic) rather than the slice coming back short and
failing to broadcast against the rest of that call's tensors.
"""
buf = torch.arange(10).reshape(1, 10, 1)
assert _wrapped_slice(buf, 2, 3).flatten().tolist() == [2, 3, 4]
assert _wrapped_slice(buf, 8, 4).flatten().tolist() == [8, 9, 0, 1]
assert _wrapped_slice(buf, 12, 3).flatten().tolist() == [2, 3, 4] # offset past the end
assert _wrapped_slice(buf, 5, 13).flatten().tolist() == [5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7]
def test_incremental_hift_empty_middle_chunk_emits_nothing():
"""An empty middle chunk (reachable from forward_streaming_batch when a
row's post-trim feat comes out empty near the lookahead boundary) must
emit zero samples and leave cache_state's phase_acc unchanged, not
re-emit the re-decoded overlap window as if it were new audio.
"""
hift = _make_hift(REAL_CONFIG)
model = _make_model(hift, window_len=64)
chunks = _chunks(total_mel=96, chunk_len=24)
empty = chunks[0][:, :, :0]
cache = None
for chunk in chunks[:2]:
_, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False)
assert cache is not None
phase_acc_before = cache["phase_acc"]
speech, cache = model._stream_hift_from_feat(empty, cache_state=cache, finalize=False)
assert speech.shape[-1] == 0
if phase_acc_before is None:
assert cache["phase_acc"] is None
else:
torch.testing.assert_close(cache["phase_acc"], phase_acc_before)
@pytest.mark.parametrize("config", CONFIGS)
@pytest.mark.parametrize("mel_frames", [1, 2, 3, 4])
def test_incremental_hift_short_finalize_chunk_does_not_crash(config, mel_frames):
"""A one-token (or otherwise sub-trim) utterance that finalizes on its very
first chunk must not crash. On finalize the F0 predictor consumes the whole
window (it does not hold back `trim` frames), so the phase-carry index must
not be inflated by `trim` -- it would point past the end of a short window.
"""
hift = _make_hift(config)
model = _make_model(hift, window_len=32)
chunk = torch.randn(1, 80, mel_frames)
speech, cache = model._stream_hift_from_feat(chunk, cache_state=None, finalize=True)
assert cache is None
assert speech.shape[0] == 1
assert speech.shape[-1] > 0
assert torch.isfinite(speech).all()
tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py
+2
View File
@@ -177,6 +177,8 @@ class _MiniDiT(nn.Module):
cnn_cache_buffer: torch.Tensor | None = None,
att_cache_buffer: torch.Tensor | None = None,
) -> torch.Tensor:
assert cnn_cache_buffer is not None
assert att_cache_buffer is not None
x = x.transpose(1, 2)
x = self.in_proj(x)
for b_idx in range(len(self.blocks)):
tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py
+54
View File
@@ -0,0 +1,54 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for MOSS audio tokenizer v2 projection modules."""
from __future__ import annotations
import pytest
import torch
import torch.nn as nn
from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
MossAudioTokenizerProjectedTransformer,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_projected_transformer_keeps_learned_projections_when_dimensions_match() -> None:
module = MossAudioTokenizerProjectedTransformer(
input_dimension=8,
output_dimension=8,
d_model=8,
module_type="transformer",
num_heads=1,
num_layers=0,
positional_embedding="rope",
)
assert isinstance(module.input_proj, nn.Linear)
assert isinstance(module.output_proj, nn.Linear)
assert module.input_proj.weight.shape == (8, 8)
assert module.output_proj.weight.shape == (8, 8)
x = torch.randn(2, 8, 5)
lengths = torch.tensor([5, 3], dtype=torch.long)
output, output_lengths = module(x, lengths)
assert output.shape == x.shape
torch.testing.assert_close(output_lengths, lengths)
def test_projected_transformer_keeps_linear_projections_when_dimensions_differ() -> None:
module = MossAudioTokenizerProjectedTransformer(
input_dimension=8,
output_dimension=6,
d_model=10,
module_type="transformer",
num_heads=1,
num_layers=0,
positional_embedding="rope",
)
assert isinstance(module.input_proj, nn.Linear)
assert isinstance(module.output_proj, nn.Linear)
tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py
+86
View File
@@ -0,0 +1,86 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from typing import Any
import pytest
import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_realtime_talker_preserves_state_across_chunked_prefill() -> None:
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
MossTTSRealtimeTalkerForGeneration,
)
talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration)
talker.n_vq = 2
talker.audio_pad_token = 1024
talker.text_vocab_size = 16
talker.text_pad_id = 2
talker._batch_state = None
talker._batch_state_spans = None
embedded_audio: list[torch.Tensor | None] = []
def build_input_embeds(input_ids, audio_codes):
embedded_audio.append(audio_codes.clone() if audio_codes is not None else None)
return torch.zeros((input_ids.shape[0], 1))
generated: list[None] = []
def generate_frame(*args, **kwargs):
generated.append(None)
return torch.ones((1, talker.n_vq), dtype=torch.long)
talker._build_input_embeds = build_input_embeds
talker.local_transformer = SimpleNamespace(generate_frame=generate_frame)
talker.local_lm_heads = []
ref_codes = torch.arange(10, dtype=torch.long).reshape(5, 2)
info: dict[str, Any] = {
"codes": {"ref": ref_codes},
"ids": {"all": [7, 8]},
"max_new_frames": [3],
}
chunks = [(0, 2, False), (2, 2, False), (4, 1, True)]
state: dict[str, Any] | None = None
for computed, length, eligible in chunks:
_, _, update = talker.preprocess(
torch.ones(length, dtype=torch.long),
None,
_omni_is_prefill=True,
_omni_num_computed_tokens=computed,
_omni_prompt_len=5,
**info,
)
info.update(update)
if state is None:
state = info["audio_state"]
else:
assert info["audio_state"] is state
output = talker.make_omni_output(
torch.ones((length, 1)),
model_intermediate_buffer=[info],
request_token_spans=[(0, length)],
request_sample_eligible=[eligible],
)
if not eligible:
assert output.multimodal_outputs == {}
logits = talker.compute_logits(torch.ones((1, 1)))
assert logits.argmax().item() == (7 if eligible else talker.text_pad_id)
assert state is not None
assert info["ref_offset"] == 5
assert state["step"] == 1
assert state["text_cursor"] == 1
assert state["remaining_text"] == [7, 8]
assert len(generated) == 1
assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 2)
for actual, expected in zip(embedded_audio, ref_codes.split((2, 2, 1)), strict=True):
assert actual is not None
torch.testing.assert_close(actual, expected)
tests/model_executor/models/moss_tts/test_realtime_generation_limit.py
+53
View File
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_realtime_talker_stops_at_max_new_frames() -> None:
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
MossTTSRealtimeTalkerForGeneration,
)
talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration)
talker.n_vq = 4
talker.audio_pad_token = 1024
talker.text_vocab_size = 8
talker.text_pad_id = 2
talker.audio_eos_id = 3
talker._batch_state = None
talker._build_input_embeds = lambda input_ids, audio_codes: torch.zeros((input_ids.shape[0], 1))
generated: list[None] = []
def generate_frame(*args, **kwargs):
generated.append(None)
return torch.ones((1, talker.n_vq), dtype=torch.long)
talker.local_transformer = SimpleNamespace(generate_frame=generate_frame)
talker.local_lm_heads = []
_, _, info = talker.preprocess(
torch.tensor([1, 2]),
None,
max_new_frames=[2],
)
assert info["audio_state"]["max_new_frames"] == 2
for _ in range(3):
output = talker.make_omni_output(
torch.ones((1, 1)),
runtime_additional_information=[info],
request_token_spans=[(0, 1)],
)
assert len(generated) == 2
assert info["audio_state"]["step"] == 2
assert info["audio_state"]["is_stopping"] is True
assert output.multimodal_outputs == {}
assert talker.compute_logits(torch.ones((1, 1))).argmax().item() == talker.audio_eos_id
tests/model_executor/models/moss_tts/test_realtime_prompt.py
+78
View File
@@ -0,0 +1,78 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import numpy as np
import pytest
import torch
from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _Tokenizer:
def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
del add_special_tokens
if text == "<|im_start|>assistant\n":
return [20, 21]
return list(range(100, 100 + len(text.split())))
class _Processor:
audio_channel_pad = 1024
audio_bos_token = 1025
def __init__(self) -> None:
self.reference_codes = None
def make_ensemble(self, prompt_audio_tokens):
self.reference_codes = prompt_audio_tokens
return np.full((3, 17), self.audio_channel_pad, dtype=np.int64)
def test_build_realtime_prompt_aligns_reference_and_prefill_text() -> None:
processor = _Processor()
reference_codes = torch.arange(80, dtype=torch.int64).reshape(5, 16)
params = build_realtime_prompt(
_Tokenizer(),
processor,
"one two three",
reference_codes,
)
assert np.array_equal(processor.reference_codes, reference_codes.numpy())
assert params["prompt_token_ids"] == [1024, 1024, 1024, 20, 21, 100, 101, 102]
audio_grid = params["codes"]["ref"]
assert audio_grid.shape == (len(params["prompt_token_ids"]), 16)
assert audio_grid[-1, 0].item() == processor.audio_bos_token
assert "ids" not in params
def test_build_realtime_prompt_streams_text_after_first_twelve_tokens() -> None:
params = build_realtime_prompt(
_Tokenizer(),
_Processor(),
" ".join(f"word{i}" for i in range(15)),
torch.zeros((2, 16), dtype=torch.int64),
)
assert params["prompt_token_ids"][-12:] == list(range(100, 112))
assert params["ids"] == {"all": [112, 113, 114]}
def test_build_realtime_prompt_rejects_wrong_processor_layout(monkeypatch: pytest.MonkeyPatch) -> None:
processor = _Processor()
monkeypatch.setattr(
processor,
"make_ensemble",
lambda prompt_audio_tokens: np.zeros((3, 16), dtype=np.int64),
)
with pytest.raises(ValueError, match=r"expected \(sequence_length, 17\)"):
build_realtime_prompt(
_Tokenizer(),
processor,
"one",
torch.zeros((2, 16), dtype=torch.int64),
)
tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py
+1 -1
View File
@@ -582,7 +582,7 @@ def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hit
def test_derive_audio_from_video_placeholders_only_pairs_true_videos():
fake_self = _fake_processor()
fake_self = _fake_qwen2_prompt_processor()
video_placeholders = [
PlaceholderFeaturesInfo(
modality="video",
tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py
+50
View File
@@ -331,6 +331,56 @@ class TestEmbedInputIds:
"Interleaved: audio positions should get audio embeddings"
)
def test_interleaved_audio_boundaries_keep_text_embeddings(self, mocker: MockerFixture):
from vllm.multimodal.processing.processor import PlaceholderFeaturesInfo
from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
Qwen2_5OmniThinkerMultiModalProcessor,
)
processor = mocker.Mock()
processor.info.get_hf_processor.return_value = mocker.Mock(audio_token="audio", video_token="video")
processor.info.get_tokenizer.return_value.get_vocab.return_value = {
"audio": AUDIO_TOKEN_ID,
"video": VIDEO_TOKEN_ID,
}
# Real audio_bos/audio_eos IDs from the failing use_audio_in_video input.
tokens = [151647, VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, 151648]
placeholder = PlaceholderFeaturesInfo(
modality="video",
item_idx=0,
start_idx=2,
tokens=tokens,
is_embed=torch.ones(len(tokens), dtype=torch.bool),
)
placeholders = Qwen2_5OmniThinkerMultiModalProcessor._derive_audio_from_video_placeholders(
processor,
{"video": [placeholder]},
{"audio": [[object()]]},
[True],
)
input_ids = torch.tensor([TEXT_TOKEN_ID] * 2 + tokens + [TEXT_TOKEN_ID])
is_multimodal = torch.zeros_like(input_ids, dtype=torch.bool)
for modality in ("audio", "video"):
derived = placeholders[modality][0].to_range()
is_multimodal[derived.offset : derived.offset + derived.length] |= derived.is_embed
model, hidden = make_mock_model(mocker, thinker_cls=_omni_qwen2_5_thinker_cls())
text_embeddings = model.get_language_model().embed_input_ids(input_ids)
result = model.embed_input_ids(
input_ids,
build_mm_embeds(audio_n=2, image_n=0, video_n=2, hidden=hidden),
is_multimodal=is_multimodal,
)
is_audio = input_ids == AUDIO_TOKEN_ID
is_video = input_ids == VIDEO_TOKEN_ID
assert torch.equal(is_multimodal, is_audio | is_video)
torch.testing.assert_close(result[is_audio], torch.full((2, hidden), 10.0))
torch.testing.assert_close(result[is_video], torch.full((2, hidden), 30.0))
torch.testing.assert_close(result[~(is_audio | is_video)], text_embeddings[~(is_audio | is_video)])
# ---------------------------------------------------------------------------
# Tests for merge_interleaved_embeddings helper
tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py
+49 -2
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for the Qwen3-Omni thinker forward return contract.
Background
@@ -41,6 +41,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import (
Qwen3OmniMoeForConditionalGeneration,
)
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
Qwen3MoeLLMModel,
Qwen3OmniMoeThinkerForConditionalGeneration,
)
@@ -54,6 +55,52 @@ def _hidden_states() -> torch.Tensor:
return torch.zeros(_TOKENS, _HIDDEN)
def test_single_rank_intermediate_capture_includes_deferred_residual(mocker):
"""Capture the Talker state on the default single-PP-rank path."""
import vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker as module
pp_group = mocker.Mock(is_first_rank=True, is_last_rank=True)
mocker.patch.object(module, "get_pp_group", return_value=pp_group)
class SplitStateLayer(nn.Module):
def forward(self, positions, hidden_states, residual):
next_residual = hidden_states if residual is None else hidden_states + residual
return hidden_states * 2, next_residual
class Combine(nn.Module):
def forward(self, hidden_states, residual):
return hidden_states + residual, None
model = object.__new__(Qwen3MoeLLMModel)
nn.Module.__init__(model)
model.layers = nn.ModuleList([SplitStateLayer(), SplitStateLayer()])
model.start_layer = 0
model.end_layer = 2
model.norm = Combine()
inputs = torch.ones(_TOKENS, _HIDDEN)
output_without_capture, _ = Qwen3MoeLLMModel.forward(
model,
input_ids=None,
positions=torch.arange(_TOKENS),
inputs_embeds=inputs,
)
output_with_capture, captured = Qwen3MoeLLMModel.forward(
model,
input_ids=None,
positions=torch.arange(_TOKENS),
inputs_embeds=inputs,
capture_layer_indices=[0, 1],
return_hidden_states=True,
)
layers = captured["hidden_states"]["layers"]
torch.testing.assert_close(layers[0], inputs)
torch.testing.assert_close(layers[1], torch.full_like(inputs, 3))
torch.testing.assert_close(output_with_capture, output_without_capture)
torch.testing.assert_close(output_with_capture, torch.full_like(inputs, 7))
class _InnerModelStub:
"""Stands in for ``Qwen3MoeLLMModel``: always returns the internal 2-tuple."""
@@ -169,7 +216,7 @@ def test_make_omni_output_accepts_bare_tensor():
def test_make_omni_output_accepts_capture_tuple():
model = _make_combined(is_staged_run=True)
captured = {"hidden_states": {"layers": {}}}
captured: dict[str, dict[str, dict[int, torch.Tensor]]] = {"hidden_states": {"layers": {}}}
out = model.make_omni_output((_hidden_states(), captured))
assert isinstance(out, OmniOutput)
assert out.multimodal_outputs == captured
tests/model_executor/models/test_minimax_h3_encoder_payload.py
+406
View File
@@ -0,0 +1,406 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from contextlib import contextmanager
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.data_entry_keys import flatten_payload
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MiniMaxH3EncoderConditioning,
MiniMaxH3EncoderMediaConditioning,
MiniMaxH3EncoderMediaInput,
)
from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_pipeline_owns_three_encoders_then_dit_and_decoders() -> None:
from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE
encoder, diffusion = MINIMAX_H3_PIPELINE.stages
assert encoder.model_arch == "MiniMaxH3Encoder"
assert encoder.custom_process_next_stage_input_func is None
assert diffusion.model_arch == "MiniMaxH3Pipeline"
assert diffusion.requires_multimodal_data is False
assert diffusion.requires_full_payload_input is False
def _patch_encoder_constructors(monkeypatch, *, rank: int):
from vllm_omni.model_executor.models.minimax_h3 import encoder as encoder_module
group = SimpleNamespace(rank_in_group=rank, world_size=2, device_group=object())
class FakeVideoVAE(torch.nn.Module):
def __init__(self, path, *, device, encode_only, trust_remote_code):
super().__init__()
self.path = path
self.parallel_args = None
self.trust_remote_code = trust_remote_code
def set_parallel_size(self, size, *, process_group):
self.parallel_args = (size, process_group)
class FakeAudioVAE(torch.nn.Module):
def __init__(self, path, *, device, encode_only, trust_remote_code):
super().__init__()
self.path = path
self.trust_remote_code = trust_remote_code
def init_backbone(self, *, vllm_config, prefix):
torch.nn.Module.__init__(self)
monkeypatch.setattr(encoder_module.MiniMaxH3TextEncoderBackbone, "__init__", init_backbone)
monkeypatch.setattr(encoder_module, "get_tp_group", lambda: group)
monkeypatch.setattr(encoder_module, "get_local_device", lambda: torch.device("cpu"))
monkeypatch.setattr(encoder_module, "MiniMaxH3VideoVAE", FakeVideoVAE)
monkeypatch.setattr(encoder_module, "MiniMaxH3AudioVAE", FakeAudioVAE)
return encoder_module, group
def _component_root(tmp_path):
root = tmp_path / "FL2VA"
for name in ("text_encoder", "video_vae", "audio_vae"):
(root / name).mkdir(parents=True)
return root
def _encoder_config(root, *, video_mode="patch", roles=None, trust_remote_code=True):
components = (
{
"text_encoder": {"parallel_mode": "tp"},
"video_vae": {"parallel_mode": video_mode},
"audio_vae": {"parallel_mode": "leader"},
}
if roles is None
else roles
)
return SimpleNamespace(
model_config=SimpleNamespace(
model=str(root / "text_encoder"),
hf_config=SimpleNamespace(minimax_h3_encoder_components=components),
trust_remote_code=trust_remote_code,
)
)
def test_encoder_requires_exactly_three_components(monkeypatch, tmp_path) -> None:
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0)
root = tmp_path / "FL2VA"
(root / "text_encoder").mkdir(parents=True)
(root / "video_vae").mkdir()
config = _encoder_config(root)
with pytest.raises(RuntimeError, match=r"requires all three encoders.*audio_vae"):
encoder_module.MiniMaxH3Encoder(vllm_config=config)
@pytest.mark.parametrize(("rank", "has_audio_vae"), [(0, True), (1, False)])
@pytest.mark.parametrize("trust_remote_code", [False, True])
def test_encoder_uses_tp_patch_and_leader_on_the_same_rank_set(
monkeypatch,
tmp_path,
rank,
has_audio_vae,
trust_remote_code,
) -> None:
encoder_module, group = _patch_encoder_constructors(monkeypatch, rank=rank)
root = _component_root(tmp_path)
config = _encoder_config(root, trust_remote_code=trust_remote_code)
model = encoder_module.MiniMaxH3Encoder(vllm_config=config)
assert model.component_config.text_parallel_mode == "tp"
assert model.component_config.video_parallel_mode == "patch"
assert model.component_config.audio_parallel_mode == "leader"
assert model.video_vae.parallel_args == (2, group.device_group)
assert model.video_vae.trust_remote_code is trust_remote_code
assert (model.audio_vae is not None) is has_audio_vae
if has_audio_vae:
assert model.audio_vae.trust_remote_code is trust_remote_code
@pytest.mark.parametrize("component", ["video", "audio"])
@pytest.mark.parametrize("trusted", [False, True])
def test_encoder_only_vae_requires_remote_code_trust(monkeypatch, component, trusted) -> None:
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
config = {"auto_map": {"AutoModel": "modeling.RemoteVAE"}, "sample_rate": 32000}
monkeypatch.setattr(vae_module, "_load_component_config", lambda _path: config)
remote = torch.nn.Module()
remote.model = torch.nn.Module()
calls = []
def load_encoder(path, config_dict):
calls.append((path, config_dict))
return remote
monkeypatch.setattr(vae_module, f"_load_{component}_vae_encoder", load_encoder)
vae_cls = vae_module.MiniMaxH3VideoVAE if component == "video" else vae_module.MiniMaxH3AudioVAE
if trusted:
model = vae_cls("unused", device=torch.device("cpu"), encode_only=True, trust_remote_code=True)
assert model.remote is remote
assert calls == [("unused", config)]
else:
with pytest.raises(ValueError, match="trust-remote-code"):
vae_cls("unused", device=torch.device("cpu"), encode_only=True)
assert calls == []
def test_encoder_requires_all_role_policies(monkeypatch, tmp_path) -> None:
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0)
root = _component_root(tmp_path)
config = _encoder_config(
root,
roles={
"text_encoder": {"parallel_mode": "tp"},
"video_vae": {"parallel_mode": "patch"},
},
)
with pytest.raises(ValueError, match=r"requires exactly.*audio_vae"):
encoder_module.MiniMaxH3Encoder(vllm_config=config)
def test_encoder_does_not_infer_role_policies(monkeypatch, tmp_path) -> None:
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0)
root = _component_root(tmp_path)
config = SimpleNamespace(
model_config=SimpleNamespace(
model=str(root / "text_encoder"),
hf_config=SimpleNamespace(),
)
)
with pytest.raises(ValueError, match="requires explicit minimax_h3_encoder_components"):
encoder_module.MiniMaxH3Encoder(vllm_config=config)
def test_video_leader_policy_skips_nonleader(monkeypatch, tmp_path) -> None:
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=1)
root = _component_root(tmp_path)
model = encoder_module.MiniMaxH3Encoder(vllm_config=_encoder_config(root, video_mode="leader"))
assert model.video_vae is None
assert model.audio_vae is None
assert (
model._encode_media(
MiniMaxH3EncoderMediaInput(
task="ref2va",
height=32,
width=32,
num_frames=48,
latent_t=1,
audio_t=80,
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
)
)
is None
)
def test_encoder_runs_video_and_audio_components_on_ar_model() -> None:
from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
class VideoVAE:
def encode_image(self, _image):
return torch.ones(1, 96)
def encode_video(self, _frames):
return torch.full((1, 96), 2.0), (1, 2, 2)
class AudioVAE:
def encode_waveform(self, waveform, _sample_rate):
return torch.ones(160, 32), waveform.shape[-1] // 800
model = MiniMaxH3Encoder.__new__(MiniMaxH3Encoder)
torch.nn.Module.__init__(model)
model._component_leader = True
model.video_vae = VideoVAE()
model.audio_vae = AudioVAE()
media = MiniMaxH3EncoderMediaInput(
task="ref2va",
height=32,
width=32,
num_frames=48,
latent_t=1,
audio_t=80,
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
videos=(torch.zeros(1, 32, 32, 3, dtype=torch.uint8),),
video_audios=((torch.zeros(64_000), 32_000),),
)
conditioning = model._encode_media(media)
assert conditioning is not None
assert conditioning.visual_condition_shapes == ((1, 2, 2), (1, 2, 2))
assert conditioning.audio_condition_lengths == (80,)
assert conditioning.ref_blocks[1]["kind"] == "video_audio"
def test_prepare_encoder_inputs_keeps_reference_audio_budgets_separate(monkeypatch) -> None:
from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
frames = torch.zeros(1, 32, 32, 3, dtype=torch.uint8).numpy()
waveform = torch.zeros(320_000)
monkeypatch.setattr(processing, "resolve_minimax_h3_shape", lambda *_args: (32, 32, 240, 60, 400))
monkeypatch.setattr(
processing,
"prepare_reference_videos",
lambda *_args, **_kwargs: [
{"prepared_path": "prepared.mp4", "original_path": "reference.mp4", "input_has_audio": True}
],
)
monkeypatch.setattr(processing, "load_video_frames", lambda _path: frames)
monkeypatch.setattr(
processing,
"sample_reference_video_frames",
lambda *_args, **_kwargs: {"frames": [frames[0]], "block_timestamps": [[0.0]]},
)
monkeypatch.setattr(processing, "load_video_audio", lambda *_args, **_kwargs: (waveform, 32_000))
prepared = processing.prepare_encoder_inputs(
{
"prompt": "reference",
"multi_modal_data": {"video": "reference.mp4", "audio": (waveform, 32_000)},
},
SimpleNamespace(extra_args={"task": "ref2va"}),
)
assert prepared.media.video_audios[0][0].shape[-1] == 320_000
assert prepared.media.audios[0][0].shape[-1] == 320_000
assert prepared.condition_labels == [("audio", 1), ("video", 1), ("audio", 2)]
@pytest.mark.parametrize(
("embedded_lengths", "standalone_lengths", "valid"),
[((400,), (400,), True), ((320, 320), (400,), False), ((400,), (320, 320), False)],
)
def test_encode_media_keeps_audio_budgets_and_component_residency_separate(
embedded_lengths, standalone_lengths, valid
) -> None:
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
active = None
scopes = []
class VideoVAE:
def encode_image(self, _image):
assert active is self
return torch.ones(1, 96)
def encode_video(self, _frames):
assert active is self
return torch.ones(1, 96), (1, 2, 2)
class AudioVAE:
def encode_waveform(self, waveform, _sample_rate):
assert active is self
length = waveform.shape[-1] // 800
return torch.ones(2 * length, 32), length
@contextmanager
def component_scope(component):
nonlocal active
assert active is None
active = component
scopes.append(component)
try:
yield
finally:
active = None
video_vae = VideoVAE()
audio_vae = AudioVAE()
media = MiniMaxH3EncoderMediaInput(
task="ref2va",
height=32,
width=32,
num_frames=360,
latent_t=90,
audio_t=600,
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
videos=tuple(torch.zeros(1, 32, 32, 3, dtype=torch.uint8) for _ in embedded_lengths),
video_audios=tuple((torch.zeros(length * 800), 32_000) for length in embedded_lengths),
audios=tuple((torch.zeros(length * 800), 32_000) for length in standalone_lengths),
)
kwargs = dict(video_vae=video_vae, audio_vae=audio_vae, emit_conditioning=True, component_scope=component_scope)
if valid:
conditioning = encode_media(media, **kwargs)
assert conditioning.audio_condition_lengths == embedded_lengths + standalone_lengths
assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video_audio", "audio"]
conditioning.to_omni_components()
else:
with pytest.raises(ValueError, match="at most 15 seconds"):
encode_media(media, **kwargs)
assert scopes == [video_vae, audio_vae]
def test_encoder_output_reuses_encoder_handoff_and_round_trips() -> None:
expected = MiniMaxH3EncoderConditioning(
hidden_states=torch.randn(3, 5120, dtype=torch.bfloat16),
token_tags=torch.tensor([1, 0, 1], dtype=torch.int64),
task="t2va",
height=256,
width=448,
num_frames=17,
latent_t=5,
audio_t=10,
)
source = SimpleNamespace(
finished=True,
request_id="request-1",
outputs=[SimpleNamespace(multimodal_output=flatten_payload(expected.to_omni_payload()))],
)
prompt = {
"prompt": "hello",
"additional_information": {"global_request_id": ["request-1"]},
}
result = encoder2diffusion([source], prompt)
assert result is not None
payload = result["additional_information"]["encoder_output"]
actual = MiniMaxH3EncoderConditioning.from_omni_payload(payload)
torch.testing.assert_close(actual.hidden_states, expected.hidden_states)
torch.testing.assert_close(actual.token_tags, expected.token_tags)
def test_encoder_releases_workspace_after_cpu_payload(monkeypatch) -> None:
from vllm_omni.model_executor.models.minimax_h3 import encoder as encoder_module
model = encoder_module.MiniMaxH3Encoder.__new__(encoder_module.MiniMaxH3Encoder)
torch.nn.Module.__init__(model)
model._component_leader = True
model._token_tags = torch.tensor([1])
info = {
"meta": {"minimax_h3_encoder_request": {}},
"_minimax_h3_encoder_media_conditioning": MiniMaxH3EncoderMediaConditioning(
task="t2va",
height=256,
width=448,
num_frames=17,
latent_t=5,
audio_t=10,
),
}
empty_cache_calls = 0
def empty_cache() -> None:
nonlocal empty_cache_calls
empty_cache_calls += 1
monkeypatch.setattr(torch.accelerator, "empty_cache", empty_cache)
output = model.make_omni_output(
torch.zeros(1, 5120),
model_intermediate_buffer=[info],
request_sample_eligible=[True],
)
assert empty_cache_calls == 1
assert output.multimodal_outputs is not None
assert output.multimodal_outputs["embed"]["embedding"][0].device.type == "cpu"
tests/model_executor/stage_input_processors/test_minimax_h3.py
+147 -218
View File
@@ -2,8 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for MiniMax H3's disaggregated text-encoder contract."""
import os
import shutil
from types import SimpleNamespace
from unittest.mock import Mock
@@ -12,39 +10,37 @@ import pytest
import torch
from PIL import Image
from vllm_omni.data_entry_keys import OmniPayload, deserialize_payload, serialize_payload
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_load_audio,
resolve_minimax_h3_diffusion_model_path,
from vllm_omni.data_entry_keys import flatten_payload
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import resolve_minimax_h3_diffusion_model_path
from vllm_omni.engine.serialization import (
deserialize_additional_information,
serialize_additional_information,
)
from vllm_omni.errors import OmniClientError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
resolve_minimax_h3_model_root,
resolve_minimax_h3_encoder_model_root,
resolve_minimax_h3_partition,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MINIMAX_H3_CONDITION_LABELS_KEY,
MINIMAX_H3_ENCODER_REQUEST_KEY,
MINIMAX_H3_PRESENTATION_TASK_KEY,
MINIMAX_H3_TEXT_CONDITIONING_SCHEMA,
MiniMaxH3TextConditioning,
MiniMaxH3EncoderConditioning,
)
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import _audio_items, _load_audio
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
minimax_h3_ref2va_presentation,
minimax_h3_ref2va_video_presentation,
)
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
deserialize_prepared_reference_videos,
)
from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
MiniMaxH3MultiModalProcessor,
_build_minimax_h3_presentation,
)
from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
_audio_items,
_diffusion_sampling_params,
prepare_text_encoder_prompt,
text_encoder2diffusion,
encoder2diffusion,
prepare_encoder_prompt,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -66,13 +62,14 @@ class _SegmentTokenizer:
return self._special_ids[token]
def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(mocker):
def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(monkeypatch):
processor = object.__new__(MiniMaxH3MultiModalProcessor)
sentinel = ([1, 2, 3], object(), True)
apply_processor = mocker.patch.object(
apply_processor = Mock(return_value=sentinel)
monkeypatch.setattr(
MiniMaxH3MultiModalProcessor,
"_apply_hf_processor",
return_value=sentinel,
apply_processor,
)
inputs = object()
timing_ctx = object()
@@ -119,9 +116,11 @@ def test_fused_audio_loader_accepts_list_waveform_pair():
def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
prompt = {
"prompt": "hello",
"additional_information": {"global_request_id": ["request-1"]},
"model_intermediate_buffer": {"private": "preserved"},
"multi_modal_data": {
"image": Image.new("RGB", (256, 256)),
"audio": [np.zeros(16), 16_000],
"audio": [np.zeros(32_000), 16_000],
},
}
sampling = OmniDiffusionSamplingParams(
@@ -130,7 +129,7 @@ def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
extra_args={"task": "ref2va"},
)
transformed = prepare_text_encoder_prompt(prompt, [sampling])
transformed = prepare_encoder_prompt(prompt, [sampling])
assert transformed["prompt"] == "hello"
assert len(transformed["multi_modal_data"]["image"]) == 1
@@ -140,223 +139,148 @@ def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
("image", 1),
("audio", 1),
]
assert transformed["model_intermediate_buffer"] == {"private": "preserved"}
runner_info = transformed["additional_information"]
for _ in range(2):
wire = serialize_additional_information(runner_info)
runner_info = deserialize_additional_information(wire)
assert runner_info["global_request_id"] == ["request-1"]
request_metadata = runner_info["meta"][MINIMAX_H3_ENCODER_REQUEST_KEY]
assert request_metadata["task"] == "ref2va"
assert isinstance(
runner_info["hidden_states"]["layers"][0],
torch.Tensor,
)
assert isinstance(
runner_info["hidden_states"]["layers"][1],
torch.Tensor,
)
from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
media = MiniMaxH3Encoder._media_input(runner_info)
assert media.task == "ref2va"
assert len(media.images) == 1
assert len(media.audios) == 1
def test_text_encoder_prompt_rejects_injected_prepared_video_descriptor():
prompt = {
"prompt": "hello",
"additional_information": {"meta": {MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY: '{"artifact_dir":"/tmp/user"}'}},
}
sampling = OmniDiffusionSamplingParams(height=256, width=448, extra_args={"task": "t2va"})
def _mock_ref2va_video_with_audio(monkeypatch, *, duration_seconds: float) -> None:
from vllm_omni.model_executor.models.minimax_h3 import encoder_processing
transformed = prepare_text_encoder_prompt(prompt, [sampling])
assert MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY not in transformed["additional_information"]["meta"]
def test_prepare_ref2va_video_uses_shared_frame_sampler_once(monkeypatch):
prepare_videos = Mock()
def prepare(_videos, *, target_frame_count, workdir, start_time_seconds):
prepare_videos(_videos, target_frame_count, workdir, start_time_seconds)
prepared_path = os.path.join(workdir, "prepared.mp4")
open(prepared_path, "wb").close()
return [
{
"original_path": "/tmp/input.mp4",
"prepared_path": prepared_path,
"input_has_audio": True,
"width": 448,
"height": 256,
"start_time_seconds": 0.0,
"duration_seconds": 5.2,
"audio_duration_seconds": 5.2,
}
]
sample_frames = Mock(return_value={"frames": [np.zeros((4, 4, 3), dtype=np.uint8)]})
monkeypatch.setattr("vllm_omni.model_executor.stage_input_processors.minimax_h3.prepare_reference_videos", prepare)
sample_rate = 16_000
frames = np.zeros((4, 32, 32, 3), dtype=np.uint8)
monkeypatch.setattr(
"vllm_omni.model_executor.stage_input_processors.minimax_h3.sample_reference_video_frames",
sample_frames,
encoder_processing,
"prepare_reference_videos",
lambda *args, **kwargs: [
{
"prepared_path": "prepared.mp4",
"original_path": "original.mp4",
"input_has_audio": True,
"duration_seconds": duration_seconds,
"audio_duration_seconds": duration_seconds,
}
],
)
monkeypatch.setattr(encoder_processing, "load_video_frames", lambda path: frames)
monkeypatch.setattr(
encoder_processing,
"sample_reference_video_frames",
lambda *args, **kwargs: {"block_timestamps": [0.0], "frames": frames[:1]},
)
monkeypatch.setattr(
encoder_processing,
"load_video_audio",
lambda *args, **kwargs: (torch.zeros(round(duration_seconds * sample_rate)), sample_rate),
)
def test_prepare_ref2va_rejects_short_embedded_video_audio(monkeypatch):
_mock_ref2va_video_with_audio(monkeypatch, duration_seconds=1.0)
sampling = OmniDiffusionSamplingParams(
height=256,
width=448,
num_frames=96,
extra_args={"task": "ref2va"},
)
prompt = {"prompt": "hello", "multi_modal_data": {"video": "original.mp4"}}
with pytest.raises(OmniClientError, match=r"duration must be in \[2, 15\] seconds, got 1\.000"):
prepare_encoder_prompt(prompt, [sampling])
def test_prepare_ref2va_keeps_embedded_and_standalone_audio_budgets_separate(monkeypatch):
duration_seconds = 8.0
sample_rate = 16_000
_mock_ref2va_video_with_audio(monkeypatch, duration_seconds=duration_seconds)
sampling = OmniDiffusionSamplingParams(
height=256,
width=448,
num_frames=240,
extra_args={"task": "ref2va"},
)
prompt = {
"prompt": "hello",
"multi_modal_data": {"video": "/tmp/input.mp4"},
"multi_modal_data": {
"video": "original.mp4",
"audio": (torch.zeros(int(duration_seconds * sample_rate)), sample_rate),
},
}
sampling = OmniDiffusionSamplingParams(height=256, width=448, num_frames=124, extra_args={"task": "ref2va"})
transformed = prepare_text_encoder_prompt(prompt, [sampling])
prepare_videos.assert_called_once()
_, target_frame_count, artifact_dir, start_time_seconds = prepare_videos.call_args.args
assert target_frame_count == 124
assert start_time_seconds is None
descriptor = transformed["additional_information"]["meta"][MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY]
described_dir, described_videos = deserialize_prepared_reference_videos(descriptor)
assert described_dir == artifact_dir
assert os.path.isfile(described_videos[0]["prepared_path"])
sample_frames.assert_called_once_with(described_videos[0]["prepared_path"])
assert prompt["multi_modal_data"] == {"video": "/tmp/input.mp4"}
transformed = prepare_encoder_prompt(prompt, [sampling])
assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
("audio", 1),
("video", 1),
("audio", 2),
]
shutil.rmtree(artifact_dir)
def test_stage_wire_rejects_multiple_text_encoder_sources():
with pytest.raises(RuntimeError, match="exactly one text-encoder source"):
text_encoder2diffusion([SimpleNamespace(), SimpleNamespace()], prompt={"prompt": "hello"})
def _encoder_output() -> dict:
return MiniMaxH3EncoderConditioning(
hidden_states=torch.randn(3, 5120, dtype=torch.bfloat16),
token_tags=torch.tensor([1, 0, 1], dtype=torch.int64),
task="t2va",
height=256,
width=448,
num_frames=17,
latent_t=5,
audio_t=10,
).to_omni_payload()
def test_stage_wire_rejects_request_id_mismatch():
source = SimpleNamespace(request_id="other", outputs=[SimpleNamespace()])
def test_encoder2diffusion_reuses_encoder_handoff() -> None:
prompt = {
"prompt": "hello",
"additional_information": {"global_request_id": ["req-1"]},
}
with pytest.raises(RuntimeError, match="request ID does not match"):
text_encoder2diffusion([source], prompt=prompt)
def test_stage_wire_rejects_multiple_completions():
source = SimpleNamespace(request_id="req-1", outputs=[SimpleNamespace(), SimpleNamespace()])
with pytest.raises(RuntimeError, match="exactly one completion"):
text_encoder2diffusion([source], prompt={"prompt": "hello"})
def _source_output(payload: dict) -> SimpleNamespace:
completion = SimpleNamespace(multimodal_output=payload)
return SimpleNamespace(request_id="request-1", outputs=[completion])
def _diffusion_prompt() -> dict:
return {
"request_id": "request-1",
"prompt": "test prompt",
"additional_information": {},
}
def _valid_text_conditioning_payload() -> OmniPayload:
return {
"hidden_states": {
"output": torch.zeros(4, 5120, dtype=torch.bfloat16),
},
"meta": {
"token_role_ids": torch.tensor([[1], [1], [0], [0]], dtype=torch.int64),
"multi_modal_data": {"image": object()},
"additional_information": {
"private": "preserved",
"meta": {MINIMAX_H3_ENCODER_REQUEST_KEY: {"task": "t2va"}},
"hidden_states": {"layers": {0: torch.zeros(1)}},
},
"model_intermediate_buffer": {"private": "encoder-only"},
}
def test_text_encoder2diffusion_reads_hidden_states_output_and_token_role_ids() -> None:
payload = _valid_text_conditioning_payload()
hidden = payload["hidden_states"]["output"]
token_role_ids = payload["meta"]["token_role_ids"]
result = text_encoder2diffusion(
[_source_output(payload)],
_diffusion_prompt(),
source = SimpleNamespace(
finished=True,
request_id="request-1",
outputs=[SimpleNamespace(multimodal_output=flatten_payload(_encoder_output()))],
)
conditioning = result["additional_information"]["text_encoder_output"]
torch.testing.assert_close(conditioning["hidden_states"], hidden)
torch.testing.assert_close(conditioning["token_tags"], token_role_ids.squeeze(-1))
result = encoder2diffusion([source], prompt)
assert result["prompt"] == "test prompt"
assert result["multi_modal_data"] is None
assert "model_intermediate_buffer" not in result
additional_information = result["additional_information"]
assert additional_information["private"] == "preserved"
assert "hidden_states" not in additional_information
assert "meta" not in additional_information
parsed = MiniMaxH3EncoderConditioning.from_omni_payload(additional_information["encoder_output"])
assert parsed.task == "t2va"
def test_text_conditioning_artifact_round_trips_through_omni_payload_wire_format() -> None:
payload = _valid_text_conditioning_payload()
wire = serialize_payload(payload)
assert wire is not None
restored = deserialize_payload(wire)
conditioning = MiniMaxH3TextConditioning.from_omni_payload(restored)
assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA == "minimax_h3.text_conditioning/v1"
torch.testing.assert_close(
conditioning.hidden_states,
payload["hidden_states"]["output"],
)
torch.testing.assert_close(
conditioning.token_tags,
payload["meta"]["token_role_ids"].squeeze(-1),
)
assert conditioning.hidden_states.dtype == torch.bfloat16
assert conditioning.hidden_states.is_contiguous()
assert conditioning.token_tags.dtype == torch.int64
assert conditioning.token_tags.is_contiguous()
@pytest.mark.parametrize(
("payload", "message"),
[
({"meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states payload"),
({"hidden_states": {}, "meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states.output tensor"),
({"hidden_states": {"output": torch.randn(4, 5120)}}, "no conditioning metadata"),
({"hidden_states": {"output": torch.randn(4, 5120)}, "meta": {}}, "no token_role_ids tensor"),
(
{
"hidden_states": {"output": torch.randn(4, 5120)},
"meta": {"token_role_ids": torch.ones(4)},
},
r"must have shape \[tokens, 1\]",
),
],
)
def test_text_encoder2diffusion_rejects_invalid_conditioning_payload(payload: dict, message: str) -> None:
with pytest.raises(RuntimeError, match=message):
text_encoder2diffusion([_source_output(payload)], _diffusion_prompt())
@pytest.mark.parametrize(
("case", "message"),
[
("hidden_width", r"hidden_states must have shape \[tokens, 5120\]"),
("hidden_dtype", "hidden_states must have dtype torch.bfloat16"),
("hidden_noncontiguous", "hidden_states must use contiguous strided layout"),
("hidden_non_strided", "hidden_states must use contiguous strided layout"),
("token_dtype", "stage-wire token_role_ids must have dtype torch.int64"),
("token_layout", "stage-wire token_role_ids must use contiguous strided layout"),
("token_count", "token_tags must align with hidden_states"),
("token_value", "token_tags must contain only 0 and 1"),
],
)
def test_text_encoder2diffusion_rejects_h3_v1_contract_mismatch(case: str, message: str) -> None:
payload = _valid_text_conditioning_payload()
if case == "hidden_width":
payload["hidden_states"]["output"] = torch.zeros(4, 5119, dtype=torch.bfloat16)
elif case == "hidden_dtype":
payload["hidden_states"]["output"] = torch.zeros(4, 5120, dtype=torch.float32)
elif case == "hidden_noncontiguous":
payload["hidden_states"]["output"] = torch.empty(5120, 4, dtype=torch.bfloat16).t()
elif case == "hidden_non_strided":
payload["hidden_states"]["output"] = torch.empty(
(4, 5120),
dtype=torch.bfloat16,
layout=torch.sparse_coo,
)
elif case == "token_dtype":
payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [0], [0]], dtype=torch.int32)
elif case == "token_layout":
payload["meta"]["token_role_ids"] = torch.tensor(
[[1], [0], [1], [0], [1], [0], [1], [0]],
dtype=torch.int64,
)[::2]
elif case == "token_count":
payload["meta"]["token_role_ids"] = torch.tensor([[1], [0], [0]], dtype=torch.int64)
elif case == "token_value":
payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [2], [0]], dtype=torch.int64)
else: # pragma: no cover - parameterization is exhaustive
raise AssertionError(f"unknown test case: {case}")
with pytest.raises(RuntimeError, match=message) as exc_info:
text_encoder2diffusion([_source_output(payload)], _diffusion_prompt())
assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA in str(exc_info.value)
def test_encoder2diffusion_waits_for_one_finished_source() -> None:
assert encoder2diffusion([SimpleNamespace(finished=False)], {"prompt": "hello"}) is None
with pytest.raises(RuntimeError, match="exactly one encoder source"):
encoder2diffusion([SimpleNamespace(), SimpleNamespace()], {"prompt": "hello"})
def test_ref2va_one_image_tokens_and_tags_match_fused_presentation():
@@ -421,15 +345,20 @@ def test_checkpoint_resolver_selects_local_partition(tmp_path):
(root / "FL2VA" / "text_encoder").mkdir(parents=True)
(root / "Ref2VA" / "text_encoder").mkdir(parents=True)
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
assert resolve_minimax_h3_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
assert resolve_minimax_h3_model_root(str(root / "Ref2VA"), None, None) == str(root / "Ref2VA" / "text_encoder")
assert resolve_minimax_h3_model_root(str(root / "FL2VA"), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
assert resolve_minimax_h3_encoder_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
assert resolve_minimax_h3_encoder_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
assert resolve_minimax_h3_encoder_model_root(str(root), None, "combined") == str(root / "FL2VA" / "text_encoder")
assert resolve_minimax_h3_encoder_model_root(str(root / "Ref2VA"), None, None) == str(
root / "Ref2VA" / "text_encoder"
)
assert resolve_minimax_h3_encoder_model_root(str(root / "FL2VA"), None, "ref2va") == str(
root / "Ref2VA" / "text_encoder"
)
def test_checkpoint_resolver_rejects_unknown_task(tmp_path):
with pytest.raises(ValueError, match="task_type must be one of"):
resolve_minimax_h3_model_root(str(tmp_path), None, "unknown")
resolve_minimax_h3_encoder_model_root(str(tmp_path), None, "unknown")
def test_partition_resolver_preserves_consumer_auto_default(tmp_path):
tests/platforms/npu/quant/test_mxfp4_step_fallback.py
+168
View File
@@ -0,0 +1,168 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Real A5 kernels: step selection, bias, shape, and reuse of prepared W4 weights."""
import pytest
import torch
pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.npu]
torch_npu = pytest.importorskip("torch_npu")
@torch.inference_mode()
def _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, *, policy, scale_alg=2):
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4Config,
NPUMxfp4LinearMethod,
NPUMxfp4OnlineLinearMethod,
)
prefix = "blocks.10.attn1.to_qkv"
steps = [1] if policy in ("step", "layer_and_step") else []
layers = [prefix] if policy in ("layer", "layer_and_step") else []
fallback = {"w4a8_fallback_steps": steps, "w4a8_fallback_layers": layers, "mxfp4_scale_alg": scale_alg}
torch.manual_seed(42)
layer = torch.nn.Linear(1024, 128, bias=with_bias, device="npu", dtype=dtype)
if mode == "single":
method = NPUMxfp4OnlineLinearMethod(DiffusionMXFP4Config(**fallback), prefix=prefix)
elif mode == "single_offline":
method = NPUMxfp4LinearMethod(
DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, **fallback), prefix=prefix
)
layer.orig_dtype = dtype
values = torch.tensor([-4.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0])
layer.weight.copy_(values[torch.randint(8, (128, 1024))].to(layer.weight))
layer.register_parameter(
"weight_scale", torch.nn.Parameter(torch.full((128, 32), 124, dtype=torch.uint8, device="npu"), False)
)
layer.register_parameter("mul_scale", torch.nn.Parameter((torch.rand(1024) * 1.5 + 0.5).npu(), False))
else:
raise ValueError(f"Unsupported single-scale test mode: {mode}")
method.process_weights_after_loading(layer)
x = torch.randn(input_shape, device="npu", dtype=dtype)
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
method.process_weights_after_loading(layer)
quantize = torch_npu.npu_dynamic_mx_quant
activation_types = []
def record_quantize(x, **kwargs):
dst = kwargs["dst_type"]
expected_alg = 0 if dst == torch_npu.float8_e4m3fn else scale_alg
assert kwargs["scale_alg"] == expected_alg
assert kwargs["axis"] == -1 and kwargs["block_size"] == 32 and kwargs["round_mode"] == "rint"
if expected_alg == 2:
assert kwargs["dst_type_max"] == 7.25
else:
assert "dst_type_max" not in kwargs
activation_types.append(dst)
return quantize(x, **kwargs)
monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", record_quantize)
results = []
for _request in range(2):
with override_forward_context(ForwardContext()):
for step in range(3):
set_forward_context_denoise_step_idx(step)
output = method.apply(layer, x, layer.bias)
torch.npu.synchronize()
assert output.shape == (*input_shape[:-1], 128)
assert output.dtype == dtype
assert torch.isfinite(output).all()
results.append(output.cpu())
expected_types = [
torch_npu.float8_e4m3fn if layers or step in steps else torch_npu.float4_e2m1fn_x2 for step in range(3)
]
assert activation_types == expected_types * 2
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
assert not hasattr(layer, "w4a8_weight")
assert not hasattr(layer, "w4a8_weight_scale")
for first, second in zip(results[:3], results[3:]):
torch.testing.assert_close(first, second, rtol=0, atol=0)
@pytest.mark.parametrize("with_bias", [False, True])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("input_shape", [(16, 1024), (2, 8, 1024)])
@pytest.mark.parametrize("mode", ["single", "single_offline"])
def test_mxfp4_step_fallback(with_bias, dtype, input_shape, mode, monkeypatch):
_exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, policy="step")
@pytest.mark.cards_1
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@torch.inference_mode()
def test_single_scale_step_fallback_matches_independent_dense_reference(dtype):
if not torch.npu.is_available() or "Ascend950" not in torch.npu.get_device_name():
pytest.skip("MXFP4 kernels require an Ascend950 NPU")
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4Config,
NPUMxfp4LinearMethod,
)
# Synthetic serialized single-scale values and E8M0 exponent bytes.
# Powers of two make both activation precisions exact: this independent
# CPU dense oracle needs no production packing, dequantization or apply.
# Real UOS-quantized checkpoint validation remains a separate checkpoint/E2E gate.
values = torch.tensor([1.0, -1.0, 2.0, -2.0]).repeat(32).unsqueeze(1).expand(128, 512)
weight = torch.cat((values, values), dim=1)
fine = torch.cat((torch.full((128, 16), 127), torch.full((128, 16), 128)), dim=1).to(torch.uint8)
smooth = torch.cat((torch.full((512,), 0.5), torch.full((512,), 2.0)))
bias = torch.tensor([0.25, -0.25]).repeat(64)
x = torch.tensor([1.0, -1.0, 0.5, -0.5]).repeat(4).unsqueeze(1).expand(16, 1024) / 64
expected = torch.nn.functional.linear(x * smooth, torch.cat((values, 2 * values), dim=1), bias)
layer = torch.nn.Linear(1024, 128, bias=True, device="npu", dtype=dtype)
layer.weight.copy_(weight)
layer.bias.copy_(bias)
layer.orig_dtype = dtype
for name, value in {"weight_scale": fine, "mul_scale": smooth}.items():
layer.register_parameter(name, torch.nn.Parameter(value.npu(), requires_grad=False))
method = NPUMxfp4LinearMethod(
DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True, w4a8_fallback_steps=[1], mxfp4_scale_alg=2
)
)
method.process_weights_after_loading(layer)
prepared_pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
assert set(layer.state_dict()) == {"weight", "bias", "weight_scale", "mul_scale"}
assert not hasattr(layer, "w4a8_weight")
assert not hasattr(layer, "w4a8_weight_scale")
x = x.to(device="npu", dtype=dtype)
with override_forward_context(ForwardContext()):
for step in range(3):
set_forward_context_denoise_step_idx(step)
actual = method.apply(layer, x, layer.bias).cpu()
# A8's documented GEMM output is BF16 before restoring the caller
# dtype; the reference does not reuse quantization or unpacking.
reference = expected.bfloat16().to(dtype) if step == 1 else expected.to(dtype)
torch.testing.assert_close(actual, reference, rtol=0, atol=0)
assert actual.dtype == dtype
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == prepared_pointers
@pytest.mark.cards_1
@pytest.mark.parametrize("mode", ["single", "single_offline"])
@pytest.mark.parametrize("policy", ["layer", "layer_and_step"])
def test_mxfp4_layer_fallback(mode, policy, monkeypatch):
"""Reuse the kernel/shape/cache assertions without multiplying the full dtype/bias matrix."""
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy=policy)
@pytest.mark.cards_1
@pytest.mark.parametrize("mode", ["single", "single_offline"])
def test_mxfp4_ocp_compatibility(mode, monkeypatch):
"""OCP compatibility: A4 -> A8 -> A4 with the same W4 across requests."""
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy="step", scale_alg=0)
vllm_omni/attention/fish_kvcache_attn.py
+8
View File
@@ -24,6 +24,14 @@ def _triton_backend() -> Any:
def is_available() -> bool:
# The Fish kvcache fast path is CUDA-only: it depends on FlashAttention's
# use_cascade metadata (absent from Ascend's metadata) and a Triton CUDA
# kernel. On other platforms report unavailable so the native attention
# backend takes over instead of crashing on attn_metadata.use_cascade.
from vllm.platforms import current_platform
if not current_platform.is_cuda():
return False
return _triton_backend().is_available()
vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py
+238 -26
View File
@@ -3,15 +3,36 @@
"""Dataset and split normalization for Omni-DuplexEval.
The loader accepts a Hugging Face dataset id, a JSON/JSONL manifest, or an
The loader accepts a Hugging Face dataset id, a local Hugging Face dataset
layout (a directory or a ``.parquet`` file), 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.
Local directory mirrors are expected to use a *single* Hugging Face
configuration whose data files are named after the benchmark splits
(``data/<SPLIT>-00000-of-00001.parquet``, with ``<SPLIT>`` one of the RTD_*/PR_*
names), so ``datasets.load_dataset`` exposes one physical split per benchmark
split. A single ``.parquet`` file and a mirror that Hugging Face collapses into
a generic ``train`` split are also accepted, but then the requested ``split`` is
applied as a filter on each row's preserved ``split``/``subset``/``config``
identity instead of as a physical split name. Rows that identify neither their
split nor their family/task raise a clear error that explains the required
layout.
A requested ``split`` that matches zero rows (for example a mistyped name) is
reported loudly: :func:`load_samples` raises a ``ValueError`` that lists the
split identities actually observed in the loaded rows instead of silently
returning an empty selection. Rows coming from a JSON/JSONL manifest or an
already materialized iterable that carry no split/subset/config identity keep
the base override behavior: ``--split`` is stamped onto them so they are kept
rather than rejected.
"""
from __future__ import annotations
import importlib
import json
from collections.abc import Iterable
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any
@@ -68,6 +89,23 @@ def family_for_split(split: str, task_type: str | None = None) -> str:
raise ValueError(f"cannot infer benchmark family from split {split!r}")
def _identity_error(split: str, sample_id: str, *, family: bool) -> str:
"""Build an actionable error for a row whose family/task cannot be derived."""
subject = "family" if family else "proactive-reminder task type"
label = f" for sample {sample_id!r}" if sample_id else ""
detail = (
"the split is not an RTD_*/PR_* name and no per-row 'task_type' is present"
if family
else "no per-row 'task_type' is present"
)
return (
f"cannot infer the Omni-DuplexEval {subject}{label} from split {split!r}: {detail}. "
"Load a local Omni-DuplexEval mirror whose data files are named after the "
"Hugging Face splits (RTD_*/PR_*) or add per-row "
"'family'/'task_type'/'split' columns."
)
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:
@@ -100,14 +138,27 @@ class DuplexSample:
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=""))
def from_row(cls, row: dict[str, Any], *, media_root: Path | None = None) -> DuplexSample:
chosen_split = str(_value(row, "split", "subset", "config", default=""))
sample_id = str(_value(row, "id", "sample_id", "uid", "name", default=""))
task_value = _value(row, "task_type", "task", "type")
family = str(_value(row, "family", default="") or family_for_split(chosen_split, task_value))
family_value = _value(row, "family", default="")
if family_value:
family = str(family_value)
else:
try:
family = family_for_split(chosen_split, task_value)
except ValueError as exc:
raise ValueError(_identity_error(chosen_split, sample_id, family=True)) from exc
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=""))
if task_value:
task = canonical_task_type(str(task_value))
else:
try:
task = task_type_for_split(chosen_split)
except ValueError as exc:
raise ValueError(_identity_error(chosen_split, sample_id, family=False)) from exc
video = _value(row, "video", "video_path", "video_file")
audio = _value(row, "question_audio", "audio", "question_wav")
if media_root:
@@ -148,6 +199,150 @@ def _read_manifest(path: Path) -> list[dict[str, Any]]:
return payload
def _has_split_identity(row: dict[str, Any]) -> bool:
"""Whether a row already identifies the benchmark split it belongs to."""
return any(row.get(key) for key in ("split", "subset", "config"))
def _stamp_split(rows: Iterable[dict[str, Any]], name: str) -> list[dict[str, Any]]:
"""Stamp a physical Hugging Face split name onto rows that lack identity.
Hugging Face collapses a config-less local mirror (a bare ``data/``
directory) or a single ``.parquet`` file into a single generic split
(``train``). A row's own split identity (``split``/``subset``/``config``)
must win so family inference is not poisoned by that generic name; the
physical split name is only stamped when the row does not identify itself.
"""
stamped: list[dict[str, Any]] = []
for row in rows:
item = dict(row)
if name and not _has_split_identity(item):
item["split"] = name
stamped.append(item)
return stamped
def _override_split(rows: list[dict[str, Any]], split: str | None) -> list[dict[str, Any]]:
"""Apply the base split-override semantics to manifest/iterable rows.
Rows that lack split/subset/config identity are stamped with the requested
split so they are kept, matching the pre-override behavior; rows that
already carry identity are left untouched so the requested split stays a
filter. Copies are returned, never mutating the caller's original dicts.
"""
if not split or split == "all":
return rows
return [row if _has_split_identity(row) else {**row, "split": split} for row in rows]
def _recast_media_features(loaded: Any) -> Any:
"""Recast Audio/Video features to decode=False so rows carry {bytes, path}.
Hugging Face ``Audio``/``Video`` features default to ``decode=True``,
producing ``{"array", "path", "sampling_rate"}`` dicts without the
``bytes`` key that :func:`materialize_media` needs. ``decode=False`` keeps
the raw parquet struct instead. Must run before row iteration
(``cast_column`` is lazy; decoding happens at row access).
"""
try:
features_mod = importlib.import_module("datasets.features")
except ImportError:
return loaded
audio_cls = getattr(features_mod, "Audio", None)
video_cls = getattr(features_mod, "Video", None)
if audio_cls is None:
return loaded
if isinstance(loaded, dict):
# A DatasetDict -- or a plain dict returned by a fake loader whose
# values are lists -- recurses one level and is returned untouched by
# the features guard below.
return {name: _recast_media_features(table) for name, table in loaded.items()}
feats = getattr(loaded, "features", None)
if feats is None:
return loaded
for column, feature in feats.items():
if isinstance(feature, audio_cls):
loaded = loaded.cast_column(column, audio_cls(decode=False))
elif video_cls is not None and isinstance(feature, video_cls):
loaded = loaded.cast_column(column, video_cls(decode=False))
return loaded
def _read_hf(load_dataset: Callable[..., Any], dataset: str, want: str | None) -> Any:
"""Call ``datasets.load_dataset`` for a local layout or a remote id.
The physical split is only requested when one was asked for; ``datasets``
then returns a single ``Dataset`` (an iterable of rows) instead of a
``DatasetDict`` (a mapping of split name to such an iterable). The result is
normalized by :func:`_rows_from_hf`.
"""
if Path(dataset).suffix.lower() == ".parquet":
# A single parquet file is always exposed as one generic ``train``
# split; the logical benchmark split is filtered from row identity
# afterwards. The ``want`` branch was unreachable (the only call site
# passes ``None``) and implied physical split support that does not
# exist, so it was removed.
return load_dataset("parquet", data_files=dataset)
if want:
return load_dataset(dataset, split=want)
return load_dataset(dataset)
def _rows_from_hf(dataset: str, *, split: str | None) -> list[dict[str, Any]]:
"""Load rows from a Hugging Face dataset id or a local dataset layout.
``datasets.load_dataset`` accepts a remote id, a local dataset directory
(``data/<SPLIT>-00000-of-00001.parquet``) and, via the ``parquet`` builder,
a single ``.parquet`` file, which keeps the benchmark usable offline with a
locally mirrored copy of the dataset.
The logical benchmark split (``RTD_OCR``, ``PR_correction``, ...) is not
forwarded blindly as a physical Hugging Face split: a single ``.parquet``
file and a collapsed local mirror only expose a generic ``train`` split, so
passing the logical name there raises ``Unknown split``. The physical
dataset is loaded instead and the requested split is applied later, on the
preserved per-row identity (see :func:`load_samples`).
"""
try:
from datasets import load_dataset
except ImportError as exc:
raise RuntimeError(
"datasets is required to load a Hugging Face dataset id or a local "
"dataset directory/parquet file; use a local JSON/JSONL manifest instead"
) from exc
requested = split if split and split != "all" else None
is_local = Path(dataset).exists()
if Path(dataset).suffix.lower() == ".parquet":
# A single parquet file is always exposed as one generic ``train``
# split; the logical benchmark split is filtered from row identity
# afterwards rather than forwarded as a physical split name.
loaded = _read_hf(load_dataset, dataset, None)
else:
try:
loaded = _read_hf(load_dataset, dataset, requested)
except ValueError as exc:
# A collapsed *local* mirror (bare ``data/`` directory) does not
# expose the logical split as a physical Hugging Face split; fall
# back to the whole physical dataset and filter by row identity.
# A remote id keeps the original error so a mistyped split still
# fails loudly instead of silently selecting nothing.
if not is_local or requested is None or "Unknown split" not in str(exc):
raise
loaded = _read_hf(load_dataset, dataset, None)
requested = None
loaded = _recast_media_features(loaded)
if isinstance(loaded, dict):
rows: list[dict[str, Any]] = []
for name, table in loaded.items():
rows.extend(_stamp_split((dict(row) for row in table), name))
return rows
return _stamp_split((dict(row) for row in loaded), requested or "")
def load_samples(
dataset: str | Path | Iterable[dict[str, Any]] = DEFAULT_DATASET,
*,
@@ -157,37 +352,54 @@ def load_samples(
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)
if isinstance(dataset, str | Path):
path = Path(str(dataset))
if path.exists():
# A directory or a ``.parquet`` file is a Hugging Face dataset
# layout (e.g. a locally mirrored Omni-DuplexEval snapshot); only
# single JSON/JSONL files are parsed as manifests. Routing by path
# shape instead of mere existence lets local Hugging Face datasets
# load through ``datasets.load_dataset``.
if path.is_dir() or path.suffix.lower() == ".parquet":
rows = _rows_from_hf(str(path), split=split)
else:
try:
rows = _override_split(_read_manifest(path), split)
except (json.JSONDecodeError, UnicodeDecodeError) as exc:
raise ValueError(
f"cannot load dataset from {str(path)!r}: expected a "
"JSON/JSONL manifest file, a local Hugging Face dataset "
"directory, or a single .parquet file"
) from exc
else:
rows = [dict(row) for row in loaded]
rows = _rows_from_hf(str(dataset), split=split)
else:
rows = _override_split(list(dataset), split)
wanted = set(str(item) for item in ids) if ids else None
root = Path(media_root).expanduser() if media_root else None
result = []
split_matched = 0
for row in rows:
sample = DuplexSample.from_row(row, split=None if split in (None, "all") else split, media_root=root)
# The requested split is a *filter* over each row's own split identity;
# never pass it as an override, which would relabel every row as the
# requested split (turning a collapsed local dataset into a full match
# instead of a filter).
sample = DuplexSample.from_row(row, media_root=root)
if split and split != "all" and sample.split != split:
continue
split_matched += 1
if family != "all" and sample.family != family:
continue
if wanted is not None and sample.id not in wanted:
continue
result.append(sample)
if split and split != "all" and split_matched == 0 and rows:
observed = sorted({str(_value(row, "split", "subset", "config", default="")) or "<none>" for row in rows})
raise ValueError(
f"--split {split!r} matched zero samples. "
f"Split identities observed in the loaded rows: {observed}. "
"Check the split name for a typo."
)
if limit is not None:
result = result[: max(0, limit)]
return result
vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py
+48 -9
View File
@@ -5,7 +5,9 @@
from __future__ import annotations
import os
import subprocess
import tempfile
import wave
from collections.abc import Iterator
from pathlib import Path
@@ -20,20 +22,57 @@ def _run_ffmpeg(args: list[str]) -> bytes:
def materialize_media(value: Any, output_dir: str | Path, stem: str, suffix: str) -> Path:
"""Resolve common Hugging Face media values to a local file."""
"""Resolve common Hugging Face media values to a local file.
``str``/``Path`` values are returned as-is. Inline byte payloads (a
``bytes``/``bytearray`` or a Hugging Face ``Audio``/``Video``/``Image``
feature mapping ``{"bytes": ..., "path": ...}``) are written to a real file
under ``output_dir`` and that existing path is returned. For such mappings
the ``path`` key is the original artifact name inside the dataset, which is
usually *not* present on the local disk, so it is only used to pick a
content extension and never treated as an existing file. A mapping or
object carrying only a ``path`` (no inline ``bytes``) is returned as an
already-materialized file reference.
"""
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
payload = value.get("bytes") if isinstance(value, dict) else getattr(value, "bytes", 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))
if payload is not None:
return _write_media_bytes(payload, output_dir, stem, _media_extension(path, suffix))
if path:
return Path(path)
raise ValueError(f"cannot materialize media value for {stem!r}")
def _media_extension(path: str | Path | None, fallback: str) -> str:
"""Pick a safe lower-case file extension from ``path`` when available."""
extension = Path(str(path)).suffix.lower() if path else ""
if extension.startswith(".") and len(extension) <= 8 and extension[1:].isalnum():
return extension
return fallback
def _write_media_bytes(payload: bytes | bytearray, output_dir: str | Path, stem: str, extension: str) -> Path:
data = bytes(payload)
root = Path(output_dir)
root.mkdir(parents=True, exist_ok=True)
destination = root / f"{stem}{extension}"
if destination.exists() and destination.stat().st_size == len(data) and destination.read_bytes() == data:
return destination
fd, temporary = tempfile.mkstemp(dir=root, prefix=f".{destination.name}.", suffix=".part")
try:
with os.fdopen(fd, "wb") as handle:
handle.write(data)
os.replace(temporary, destination)
except OSError:
try:
os.unlink(temporary)
except OSError:
pass
raise
return destination
vllm_omni/benchmarks/patch/patch.py
+8
View File
@@ -3078,6 +3078,14 @@ async def benchmark(
"max_concurrent_requests": metrics.max_concurrent_requests,
"rtfx": metrics.rtfx,
}
for sample_count in (
"num_ttft_samples",
"num_tpot_samples",
"num_itl_samples",
"num_audio_ttfp_samples",
"num_audio_rtf_samples",
):
result[sample_count] = getattr(metrics, sample_count)
else:
result = {
"duration": benchmark_duration,
vllm_omni/clients/duplex.py
+6 -1
View File
@@ -645,6 +645,7 @@ class DuplexClientBase(ABC):
self._closing = False
self._closed = asyncio.Event()
self._closed_marker: _ClosedMarker | None = None
self._last_error_reason: str | None = None
# -- transport hooks ------------------------------------------------------
@@ -1061,12 +1062,16 @@ class DuplexClientBase(ABC):
# on the response path too so a consumer waiting in responses()
# is not left waiting forever (see the responses() docstring).
_put_drop_oldest(self._response_queue, _ErrorMarker(event))
# Keep the text for the close that usually follows: a client that
# is streaming input is inside send(), not draining responses(),
# and would otherwise only ever see "closed".
self._last_error_reason = f"{event.code}: {event.message}" if event.code else event.message
for queue in list(self._subscribers):
_put_drop_oldest(queue, event)
if isinstance(event, SessionClosed):
self._finalize("closed", expected=True)
self._finalize(self._last_error_reason or "closed", expected=True)
elif isinstance(event, SessionExpired):
reason = data.get("reason")
self._finalize(f"expired: {reason}" if reason else "expired", expected=False)
vllm_omni/config/omni_config.py
+10 -2
View File
@@ -50,6 +50,8 @@ from vllm_omni.config.stage_config import (
merge_sampling_constraints,
normalize_pipeline_cli_overrides,
reconcile_diffusion_attention_overrides,
resolve_stage_async_chunk,
validate_stage_async_chunk_edges,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -804,6 +806,7 @@ class _DiffusionConfigProjection:
"transformer": True,
"vae": True,
"text_encoder": True,
"vae_encoder": True,
}
)
override_transformer_cls_name: str | None = None
@@ -814,6 +817,8 @@ class _DiffusionConfigProjection:
enable_stage_verification: bool = True
prompt_file_path: str | None = None
quantization_config: _QuantizationConfigType = None
# Internal provenance, retained across config projection and worker transport.
quantization_config_is_auto_detected: bool = False
extras: dict[str, Any] = field(default_factory=dict)
@field_validator("kv_transfer_config", mode="before")
@@ -961,6 +966,8 @@ class _DiffusionConfigProjection:
or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
)
if should_use_checkpoint_config:
if self.quantization_config is None:
self.quantization_config_is_auto_detected = True
self.quantization_config = quant_config
if quant_method is not None:
self.additional_config.setdefault("auto_detected_quant_method", quant_method)
@@ -1519,7 +1526,7 @@ def _build_common_stage_config_kwargs(
*,
model: str | None,
) -> tuple[dict[str, Any], str | None, str | None]:
input_proc, next_stage_proc = _select_processor_funcs(topology, bool(deploy.async_chunk))
input_proc, next_stage_proc = _select_processor_funcs(topology, resolve_stage_async_chunk(deploy, stage_deploy))
quantization_config = _build_quantization_config(deploy, engine.quantization)
parallel_config = _build_parallel_config(deploy, engine.parallel, parallel_config_cls)
@@ -1834,7 +1841,7 @@ def _build_connector_config(
output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None
input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None
return cast(Any, OmniStageConnectorConfig)(
async_chunk=bool(deploy.async_chunk),
async_chunk=resolve_stage_async_chunk(deploy, stage_deploy),
omni_kv_config=_copy_value(engine.get("omni_kv_config")),
output_connectors=_copy_value(output_connectors) if output_connectors else None,
input_connectors=_copy_value(input_connectors) if input_connectors else None,
@@ -1965,6 +1972,7 @@ class VllmOmniConfig:
if len(pipeline_cfg.stages) <= 1:
deploy.async_chunk = False
_validate_async_chunk_support(pipeline_cfg, deploy)
validate_stage_async_chunk_edges(pipeline_cfg, deploy)
deploy_by_id = {stage.stage_id: stage for stage in deploy.stages}
model = cli_overrides.get("model")
Loaded 100 of 155 files, more files were not shown because too many files have changed in this diff. Show more