diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md
index 8f445635d..f9ee3b0c1 100644
--- a/apps/ComfyUI-vLLM-Omni/README.md
+++ b/apps/ComfyUI-vLLM-Omni/README.md
@@ -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**)
+
+
+
+
+
+
+
+
### Chaining multiple model services
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
diff --git a/apps/ComfyUI-vLLM-Omni/__init__.py b/apps/ComfyUI-vLLM-Omni/__init__.py
index 55b003ef1..8c1b03d87 100644
--- a/apps/ComfyUI-vLLM-Omni/__init__.py
+++ b/apps/ComfyUI-vLLM-Omni/__init__.py
@@ -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 ===
diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py
index 47bdd9a2c..64020b797 100644
--- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py
+++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py
@@ -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):
diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg
new file mode 100644
index 000000000..1c8130fbd
Binary files /dev/null and b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-music-generation.jpg differ
diff --git a/apps/ComfyUI-vLLM-Omni/docs/minimax-h3-t2v.md b/apps/ComfyUI-vLLM-Omni/docs/minimax-h3-t2v.md
new file mode 100644
index 000000000..790aa0f53
--- /dev/null
+++ b/apps/ComfyUI-vLLM-Omni/docs/minimax-h3-t2v.md
@@ -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.
diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/MiniMax_H3_Text_to_Video.json b/apps/ComfyUI-vLLM-Omni/example_workflows/MiniMax_H3_Text_to_Video.json
new file mode 100644
index 000000000..137ee8c2d
--- /dev/null
+++ b/apps/ComfyUI-vLLM-Omni/example_workflows/MiniMax_H3_Text_to_Video.json
@@ -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
+}
diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json
new file mode 100644
index 000000000..9dc7370c0
--- /dev/null
+++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Music Generation.json
@@ -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
+}
diff --git a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py
index ee69f7e59..2eaa097bd 100644
--- a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py
+++ b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py
@@ -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(
diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py
index 1e5aaf29c..6bf64ec1d 100644
--- a/benchmarks/diffusion/diffusion_benchmark_serving.py
+++ b/benchmarks/diffusion/diffusion_benchmark_serving.py
@@ -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))
diff --git a/benchmarks/kernels/bench_qwen_select01_modulation.py b/benchmarks/kernels/bench_qwen_select01_modulation.py
new file mode 100644
index 000000000..9a97c8de3
--- /dev/null
+++ b/benchmarks/kernels/bench_qwen_select01_modulation.py
@@ -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())
diff --git a/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py b/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py
new file mode 100644
index 000000000..fb45b1743
--- /dev/null
+++ b/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py
@@ -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()
diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md
index aa81d42c6..199241cf6 100644
--- a/docs/configuration/stage_configs.md
+++ b/docs/configuration/stage_configs.md
@@ -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). |
diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md
index 58722ffa0..862516831 100644
--- a/docs/design/architecture_overview.md
+++ b/docs/design/architecture_overview.md
@@ -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 |
diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md
index 4646b85cd..41a050a47 100644
--- a/docs/design/fullduplex.md
+++ b/docs/design/fullduplex.md
@@ -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.
diff --git a/docs/features/comfyui.md b/docs/features/comfyui.md
index b4d6fa101..dda68b793 100644
--- a/docs/features/comfyui.md
+++ b/docs/features/comfyui.md
@@ -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.
diff --git a/docs/serving/full_duplex_api.md b/docs/serving/full_duplex_api.md
index 92511e84f..9b3eceb58 100644
--- a/docs/serving/full_duplex_api.md
+++ b/docs/serving/full_duplex_api.md
@@ -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
diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md
index cc10d99c9..c18b416c7 100644
--- a/docs/serving/realtime_duplex_api.md
+++ b/docs/serving/realtime_duplex_api.md
@@ -32,33 +32,19 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
mounts `ws://:8099/v1/realtime?duplex=1` (this page; `ws://: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).
diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md
index b18b3e7b3..00ff97c01 100644
--- a/docs/user_guide/quantization/mxfp4.md
+++ b/docs/user_guide/quantization/mxfp4.md
@@ -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="",
+ 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 --quantization mxfp4
+python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4
# Online serving
vllm serve --omni --quantization mxfp4
@@ -113,19 +299,19 @@ vllm serve --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="", quantization="mxfp4_dualscale")
```
```bash
-python text_to_video.py --model --quantization mxfp4_dualscale
+python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4_dualscale
```
If accuracy debugging identifies additional precision-sensitive layers, they
@@ -136,21 +322,22 @@ omni = Omni(
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.
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index 06179ef9e..7b7bcb890 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -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
diff --git a/pyproject.toml b/pyproject.toml
index 434c4305c..c463d3917 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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)
diff --git a/recipes/BreezeBlue/Breeze-TTS-2.md b/recipes/BreezeBlue/Breeze-TTS-2.md
index 50e8858af..fb26ce91c 100644
--- a/recipes/BreezeBlue/Breeze-TTS-2.md
+++ b/recipes/BreezeBlue/Breeze-TTS-2.md
@@ -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) |
diff --git a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md
index 62f744ee5..54ea616af 100644
--- a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md
+++ b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md
@@ -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.
diff --git a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md
index efd88c068..f4304cd4a 100644
--- a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md
+++ b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md
@@ -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.
diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md
index 891332cdb..b48a1f472 100644
--- a/recipes/MiniMaxAI/MiniMax-H3.md
+++ b/recipes/MiniMaxAI/MiniMax-H3.md
@@ -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.
diff --git a/requirements/common.txt b/requirements/common.txt
index 855fb6b0a..d7b845e32 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -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
diff --git a/tests/attention/test_fish_kvcache_attn.py b/tests/attention/test_fish_kvcache_attn.py
index 4a8ea17e8..73d44295a 100644
--- a/tests/attention/test_fish_kvcache_attn.py
+++ b/tests/attention/test_fish_kvcache_attn.py
@@ -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)
diff --git a/tests/benchmarks/duplex/test_omni_duplex_eval.py b/tests/benchmarks/duplex/test_omni_duplex_eval.py
index 5f77aacc0..f5d3ebcd8 100644
--- a/tests/benchmarks/duplex/test_omni_duplex_eval.py
+++ b/tests/benchmarks/duplex/test_omni_duplex_eval.py
@@ -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
diff --git a/tests/benchmarks/duplex/test_omni_duplex_eval_media.py b/tests/benchmarks/duplex/test_omni_duplex_eval_media.py
new file mode 100644
index 000000000..28a502428
--- /dev/null
+++ b/tests/benchmarks/duplex/test_omni_duplex_eval_media.py
@@ -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")
diff --git a/tests/benchmarks/patch/test_benchmark_result.py b/tests/benchmarks/patch/test_benchmark_result.py
new file mode 100644
index 000000000..1e9619fe3
--- /dev/null
+++ b/tests/benchmarks/patch/test_benchmark_result.py
@@ -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
diff --git a/tests/benchmarks/test_diffusion_benchmark_warmups.py b/tests/benchmarks/test_diffusion_benchmark_warmups.py
new file mode 100644
index 000000000..9d72189b7
--- /dev/null
+++ b/tests/benchmarks/test_diffusion_benchmark_warmups.py
@@ -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
diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py
index 43b7547cb..75ef06e3f 100644
--- a/tests/clients/test_duplex_client.py
+++ b/tests/clients/test_duplex_client.py
@@ -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
diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py
index 3699eeda7..7940e2176 100644
--- a/tests/config/test_config_factory.py
+++ b/tests/config/test_config_factory.py
@@ -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():
diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py
index 73a7bcda2..65c315229 100644
--- a/tests/config/test_forced_aligner_injection.py
+++ b/tests/config/test_forced_aligner_injection.py
@@ -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 == []
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index a8a2056b6..1ed62f7f1 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -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])
diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py
index 604904dd0..810c506c7 100644
--- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py
+++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py
@@ -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
diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py
index bca7d095c..154644fb2 100644
--- a/tests/core/sched/test_omni_ar_scheduler_streaming.py
+++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py
@@ -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
diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py
index df99f090d..a0f966121 100644
--- a/tests/dfx/perf/tests/test_runner_metadata.py
+++ b/tests/dfx/perf/tests/test_runner_metadata.py
@@ -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,
)
diff --git a/tests/diffusion/layers/test_qwen_select01_modulation.py b/tests/diffusion/layers/test_qwen_select01_modulation.py
new file mode 100644
index 000000000..050da4717
--- /dev/null
+++ b/tests/diffusion/layers/test_qwen_select01_modulation.py
@@ -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)
diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py
index 79954f1db..32061f648 100644
--- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py
+++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py
@@ -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.`` -> ``.to_out.0.`` (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)),
]
)
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py
new file mode 100644
index 000000000..9d44bf445
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py
@@ -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")
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
index 67498abc5..4b6fc0d88 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
@@ -27,174 +27,34 @@ def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _
return result
-def _turbo_spec(filename: str) -> TurboSpec:
- """Build the spec a loaded artifact of this name would carry."""
- spec = parse_turbo_filename(filename)
- assert spec is not None
- return spec
+def _encoder_payload(
+ *,
+ task: str = "t2va",
+ height: int = 768,
+ width: int = 1344,
+ num_frames: int = 124,
+) -> dict[str, Any]:
+ from vllm_omni.diffusion.models.minimax_h3.time_request import MINIMAX_H3_SHAPE_PLANNER
+ from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderConditioning
-
-def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
- from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
-
- class FakeEncoder:
- instances = []
-
- def __init__(self, **kwargs):
- self.pushes = []
- self.kwargs = kwargs
- self.__class__.instances.append(self)
-
- def push(self, frames):
- self.pushes.append(np.array(frames, copy=True))
-
- def finish(self):
- return b"mp4"
-
- def abort(self):
- raise AssertionError("unexpected abort")
-
- class FakeAudioVAE:
- def decode_latent(self, latent):
- return torch.zeros(1, 1, 2)
-
- class FakeVideoVAE:
- def decode_with_chunks(self, latent, *, on_chunk):
- for value in (0.0, 0.25, 0.5):
- on_chunk(torch.full((1, 3, 3, 2, 2), value))
-
- pipeline = object.__new__(mod.MiniMaxH3Pipeline)
- torch.nn.Module.__init__(pipeline)
- pipeline.audio_vae = FakeAudioVAE()
- pipeline.video_vae = FakeVideoVAE()
- pipeline.device = torch.device("cpu")
- monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
- monkeypatch.setattr(
- "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder",
- FakeEncoder,
+ conditioning = MiniMaxH3EncoderConditioning(
+ hidden_states=torch.ones(1, 5120, dtype=torch.bfloat16),
+ token_tags=torch.ones(1, dtype=torch.int64),
+ task=task,
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ latent_t=MINIMAX_H3_SHAPE_PLANNER.video_latent_t(num_frames),
+ audio_t=MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(num_frames / 24),
)
-
- output = pipeline.decode_to_mp4(
- torch.zeros(1),
- torch.zeros(1),
- height=2,
- width=2,
- batch_frames=4,
- )
-
- assert output == b"mp4"
- encoder = FakeEncoder.instances[-1]
- assert [frames.shape[0] for frames in encoder.pushes] == [6, 3]
- assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8))
- assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8))
- assert encoder.pushes[1].shape == (3, 2, 2, 3)
+ return conditioning.to_omni_payload()
-def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch):
- """A client's encoder options must survive the worker-side pre-encode path."""
- from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
-
- class FakeEncoder:
- instances = []
-
- def __init__(self, **kwargs):
- self.kwargs = kwargs
- self.__class__.instances.append(self)
-
- def push(self, frames):
- del frames
-
- def finish(self):
- return b"mp4"
-
- def abort(self):
- raise AssertionError("unexpected abort")
-
- class FakeAudioVAE:
- def decode_latent(self, latent):
- return torch.zeros(1, 1, 2)
-
- class FakeVideoVAE:
- def decode_with_chunks(self, latent, *, on_chunk):
- on_chunk(torch.zeros(1, 3, 1, 2, 2))
-
- pipeline = object.__new__(mod.MiniMaxH3Pipeline)
- torch.nn.Module.__init__(pipeline)
- pipeline.audio_vae = FakeAudioVAE()
- pipeline.video_vae = FakeVideoVAE()
- pipeline.device = torch.device("cpu")
- monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
- monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder)
-
- pipeline.decode_to_mp4(
- torch.zeros(1),
- torch.zeros(1),
- height=2,
- width=2,
- video_codec_options={"preset": "ultrafast"},
- )
-
- assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"}
-
-
-@pytest.mark.parametrize(
- ("codec_extra", "expected"),
- [
- ({}, {"preset": "ultrafast", "threads": "0"}),
- ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}),
- ({"video_codec_options": {}}, {}),
- ({"video_codec_options": None}, None),
- ],
-)
-@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)])
-def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames):
- from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
- from vllm_omni.diffusion.request import OmniDiffusionRequest
- from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
- from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
- pipeline = object.__new__(MiniMaxH3Pipeline)
- torch.nn.Module.__init__(pipeline)
- pipeline.partition = "fl2va"
- pipeline.supported_tasks = frozenset({"t2va"})
- pipeline.default_video_shift = 12.0
- pipeline.default_audio_shift = 3.0
- pipeline.device = torch.device("cpu")
- pipeline.od_config = SimpleNamespace()
- pipeline.text_encoder = object()
- pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long)))
- pipeline._quality_policy = Mock()
- pipeline._cache_dit_runtime = Mock()
- pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
- pipeline.decode_to_mp4 = Mock(return_value=b"mp4")
- sampling = OmniDiffusionSamplingParams(
- width=1344,
- height=768,
- fps=24,
- num_frames=124,
- num_inference_steps=50,
- extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra},
- )
- batch = DiffusionRequestBatch(
- [OmniDiffusionRequest(prompt="test", sampling_params=sampling, request_id="codec-defaults")]
- )
-
- pipeline.forward(batch)
-
- assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected
- assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames
-
-
-def test_video_codec_options_are_normalized_for_the_encoder():
- from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options
-
- assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == {
- "preset": "ultrafast",
- "threads": "0",
+def _encoder_prompt(text: str = "test") -> dict[str, Any]:
+ return {
+ "prompt": text,
+ "additional_information": {"encoder_output": _encoder_payload()},
}
- assert normalize_video_codec_options(None) is None
- with pytest.raises(ValueError, match="video_codec_options"):
- normalize_video_codec_options("preset=ultrafast")
def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
@@ -205,25 +65,22 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
+ pipeline.load_text_encoder = False
+ pipeline.load_vae_encoder = False
pipeline.partition = "fl2va"
pipeline.supported_tasks = frozenset({"t2va"})
pipeline.default_video_shift = 12.0
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
- pipeline.text_encoder = object()
+ pipeline.load_text_encoder = False
+ pipeline.load_vae_encoder = False
cache_spec = object()
quality_plan = SimpleNamespace(cache_dit=cache_spec)
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = quality_plan
events = []
pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: events.append(("prepare", spec)))
- pipeline.encode_prompt = Mock(
- return_value=(
- torch.ones(1, 2),
- torch.ones(1, dtype=torch.long),
- )
- )
def diffuse(**kwargs):
events.append(("diffuse", kwargs))
@@ -244,7 +101,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
batch = DiffusionRequestBatch(
[
OmniDiffusionRequest(
- prompt="quality boundary",
+ prompt=_encoder_prompt("quality boundary"),
sampling_params=sampling,
request_id="quality-boundary",
)
@@ -287,7 +144,7 @@ def test_pipeline_import_registry_and_component_discovery():
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
-def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
+def test_diffusion_stage_has_only_dit_and_decoder_vaes_during_discovery():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
@@ -296,7 +153,6 @@ def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
pipeline.transformer = torch.nn.Linear(1, 1)
pipeline.video_vae = torch.nn.Linear(1, 1)
pipeline.audio_vae = torch.nn.Linear(1, 1)
- pipeline.text_encoder = None
pipeline._dit_modules = ["transformer"]
pipeline._encoder_modules = []
pipeline._vae_modules = ["video_vae", "audio_vae"]
@@ -308,6 +164,14 @@ def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
assert discovered.vae_names == ["video_vae", "audio_vae"]
+def test_diffusion_stage_requires_the_encoder_payload():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.errors import OmniClientError
+
+ with pytest.raises(OmniClientError, match="requires encoder conditioning"):
+ MiniMaxH3Pipeline._extract_encoder_conditioning(None)
+
+
def test_encoder_free_stage_skips_missing_component_during_dlo_registration():
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
@@ -339,10 +203,11 @@ def _write_partition_index(path, *, partition, tasks):
)
+@pytest.mark.skip(reason="upstream utility was removed from this branch")
def test_modular_diffusers_index_is_resolved_generically(tmp_path):
- from vllm_omni.config.resolver import resolve_omni_config
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
+ from vllm_omni.entrypoints.utils import resolve_model_config_path
(tmp_path / "modular_model_index.json").write_text(
json.dumps(
@@ -356,16 +221,7 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path):
assert is_diffusion_model(str(tmp_path))
assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline"
- resolved = resolve_omni_config(
- str(tmp_path),
- trust_remote_code=False,
- deploy_config_path=None,
- cli_overrides=None,
- stage_overrides=None,
- strategy_config_path=None,
- )
- assert resolved.config_path is None
- assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline"
+ assert resolve_model_config_path(str(tmp_path)) is None
config = OmniDiffusionConfig(model=str(tmp_path))
config.enrich_config()
@@ -412,7 +268,7 @@ def test_startup_task_rejects_unsupported_value():
_minimax_h3_partition_for_task("unsupported")
-def test_combined_task_inference_and_transformer_routing():
+def test_encoder_task_validation_and_transformer_routing():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.errors import OmniClientError
@@ -434,7 +290,7 @@ def test_combined_task_inference_and_transformer_routing():
assert pipeline._resolve_task("ref2va", {"image": object()}, turbo_spec=ref2v_spec) == "ref2va"
pipeline.partition = "ref2va"
pipeline.supported_tasks = frozenset({"ref2va"})
- assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
+ assert pipeline._resolve_task("ref2va") == "ref2va"
pipeline.partition = "combined"
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
@@ -450,15 +306,14 @@ def test_combined_task_inference_and_transformer_routing():
"component_partition",
"source_partitions",
"expected_tasks",
- "load_text_encoder",
),
[
- (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}, True),
- ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, True),
- ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}, True),
- ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, False),
+ (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}),
+ ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}),
+ ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}),
],
)
+@pytest.mark.parametrize("load_text_encoder", [None, True, False])
@pytest.mark.parametrize("trust_remote_code", [False, True])
def test_pipeline_loads_task_selected_components_with_encoder_ownership(
monkeypatch,
@@ -493,16 +348,13 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
for partition_name in ("FL2VA", "Ref2VA"):
for component in (
"transformer",
- "tokenizer",
- "processor",
- "text_encoder",
"video_vae",
"audio_vae",
):
(tmp_path / partition_name / component).mkdir()
- created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
- created_kwargs: dict[str, list[Any]] = {"text_encoder": [], "video_vae": [], "audio_vae": []}
+ created: dict[str, list[Any]] = {"dit": [], "video_vae": [], "audio_vae": [], "text_encoder": []}
+ component_options = {}
class FakeModule(torch.nn.Module):
def __init__(self, *args, **kwargs):
@@ -516,27 +368,17 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
def component_factory(name):
def create(path, *args, **kwargs):
created[name].append(str(path))
- created_kwargs[name].append(kwargs)
+ component_options[name] = kwargs
return FakeModule()
return create
- tokenizer_calls: list[Any] = []
- processor_calls: list[Any] = []
- tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast)
- tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
- tokenizer_calls, (args, kwargs), object()
- )
- processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor)
- processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
- processor_calls, (args, kwargs), object()
- )
monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT)
- monkeypatch.setattr(
- pipeline_module,
- "MiniMaxH3Qwen3VLEncoder",
- component_factory("text_encoder"),
- )
+ monkeypatch.setattr(pipeline_module, "MiniMaxH3Qwen3VLEncoder", component_factory("text_encoder"))
+ tokenizer_loader = Mock()
+ processor_loader = Mock()
+ monkeypatch.setattr(pipeline_module.Qwen2TokenizerFast, "from_pretrained", tokenizer_loader)
+ monkeypatch.setattr(pipeline_module.Qwen3VLProcessor, "from_pretrained", processor_loader)
monkeypatch.setattr(
pipeline_module,
"MiniMaxH3VideoVAE",
@@ -547,21 +389,6 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
"MiniMaxH3AudioVAE",
component_factory("audio_vae"),
)
- monkeypatch.setattr(
- pipeline_module,
- "Qwen2TokenizerFast",
- tokenizer_cls,
- )
- monkeypatch.setattr(
- pipeline_module,
- "Qwen3VLProcessor",
- processor_cls,
- )
- monkeypatch.setattr(
- pipeline_module,
- "_dit_rank_world",
- lambda: (None, 0, 1),
- )
monkeypatch.setattr(
pipeline_module,
"get_local_device",
@@ -581,21 +408,17 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
od_config = OmniDiffusionConfig(
model="MiniMaxAI/MiniMax-H3",
+ trust_remote_code=trust_remote_code,
revision=None,
task_type=task_type,
quantization_config=None,
- model_loaded={
- "transformer": True,
- "vae": True,
- "text_encoder": load_text_encoder,
- },
- parallel_config=DiffusionParallelConfig(
- cfg_parallel_size=1,
- text_encoder_tp_size=1,
- ),
- trust_remote_code=trust_remote_code,
+ parallel_config=DiffusionParallelConfig(cfg_parallel_size=1),
)
+ if load_text_encoder is not None:
+ od_config.model_loaded["text_encoder"] = load_text_encoder
+ od_config.model_loaded["vae_encoder"] = load_text_encoder
pipeline = pipeline_module.MiniMaxH3Pipeline(od_config=od_config)
+ expect_encoder = load_text_encoder is not False
assert pipeline.partition == expected_partition
assert pipeline.supported_tasks == expected_tasks
@@ -603,59 +426,47 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
component_path = tmp_path / component_partition
assert created["video_vae"] == [str(component_path / "video_vae")]
assert created["audio_vae"] == [str(component_path / "audio_vae")]
- # The checkpoint VAEs execute the modeling code shipped with the weights, so
- # the pipeline must hand them the engine's trust_remote_code decision. Both
- # VAEs are created for every parametrization, so this stays unconditional.
- for component in ("video_vae", "audio_vae"):
- assert [kwargs["trust_remote_code"] for kwargs in created_kwargs[component]] == [trust_remote_code]
- # The tokenizer/processor "loaded once" assertions are asserted per
- # load_text_encoder branch below, matching upstream's own expectations.
+ assert created["text_encoder"] == ([str(component_path / "text_encoder")] if expect_encoder else [])
+ assert tokenizer_loader.call_count == processor_loader.call_count == int(expect_encoder)
+ assert pipeline.load_text_encoder is expect_encoder
+ assert pipeline._encoder_modules == (["text_encoder"] if expect_encoder else [])
+ assert ("text_encoder" in pipeline._offload_plan.encoder_block_attrs) is expect_encoder
+ assert ("text_encoder" in pipeline._offload_plan.on_demand_component_paths) is expect_encoder
+ assert ("encode_prompt" in pipeline._PROFILER_TARGETS) is expect_encoder
+ assert ("_encode_local_media" in pipeline._PROFILER_TARGETS) is expect_encoder
+ assert component_options["video_vae"]["decode_only"] is not expect_encoder
+ assert component_options["audio_vae"]["decode_only"] is not expect_encoder
+ assert component_options["video_vae"]["trust_remote_code"] is trust_remote_code
+ assert component_options["audio_vae"]["trust_remote_code"] is trust_remote_code
+ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+ assert ModuleDiscovery.discover(pipeline).encoder_names == (["text_encoder"] if expect_encoder else [])
expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"]
assert pipeline._dit_modules == expected_dit_modules
assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2)
- if load_text_encoder:
- assert created["text_encoder"] == [str(component_path / "text_encoder")]
- assert len(tokenizer_calls) == 1
- assert len(processor_calls) == 1
- assert pipeline.text_encoder is not None
- assert pipeline.tokenizer is not None
- assert pipeline.processor is not None
- assert pipeline._encoder_modules == ["text_encoder"]
- assert pipeline.text_encoder_tp_size == 1
- else:
- assert created["text_encoder"] == []
- assert tokenizer_calls == []
- assert processor_calls == []
- assert pipeline.text_encoder is None
- assert pipeline.tokenizer is None
- assert pipeline.processor is None
- assert pipeline._encoder_modules == []
- assert pipeline.text_encoder_tp_size == 0
- assert pipeline.text_encoder_group is None
if expected_partition == "ref2va":
assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
assert [source.model_or_path for source in pipeline.weights_sources] == [
*(str(tmp_path / partition_name) for partition_name in source_partitions),
- *((str(component_path),) if load_text_encoder else ()),
+ *([str(component_path)] if expect_encoder else []),
]
assert [source.subfolder for source in pipeline.weights_sources] == [
*("transformer" for _ in source_partitions),
- *(("text_encoder",) if load_text_encoder else ()),
+ *(["text_encoder"] if expect_encoder else []),
]
expected_prefixes = ["transformer."]
if expected_dits == 2:
expected_prefixes.append("transformers_ref.")
- if load_text_encoder:
+ if expect_encoder:
expected_prefixes.append("text_encoder.")
assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes
- if load_text_encoder:
+ expected_patterns = pipeline_module.MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[expected_partition]
+ if expect_encoder:
expected_patterns = (
pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
if expected_partition == "combined"
else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition]
)
- else:
- expected_patterns = pipeline_module.MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[expected_partition]
assert download_calls == [
{
"model_name_or_path": "MiniMaxAI/MiniMax-H3",
@@ -687,7 +498,6 @@ def test_combined_weight_loader_routes_each_contiguous_partition():
torch.nn.Module.__init__(pipeline)
pipeline.transformer = FakeTransformer()
pipeline.transformers_ref = FakeTransformer()
- pipeline.text_encoder = torch.nn.Identity()
pipeline.video_vae = torch.nn.Identity()
pipeline.audio_vae = torch.nn.Identity()
loaded = pipeline.load_weights(
@@ -740,26 +550,6 @@ def test_joint_postprocess_is_multiprocessing_picklable():
assert result["fps"] == 24
-def test_joint_video_output_is_quantized_before_transfer():
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- _minimax_h3_post_process,
- _prepare_minimax_h3_video_output,
- )
-
- video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5)
- expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1)
- prepared = _prepare_minimax_h3_video_output(video)
-
- assert prepared.dtype == torch.uint8
- assert prepared.is_contiguous()
- assert prepared.shape == (2, 2, 4, 5, 3)
- torch.testing.assert_close(prepared, expected, rtol=0, atol=0)
-
- result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6)))
- assert result["video"][0].dtype == np.uint8
- np.testing.assert_array_equal(result["video"][0], expected[0].numpy())
-
-
def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -865,17 +655,12 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
- pipeline.text_encoder = Mock()
pipeline._base_schedule_by_partition = schedules
+ pipeline.load_text_encoder = False
+ pipeline.load_vae_encoder = False
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: None)
- pipeline.encode_prompt = Mock(
- return_value=(
- torch.ones(1, 2),
- torch.ones(1, dtype=torch.long),
- )
- )
def diffuse(**kwargs):
diffuse_calls.append(kwargs)
@@ -903,7 +688,7 @@ def _t2va_batch(num_inference_steps=None):
return DiffusionRequestBatch(
[
OmniDiffusionRequest(
- prompt="distilled schedule",
+ prompt=_encoder_prompt("distilled schedule"),
sampling_params=sampling,
request_id="distilled",
)
@@ -1024,10 +809,6 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
def test_packed_attention_skips_mask_for_packed_mask_free_backend():
- from vllm_omni.diffusion.attention.backends.abstract import (
- VideoTokenLayout,
- VideoTokenSpan,
- )
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention,
)
@@ -1054,10 +835,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
torch.nn.Module.__init__(attention)
attention.attention = FakeAttention()
q = torch.randn(8, 2, 4)
- video_layout = VideoTokenLayout(
- used_len=5,
- video_spans=(VideoTokenSpan(start=3, latent_grid=(1, 1, 2), role="target"),),
- )
attention._run_packed_attention(
q,
@@ -1066,9 +843,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
max_seqlen=5,
packed_total=8,
- video_layout=video_layout,
- vsa_prefix_segments=(1, 2),
- gate_compress=torch.ones_like(q),
)
metadata = attention.attention.metadata
@@ -1076,11 +850,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
assert metadata.attn_mask is None
assert metadata.extra["valid_kv_length"] == 5
assert metadata.extra["npu_attn_varlen"] is True
- assert metadata.video_layout is video_layout
- assert metadata.extra["vsa_h3_prefix_segments"] == (1, 2)
- assert "vsa_h3_video_shape" not in metadata.extra
- assert "vsa_h3_target_start" not in metadata.extra
- assert metadata.extra["gate_compress"].shape == (1, *q.shape)
packed_padding = metadata.packed_padding
assert packed_padding is not None
assert packed_padding.q_length == 5
@@ -1275,20 +1044,20 @@ def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_
def test_reference_image_resize_contract():
from PIL import Image
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- _reference_image_shape,
+ from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
+ resolve_minimax_h3_reference_image_shape,
)
- assert _reference_image_shape(Image.new("RGB", (1080, 1440))) == (
+ assert resolve_minimax_h3_reference_image_shape(Image.new("RGB", (1080, 1440))) == (
2048,
2720,
)
with pytest.raises(ValueError, match="aspect ratio"):
- _reference_image_shape(Image.new("RGB", (100, 501)))
+ resolve_minimax_h3_reference_image_shape(Image.new("RGB", (100, 501)))
def test_fl2va_supports_first_last_and_explicit_frame_index_contracts():
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
_resolve_fl2va_keyframe_indices,
)
@@ -1301,21 +1070,20 @@ def test_fl2va_supports_first_last_and_explicit_frame_index_contracts():
def test_minimax_h3_uses_the_official_output_canvas_policy():
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- _resolve_output_canvas,
+ from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
+ resolve_minimax_h3_output_canvas,
)
- assert _resolve_output_canvas(21 / 9, 768) == (672, 1536)
- assert _resolve_output_canvas(16 / 9, 768) == (768, 1344)
- assert _resolve_output_canvas(9 / 16, 768) == (1344, 768)
+ assert resolve_minimax_h3_output_canvas(21 / 9, 768) == (672, 1536)
+ assert resolve_minimax_h3_output_canvas(16 / 9, 768) == (768, 1344)
+ assert resolve_minimax_h3_output_canvas(9 / 16, 768) == (1344, 768)
with pytest.raises(ValueError, match="short_edge"):
- _resolve_output_canvas(16 / 9, 720)
+ resolve_minimax_h3_output_canvas(16 / 9, 720)
def test_minimax_h3_accepts_sglang_auto_aspect_ratio_alias():
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape
- pipeline = object.__new__(MiniMaxH3Pipeline)
sampling = SimpleNamespace(
fps=24,
num_frames=1,
@@ -1323,7 +1091,7 @@ def test_minimax_h3_accepts_sglang_auto_aspect_ratio_alias():
width=None,
extra_args={"duration": 5.0, "aspect_ratio": "auto"},
)
- height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None)
+ height, width, *_ = resolve_minimax_h3_shape("ref2va", sampling, None)
assert (height, width) == (768, 1344)
@@ -1476,12 +1244,6 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
from vllm_omni.diffusion.models.minimax_h3.encoder import (
MiniMaxH3Qwen3VLEncoder,
)
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- MiniMaxH3Pipeline,
- )
- from vllm_omni.diffusion.offloader.component_utils import (
- get_encoder_block_groups,
- )
encoder = MiniMaxH3Qwen3VLEncoder(
"/nonexistent/text_encoder",
@@ -1494,25 +1256,6 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
# The stub has no parameters, so it never contributes to the runner's
# strict missing-parameter check on non-encoder ranks.
assert list(encoder.named_parameters()) == []
- assert (
- get_encoder_block_groups(
- encoder,
- "text_encoder",
- MiniMaxH3Pipeline._offload_plan,
- strict=True,
- )
- == []
- )
-
- # An arbitrary empty module is not silently treated as a distributed
- # stub; explicit plans still validate their declared paths.
- with pytest.raises(ValueError, match=r"text_encoder\.vision\.blocks was not found"):
- get_encoder_block_groups(
- nn.Module(),
- "text_encoder",
- MiniMaxH3Pipeline._offload_plan,
- strict=True,
- )
def test_global_quant_config_is_shared_by_dit_and_encoder():
@@ -1575,149 +1318,22 @@ def test_text_encoder_linear_delegates_quantization_to_vllm_factory():
quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix)
-def test_no_offload_keeps_text_encoder_resident():
+def test_distributed_layerwise_offload_stages_vae_component():
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=False,
enable_layerwise_offload=False,
- )
- pipeline.text_encoder = Mock()
- expected = torch.ones(2, 3)
- pipeline.text_encoder.encode_ids.return_value = expected
- input_ids = torch.tensor([1, 2])
-
- actual = pipeline._encode_text_hidden(input_ids, {})
-
- assert actual is expected
- pipeline.text_encoder.load_to_device.assert_called_once_with()
- pipeline.text_encoder.offload_to_cpu.assert_not_called()
-
-
-def test_model_offload_uses_hooked_text_encoder_call():
- 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,
- )
- expected = torch.ones(2, 3)
- pipeline.text_encoder = Mock(return_value=expected)
- pipeline._model_cpu_offload_modules = [pipeline.text_encoder]
- input_ids = torch.tensor([1, 2])
- vision_kwargs = {"pixel_values": torch.ones(1, 4)}
-
- actual = pipeline._encode_text_hidden(input_ids, vision_kwargs)
-
- assert actual is expected
- pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs)
- pipeline.text_encoder.load_to_device.assert_not_called()
-
-
-@pytest.mark.parametrize(
- ("enable_layerwise", "enable_distributed"),
- [(True, False), (False, True)],
-)
-def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed):
- from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
- pipeline = object.__new__(MiniMaxH3Pipeline)
- torch.nn.Module.__init__(pipeline)
- pipeline.od_config = SimpleNamespace(
- diffusion_offload_config=None,
- enable_cpu_offload=False,
- enable_layerwise_offload=enable_layerwise,
- enable_distributed_layerwise_offload=enable_distributed,
- )
- 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.offload_to_cpu.assert_called_once_with()
-
- vae = Mock()
- with pipeline._component_on_device(vae):
- pass
- vae.load_to_device.assert_called_once_with()
- vae.offload_to_cpu.assert_called_once_with()
-
-
-def test_layerwise_encoder_selection_releases_text_encoder():
- 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=False,
- enable_layerwise_offload=False,
- enable_distributed_layerwise_offload=False,
- diffusion_offload_config={
- "mode": "layer",
- "components": ["text_encoder"],
- },
- )
- 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.offload_to_cpu.assert_called_once_with()
-
-
-def test_layerwise_dit_only_keeps_text_encoder_resident():
- 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=False,
- enable_layerwise_offload=False,
- enable_distributed_layerwise_offload=False,
- diffusion_offload_config={"mode": "layer", "components": ["dit"]},
- )
- 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.offload_to_cpu.assert_not_called()
-
-
-def test_dit_encoder_selection_keeps_vae_resident():
- 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=False,
- enable_layerwise_offload=False,
- enable_distributed_layerwise_offload=False,
- diffusion_offload_config={
- "mode": "layer",
- "components": ["dit", "text_encoder"],
- },
+ enable_distributed_layerwise_offload=True,
)
component = Mock()
with pipeline._component_on_device(component):
- pass
+ component.load_to_device.assert_called_once_with()
+ component.offload_to_cpu.assert_not_called()
- component.load_to_device.assert_not_called()
- component.offload_to_cpu.assert_not_called()
+ component.offload_to_cpu.assert_called_once_with()
def test_distributed_layerwise_resident_blocks_are_stage_scoped():
@@ -1766,6 +1382,58 @@ def test_distributed_layerwise_resident_blocks_can_be_skipped():
controller.offload_resident_layers.assert_not_called()
+@pytest.mark.skip(reason="encoder layerwise API is not part of this branch")
+def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3.encoder import (
+ MiniMaxH3Qwen3VLEncoder,
+ )
+
+ class Stack(torch.nn.Module):
+ def __init__(self, count):
+ super().__init__()
+ self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
+
+ class TextStack(torch.nn.Module):
+ def __init__(self, count):
+ super().__init__()
+ self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
+
+ hooks = []
+
+ def fake_apply(block, next_block, device, stream, pin_memory):
+ hook = SimpleNamespace(
+ block=block,
+ next_block=next_block,
+ device=device,
+ stream=stream,
+ pin_memory=pin_memory,
+ _prev_hook=None,
+ offload_layer=Mock(),
+ )
+ hooks.append(hook)
+ return hook
+
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook",
+ fake_apply,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.platforms.current_omni_platform.Stream",
+ Mock(return_value="copy-stream"),
+ )
+ encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+ torch.nn.Module.__init__(encoder)
+ encoder.device_target = torch.device("cpu")
+ encoder.vision = Stack(2)
+ encoder.text_model = TextStack(3)
+
+ encoder.enable_omni_layerwise_offload(pin_memory=False)
+
+ assert len(hooks) == 5
+ assert all(hook._prev_hook is not None for hook in hooks)
+ assert all(hook.device == torch.device("cpu") for hook in hooks)
+
+
def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
@@ -1797,97 +1465,6 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
assert next(video_vae.parameters()).dtype == torch.float32
-def test_keyframe_encode_pins_and_restores_cudnn_settings():
- from vllm_omni.diffusion.models.minimax_h3.vae import (
- _minimax_h3_keyframe_encode_context,
- )
-
- cudnn = torch.backends.cudnn
- original = (
- cudnn.enabled,
- cudnn.benchmark,
- cudnn.deterministic,
- cudnn.allow_tf32,
- )
- try:
- cudnn.enabled = False
- cudnn.benchmark = True
- cudnn.deterministic = False
- cudnn.allow_tf32 = False
-
- with _minimax_h3_keyframe_encode_context(torch.device("cuda")):
- assert cudnn.enabled
- assert not cudnn.benchmark
- assert cudnn.deterministic
- assert cudnn.allow_tf32
-
- assert not cudnn.enabled
- assert cudnn.benchmark
- assert not cudnn.deterministic
- assert not cudnn.allow_tf32
- finally:
- (
- cudnn.enabled,
- cudnn.benchmark,
- cudnn.deterministic,
- cudnn.allow_tf32,
- ) = original
-
-
-@pytest.mark.parametrize("fail_encode", [False, True])
-def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
- from PIL import Image
-
- from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
-
- events = []
- input_image = Image.new("RGB", (32, 32))
-
- @contextmanager
- def track_context(device):
- assert device == torch.device("cpu")
- events.append("enter")
- try:
- yield
- finally:
- events.append("exit")
-
- class FakeModel(torch.nn.Module):
- def __init__(self):
- super().__init__()
- self.weight = torch.nn.Parameter(torch.ones(1))
- self.parallel_tiling = True
-
- def encode_images(self, image, *, use_fp16_latent):
- assert events == ["enter"]
- assert image is input_image
- assert use_fp16_latent
- events.append("encode")
- if fail_encode:
- raise RuntimeError("encode failed")
- return [torch.ones(1, 1, 2, 2)]
-
- video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
- torch.nn.Module.__init__(video_vae)
- video_vae.model = FakeModel()
- video_vae.config_dict = {
- "latent_channels": 1,
- "latents_mean": [0.0],
- "latents_std": [1.0],
- }
- monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context)
-
- if fail_encode:
- with pytest.raises(RuntimeError, match="encode failed"):
- video_vae.encode_image(input_image)
- else:
- rows = video_vae.encode_image(input_image)
- torch.testing.assert_close(rows, torch.ones(1, 4))
-
- assert events == ["enter", "encode", "exit"]
- assert video_vae.model.parallel_tiling
-
-
def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
@@ -2115,69 +1692,6 @@ def test_video_vae_encode_uses_configured_parallel_tiling():
assert shape == (2, 2, 2)
-def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
- from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
- from vllm_omni.diffusion.models.minimax_h3 import (
- pipeline_minimax_h3 as pipeline_module,
- )
-
- prepared = [
- {"prepared_path": "video-1.mp4"},
- {"prepared_path": "video-2.mp4"},
- ]
-
- class FakeVideoVAE:
- def __init__(self):
- self.calls = []
-
- def is_distributed_enabled(self):
- return True
-
- def encode_video(self, frames):
- self.calls.append(frames)
- index = len(self.calls)
- return torch.full((1, 2), index, dtype=torch.float32), (index, 2, 3)
-
- pipeline = object.__new__(MiniMaxH3Pipeline)
- torch.nn.Module.__init__(pipeline)
- pipeline.device = torch.device("cpu")
- pipeline.video_vae = FakeVideoVAE()
-
- monkeypatch.setattr(
- pipeline_module,
- "_dit_rank_world",
- lambda: ("dit-group", 1, 4),
- )
-
- def fake_broadcast_object_list(values, *, src, group, device):
- assert values == [None]
- assert (src, group, device) == (0, "dit-group", torch.device("cpu"))
- values[0] = prepared
-
- monkeypatch.setattr(
- pipeline_module.dist,
- "broadcast_object_list",
- fake_broadcast_object_list,
- )
- monkeypatch.setattr(
- pipeline_module,
- "load_video_frames",
- lambda path: f"frames:{path}",
- )
-
- rows, shapes = pipeline._encode_video_conditions_resident(None, count=2)
-
- assert pipeline.video_vae.calls == [
- "frames:video-1.mp4",
- "frames:video-2.mp4",
- ]
- torch.testing.assert_close(
- rows,
- torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
- )
- assert shapes == [(1, 2, 3), (2, 2, 3)]
-
-
@pytest.mark.parametrize(
("case", "extra", "image_count", "expected"),
[
@@ -2186,7 +1700,7 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
],
)
def test_f1_f2_official_fl2va_keyframe_matrix(case, extra, image_count, expected):
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
_resolve_fl2va_keyframe_indices,
)
@@ -2206,7 +1720,7 @@ def test_f1_f2_official_fl2va_keyframe_matrix(case, extra, image_count, expected
],
)
def test_r1_r6_ref2va_reference_count_matrix(case, counts):
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
_validate_ref2va_reference_counts,
)
@@ -2215,10 +1729,10 @@ def test_r1_r6_ref2va_reference_count_matrix(case, counts):
def test_ref2va_reference_count_validation_preserves_client_error_metadata():
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ from vllm_omni.errors import OmniClientError
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
_validate_ref2va_reference_counts,
)
- from vllm_omni.errors import OmniClientError
with pytest.raises(OmniClientError, match="at least one image or video"):
_validate_ref2va_reference_counts(0, 0, 0)
@@ -2322,32 +1836,6 @@ def test_prepared_reference_video_descriptor_rejects_unknown_fields():
deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}')
-def test_stage_one_reuses_existing_prepared_reference_video(tmp_path):
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- _reuse_prepared_reference_videos,
- )
-
- prepared_path = tmp_path / "prepared.mp4"
- prepared_path.touch()
- prepared = [{"prepared_path": str(prepared_path)}]
-
- assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared
-
-
-def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path):
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- _reuse_prepared_reference_videos,
- )
-
- with pytest.raises(ValueError, match="count does not match"):
- _reuse_prepared_reference_videos([], expected_count=1)
- with pytest.raises(ValueError, match="prepared reference video is unavailable"):
- _reuse_prepared_reference_videos(
- [{"prepared_path": str(tmp_path / "missing.mp4")}],
- expected_count=1,
- )
-
-
def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
@@ -2667,9 +2155,8 @@ def test_ref2va_two_video_recipe_rejects_real_duration_overflow(monkeypatch, tmp
@pytest.mark.parametrize("duration", [4.0, 15.0])
def test_g2_output_duration_accepts_official_boundaries(duration):
- from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape
- pipeline = object.__new__(MiniMaxH3Pipeline)
sampling = SimpleNamespace(
fps=24,
num_frames=1,
@@ -2678,15 +2165,14 @@ def test_g2_output_duration_accepts_official_boundaries(duration):
extra_args={"duration": duration},
)
- height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None)
+ height, width, *_ = resolve_minimax_h3_shape("ref2va", sampling, None)
assert height > 0 and width > 0
@pytest.mark.parametrize("duration", [3.99, 15.01, float("nan"), "not-a-duration"])
def test_g2_output_duration_rejects_out_of_contract_values(duration):
- from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape
- pipeline = object.__new__(MiniMaxH3Pipeline)
sampling = SimpleNamespace(
fps=24,
num_frames=1,
@@ -2695,7 +2181,7 @@ def test_g2_output_duration_rejects_out_of_contract_values(duration):
extra_args={"duration": duration},
)
with pytest.raises(ValueError, match="duration"):
- pipeline._resolve_shape("ref2va", sampling, None)
+ resolve_minimax_h3_shape("ref2va", sampling, None)
def test_g1_fanout_uses_incrementing_output_seeds():
@@ -2713,26 +2199,25 @@ def test_g1_fanout_uses_incrementing_output_seeds():
def test_g3_task_specific_aspect_ratio_policy():
from PIL import Image
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- _resolve_minimax_h3_aspect_ratio,
+ from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
+ resolve_minimax_h3_aspect_ratio,
)
image = Image.new("RGB", (1280, 720))
- assert _resolve_minimax_h3_aspect_ratio("fl2va", "9:16", image) == pytest.approx(16 / 9)
- assert _resolve_minimax_h3_aspect_ratio("ref2va", None, None) == pytest.approx(16 / 9)
- assert _resolve_minimax_h3_aspect_ratio("ref2va", "auto", None) == pytest.approx(16 / 9)
+ assert resolve_minimax_h3_aspect_ratio("fl2va", "9:16", image) == pytest.approx(16 / 9)
+ assert resolve_minimax_h3_aspect_ratio("ref2va", None, None) == pytest.approx(16 / 9)
+ assert resolve_minimax_h3_aspect_ratio("ref2va", "auto", None) == pytest.approx(16 / 9)
with pytest.raises(ValueError, match="requires an explicit"):
- _resolve_minimax_h3_aspect_ratio("t2va", None, None)
+ resolve_minimax_h3_aspect_ratio("t2va", None, None)
with pytest.raises(ValueError, match="one of"):
- _resolve_minimax_h3_aspect_ratio("t2va", "2:1", None)
+ resolve_minimax_h3_aspect_ratio("t2va", "2:1", None)
def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override():
from PIL import Image
- from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape
- pipeline = object.__new__(MiniMaxH3Pipeline)
t2va_sampling = SimpleNamespace(
fps=24,
num_frames=1,
@@ -2741,7 +2226,7 @@ def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override():
extra_args={"duration": 5.0},
)
with pytest.raises(ValueError, match="requires an explicit aspect_ratio"):
- pipeline._resolve_shape("t2va", t2va_sampling, None)
+ resolve_minimax_h3_shape("t2va", t2va_sampling, None)
fl2va_sampling = SimpleNamespace(
fps=24,
@@ -2750,14 +2235,14 @@ def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override():
width=None,
extra_args={"duration": 5.0, "aspect_ratio": "9:16"},
)
- height, width, *_ = pipeline._resolve_shape("fl2va", fl2va_sampling, Image.new("RGB", (1280, 720)))
+ height, width, *_ = resolve_minimax_h3_shape("fl2va", fl2va_sampling, Image.new("RGB", (1280, 720)))
assert (height, width) == (768, 1344)
def test_g4_reference_image_boundaries_and_aspect_ratio():
from PIL import Image
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
_validate_reference_image,
)
@@ -2772,18 +2257,18 @@ def test_g4_reference_image_boundaries_and_aspect_ratio():
def test_g4_reference_image_file_format_and_size_contract(tmp_path):
from PIL import Image
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
- _load_image,
+ from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
+ load_minimax_h3_images,
)
valid_path = tmp_path / "reference.png"
Image.new("RGB", (256, 256)).save(valid_path)
- assert _load_image(valid_path).size == (256, 256)
+ assert load_minimax_h3_images(valid_path)[0].size == (256, 256)
invalid_path = tmp_path / "reference.bmp"
Image.new("RGB", (256, 256)).save(invalid_path)
with pytest.raises(ValueError, match="must use"):
- _load_image(invalid_path)
+ load_minimax_h3_images(invalid_path)
def test_g4_standalone_audio_duration_and_total_duration_contract():
@@ -2809,97 +2294,6 @@ def test_g4_standalone_audio_duration_and_total_duration_contract():
)
-def test_ref2va_video_soundtrack_does_not_consume_standalone_audio_budget():
- from PIL import Image
-
- from vllm_omni.diffusion.cache.cachedit.runtime import RequestScopedCacheDiTRuntime
- from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
- from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
- from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
- pipeline = object.__new__(MiniMaxH3Pipeline)
- torch.nn.Module.__init__(pipeline)
- pipeline.partition = "ref2va"
- pipeline.supported_tasks = frozenset({"ref2va"})
- pipeline.device = torch.device("cpu")
- pipeline.default_video_shift = 12.0
- pipeline.default_audio_shift = 3.0
- pipeline._resolve_sigma_positions = Mock(return_value=(None, 50))
- pipeline._quality_policy = MiniMaxH3QualityPolicy(None)
- pipeline._cache_dit_runtime = Mock(spec=RequestScopedCacheDiTRuntime)
- pipeline.text_encoder = object()
- pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long)))
- pipeline._prepare_reference_videos = Mock(return_value=[{"input_has_audio": True}])
- pipeline._encode_visual_conditions = Mock(return_value=(torch.ones(1, 96), [(1, 16, 16), (17, 48, 84)]))
- # A 2.35-second video soundtrack and a 15-second standalone reference
- # each satisfy their modality's separate 15-second duration budget.
- embedded = torch.full((188, 32), 1.0)
- standalone = torch.full((1200, 32), 2.0)
- pipeline._encode_reference_audio_conditions = Mock(return_value=(embedded, [94], standalone, [600]))
- sampling = OmniDiffusionSamplingParams(
- width=1344,
- height=768,
- fps=24,
- num_inference_steps=50,
- extra_args={"task": "ref2va", "duration": 15.0},
- )
-
- context = pipeline._prepare_request_inputs(
- prompt="A person carrying books down a street",
- multi_modal_data={
- "image": Image.new("RGB", (256, 256)),
- "video": ["reference.mp4"],
- "audio": (torch.zeros(1, 15 * 16000), 16000),
- },
- sampling=sampling,
- )
-
- assert context["audio_condition_lengths"] == [94, 600]
- assert [block["kind"] for block in context["ref_blocks"]] == ["image", "video_audio", "audio"]
- torch.testing.assert_close(context["audio_condition"], torch.cat([embedded, standalone]))
-
-
-def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch):
- from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
-
- pipeline = object.__new__(pipeline_module.MiniMaxH3Pipeline)
- pipeline.device = torch.device("cpu")
- monkeypatch.setattr(pipeline_module, "_dit_rank_world", lambda: (None, 1, 2))
-
- waveform = torch.zeros(1, 10)
- with pytest.raises(ValueError, match="max_duration_seconds must be positive"):
- pipeline._encode_audio_conditions_resident(
- [(waveform, 10)],
- max_duration_seconds=0,
- )
-
-
-def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
-
- pipeline = object.__new__(MiniMaxH3Pipeline)
- pipeline.device = torch.device("cpu")
- observed = []
-
- class FakeAudioVAE:
- def encode_waveform(self, waveform, sample_rate):
- observed.append((waveform.clone(), sample_rate))
- return waveform, waveform.shape[-1]
-
- pipeline.audio_vae = FakeAudioVAE()
- waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40)
-
- encoded, lengths = pipeline._encode_audio_conditions_resident(
- [(waveform, 10)],
- max_duration_seconds=2.5,
- )
-
- assert observed[0][0].shape == (1, 25)
- assert observed[0][1] == 10
- assert encoded.shape == (1, 25)
- assert lengths == [25]
-
-
@pytest.mark.parametrize(
("field", "value", "message"),
[
@@ -2963,9 +2357,8 @@ def _one_layer_text_encoder():
MiniMaxH3Qwen3VLQKVParallelLinear,
MiniMaxH3Qwen3VLRMSNorm,
)
- from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _SingleRankEncoderGroup
- group = _SingleRankEncoderGroup(0)
+ group = SimpleNamespace(rank_in_group=0, world_size=1, device_group=None, ranks=[0])
layer = nn.Module()
layer.self_attn = nn.Module()
layer.mlp = nn.Module()
@@ -3103,7 +2496,7 @@ def _pad_seq_len_batch(pad_seq_len):
return DiffusionRequestBatch(
[
OmniDiffusionRequest(
- prompt="pinned packed length",
+ prompt=_encoder_prompt("pinned packed length"),
sampling_params=sampling,
request_id="pad-seq-len",
)
@@ -3189,3 +2582,431 @@ def test_pad_seq_len_request_validation_accepts_aligned_values():
assert _resolve_pad_seq_len(None) is None
assert _resolve_pad_seq_len(54080) == 54080
assert _resolve_pad_seq_len(MINIMAX_H3_MAX_PAD_SEQ_LEN) == MINIMAX_H3_MAX_PAD_SEQ_LEN
+
+
+def _turbo_spec(filename: str) -> TurboSpec:
+ """Build the spec a loaded artifact of this name would carry."""
+ spec = parse_turbo_filename(filename)
+ assert spec is not None
+ return spec
+
+
+def test_joint_video_output_is_quantized_before_transfer():
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ _minimax_h3_post_process,
+ _prepare_minimax_h3_video_output,
+ )
+
+ video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5)
+ expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1)
+ prepared = _prepare_minimax_h3_video_output(video)
+
+ assert prepared.dtype == torch.uint8
+ assert prepared.is_contiguous()
+ assert prepared.shape == (2, 2, 4, 5, 3)
+ torch.testing.assert_close(prepared, expected, rtol=0, atol=0)
+
+ result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6)))
+ assert result["video"][0].dtype == np.uint8
+ np.testing.assert_array_equal(result["video"][0], expected[0].numpy())
+
+
+def test_no_offload_keeps_text_encoder_resident():
+ 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=False,
+ enable_layerwise_offload=False,
+ )
+ pipeline.text_encoder = Mock()
+ expected = torch.ones(2, 3)
+ pipeline.text_encoder.encode_ids.return_value = expected
+ input_ids = torch.tensor([1, 2])
+
+ actual = pipeline._encode_text_hidden(input_ids, {})
+
+ assert actual is expected
+ pipeline.text_encoder.load_to_device.assert_called_once_with()
+ pipeline.text_encoder.offload_to_cpu.assert_not_called()
+
+
+def test_model_offload_uses_hooked_text_encoder_call():
+ 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,
+ )
+ expected = torch.ones(2, 3)
+ pipeline.text_encoder = Mock(return_value=expected)
+ pipeline._model_cpu_offload_modules = [pipeline.text_encoder]
+ input_ids = torch.tensor([1, 2])
+ vision_kwargs = {"pixel_values": torch.ones(1, 4)}
+
+ actual = pipeline._encode_text_hidden(input_ids, vision_kwargs)
+
+ assert actual is expected
+ pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs)
+ pipeline.text_encoder.load_to_device.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ ("enable_layerwise", "enable_distributed"),
+ [(True, False), (False, True)],
+)
+def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.od_config = SimpleNamespace(
+ diffusion_offload_config=None,
+ enable_cpu_offload=False,
+ enable_layerwise_offload=enable_layerwise,
+ enable_distributed_layerwise_offload=enable_distributed,
+ )
+ 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.offload_to_cpu.assert_called_once_with()
+
+ vae = Mock()
+ with pipeline._component_on_device(vae):
+ pass
+ vae.load_to_device.assert_called_once_with()
+ vae.offload_to_cpu.assert_called_once_with()
+
+
+def test_layerwise_encoder_selection_releases_text_encoder():
+ 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=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
+ diffusion_offload_config={
+ "mode": "layer",
+ "components": ["text_encoder"],
+ },
+ )
+ 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.offload_to_cpu.assert_called_once_with()
+
+
+def test_layerwise_dit_only_keeps_text_encoder_resident():
+ 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=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
+ diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+ )
+ 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.offload_to_cpu.assert_not_called()
+
+
+def test_dit_encoder_selection_keeps_vae_resident():
+ 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=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
+ diffusion_offload_config={
+ "mode": "layer",
+ "components": ["dit", "text_encoder"],
+ },
+ )
+ component = Mock()
+
+ with pipeline._component_on_device(component):
+ pass
+
+ component.load_to_device.assert_not_called()
+ component.offload_to_cpu.assert_not_called()
+
+
+def test_keyframe_encode_pins_and_restores_cudnn_settings():
+ from vllm_omni.diffusion.models.minimax_h3.vae import (
+ _minimax_h3_keyframe_encode_context,
+ )
+
+ cudnn = torch.backends.cudnn
+ original = (
+ cudnn.enabled,
+ cudnn.benchmark,
+ cudnn.deterministic,
+ cudnn.allow_tf32,
+ )
+ try:
+ cudnn.enabled = False
+ cudnn.benchmark = True
+ cudnn.deterministic = False
+ cudnn.allow_tf32 = False
+
+ with _minimax_h3_keyframe_encode_context(torch.device("cuda")):
+ assert cudnn.enabled
+ assert not cudnn.benchmark
+ assert cudnn.deterministic
+ assert cudnn.allow_tf32
+
+ assert not cudnn.enabled
+ assert cudnn.benchmark
+ assert not cudnn.deterministic
+ assert not cudnn.allow_tf32
+ finally:
+ (
+ cudnn.enabled,
+ cudnn.benchmark,
+ cudnn.deterministic,
+ cudnn.allow_tf32,
+ ) = original
+
+
+@pytest.mark.parametrize("fail_encode", [False, True])
+def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
+ from PIL import Image
+
+ from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+ events = []
+ input_image = Image.new("RGB", (32, 32))
+
+ @contextmanager
+ def track_context(device):
+ assert device == torch.device("cpu")
+ events.append("enter")
+ try:
+ yield
+ finally:
+ events.append("exit")
+
+ class FakeModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.ones(1))
+ self.parallel_tiling = True
+
+ def encode_images(self, image, *, use_fp16_latent):
+ assert events == ["enter"]
+ assert image is input_image
+ assert use_fp16_latent
+ events.append("encode")
+ if fail_encode:
+ raise RuntimeError("encode failed")
+ return [torch.ones(1, 1, 2, 2)]
+
+ video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
+ torch.nn.Module.__init__(video_vae)
+ video_vae.model = FakeModel()
+ video_vae.config_dict = {
+ "latent_channels": 1,
+ "latents_mean": [0.0],
+ "latents_std": [1.0],
+ }
+ monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context)
+
+ if fail_encode:
+ with pytest.raises(RuntimeError, match="encode failed"):
+ video_vae.encode_image(input_image)
+ else:
+ rows = video_vae.encode_image(input_image)
+ torch.testing.assert_close(rows, torch.ones(1, 4))
+
+ assert events == ["enter", "encode", "exit"]
+ assert video_vae.model.parallel_tiling
+
+
+def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+
+ class FakeEncoder:
+ instances = []
+
+ def __init__(self, **kwargs):
+ self.pushes = []
+ self.kwargs = kwargs
+ self.__class__.instances.append(self)
+
+ def push(self, frames):
+ self.pushes.append(np.array(frames, copy=True))
+
+ def finish(self):
+ return b"mp4"
+
+ def abort(self):
+ raise AssertionError("unexpected abort")
+
+ class FakeAudioVAE:
+ def decode_latent(self, latent):
+ return torch.zeros(1, 1, 2)
+
+ class FakeVideoVAE:
+ def decode_with_chunks(self, latent, *, on_chunk):
+ for value in (0.0, 0.25, 0.5):
+ on_chunk(torch.full((1, 3, 3, 2, 2), value))
+
+ pipeline = object.__new__(mod.MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.audio_vae = FakeAudioVAE()
+ pipeline.video_vae = FakeVideoVAE()
+ pipeline.device = torch.device("cpu")
+ monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder",
+ FakeEncoder,
+ )
+
+ output = pipeline.decode_to_mp4(
+ torch.zeros(1),
+ torch.zeros(1),
+ height=2,
+ width=2,
+ batch_frames=4,
+ )
+
+ assert output == b"mp4"
+ encoder = FakeEncoder.instances[-1]
+ assert [frames.shape[0] for frames in encoder.pushes] == [6, 3]
+ assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8))
+ assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8))
+ assert encoder.pushes[1].shape == (3, 2, 2, 3)
+
+
+def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch):
+ """A client's encoder options must survive the worker-side pre-encode path."""
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+
+ class FakeEncoder:
+ instances = []
+
+ def __init__(self, **kwargs):
+ self.kwargs = kwargs
+ self.__class__.instances.append(self)
+
+ def push(self, frames):
+ del frames
+
+ def finish(self):
+ return b"mp4"
+
+ def abort(self):
+ raise AssertionError("unexpected abort")
+
+ class FakeAudioVAE:
+ def decode_latent(self, latent):
+ return torch.zeros(1, 1, 2)
+
+ class FakeVideoVAE:
+ def decode_with_chunks(self, latent, *, on_chunk):
+ on_chunk(torch.zeros(1, 3, 1, 2, 2))
+
+ pipeline = object.__new__(mod.MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.audio_vae = FakeAudioVAE()
+ pipeline.video_vae = FakeVideoVAE()
+ pipeline.device = torch.device("cpu")
+ monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
+ monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder)
+
+ pipeline.decode_to_mp4(
+ torch.zeros(1),
+ torch.zeros(1),
+ height=2,
+ width=2,
+ video_codec_options={"preset": "ultrafast"},
+ )
+
+ assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"}
+
+
+@pytest.mark.parametrize(
+ ("codec_extra", "expected"),
+ [
+ ({}, {"preset": "ultrafast", "threads": "0"}),
+ ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}),
+ ({"video_codec_options": {}}, {}),
+ ({"video_codec_options": None}, None),
+ ],
+)
+@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)])
+def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.request import OmniDiffusionRequest
+ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline.supported_tasks = frozenset({"t2va"})
+ pipeline.default_video_shift = 12.0
+ pipeline.default_audio_shift = 3.0
+ pipeline.device = torch.device("cpu")
+ pipeline.od_config = SimpleNamespace()
+ pipeline.load_text_encoder = False
+ pipeline.load_vae_encoder = False
+ pipeline.text_encoder = object()
+ pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long)))
+ pipeline._quality_policy = Mock()
+ pipeline._cache_dit_runtime = Mock()
+ pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+ pipeline.decode_to_mp4 = Mock(return_value=b"mp4")
+ sampling = OmniDiffusionSamplingParams(
+ width=1344,
+ height=768,
+ fps=24,
+ num_frames=124,
+ num_inference_steps=50,
+ extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra},
+ )
+ batch = DiffusionRequestBatch(
+ [OmniDiffusionRequest(prompt=_encoder_prompt("test"), sampling_params=sampling, request_id="codec-defaults")]
+ )
+
+ pipeline.forward(batch)
+
+ assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected
+ assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames
+
+
+def test_video_codec_options_are_normalized_for_the_encoder():
+ from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options
+
+ assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == {
+ "preset": "ultrafast",
+ "threads": "0",
+ }
+ assert normalize_video_codec_options(None) is None
+ with pytest.raises(ValueError, match="video_codec_options"):
+ normalize_video_codec_options("preset=ultrafast")
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py
new file mode 100644
index 000000000..510a7b33e
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py
@@ -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]
+ )
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py
index 31a8916a6..24206952d 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py
@@ -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()
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
index c7e53ca62..87b04613a 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
@@ -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)
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py
index 1b123929c..e202c0d0e 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py
@@ -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)
diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py
index 2f10ea33e..16dc85115 100644
--- a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py
+++ b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py
@@ -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)
diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py
index 4f3a81491..3fd141dc1 100644
--- a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py
+++ b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py
@@ -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)
diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py b/tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py
new file mode 100644
index 000000000..1a3acbb48
--- /dev/null
+++ b/tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py
@@ -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
diff --git a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py
index bb95363fd..b01f275eb 100644
--- a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py
+++ b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py
@@ -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
diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
index 526eeeb15..58a5e961b 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
@@ -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]] = []
diff --git a/tests/diffusion/quantization/test_mxfp4_checkpoint.py b/tests/diffusion/quantization/test_mxfp4_checkpoint.py
new file mode 100644
index 000000000..e183df078
--- /dev/null
+++ b/tests/diffusion/quantization/test_mxfp4_checkpoint.py
@@ -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
diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py
index 240ee2509..c51b2013b 100644
--- a/tests/diffusion/quantization/test_mxfp4_config.py
+++ b/tests/diffusion/quantization/test_mxfp4_config.py
@@ -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)
diff --git a/tests/diffusion/quantization/test_wan_mxfp4_loading.py b/tests/diffusion/quantization/test_wan_mxfp4_loading.py
new file mode 100644
index 000000000..fc7a43d43
--- /dev/null
+++ b/tests/diffusion/quantization/test_wan_mxfp4_loading.py
@@ -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
diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py
index 23ff796e2..6f37f823d 100644
--- a/tests/diffusion/test_diffusion_config_propagation.py
+++ b/tests/diffusion/test_diffusion_config_propagation.py
@@ -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.
diff --git a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py
index ad659541d..acf717e42 100644
--- a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py
+++ b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py
@@ -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 = {
diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py
index fb2f8cbc8..1fca04f99 100644
--- a/tests/e2e/accuracy/test_qwen_image.py
+++ b/tests/e2e/accuracy/test_qwen_image.py
@@ -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"
diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py
index fe42dd9bf..480381f0a 100644
--- a/tests/e2e/features/comfyui/test_comfyui_integration.py
+++ b/tests/e2e/features/comfyui/test_comfyui_integration.py
@@ -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)
diff --git a/tests/e2e/offline_inference/test_moss_tts_expansion.py b/tests/e2e/offline_inference/test_moss_tts_expansion.py
index 81a64fcd9..b899252b4 100644
--- a/tests/e2e/offline_inference/test_moss_tts_expansion.py
+++ b/tests/e2e/offline_inference/test_moss_tts_expansion.py
@@ -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 = [
diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py
index 81d97cd74..0f0bf59c7 100644
--- a/tests/e2e/offline_inference/test_moss_tts_realtime.py
+++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py
@@ -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, "你好，这是语音合成测试。")
diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py
index b84e5a0ff..678fc7e55 100644
--- a/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py
+++ b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py
@@ -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)
diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py
index 5417d2b16..a1d045b72 100644
--- a/tests/e2e/online_serving/test_minicpmo_4_5.py
+++ b/tests/e2e/online_serving/test_minicpmo_4_5.py
@@ -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.
diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py
index 598326a1e..bd38a7fc7 100644
--- a/tests/e2e/online_serving/test_moss_tts_expansion.py
+++ b/tests/e2e/online_serving/test_moss_tts_expansion.py
@@ -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,
diff --git a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
index 1e5e59aed..389b0b6ba 100644
--- a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
+++ b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
@@ -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)
diff --git a/tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py b/tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py
new file mode 100644
index 000000000..d8bd6302e
--- /dev/null
+++ b/tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py
@@ -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"]
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index fb2c01df2..acac20329 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/tests/engine/test_async_omni_engine_stage_init.py
@@ -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
diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py
index beddd2106..3b6ee7e7b 100644
--- a/tests/engine/test_orchestrator.py
+++ b/tests/engine/test_orchestrator.py
@@ -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()
diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py
index ca2a26b70..25e662d27 100644
--- a/tests/engine/test_orchestrator_kv_sender_info.py
+++ b/tests/engine/test_orchestrator_kv_sender_info.py
@@ -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]
diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py
index 8618061d8..419be9305 100644
--- a/tests/engine/test_output_processor.py
+++ b/tests/engine/test_output_processor.py
@@ -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]
diff --git a/tests/entrypoints/openai_api/test_duplex_api_server.py b/tests/entrypoints/openai_api/test_duplex_api_server.py
index b528cbc2a..9cf58b114 100644
--- a/tests/entrypoints/openai_api/test_duplex_api_server.py
+++ b/tests/entrypoints/openai_api/test_duplex_api_server.py
@@ -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
# --------------------------------------------------------------------------- #
diff --git a/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py b/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py
new file mode 100644
index 000000000..d69597634
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py
@@ -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())
diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py
index aebbb646c..c92d876d9 100644
--- a/tests/entrypoints/openai_api/test_serving_speech.py
+++ b/tests/entrypoints/openai_api/test_serving_speech.py
@@ -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)
diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py
index 39c9e66d9..1d34b6e0b 100644
--- a/tests/entrypoints/openai_api/test_serving_speech_stream.py
+++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py
@@ -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
diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py
index 32d7d9988..ea57d07ab 100644
--- a/tests/entrypoints/test_omni_entrypoints.py
+++ b/tests/entrypoints/test_omni_entrypoints.py
@@ -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)
diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py
index c588b7a38..e0c381bf3 100644
--- a/tests/examples/offline_inference/test_text_to_video_prompts.py
+++ b/tests/examples/offline_inference/test_text_to_video_prompts.py
@@ -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")],
diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py
index 7dc7a6ab3..5df9dfa7e 100644
--- a/tests/helpers/assertions.py
+++ b/tests/helpers/assertions.py
@@ -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:
diff --git a/tests/helpers/client.py b/tests/helpers/client.py
index ed5158b92..e1042321a 100644
--- a/tests/helpers/client.py
+++ b/tests/helpers/client.py
@@ -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]
diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py
index e9b1a771c..55956a6f2 100644
--- a/tests/helpers/tests/test_assertions.py
+++ b/tests/helpers/tests/test_assertions.py
@@ -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)
diff --git a/tests/model_executor/models/breeze_tts_2/test_instruction.py b/tests/model_executor/models/breeze_tts_2/test_instruction.py
index 3817ccb30..c7f2fef66 100644
--- a/tests/model_executor/models/breeze_tts_2/test_instruction.py
+++ b/tests/model_executor/models/breeze_tts_2/test_instruction.py
@@ -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
diff --git a/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py b/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py
new file mode 100644
index 000000000..0edaabc6f
--- /dev/null
+++ b/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py
@@ -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
diff --git a/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py b/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py
new file mode 100644
index 000000000..8cd389d17
--- /dev/null
+++ b/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py
@@ -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
diff --git a/tests/model_executor/models/breeze_tts_2/test_streaming.py b/tests/model_executor/models/breeze_tts_2/test_streaming.py
index 045444ec3..a630b90bc 100644
--- a/tests/model_executor/models/breeze_tts_2/test_streaming.py
+++ b/tests/model_executor/models/breeze_tts_2/test_streaming.py
@@ -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 == {}
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
index 7dfbe6e74..74fc5059b 100644
--- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
@@ -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)
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py
new file mode 100644
index 000000000..1e8eff599
--- /dev/null
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py
@@ -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()
diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py
index 959156b8f..555e72e46 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py
@@ -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)):
diff --git a/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py
new file mode 100644
index 000000000..f5d799dd7
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py
@@ -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)
diff --git a/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py b/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py
new file mode 100644
index 000000000..916213318
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py
@@ -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)
diff --git a/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py b/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py
new file mode 100644
index 000000000..1bbfaeeef
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py
@@ -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
diff --git a/tests/model_executor/models/moss_tts/test_realtime_prompt.py b/tests/model_executor/models/moss_tts/test_realtime_prompt.py
new file mode 100644
index 000000000..cb36f6dd8
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_realtime_prompt.py
@@ -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),
+ )
diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py
index 54efc99e4..9ebad492b 100644
--- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py
+++ b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py
@@ -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",
diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py
index 7b624ba11..0d7cd8175 100644
--- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py
+++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py
@@ -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
diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py
index 0055d1704..f4635a81f 100644
--- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py
+++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py
@@ -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
diff --git a/tests/model_executor/models/test_minimax_h3_encoder_payload.py b/tests/model_executor/models/test_minimax_h3_encoder_payload.py
new file mode 100644
index 000000000..4c9a7137b
--- /dev/null
+++ b/tests/model_executor/models/test_minimax_h3_encoder_payload.py
@@ -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"
diff --git a/tests/model_executor/stage_input_processors/test_minimax_h3.py b/tests/model_executor/stage_input_processors/test_minimax_h3.py
index 9fda89953..aac16d29f 100644
--- a/tests/model_executor/stage_input_processors/test_minimax_h3.py
+++ b/tests/model_executor/stage_input_processors/test_minimax_h3.py
@@ -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):
diff --git a/tests/platforms/npu/quant/test_mxfp4_step_fallback.py b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py
new file mode 100644
index 000000000..8431ebf17
--- /dev/null
+++ b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py
@@ -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)
diff --git a/vllm_omni/attention/fish_kvcache_attn.py b/vllm_omni/attention/fish_kvcache_attn.py
index e7235168c..6e402ca54 100644
--- a/vllm_omni/attention/fish_kvcache_attn.py
+++ b/vllm_omni/attention/fish_kvcache_attn.py
@@ -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()
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py
index 9d6013e80..349d80fc3 100644
--- a/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py
@@ -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/-00000-of-00001.parquet``, with ```` 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/-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 "" 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
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py
index 629d8e877..b45686fb8 100644
--- a/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py
@@ -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
diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py
index bee910c90..5a5e74b0d 100644
--- a/vllm_omni/benchmarks/patch/patch.py
+++ b/vllm_omni/benchmarks/patch/patch.py
@@ -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,
diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py
index 52266b64e..17358906f 100644
--- a/vllm_omni/clients/duplex.py
+++ b/vllm_omni/clients/duplex.py
@@ -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)
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 0b8723c0e..c13935489 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -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")
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index c96f514e8..24453ae88 100644
--- a/vllm_omni/config/stage_config.py
+++ b/vllm_omni/config/stage_config.py
@@ -309,10 +309,8 @@ class PipelineConfig:
diffusers_class_name: str | None = None
diffusers_class_aliases: tuple[str, ...] = ()
endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
- # Dotted path of the model's ``DuplexModelPlugin``. A pipeline is a duplex
- # model iff this is set: ``vllm-omni serve`` then always runs it through
- # ``DuplexOmni`` (every served surface runs on a duplex session) and the engine hosts one
- # ``DuplexOrchestrator`` with the plugin loaded.
+ # Dotted path of the model's ``DuplexModelPlugin``. Online serving uses
+ # DuplexOmni only when the deploy configuration selects session_mode: duplex.
duplex_plugin: str | None = None
# Legacy duplex wiring of the models that are not ported to the plugin
# framework yet (PersonaPlex, Nemotron VoiceChat). Nothing reads them: a
@@ -390,6 +388,9 @@ class StageDeployConfig:
num_replicas: int = 1
env: dict[str, Any] | None = None
+ # False opts this stage out of pipeline-wide async chunking.
+ async_chunk: bool | None = None
+
# Inter-stage connector wiring and request defaults.
output_connectors: dict[str, str] | None = None
input_connectors: dict[str, str] | None = None
@@ -587,6 +588,7 @@ class DeployConfig:
_STAGE_RESERVED_KEYS = frozenset(
{
+ "async_chunk",
"stage_id",
"devices",
"num_replicas",
@@ -645,6 +647,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
if name in flat_args:
kwargs[name] = flat_args.pop(name)
+ kwargs["async_chunk"] = stage_data.get("async_chunk")
kwargs["output_connectors"] = stage_data.get("output_connectors")
kwargs["input_connectors"] = stage_data.get("input_connectors")
kwargs["default_sampling_params"] = stage_data.get("default_sampling_params")
@@ -888,6 +891,28 @@ def _resolve_execution_mode(
return _EXECUTION_TYPE_TO_STAGE_WORKER.get(execution_type, (StageType.LLM, None))
+def resolve_stage_async_chunk(deploy: DeployConfig, stage: StageDeployConfig | None) -> bool:
+ return bool(deploy.async_chunk and (stage is None or stage.async_chunk is not False))
+
+
+def validate_stage_async_chunk_edges(pipeline: PipelineConfig, deploy: DeployConfig) -> None:
+ stages = {stage.stage_id: stage for stage in pipeline.stages}
+ deploy_by_id = {stage.stage_id: stage for stage in deploy.stages}
+ for consumer in pipeline.stages:
+ for source_id in consumer.input_sources:
+ producer = stages[source_id]
+ if not producer.async_chunk_process_next_stage_input_func:
+ continue
+ producer_async = resolve_stage_async_chunk(deploy, deploy_by_id.get(source_id))
+ consumer_async = resolve_stage_async_chunk(deploy, deploy_by_id.get(consumer.stage_id))
+ if producer_async != consumer_async:
+ raise ValueError(
+ f"Pipeline {pipeline.model_type!r} has incompatible async_chunk settings on "
+ f"connector edge {source_id} -> {consumer.stage_id}. "
+ "Set the same async_chunk mode on both stages or disable pipeline-wide async_chunk."
+ )
+
+
def _select_processor_funcs(
ps: StagePipelineConfig,
async_chunk: bool,
@@ -964,9 +989,7 @@ def _build_engine_args(
continue
engine_args[k] = v
engine_args.update(ds.engine_extras)
- # Materialize the resolved pipeline-wide async_chunk value into every
- # stage so explicit False overrides do not get lost downstream.
- engine_args["async_chunk"] = bool(deploy.async_chunk)
+ engine_args["async_chunk"] = resolve_stage_async_chunk(deploy, ds)
engine_args["session_mode"] = deploy.session_mode
if deploy.session_mode == "duplex":
# The engine admission limit is also the authoritative capacity for
@@ -1060,11 +1083,12 @@ def merge_pipeline_deploy(
"Either set async_chunk=False or implement an async-chunk producer on the pipeline."
)
+ validate_stage_async_chunk_edges(pipeline, deploy)
result: list[StageConfig] = []
for ps in pipeline.stages:
ds = deploy_by_id.get(ps.stage_id)
stage_type, worker_type = _resolve_execution_mode(ps.execution_type)
- input_proc, next_stage_proc = _select_processor_funcs(ps, deploy.async_chunk)
+ input_proc, next_stage_proc = _select_processor_funcs(ps, resolve_stage_async_chunk(deploy, ds))
engine_args = _build_engine_args(ps, ds, pipeline, deploy, next_stage_proc)
# Downstream stages may share a multimodal wrapper class without owning
# an encoder. Do not make vLLM profile dummy multimodal inputs for them.
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 329e8ad81..3d866fda7 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -16,7 +16,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMSchedu
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.core.sched.request_queue import create_request_queue
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason
from vllm.v1.metrics.perf import PerfStats
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
@@ -116,6 +116,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
self._init_omni_io_scheduling_state()
# Snapshot prompt length for each streaming input update
self._new_prompt_len_snapshot: dict[str, int] = {}
+ # Streaming sessions finished because their next prompt extension
+ # would exceed max_model_len: request_id -> (client_index, reason).
+ # Drained into an explicit FinishReason.ERROR output on the next
+ # update_from_output so the client learns why the session ended.
+ self._streaming_context_overflow: dict[str, tuple[int, str]] = {}
def _get_confirmed_num_computed_tokens(self, request: Request) -> int:
"""num_computed_tokens minus async placeholders (KV actually on GPU)."""
@@ -704,6 +709,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
failed_kv_load_req_ids,
outputs,
)
+ self._emit_streaming_context_overflow_outputs(outputs)
if self.chunk_transfer_adapter is not None:
for request in failed_requests:
self.chunk_transfer_adapter.cleanup_receiver(request.request_id)
@@ -883,11 +889,143 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
self._release_replaced_streaming_prompt_cache(session)
self._replace_streaming_session(session, update)
return
+ if self._streaming_update_overflows(session, update):
+ return
session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
super()._update_request_as_session(session, update)
if hasattr(update, "model_intermediate_buffer"):
session.model_intermediate_buffer = update.model_intermediate_buffer
+ # Prefix of the stop_reason carried by the FinishReason.ERROR output, so the
+ # serving side can map it to a stable error code.
+ STREAMING_CONTEXT_OVERFLOW_STOP_REASON = "context_length_exceeded"
+
+ def _streaming_update_overflows(self, session: Request, update: StreamingUpdate) -> bool:
+ """Finish a streaming session whose next extension cannot fit the model.
+
+ Upstream ``_update_request_as_session`` appends the update to the
+ session prompt without checking ``max_model_len``. The worker's input
+ batch then fails to copy the prompt (``could not broadcast input array
+ from shape (N,) into shape (max_model_len,)``) and the EngineCore dies,
+ taking every session on the replica with it. A native duplex session
+ grows by tens to hundreds of tokens per second of input, so long
+ sessions reach this point in normal use.
+
+ A prompt that fills the model exactly is over the line as well: the
+ session samples at least one listen/speak token after every append,
+ and upstream's running-request budget
+ ``max_model_len - num_computed_tokens - num_sampled_tokens_per_step``
+ then goes negative, which the ``num_new_tokens == 0`` guard in
+ ``schedule()`` does not catch (``allocate_slots`` dies, or the worker
+ asserts ``max_model_len + 1`` sampled positions). So the extended
+ prompt must leave room for the tokens sampled in one step.
+
+ The update is dropped and only this request is finished, right here:
+ a parked session does not make the engine schedule, so deferring the
+ finish to the next ``schedule()`` would leave the client waiting. The
+ reason is emitted with the terminal output (see
+ :meth:`_emit_streaming_context_overflow_outputs`).
+ """
+ max_model_len = getattr(self, "max_model_len", None)
+ if max_model_len is None:
+ model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
+ max_model_len = getattr(model_config, "max_model_len", None)
+ if not max_model_len:
+ return False
+ new_tokens = len(update.prompt_token_ids or ())
+ # The extended prompt is the current prompt plus the computed output
+ # tokens upstream keeps, then the update: num_computed_tokens covers
+ # both when the prompt was fully computed.
+ projected = max(int(session.num_prompt_tokens), int(session.num_computed_tokens)) + new_tokens
+ # Room for the tokens one step samples on top of the prompt (1 without
+ # speculative decoding). __new__-built test schedulers carry no
+ # num_sampled_tokens_per_step.
+ sample_room = max(1, int(getattr(self, "num_sampled_tokens_per_step", 1) or 1))
+ if projected + sample_room <= int(max_model_len):
+ return False
+ reason = (
+ f"{self.STREAMING_CONTEXT_OVERFLOW_STOP_REASON}: streaming session prompt would grow to "
+ f"{projected} tokens, leaving no room to sample within max_model_len {int(max_model_len)}"
+ )
+ logger.error(
+ "[Omni] %s: %s; finishing the request instead of extending it",
+ session.request_id,
+ reason,
+ )
+ overflow = getattr(self, "_streaming_context_overflow", None)
+ if overflow is None:
+ overflow = self._streaming_context_overflow = {}
+ overflow[session.request_id] = (int(getattr(session, "client_index", 0) or 0), reason)
+ if session.is_finished():
+ # Reached from ``_handle_stopped_request`` with a queued update.
+ # ``update_from_output`` frees every request that call reports as
+ # finished, so finishing the session here too would free it twice.
+ # ``_handle_stopped_request`` below takes it out of admission and
+ # leaves the single free to the caller.
+ return True
+ self.finish_requests((session.request_id,), RequestStatus.FINISHED_ERROR)
+ return True
+
+ def _handle_stopped_request(self, request: Request) -> bool:
+ """Do not resume a session whose queued update overflowed the model.
+
+ Upstream pops one queued ``StreamingUpdate``, applies it through
+ ``_update_request_as_session`` and then re-enqueues the request
+ unconditionally. When that update overflows, the request must not go
+ back into the waiting queue: it is terminal, and admission raises
+ ``RuntimeError: Invalid request status`` on anything that is neither
+ WAITING nor PREEMPTED, which would kill the EngineCore this guard
+ exists to keep alive. The same holds for a session whose overflow was
+ recorded before this call and that upstream still reports as resumed.
+ """
+ finished = super()._handle_stopped_request(request)
+ if finished:
+ return True
+ overflow = getattr(self, "_streaming_context_overflow", None)
+ if not overflow or request.request_id not in overflow:
+ return False
+ # Whether the overflow was recorded by this call's queued update or
+ # earlier makes no difference: a session in the overflow map is
+ # terminal, and upstream has just put it back into admission.
+ self._finish_overflowed_streaming_session(request)
+ return True
+
+ def _finish_overflowed_streaming_session(self, request: Request) -> None:
+ """Take a terminal session back out of admission.
+
+ Queues and status only. ``update_from_output`` frees every request
+ ``_handle_stopped_request`` reports as finished, so freeing here as
+ well deletes it from ``self.requests`` twice (``KeyError`` in
+ ``_free_blocks``) and skips the caller's input-coordinator cleanup.
+ """
+ self.waiting.remove_requests((request,))
+ self.skipped_waiting.remove_requests((request,))
+ if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
+ self.num_waiting_for_streaming_input -= 1
+ request.status = RequestStatus.FINISHED_ERROR
+ request.resumable = False
+
+ def _emit_streaming_context_overflow_outputs(self, outputs: dict[int, list[EngineCoreOutput]]) -> None:
+ """Turn recorded context overflows into explicit error outputs.
+
+ Without this the finished session would only get the synthesized
+ ``FinishReason.ABORT`` output, which a client cannot tell apart from
+ its own cancel.
+ """
+ overflow = getattr(self, "_streaming_context_overflow", None)
+ if not overflow:
+ return
+ for request_id, (client_index, reason) in list(overflow.items()):
+ outputs.setdefault(client_index, []).append(
+ OmniEngineCoreOutput(
+ request_id=request_id,
+ new_token_ids=[],
+ finish_reason=FinishReason.ERROR,
+ stop_reason=reason,
+ )
+ )
+ overflow.clear()
+
def _free_request(
self, request: Request, delay_free_blocks: bool = False
) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
diff --git a/vllm_omni/deploy/breeze_tts_2.yaml b/vllm_omni/deploy/breeze_tts_2.yaml
index 1d1b85391..fd75dd69e 100644
--- a/vllm_omni/deploy/breeze_tts_2.yaml
+++ b/vllm_omni/deploy/breeze_tts_2.yaml
@@ -32,8 +32,12 @@ stages:
temperature: 0.0
top_p: 1.0
top_k: -1
+ # Match upstream's penalty on generated codebook-0 history.
+ repetition_penalty: 1.1
max_tokens: 2048
- min_tokens: 1
+ # min_tokens also masks the tokenizer's EOS 1 despite ignore_eos.
+ # Keep valid codec token 1 and an immediate backbone EOS available.
+ min_tokens: 0
- stage_id: 1
max_num_seqs: 8
diff --git a/vllm_omni/deploy/fish_qwen3_omni.yaml b/vllm_omni/deploy/fish_qwen3_omni.yaml
index f10ed1656..0f4866a1f 100644
--- a/vllm_omni/deploy/fish_qwen3_omni.yaml
+++ b/vllm_omni/deploy/fish_qwen3_omni.yaml
@@ -64,6 +64,15 @@ stages:
repetition_penalty: 1.0
platforms:
+ npu:
+ # Ascend paged attention only accepts 128-token cache blocks
+ # (AscendAttentionBackend.get_supported_kernel_block_sizes() == [128]),
+ # while vLLM's CacheConfig default is 16. Without this override stage 0
+ # aborts during KV-cache init with:
+ # ValueError: No common block size for 16.
+ stages:
+ - stage_id: 0
+ block_size: 128
xpu:
stages:
# Verified on Intel Arc B60.
diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml
index 5c26fc2c6..a04cff400 100644
--- a/vllm_omni/deploy/minicpmo_4_5.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5.yaml
@@ -1,8 +1,7 @@
# MiniCPM-o 4.5 three-stage deploy on one large-memory GPU.
pipeline: minicpmo_4_5
async_chunk: true
-# session_mode: duplex makes vllm-omni serve run this model duplex-only
-# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default.
+# Use minicpmo_4_5_turn.yaml to select the ordinary turn-based serving engine.
session_mode: duplex
# Keep the active window aligned with the session and stage capacities below.
# A smaller window serializes first-packet generation for concurrent chat turns.
diff --git a/vllm_omni/deploy/minicpmo_4_5_turn.yaml b/vllm_omni/deploy/minicpmo_4_5_turn.yaml
new file mode 100644
index 000000000..8dada1257
--- /dev/null
+++ b/vllm_omni/deploy/minicpmo_4_5_turn.yaml
@@ -0,0 +1,3 @@
+# MiniCPM-o 4.5 using the turn-based engine and ordinary HTTP serving.
+base_config: minicpmo_4_5.yaml
+session_mode: turn
diff --git a/vllm_omni/deploy/minimax_h3_disaggregated.yaml b/vllm_omni/deploy/minimax_h3_disaggregated.yaml
index 1c1ee6c40..25a561c27 100644
--- a/vllm_omni/deploy/minimax_h3_disaggregated.yaml
+++ b/vllm_omni/deploy/minimax_h3_disaggregated.yaml
@@ -1,4 +1,4 @@
-# MiniMax H3: vLLM Qwen3-VL text encoder + diffusion video/audio generation.
+# MiniMax H3: unified text/video/audio encoder stage + diffusion generation stage.
pipeline: minimax_h3_disaggregated
async_chunk: false
trust_remote_code: true
@@ -8,7 +8,16 @@ stages:
- stage_id: 0
devices: "0,1"
max_num_seqs: 1
+ gpu_memory_utilization: 0.6
tensor_parallel_size: 2
+ hf_overrides:
+ minimax_h3_encoder_components:
+ text_encoder:
+ parallel_mode: tp
+ video_vae:
+ parallel_mode: patch
+ audio_vae:
+ parallel_mode: leader
default_sampling_params:
max_tokens: 1
temperature: 0.0
@@ -19,6 +28,7 @@ stages:
max_num_seqs: 1
model_loaded:
text_encoder: false
+ vae_encoder: false
parallel_config:
tensor_parallel_size: 1
ulysses_degree: 4
diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
index d7e84db42..ec11796b6 100644
--- a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
+++ b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
@@ -8,7 +8,16 @@ stages:
- stage_id: 0
devices: "0,1"
max_num_seqs: 1
+ gpu_memory_utilization: 0.6
tensor_parallel_size: 2
+ hf_overrides:
+ minimax_h3_encoder_components:
+ text_encoder:
+ parallel_mode: tp
+ video_vae:
+ parallel_mode: patch
+ audio_vae:
+ parallel_mode: leader
default_sampling_params:
max_tokens: 1
temperature: 0.0
@@ -19,6 +28,7 @@ stages:
max_num_seqs: 1
model_loaded:
text_encoder: false
+ vae_encoder: false
parallel_config:
tensor_parallel_size: 1
ulysses_degree: 4
diff --git a/vllm_omni/deploy/personaplex.yaml b/vllm_omni/deploy/personaplex.yaml
index 9f1da7e63..c4cc472d1 100644
--- a/vllm_omni/deploy/personaplex.yaml
+++ b/vllm_omni/deploy/personaplex.yaml
@@ -25,12 +25,8 @@ connectors:
extra:
shm_threshold_bytes: 65536
codec_streaming: true
- connector_get_sleep_s: 0.01
- connector_get_max_wait_first_chunk: 3000
- connector_get_max_wait: 300
# Mimi runs at 12.5 Hz (80 ms / frame). Stream a few frames per chunk.
codec_chunk_frames: 5
- codec_left_context_frames: 10
# Emit only the first audio chunk early, then return to codec_chunk_frames.
initial_codec_chunk_frames: 1
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 10f8642c6..0ffd0c54e 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -984,6 +984,7 @@ class OmniDiffusionConfig:
"transformer": True,
"vae": True,
"text_encoder": True,
+ "vae_encoder": True,
}
)
override_transformer_cls_name: str | None = None
@@ -1020,6 +1021,8 @@ class OmniDiffusionConfig:
# str is resolved to {"method": } internally.
# Per-component: {"transformer": {"method": "fp8"}, "vae": None}
quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
+ # Internal provenance, retained across config projection and worker transport.
+ quantization_config_is_auto_detected: bool = False
# Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This
# does not enable FP8 by itself; it only selects CUTLASS once the checkpoint
# has already resolved to vLLM's ModelOpt FP8 linear method.
@@ -1357,6 +1360,8 @@ class OmniDiffusionConfig:
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 = tf_config.quant_config
logger.info(
"Auto-detected quantization '%s' from model config",
diff --git a/vllm_omni/diffusion/layers/qwen_select01_modulation.py b/vllm_omni/diffusion/layers/qwen_select01_modulation.py
new file mode 100644
index 000000000..4fe1e3e0e
--- /dev/null
+++ b/vllm_omni/diffusion/layers/qwen_select01_modulation.py
@@ -0,0 +1,491 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+
+@triton.jit
+def _layernorm_select01_kernel(
+ output_ptr,
+ gate_out_ptr,
+ x_ptr,
+ weight_ptr,
+ bias_ptr,
+ scale0_ptr,
+ shift0_ptr,
+ gate0_ptr,
+ scale1_ptr,
+ shift1_ptr,
+ gate1_ptr,
+ index_ptr,
+ inner_dim,
+ seq_len,
+ stride_x_row,
+ stride_out_row,
+ stride_gate_out_row,
+ stride_w,
+ stride_b,
+ stride_s0_b,
+ stride_s0_c,
+ stride_sh0_b,
+ stride_sh0_c,
+ stride_g0_b,
+ stride_g0_c,
+ stride_s1_b,
+ stride_s1_c,
+ stride_sh1_b,
+ stride_sh1_c,
+ stride_g1_b,
+ stride_g1_c,
+ stride_i_b,
+ stride_i_l,
+ eps,
+ has_weight: tl.constexpr,
+ has_bias: tl.constexpr,
+ block_n: tl.constexpr,
+):
+ row = tl.program_id(0)
+ cols = tl.arange(0, block_n)
+ mask = cols < inner_dim
+
+ x_row_ptr = x_ptr + row * stride_x_row
+ x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+
+ mean = tl.sum(x, axis=0) / inner_dim
+ xbar = tl.where(mask, x - mean, 0.0)
+ var = tl.sum(xbar * xbar, axis=0) / inner_dim
+ x_hat = (x - mean) * tl.rsqrt(var + eps)
+
+ if has_weight:
+ weight = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+ x_hat = x_hat * weight
+ if has_bias:
+ bias = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+ x_hat = x_hat + bias
+
+ batch_idx = row // seq_len
+ seq_idx = row % seq_len
+ idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+ scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
+ shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
+ gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
+ scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
+ shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
+ gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
+
+ scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs)
+ shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs)
+ gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs)
+
+ scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32)
+ shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32)
+ gate = tl.load(gate_ptrs, mask=mask, other=0.0)
+
+ y = x_hat * (1.0 + scale) + shift
+ tl.store(output_ptr + row * stride_out_row + cols, y, mask=mask)
+ tl.store(gate_out_ptr + row * stride_gate_out_row + cols, gate, mask=mask)
+
+
+@triton.jit
+def _residual_layernorm_select01_kernel(
+ output_ptr,
+ residual_out_ptr,
+ gate_out_ptr,
+ x_ptr,
+ residual_ptr,
+ residual_gate_ptr,
+ weight_ptr,
+ bias_ptr,
+ scale0_ptr,
+ shift0_ptr,
+ gate0_ptr,
+ scale1_ptr,
+ shift1_ptr,
+ gate1_ptr,
+ index_ptr,
+ inner_dim,
+ seq_len,
+ stride_x_row,
+ stride_res_row,
+ stride_res_gate_row,
+ stride_out_row,
+ stride_res_out_row,
+ stride_gate_out_row,
+ stride_w,
+ stride_b,
+ stride_s0_b,
+ stride_s0_c,
+ stride_sh0_b,
+ stride_sh0_c,
+ stride_g0_b,
+ stride_g0_c,
+ stride_s1_b,
+ stride_s1_c,
+ stride_sh1_b,
+ stride_sh1_c,
+ stride_g1_b,
+ stride_g1_c,
+ stride_i_b,
+ stride_i_l,
+ eps,
+ has_weight: tl.constexpr,
+ has_bias: tl.constexpr,
+ block_n: tl.constexpr,
+):
+ row = tl.program_id(0)
+ cols = tl.arange(0, block_n)
+ mask = cols < inner_dim
+
+ x = tl.load(x_ptr + row * stride_x_row + cols, mask=mask, other=0.0).to(tl.float32)
+ residual = tl.load(residual_ptr + row * stride_res_row + cols, mask=mask, other=0.0).to(tl.float32)
+ residual_gate = tl.load(
+ residual_gate_ptr + row * stride_res_gate_row + cols,
+ mask=mask,
+ other=0.0,
+ ).to(tl.float32)
+
+ residual_out = residual + residual_gate * x
+ tl.store(residual_out_ptr + row * stride_res_out_row + cols, residual_out, mask=mask)
+
+ mean = tl.sum(residual_out, axis=0) / inner_dim
+ xbar = tl.where(mask, residual_out - mean, 0.0)
+ var = tl.sum(xbar * xbar, axis=0) / inner_dim
+ x_hat = (residual_out - mean) * tl.rsqrt(var + eps)
+
+ if has_weight:
+ weight = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+ x_hat = x_hat * weight
+ if has_bias:
+ bias = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+ x_hat = x_hat + bias
+
+ batch_idx = row // seq_len
+ seq_idx = row % seq_len
+ idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+ scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
+ shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
+ gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
+ scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
+ shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
+ gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
+
+ scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs)
+ shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs)
+ gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs)
+
+ scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32)
+ shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32)
+ gate = tl.load(gate_ptrs, mask=mask, other=0.0)
+
+ y = x_hat * (1.0 + scale) + shift
+ tl.store(output_ptr + row * stride_out_row + cols, y, mask=mask)
+ tl.store(gate_out_ptr + row * stride_gate_out_row + cols, gate, mask=mask)
+
+
+def split_select01_mod_params(mod_params: torch.Tensor) -> tuple[torch.Tensor, ...]:
+ shift, scale, gate = mod_params.chunk(3, dim=-1)
+ actual_batch = shift.size(0) // 2
+ if actual_batch == 0 or shift.size(0) != actual_batch * 2:
+ raise ValueError("select01 modulation expects mod_params batch dim to be 2 * batch")
+ return (
+ scale[:actual_batch],
+ shift[:actual_batch],
+ gate[:actual_batch],
+ scale[actual_batch:],
+ shift[actual_batch:],
+ gate[actual_batch:],
+ )
+
+
+def select01_modulation_native(
+ mod_params: torch.Tensor,
+ index: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params)
+ idx = index.to(dtype=torch.bool).unsqueeze(-1)
+ scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
+ shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
+ gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
+ return scale, shift, gate
+
+
+def _validate_select01_inputs(
+ x: torch.Tensor,
+ index: torch.Tensor,
+ tensors: tuple[torch.Tensor, ...],
+) -> None:
+ if x.dim() != 3:
+ raise ValueError("x must be 3D [B, L, C]")
+ if index.dim() != 2 or index.shape[1] != x.shape[1] or index.shape[0] not in (1, x.shape[0]):
+ raise ValueError("index must be 2D [B or 1, L] and match x batch/sequence")
+ batch_size, _, hidden_size = x.shape
+ for tensor in tensors:
+ if tensor.dim() != 2 or tensor.shape != (batch_size, hidden_size):
+ raise ValueError("scale/shift/gate tensors must be 2D [B, C]")
+
+
+def _maybe_contiguous(x: torch.Tensor) -> torch.Tensor:
+ return x if x.is_contiguous() else x.contiguous()
+
+
+def _launch_layernorm_select01(
+ x: torch.Tensor,
+ weight: torch.Tensor | None,
+ bias: torch.Tensor | None,
+ scale0: torch.Tensor,
+ shift0: torch.Tensor,
+ gate0: torch.Tensor,
+ scale1: torch.Tensor,
+ shift1: torch.Tensor,
+ gate1: torch.Tensor,
+ index: torch.Tensor,
+ eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ _validate_select01_inputs(x, index, (scale0, shift0, gate0, scale1, shift1, gate1))
+ batch_size, seq_len, hidden_size = x.shape
+ output = torch.empty_like(x)
+ gate_out = torch.empty_like(x)
+
+ x_2d = _maybe_contiguous(x).view(batch_size * seq_len, hidden_size)
+ output_2d = output.view(batch_size * seq_len, hidden_size)
+ gate_out_2d = gate_out.view(batch_size * seq_len, hidden_size)
+ index = _maybe_contiguous(index)
+ scale0 = _maybe_contiguous(scale0)
+ shift0 = _maybe_contiguous(shift0)
+ gate0 = _maybe_contiguous(gate0)
+ scale1 = _maybe_contiguous(scale1)
+ shift1 = _maybe_contiguous(shift1)
+ gate1 = _maybe_contiguous(gate1)
+
+ weight_arg = _maybe_contiguous(weight) if weight is not None else x_2d
+ bias_arg = _maybe_contiguous(bias) if bias is not None else x_2d
+ block_n = min(65536 // x_2d.element_size(), triton.next_power_of_2(hidden_size))
+ if hidden_size > block_n:
+ raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+
+ _layernorm_select01_kernel[(batch_size * seq_len,)](
+ output_2d,
+ gate_out_2d,
+ x_2d,
+ weight_arg,
+ bias_arg,
+ scale0,
+ shift0,
+ gate0,
+ scale1,
+ shift1,
+ gate1,
+ index,
+ hidden_size,
+ seq_len,
+ x_2d.stride(0),
+ output_2d.stride(0),
+ gate_out_2d.stride(0),
+ weight_arg.stride(0) if weight is not None else 0,
+ bias_arg.stride(0) if bias is not None else 0,
+ scale0.stride(0),
+ scale0.stride(1),
+ shift0.stride(0),
+ shift0.stride(1),
+ gate0.stride(0),
+ gate0.stride(1),
+ scale1.stride(0),
+ scale1.stride(1),
+ shift1.stride(0),
+ shift1.stride(1),
+ gate1.stride(0),
+ gate1.stride(1),
+ index.stride(0) if index.shape[0] > 1 else 0,
+ index.stride(1),
+ eps,
+ has_weight=weight is not None,
+ has_bias=bias is not None,
+ block_n=block_n,
+ num_warps=4,
+ num_stages=4,
+ )
+ return output, gate_out
+
+
+def _launch_residual_layernorm_select01(
+ x: torch.Tensor,
+ residual: torch.Tensor,
+ residual_gate: torch.Tensor,
+ weight: torch.Tensor | None,
+ bias: torch.Tensor | None,
+ scale0: torch.Tensor,
+ shift0: torch.Tensor,
+ gate0: torch.Tensor,
+ scale1: torch.Tensor,
+ shift1: torch.Tensor,
+ gate1: torch.Tensor,
+ index: torch.Tensor,
+ eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ _validate_select01_inputs(x, index, (scale0, shift0, gate0, scale1, shift1, gate1))
+ if residual.shape != x.shape:
+ raise ValueError("residual must have the same shape as x")
+ if residual_gate.shape != x.shape:
+ raise ValueError("residual_gate must have the same shape as x")
+
+ batch_size, seq_len, hidden_size = x.shape
+ output = torch.empty_like(x)
+ residual_out = torch.empty_like(x)
+ gate_out = torch.empty_like(x)
+
+ x_2d = _maybe_contiguous(x).view(batch_size * seq_len, hidden_size)
+ residual_2d = _maybe_contiguous(residual).view(batch_size * seq_len, hidden_size)
+ residual_gate_2d = _maybe_contiguous(residual_gate).view(batch_size * seq_len, hidden_size)
+ output_2d = output.view(batch_size * seq_len, hidden_size)
+ residual_out_2d = residual_out.view(batch_size * seq_len, hidden_size)
+ gate_out_2d = gate_out.view(batch_size * seq_len, hidden_size)
+ index = _maybe_contiguous(index)
+ scale0 = _maybe_contiguous(scale0)
+ shift0 = _maybe_contiguous(shift0)
+ gate0 = _maybe_contiguous(gate0)
+ scale1 = _maybe_contiguous(scale1)
+ shift1 = _maybe_contiguous(shift1)
+ gate1 = _maybe_contiguous(gate1)
+
+ weight_arg = _maybe_contiguous(weight) if weight is not None else x_2d
+ bias_arg = _maybe_contiguous(bias) if bias is not None else x_2d
+ block_n = min(65536 // x_2d.element_size(), triton.next_power_of_2(hidden_size))
+ if hidden_size > block_n:
+ raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+
+ _residual_layernorm_select01_kernel[(batch_size * seq_len,)](
+ output_2d,
+ residual_out_2d,
+ gate_out_2d,
+ x_2d,
+ residual_2d,
+ residual_gate_2d,
+ weight_arg,
+ bias_arg,
+ scale0,
+ shift0,
+ gate0,
+ scale1,
+ shift1,
+ gate1,
+ index,
+ hidden_size,
+ seq_len,
+ x_2d.stride(0),
+ residual_2d.stride(0),
+ residual_gate_2d.stride(0),
+ output_2d.stride(0),
+ residual_out_2d.stride(0),
+ gate_out_2d.stride(0),
+ weight_arg.stride(0) if weight is not None else 0,
+ bias_arg.stride(0) if bias is not None else 0,
+ scale0.stride(0),
+ scale0.stride(1),
+ shift0.stride(0),
+ shift0.stride(1),
+ gate0.stride(0),
+ gate0.stride(1),
+ scale1.stride(0),
+ scale1.stride(1),
+ shift1.stride(0),
+ shift1.stride(1),
+ gate1.stride(0),
+ gate1.stride(1),
+ index.stride(0) if index.shape[0] > 1 else 0,
+ index.stride(1),
+ eps,
+ has_weight=weight is not None,
+ has_bias=bias is not None,
+ block_n=block_n,
+ num_warps=4,
+ num_stages=4,
+ )
+ return output, residual_out, gate_out
+
+
+def can_use_qwen_select01_triton(x: torch.Tensor) -> bool:
+ return HAS_TRITON and current_platform.is_cuda() and x.is_cuda
+
+
+def fused_layernorm_select01(
+ x: torch.Tensor,
+ mod_params: torch.Tensor,
+ index: torch.Tensor,
+ eps: float,
+ weight: torch.Tensor | None = None,
+ bias: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params)
+ is_compiling = torch.compiler.is_compiling()
+ if can_use_qwen_select01_triton(x) and not is_compiling:
+ return _launch_layernorm_select01(
+ x,
+ weight,
+ bias,
+ scale0,
+ shift0,
+ gate0,
+ scale1,
+ shift1,
+ gate1,
+ index,
+ eps,
+ )
+
+ scale, shift, gate = select01_modulation_native(mod_params, index)
+ out = F.layer_norm(
+ x.float(),
+ (x.shape[-1],),
+ weight=weight.float() if weight is not None else None,
+ bias=bias.float() if bias is not None else None,
+ eps=eps,
+ ).to(x.dtype)
+ return out * (1 + scale) + shift, gate
+
+
+def fused_residual_layernorm_select01(
+ x: torch.Tensor,
+ residual: torch.Tensor,
+ residual_gate: torch.Tensor,
+ mod_params: torch.Tensor,
+ index: torch.Tensor,
+ eps: float,
+ weight: torch.Tensor | None = None,
+ bias: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params)
+ is_compiling = torch.compiler.is_compiling()
+ if can_use_qwen_select01_triton(x) and not is_compiling:
+ return _launch_residual_layernorm_select01(
+ x,
+ residual,
+ residual_gate,
+ weight,
+ bias,
+ scale0,
+ shift0,
+ gate0,
+ scale1,
+ shift1,
+ gate1,
+ index,
+ eps,
+ )
+
+ 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],),
+ weight=weight.float() if weight is not None else None,
+ bias=bias.float() if bias is not None else None,
+ eps=eps,
+ ).to(residual_out.dtype)
+ return out * (1 + scale) + shift, residual_out, gate
diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py
index 78247e7b9..834887f41 100644
--- a/vllm_omni/diffusion/model_loader/diffusers_loader.py
+++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py
@@ -7,7 +7,7 @@ import json
import os
import re
import time
-from collections.abc import Generator, Iterable, Sequence
+from collections.abc import Callable, Generator, Iterable, Sequence
from pathlib import Path
from typing import cast
@@ -481,8 +481,14 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
)
def _get_expected_parameter_names(self, model: nn.Module) -> set[str]:
- """Return parameter names that should be covered by strict load checks."""
- all_parameter_names = {name for name, _ in model.named_parameters()}
+ """Return parameter names that should be covered by strict load checks.
+
+ A parameter with an initialized checkpoint default can explicitly set
+ ``is_checkpoint_optional``. It is still loaded when present on disk.
+ """
+ all_parameter_names = {
+ name for name, param in model.named_parameters() if not getattr(param, "is_checkpoint_optional", False)
+ }
sources = self._get_weight_sources(model)
# Keep strict behavior if no source metadata exists.
@@ -853,9 +859,24 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
del model
return self.load_fresh_canonical_model()
raise
+ self._log_w4a8_fallback_load_summaries(model)
self._attach_offload_startup_state(model)
return model
+ @staticmethod
+ def _log_w4a8_fallback_load_summaries(model: nn.Module) -> None:
+ """Ask discovered DiTs to report W4A8 state at the common load exit.
+
+ Each DiT derives readiness from its processed layers, so a deferred
+ weight plan cannot be mistaken for the ordinary or HSDP load path.
+ """
+ components = ModuleDiscovery.discover(model)
+ for component_name, dit in zip(components.dit_names, components.dits):
+ candidate = getattr(dit, "_log_w4a8_fallback_load_summary", None)
+ if callable(candidate):
+ reporter = cast(Callable[[str], None], candidate)
+ reporter(component_name)
+
@staticmethod
def _request_offload_after_quant(model: nn.Module) -> int:
"""Ask online-quant layers to return to host memory once quantized.
@@ -1065,6 +1086,16 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
# that have loaded weights tracking currently.
if loaded_weights is not None:
weights_not_loaded = weights_to_load - loaded_weights
+ # Offline formats can require scales that older online loaders
+ # were allowed to synthesize. Do not apply that legacy tolerance
+ # to an explicitly required checkpoint tensor.
+ required_missing = {
+ name
+ for name, param in model.named_parameters()
+ if name in weights_not_loaded and getattr(param, "is_checkpoint_required", False)
+ }
+ if required_missing:
+ raise ValueError(f"Required weights were not initialized from checkpoint: {required_missing}")
# NOTE: if the model is quantized, ignore not_loaded check for scale
# weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a
# static activation `input_scale`, which the quant method may
diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py
index 93d52b960..f3fdeedc8 100644
--- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py
+++ b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py
@@ -15,6 +15,7 @@
# inference caches (TeaCache/TaylorSeer) are not ported.
import itertools
+from collections import defaultdict
from collections.abc import Iterable
from typing import TYPE_CHECKING
@@ -24,9 +25,9 @@ import torch.nn.functional as F
from diffusers.models.embeddings import Timesteps, get_1d_rotary_pos_embed
from einops import rearrange, repeat
from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
- ColumnParallelLinear,
+ MergedColumnParallelLinear,
+ QKVParallelLinear,
ReplicatedLinear,
RowParallelLinear,
)
@@ -42,6 +43,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
fused_qk_norm_rope,
fused_qk_norm_rope_min_tokens,
)
+from vllm_omni.diffusion.layers.norm import RMSNorm
from vllm_omni.diffusion.models.utils import make_attention_mask
from vllm_omni.platforms import current_omni_platform
@@ -260,7 +262,14 @@ class LuminaLayerNormContinuous(nn.Module):
class LuminaFeedForward(nn.Module):
- """SwiGLU feed-forward with tensor-parallel projections."""
+ """SwiGLU feed-forward with tensor-parallel projections.
+
+ The ``gate`` and ``input`` projections are fused into a single
+ ``MergedColumnParallelLinear`` (one GEMM instead of two), matching the
+ ``gate_up_proj`` convention used by the Hunyuan Image 3 port. The SwiGLU
+ activation keeps its float32 computation for numerical parity with
+ upstream.
+ """
def __init__(
self,
@@ -276,20 +285,13 @@ class LuminaFeedForward(nn.Module):
inner_dim = int(ffn_dim_multiplier * inner_dim)
inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
- self.linear_1 = ColumnParallelLinear(
- dim,
- inner_dim,
+ self.gate_up_proj = MergedColumnParallelLinear(
+ input_size=dim,
+ output_sizes=[inner_dim, inner_dim],
bias=False,
quant_config=quant_config,
- prefix=_join_prefix(prefix, "linear_1"),
- ) # gate
- self.linear_3 = ColumnParallelLinear(
- dim,
- inner_dim,
- bias=False,
- quant_config=quant_config,
- prefix=_join_prefix(prefix, "linear_3"),
- ) # input
+ prefix=_join_prefix(prefix, "gate_up_proj"),
+ )
self.linear_2 = RowParallelLinear(
inner_dim,
dim,
@@ -299,8 +301,8 @@ class LuminaFeedForward(nn.Module):
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
- h1, _ = self.linear_1(x)
- h2, _ = self.linear_3(x)
+ gate_up, _ = self.gate_up_proj(x)
+ h1, h2 = gate_up.chunk(2, dim=-1)
out, _ = self.linear_2(swiglu(h1, h2))
return out
@@ -593,28 +595,15 @@ class BooguImageSelfAttention(nn.Module):
) -> None:
super().__init__()
self.head_dim = dim // num_attention_heads
- kv_dim = self.head_dim * num_kv_heads
- self.to_q = ColumnParallelLinear(
- dim,
- dim,
+ self.to_qkv = QKVParallelLinear(
+ hidden_size=dim,
+ head_size=self.head_dim,
+ total_num_heads=num_attention_heads,
+ total_num_kv_heads=num_kv_heads,
bias=False,
quant_config=quant_config,
- prefix=_join_prefix(prefix, "to_q"),
- )
- self.to_k = ColumnParallelLinear(
- dim,
- kv_dim,
- bias=False,
- quant_config=quant_config,
- prefix=_join_prefix(prefix, "to_k"),
- )
- self.to_v = ColumnParallelLinear(
- dim,
- kv_dim,
- bias=False,
- quant_config=quant_config,
- prefix=_join_prefix(prefix, "to_v"),
+ prefix=_join_prefix(prefix, "to_qkv"),
)
self.norm_q = RMSNorm(self.head_dim, eps=1e-5)
self.norm_k = RMSNorm(self.head_dim, eps=1e-5)
@@ -626,8 +615,10 @@ class BooguImageSelfAttention(nn.Module):
prefix=_join_prefix(prefix, "to_out"),
)
- self.num_local_heads = self.to_q.output_size_per_partition // self.head_dim
- self.num_local_kv_heads = self.to_k.output_size_per_partition // self.head_dim
+ self.num_local_heads = self.to_qkv.num_heads
+ self.num_local_kv_heads = self.to_qkv.num_kv_heads
+ self.q_size = self.num_local_heads * self.head_dim
+ self.kv_size = self.num_local_kv_heads * self.head_dim
self.attn = Attention(
num_heads=self.num_local_heads,
@@ -645,9 +636,8 @@ class BooguImageSelfAttention(nn.Module):
) -> torch.Tensor:
dtype = hidden_states.dtype
- query, _ = self.to_q(hidden_states)
- key, _ = self.to_k(hidden_states)
- value, _ = self.to_v(hidden_states)
+ qkv, _ = self.to_qkv(hidden_states)
+ query, key, value = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
query = query.unflatten(-1, (self.num_local_heads, self.head_dim))
key = key.unflatten(-1, (self.num_local_kv_heads, self.head_dim))
@@ -683,50 +673,24 @@ class BooguImageJointAttention(nn.Module):
) -> None:
super().__init__()
self.head_dim = dim // num_attention_heads
- kv_dim = self.head_dim * num_kv_heads
- self.img_to_q = ColumnParallelLinear(
- dim,
- dim,
+ self.img_to_qkv = QKVParallelLinear(
+ hidden_size=dim,
+ head_size=self.head_dim,
+ total_num_heads=num_attention_heads,
+ total_num_kv_heads=num_kv_heads,
bias=False,
quant_config=quant_config,
- prefix=_join_prefix(prefix, "img_to_q"),
+ prefix=_join_prefix(prefix, "img_to_qkv"),
)
- self.img_to_k = ColumnParallelLinear(
- dim,
- kv_dim,
+ self.instruct_to_qkv = QKVParallelLinear(
+ hidden_size=dim,
+ head_size=self.head_dim,
+ total_num_heads=num_attention_heads,
+ total_num_kv_heads=num_kv_heads,
bias=False,
quant_config=quant_config,
- prefix=_join_prefix(prefix, "img_to_k"),
- )
- self.img_to_v = ColumnParallelLinear(
- dim,
- kv_dim,
- bias=False,
- quant_config=quant_config,
- prefix=_join_prefix(prefix, "img_to_v"),
- )
-
- self.instruct_to_q = ColumnParallelLinear(
- dim,
- dim,
- bias=False,
- quant_config=quant_config,
- prefix=_join_prefix(prefix, "instruct_to_q"),
- )
- self.instruct_to_k = ColumnParallelLinear(
- dim,
- kv_dim,
- bias=False,
- quant_config=quant_config,
- prefix=_join_prefix(prefix, "instruct_to_k"),
- )
- self.instruct_to_v = ColumnParallelLinear(
- dim,
- kv_dim,
- bias=False,
- quant_config=quant_config,
- prefix=_join_prefix(prefix, "instruct_to_v"),
+ prefix=_join_prefix(prefix, "instruct_to_qkv"),
)
self.norm_q = RMSNorm(self.head_dim, eps=1e-5)
@@ -757,8 +721,10 @@ class BooguImageJointAttention(nn.Module):
prefix=_join_prefix(prefix, "to_out"),
)
- self.num_local_heads = self.img_to_q.output_size_per_partition // self.head_dim
- self.num_local_kv_heads = self.img_to_k.output_size_per_partition // self.head_dim
+ self.num_local_heads = self.img_to_qkv.num_heads
+ self.num_local_kv_heads = self.img_to_qkv.num_kv_heads
+ self.q_size = self.num_local_heads * self.head_dim
+ self.kv_size = self.num_local_kv_heads * self.head_dim
self.attn = Attention(
num_heads=self.num_local_heads,
@@ -780,13 +746,13 @@ class BooguImageJointAttention(nn.Module):
dtype = img_hidden_states.dtype
batch_size = img_hidden_states.shape[0]
- img_query, _ = self.img_to_q(img_hidden_states)
- img_key, _ = self.img_to_k(img_hidden_states)
- img_value, _ = self.img_to_v(img_hidden_states)
+ img_qkv, _ = self.img_to_qkv(img_hidden_states)
+ img_query, img_key, img_value = img_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
- instruct_query, _ = self.instruct_to_q(instruct_hidden_states)
- instruct_key, _ = self.instruct_to_k(instruct_hidden_states)
- instruct_value, _ = self.instruct_to_v(instruct_hidden_states)
+ instruct_qkv, _ = self.instruct_to_qkv(instruct_hidden_states)
+ instruct_query, instruct_key, instruct_value = instruct_qkv.split(
+ [self.q_size, self.kv_size, self.kv_size], dim=-1
+ )
query, key, value = _concat_instruction_image_features(
[img_query, img_key, img_value],
@@ -1135,6 +1101,44 @@ def _cal_preprocessed_instruction_feat_dim(instruction_feature_configs: dict) ->
raise ValueError(f"Invalid reduce_type: {reduce_type}")
+# Packed (fused) projections vs. the logical sub-projections the diffusers
+# checkpoint stores, in the ``(param_name, shard_name, shard_id)`` form used
+# across vllm-omni. ``load_weights`` consumes it directly, and loader consumers
+# that discover it via ``stacked_params_mapping`` (LoRA, quantized loaders) get
+# the packed -> sub-layer relationship for free.
+#
+# The QKV entries stay leaf-scoped: each fused source (``to_q`` on the
+# self-attention module, ``img_to_q`` / ``instruct_to_q`` on the joint
+# attention) has a distinct leaf, and ``.to_q.`` cannot collide with
+# ``.img_to_q.`` because the preceding character is ``_`` rather than ``.``.
+# The FFN entries are path-qualified instead: ``linear_1`` also names the
+# timestep embedder and the ``norm_out`` projections, which must load 1:1 and
+# are not fused.
+_BOOGU_STACKED_PARAMS_MAPPING = (
+ # self-attention: noise / reference-image / context refiners, the
+ # single-stream blocks, and the double-stream image self-attention.
+ (".to_qkv.", ".to_q.", "q"),
+ (".to_qkv.", ".to_k.", "k"),
+ (".to_qkv.", ".to_v.", "v"),
+ # joint (instruction + image) attention of the double-stream blocks.
+ (".img_to_qkv.", ".img_to_q.", "q"),
+ (".img_to_qkv.", ".img_to_k.", "k"),
+ (".img_to_qkv.", ".img_to_v.", "v"),
+ (".instruct_to_qkv.", ".instruct_to_q.", "q"),
+ (".instruct_to_qkv.", ".instruct_to_k.", "k"),
+ (".instruct_to_qkv.", ".instruct_to_v.", "v"),
+ # feed-forward gate/up: ``linear_1`` is the gate (shard 0) and
+ # ``linear_3`` the input (shard 1), matching the packed order the SwiGLU
+ # activation consumes.
+ (".feed_forward.gate_up_proj.", ".feed_forward.linear_1.", 0),
+ (".feed_forward.gate_up_proj.", ".feed_forward.linear_3.", 1),
+ (".img_feed_forward.gate_up_proj.", ".img_feed_forward.linear_1.", 0),
+ (".img_feed_forward.gate_up_proj.", ".img_feed_forward.linear_3.", 1),
+ (".instruct_feed_forward.gate_up_proj.", ".instruct_feed_forward.linear_1.", 0),
+ (".instruct_feed_forward.gate_up_proj.", ".instruct_feed_forward.linear_3.", 1),
+)
+
+
class BooguImageTransformer2DModel(nn.Module):
"""Boogu-Image transformer with mixed stream topology.
@@ -1160,6 +1164,10 @@ class BooguImageTransformer2DModel(nn.Module):
prefix: str = "",
) -> None:
super().__init__()
+ # Published here rather than in ``load_weights`` so loader consumers
+ # that only inspect the module tree (LoRA discovery, quantized weight
+ # loaders) always see the packed -> sub-layer relationship.
+ self.stacked_params_mapping = list(_BOOGU_STACKED_PARAMS_MAPPING)
self.od_config = od_config
cfg = od_config.tf_model_config
@@ -1710,7 +1718,7 @@ class BooguImageTransformer2DModel(nn.Module):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Load diffusers-named checkpoint weights into the native module.
- Two name promotions relative to upstream (see step 8/10 findings):
+ Name promotions relative to upstream (see step 8/10 findings):
- ``*.img_instruct_attn.processor.{img,instruct}_{to_q,to_k,to_v}`` /
``{instruct,img}_out`` -> drop ``.processor`` (upstream keeps the
@@ -1719,9 +1727,29 @@ class BooguImageTransformer2DModel(nn.Module):
- ``*.to_out.0.weight`` -> ``*.to_out.weight`` (diffusers wraps the
output projection in a ``ModuleList``; the native module uses a plain
linear).
+
+ Packed Q/K/V and FFN gate/input matrices are folded onto the fused
+ ``QKVParallelLinear`` / ``MergedColumnParallelLinear`` parameters using
+ :attr:`stacked_params_mapping` — the same mapping loader consumers such
+ as LoRA discovery and quantized weight loaders read.
+
+ A fused parameter is only reported as loaded once *every* one of its
+ source matrices has arrived. The caller compares parameter names, so
+ reporting e.g. ``to_qkv`` complete after the first shard would let a
+ checkpoint carrying only ``to_q`` start up with the ``k``/``v`` slices
+ left uninitialized.
"""
+ stacked_params_mapping = self.stacked_params_mapping
+ # The shard ids each fused parameter is assembled from, keyed by the
+ # mapping entry that produces it.
+ expected_shards: dict[str, set[str | int]] = defaultdict(set)
+ for param_name, _weight_name, packed_shard_id in stacked_params_mapping:
+ expected_shards[param_name].add(packed_shard_id)
+
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
+ # Fused parameter -> (mapping entry that produced it, shards received).
+ fused_params: dict[str, tuple[str, set[str | int]]] = {}
for name, loaded_weight in weights:
original_name = name
@@ -1730,15 +1758,36 @@ class BooguImageTransformer2DModel(nn.Module):
if ".to_out.0." in name:
name = name.replace(".to_out.0.", ".to_out.")
+ shard_id: str | int | None = None
+ matched_entry: str | None = None
+ for param_name, weight_name, packed_shard_id in stacked_params_mapping:
+ if weight_name in name:
+ name = name.replace(weight_name, param_name)
+ shard_id = packed_shard_id
+ matched_entry = param_name
+ break
+
if name not in params_dict:
logger.warning("Skipping unexpected checkpoint weight %s", original_name)
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
- weight_loader(param, loaded_weight)
+ if matched_entry is None:
+ weight_loader(param, loaded_weight)
+ else:
+ weight_loader(param, loaded_weight, shard_id)
+ _, received = fused_params.setdefault(name, (matched_entry, set()))
+ received.add(shard_id)
loaded_params.add(name)
+ # Withdraw any fused parameter that only received part of its source
+ # matrices, so the caller's "not initialized from checkpoint" check
+ # reports it as missing instead of accepting a half-filled parameter.
+ for name, (matched_entry, received) in fused_params.items():
+ if received != expected_shards[matched_entry]:
+ loaded_params.discard(name)
+
unloaded_params = sorted(params_dict.keys() - loaded_params)
if unloaded_params:
logger.warning("Model parameters not loaded from checkpoint: %s", unloaded_params)
diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
index 8f814a27a..1b8bd7fea 100644
--- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
+++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
@@ -7,9 +7,9 @@ from __future__ import annotations
import json
import math
import os
-import tempfile
from collections.abc import Iterable, Mapping, Sequence
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
+from dataclasses import fields, replace
from itertools import groupby
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar
@@ -18,7 +18,6 @@ import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
-from PIL import Image
from transformers import Qwen2TokenizerFast, Qwen3VLProcessor
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -29,10 +28,7 @@ from vllm_omni.diffusion.cache.cachedit import (
RequestScopedCacheDiTRuntime,
)
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.parallel_state import (
- get_world_group,
- init_world_group,
-)
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
from vllm_omni.diffusion.model_loader.diffusers_loader import (
@@ -74,36 +70,20 @@ from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
resolve_minimax_h3_partition,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
+ MiniMaxH3EncoderConditioning,
+ MiniMaxH3EncoderMediaConditioning,
+ MiniMaxH3EncoderMediaInput,
MiniMaxH3TextConditioning,
)
-from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
- load_minimax_h3_images as _load_images,
-)
-from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
- minimax_h3_multi_image_presentation,
- minimax_h3_ref2va_presentation,
- minimax_h3_ref2va_video_presentation,
- minimax_h3_text_only_ids,
-)
-from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
- resolve_minimax_h3_aspect_ratio as _resolve_minimax_h3_aspect_ratio,
-)
-from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
- resolve_minimax_h3_output_canvas as _resolve_output_canvas,
-)
-from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
- resolve_minimax_h3_reference_image_shape as _reference_image_shape,
+from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
+ PreparedEncoderInputs,
+ encode_media,
+ prepare_encoder_inputs,
)
+from vllm_omni.model_executor.models.minimax_h3.preprocessing import build_minimax_h3_presentation
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
deserialize_prepared_reference_videos,
- load_audio_file,
- load_video_audio,
- load_video_frames,
- prepare_reference_videos,
- sample_reference_video_frames,
- validate_reference_audio_files,
- validate_reference_audio_waveforms,
)
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization import (
@@ -154,7 +134,6 @@ from .scheduling_minimax_h3_euler_ancestral import (
)
from .time_request import (
MINIMAX_H3_SHAPE_PLANNER,
- minimax_h3_align_frame_count,
minimax_h3_time_shift_sigmas,
)
from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE
@@ -175,22 +154,6 @@ MINIMAX_H3_FPS = 24
MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0
-MINIMAX_H3_OUTPUT_SHORT_EDGE = 768
-MINIMAX_H3_OUTPUT_MAX_PIXELS = 768 * 1344
-MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048
-MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32
-MINIMAX_H3_SUPPORTED_ASPECT_RATIOS = {
- "21:9": 21.0 / 9.0,
- "16:9": 16.0 / 9.0,
- "4:3": 4.0 / 3.0,
- "1:1": 1.0,
- "3:4": 3.0 / 4.0,
- "9:16": 9.0 / 16.0,
-}
-MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024
-MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"})
-MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0
-MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0
MINIMAX_H3_DOWNLOAD_PATTERNS = [
"FL2VA/**",
"Ref2VA/model_index.json",
@@ -267,7 +230,6 @@ def _resolve_minimax_h3_model_root(
)
-# Keys of ``_prepare_request_inputs`` that feed ``diffuse`` / ``_build_denoise_inputs``.
_MINIMAX_H3_DENOISE_INPUT_KEYS = (
"task",
"text_embeddings",
@@ -427,114 +389,6 @@ def get_minimax_h3_post_process_func(
return _minimax_h3_post_process
-def _align_multiple(value: float, multiple: int = 32) -> int:
- return max(multiple, int(round(float(value) / multiple)) * multiple)
-
-
-def _load_image(value: Any) -> Image.Image:
- images = _load_images(value)
- if len(images) != 1:
- raise OmniClientError(f"MiniMax H3 expected one image, got {len(images)}")
- return images[0]
-
-
-def _load_audio(value: Any) -> tuple[torch.Tensor, int]:
- if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))):
- audios = _load_audios(value)
- if len(audios) != 1:
- raise OmniClientError(f"MiniMax H3 expected one audio, got {len(audios)}")
- return audios[0]
- if isinstance(value, (str, os.PathLike)):
- return load_audio_file(str(value))
- if isinstance(value, (list, tuple)) and len(value) == 2:
- waveform, sample_rate = value
- waveform = torch.as_tensor(waveform).float()
- return waveform, int(sample_rate)
- if isinstance(value, dict):
- waveform = value.get("waveform", value.get("array"))
- sample_rate = value.get("sample_rate", value.get("sampling_rate"))
- if waveform is not None and sample_rate is not None:
- return torch.as_tensor(waveform).float(), int(sample_rate)
- raise OmniClientError("MiniMax H3 audio input must be a path, (waveform, sample_rate), or a waveform mapping")
-
-
-def _load_audios(value: Any) -> list[tuple[torch.Tensor, int]]:
- if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))):
- if not value:
- raise OmniClientError("MiniMax H3 audio input must not be empty")
- return [_load_audio(item) for item in value]
- return [_load_audio(value)]
-
-
-def _as_int_list(value: Any, *, name: str) -> list[int]:
- if isinstance(value, bool):
- raise OmniClientError(f"{name} must be an integer or a list of integers")
- if isinstance(value, (int, np.integer)):
- return [int(value)]
- if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
- result = list(value)
- if not result:
- raise OmniClientError(f"{name} must not be empty")
- if any(isinstance(item, bool) or not isinstance(item, (int, np.integer)) for item in result):
- raise OmniClientError(f"{name} must contain only integers")
- return [int(item) for item in result]
- raise OmniClientError(f"{name} must be an integer or a list of integers")
-
-
-def _resolve_fl2va_keyframe_indices(extra: Mapping[str, Any], image_count: int) -> list[int]:
- target = extra.get("target")
- target = target if isinstance(target, Mapping) else {}
- raw = extra.get("frame_indices", extra.get("frame_index"))
- if raw is None:
- raw = target.get("frame_indices", target.get("frame_index"))
- if raw is None:
- raw_indices = [0] if image_count == 1 else [0, -1]
- else:
- raw_indices = _as_int_list(raw, name="frame_indices")
- if len(raw_indices) != image_count:
- raise OmniClientError(
- f"MiniMax H3 FL2VA requires one frame index per image: got {raw_indices!r} for {image_count} image(s)"
- )
- if tuple(raw_indices) not in ((0,), (-1,), (0, -1)):
- raise OmniClientError("MiniMax H3 FL2VA frame_indices must be [0], [-1], or [0, -1]")
- return raw_indices
-
-
-def _reuse_prepared_reference_videos(
- prepared: list[dict[str, Any]] | None,
- *,
- expected_count: int,
-) -> list[dict[str, Any]] | None:
- if prepared is None:
- return None
- if len(prepared) != expected_count:
- raise OmniClientError("MiniMax H3 prepared-reference-video count does not match the request")
- for item in prepared:
- if not os.path.isfile(item["prepared_path"]):
- raise OmniClientError(f"MiniMax H3 prepared reference video is unavailable: {item['prepared_path']}")
- return prepared
-
-
-def _validate_ref2va_reference_counts(
- image_count: int,
- video_count: int,
- audio_count: int,
-) -> None:
- """Validate the official Ref2VA reference-count contract."""
- if image_count < 0 or video_count < 0 or audio_count < 0:
- raise OmniClientError("MiniMax H3 reference counts must be non-negative")
- if image_count + video_count == 0:
- raise OmniClientError("ref2va requires at least one image or video reference")
- if image_count > 9:
- raise OmniClientError("ref2va accepts at most 9 image references")
- if video_count > 3:
- raise OmniClientError("ref2va accepts at most 3 video references")
- if audio_count > 3:
- raise OmniClientError("ref2va accepts at most 3 standalone audio references")
- if image_count + video_count + audio_count > 12:
- raise OmniClientError("ref2va accepts at most 12 total references")
-
-
def _resolve_minimax_h3_num_outputs(value: Any) -> int:
if value is None:
return 1
@@ -572,17 +426,6 @@ def _minimax_h3_output_seeds(seed: int, num_outputs: int) -> list[int]:
return [int(seed) + output_index for output_index in range(int(num_outputs))]
-def _validate_reference_image(image: Image.Image) -> None:
- width, height = image.size
- if min(width, height) < 256 or max(width, height) > 5760:
- raise OmniClientError(
- f"MiniMax H3 reference image dimensions must be in [256, 5760] pixels, got {width}x{height}"
- )
- ratio = width / height
- if not 0.4 <= ratio <= 2.5:
- raise OmniClientError(f"MiniMax H3 reference image aspect ratio must be in [0.4, 2.5], got {width}x{height}")
-
-
def _dit_rank_world() -> tuple[Any, int, int]:
if not dist.is_initialized():
return None, 0, 1
@@ -716,10 +559,8 @@ class MiniMaxH3Pipeline(
on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
)
_PROFILER_TARGETS: ClassVar[list[str]] = [
- "_prepare_reference_videos",
"encode_prompt",
- "_encode_visual_conditions",
- "_encode_reference_audio_conditions",
+ "_encode_local_media",
"diffuse",
"decode",
"video_vae.decode_latent",
@@ -809,6 +650,18 @@ class MiniMaxH3Pipeline(
f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}"
)
+ def _active_turbo_spec(self, sampling: Any) -> TurboSpec | None:
+ """Return the spec of the Turbo adapter this request actually applies.
+
+ A recognized adapter at scale 0 contributes nothing, so it neither
+ constrains the task nor imposes its sampler contract.
+ """
+
+ lora_request = sampling.lora_request
+ if lora_request is None or math.isclose(0.0, float(sampling.lora_scale)):
+ return None
+ return self._turbo_lora_specs.get(lora_request.lora_int_id)
+
def _has_active_native_lora(self, sampling: Any) -> bool:
lora_request = sampling.lora_request
return (
@@ -868,18 +721,6 @@ class MiniMaxH3Pipeline(
return adapter_schedule
return self._base_schedule_for_task(task)
- def _active_turbo_spec(self, sampling: Any) -> TurboSpec | None:
- """Return the spec of the Turbo adapter this request actually applies.
-
- A recognized adapter at scale 0 contributes nothing, so it neither
- constrains the task nor imposes its sampler contract.
- """
-
- lora_request = sampling.lora_request
- if lora_request is None or math.isclose(0.0, float(sampling.lora_scale)):
- return None
- return self._turbo_lora_specs.get(lora_request.lora_int_id)
-
def _validate_turbo_sampling(self, sampling: Any, spec: TurboSpec) -> None:
"""Hold a request to the contract of the artifact that is loaded.
@@ -935,7 +776,30 @@ class MiniMaxH3Pipeline(
if int(self.parallel_config.cfg_parallel_size) != 1:
raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1")
self.device = get_local_device()
- self.load_text_encoder = od_config.model_loaded["text_encoder"]
+ self.load_text_encoder = od_config.model_loaded.get("text_encoder", True)
+ self.load_vae_encoder = od_config.model_loaded.get("vae_encoder", True)
+ if self.load_vae_encoder is False and self.load_text_encoder is True:
+ raise ValueError(
+ "MiniMax H3 does not support local text encoding with external media conditioning; "
+ "set text_encoder=false or enable vae_encoder"
+ )
+ self._encoder_modules = ["text_encoder"] if self.load_text_encoder else []
+ self._PROFILER_TARGETS = list(type(self)._PROFILER_TARGETS)
+ encoder_block_attrs = dict(self._offload_plan.encoder_block_attrs)
+ if not self.load_text_encoder:
+ encoder_block_attrs.pop("text_encoder", None)
+ on_demand_component_paths = set(self._offload_plan.on_demand_component_paths)
+ if not self.load_text_encoder:
+ on_demand_component_paths.discard("text_encoder")
+ self._offload_plan = replace(
+ self._offload_plan,
+ encoder_block_attrs=encoder_block_attrs,
+ on_demand_component_paths=frozenset(on_demand_component_paths),
+ )
+ if not self.load_text_encoder:
+ self._PROFILER_TARGETS.remove("encode_prompt")
+ if not self.load_vae_encoder:
+ self._PROFILER_TARGETS.remove("_encode_local_media")
self.partition = _minimax_h3_partition_for_task(
getattr(od_config, "task_type", None),
str(od_config.model),
@@ -1102,12 +966,14 @@ class MiniMaxH3Pipeline(
os.path.join(model_path, "video_vae"),
device=self.device,
load_device=component_load_device,
+ decode_only=not self.load_vae_encoder,
trust_remote_code=od_config.trust_remote_code,
)
self.audio_vae = MiniMaxH3AudioVAE(
os.path.join(model_path, "audio_vae"),
device=self.device,
load_device=component_load_device,
+ decode_only=not self.load_vae_encoder,
trust_remote_code=od_config.trust_remote_code,
)
# Registry-side VAE patch-parallel discovery uses ``pipeline.vae``.
@@ -1154,6 +1020,8 @@ class MiniMaxH3Pipeline(
raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous")
loaded_prefixes.add(prefix)
component = getattr(self, prefix.removesuffix("."))
+ if component is None:
+ raise ValueError(f"MiniMax-H3 component {prefix.removesuffix('.')!r} is disabled in this deployment")
stream = ((name[len(prefix) :], tensor) for name, tensor in grouped_weights)
if prefix == "transformer." and self._fasth3 is not None:
# Fuse before the model shards anything, which is also the only
@@ -1227,11 +1095,12 @@ class MiniMaxH3Pipeline(
def _resolve_task(
self,
requested: str | None,
- multi_modal_data: dict[str, Any],
+ multi_modal_data: dict[str, Any] | None = None,
*,
turbo_spec: TurboSpec | None = None,
has_native_lora: bool = False,
) -> str:
+ multi_modal_data = multi_modal_data or {}
if requested is None:
# A Ref2VA-only startup has no FL2VA transformer; preserve its
# historical implicit default even for image-only references.
@@ -1259,217 +1128,6 @@ class MiniMaxH3Pipeline(
self._fasth3.check_task(task)
return task
- def _resolve_shape(
- self,
- task: str,
- sampling: Any,
- image: Image.Image | None,
- ) -> tuple[int, int, int, int, int]:
- fps = int(sampling.fps or MINIMAX_H3_FPS)
- if fps != MINIMAX_H3_FPS:
- raise OmniClientError(f"MiniMax H3 output fps is fixed at {MINIMAX_H3_FPS}")
- extra = sampling.extra_args or {}
- target = extra.get("target")
- if target is not None and not isinstance(target, Mapping):
- raise OmniClientError("MiniMax H3 extra_args['target'] must be an object")
- target = target if isinstance(target, Mapping) else {}
- duration = target.get("duration_seconds", extra.get("duration_seconds", extra.get("duration")))
- if duration is not None:
- if isinstance(duration, bool):
- raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}")
- try:
- duration = float(duration)
- except (TypeError, ValueError) as exc:
- raise OmniClientError(
- f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}"
- ) from exc
- if (
- not math.isfinite(duration)
- or not MINIMAX_H3_MIN_OUTPUT_SECONDS <= duration <= MINIMAX_H3_MAX_OUTPUT_SECONDS
- ):
- raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration}")
- requested_frames = int(round(duration * fps))
- elif int(sampling.num_frames or 1) > 1:
- requested_frames = int(sampling.num_frames)
- else:
- requested_frames = 124 if task == "ref2va" else 209
- duration = requested_frames / fps
- if not MINIMAX_H3_MIN_OUTPUT_SECONDS <= requested_frames / fps <= MINIMAX_H3_MAX_OUTPUT_SECONDS:
- raise OmniClientError(
- f"MiniMax H3 output duration must be in [4, 15] seconds, got {requested_frames / fps:.3f}"
- )
- num_frames = minimax_h3_align_frame_count(requested_frames)
-
- height = sampling.height
- width = sampling.width
- aspect_ratio = target.get("aspect_ratio", extra.get("aspect_ratio"))
- raw_short_edge = target.get("short_edge", extra.get("short_edge", MINIMAX_H3_OUTPUT_SHORT_EDGE))
- if isinstance(raw_short_edge, bool) or not isinstance(raw_short_edge, (int, np.integer)):
- raise OmniClientError(
- f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {raw_short_edge!r}"
- )
- short_edge = int(raw_short_edge)
-
- aspect_ratio = _resolve_minimax_h3_aspect_ratio(
- task,
- aspect_ratio,
- image,
- )
- if not 0.25 <= aspect_ratio <= 4.0:
- raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be in [1:4, 4:1], got {aspect_ratio}")
-
- if height is None or width is None:
- height, width = _resolve_output_canvas(aspect_ratio, short_edge)
- height = int(height) // 32 * 32
- width = int(width) // 32 * 32
- if min(height, width) <= 0:
- raise OmniClientError(f"invalid MiniMax H3 canvas {width}x{height}")
- if width > 4 * height or height > 4 * width:
- raise OmniClientError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]")
-
- latent_t = MINIMAX_H3_SHAPE_PLANNER.video_latent_t(num_frames)
- audio_t = MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(num_frames / fps)
- return height, width, num_frames, latent_t, audio_t
-
- def encode_prompt(
- self,
- *,
- task: str,
- prompt: str,
- image: Image.Image | None = None,
- images: list[Image.Image] | None = None,
- prepared_videos: list[dict[str, Any]] | None = None,
- condition_labels: list[tuple[str, int]] | None = None,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- _, rank, _ = _dit_rank_world()
- hidden = None
- tags = None
- ids = None
- vision_kwargs: dict[str, torch.Tensor] = {}
- images = list(images) if images is not None else ([image] if image is not None else [])
- if rank == 0:
- if task == "t2va":
- ids = minimax_h3_text_only_ids(self.tokenizer, prompt)
- tags = torch.ones(ids.shape[0], dtype=torch.long)
- vision_kwargs = {}
- else:
- image_token_counts: list[int] = []
- if images:
- vision = self.processor.image_processor(
- images=images,
- return_tensors="pt",
- )
- image_grid = vision["image_grid_thw"]
- merge = int(self.processor.image_processor.merge_size) ** 2
- image_token_counts = [int(grid.prod().item()) // merge for grid in image_grid]
- vision_kwargs.update(
- {
- "pixel_values": vision["pixel_values"],
- "image_grid_thw": image_grid,
- }
- )
-
- video_block_counts: list[list[int]] = []
- video_block_timestamps: list[list[float]] = []
- if prepared_videos:
- videos = []
- sampled_videos = []
- for index, item in enumerate(prepared_videos):
- sampled = sample_reference_video_frames(item["prepared_path"])
- videos.append(np.stack(sampled["frames"]))
- sampled_videos.append(sampled)
- vision = self.processor.video_processor(
- videos=videos,
- do_sample_frames=False,
- return_tensors="pt",
- )
- video_grid = vision["video_grid_thw"]
- merge = int(self.processor.image_processor.merge_size) ** 2
- for index, sampled in enumerate(sampled_videos):
- blocks = int(video_grid[index, 0])
- per_block = int(video_grid[index, 1]) * int(video_grid[index, 2]) // merge
- timestamps = sampled["block_timestamps"]
- if len(timestamps) != blocks:
- raise ValueError(
- f"video block count mismatch: processor={blocks}, timestamps={len(timestamps)}"
- )
- video_block_counts.append([per_block] * blocks)
- video_block_timestamps.append(timestamps)
- vision_kwargs.update(
- {
- "pixel_values_videos": vision["pixel_values_videos"],
- "video_grid_thw": video_grid,
- }
- )
-
- if not images and not prepared_videos:
- raise OmniClientError(f"{task} requires an image or video condition")
- if condition_labels is None:
- condition_labels = []
- for image_index in range(1, len(images) + 1):
- condition_labels.append(("image", image_index))
- audio_index = 0
- for video_index, item in enumerate(prepared_videos or (), start=1):
- if item["input_has_audio"]:
- audio_index += 1
- condition_labels.append(("audio", audio_index))
- condition_labels.append(("video", video_index))
-
- if task == "fl2va":
- if prepared_videos:
- raise OmniClientError("fl2va does not accept video conditions")
- ids, tags = minimax_h3_multi_image_presentation(
- self.tokenizer,
- prompt=prompt,
- image_token_counts=image_token_counts,
- )
- elif prepared_videos:
- ids, tags = minimax_h3_ref2va_video_presentation(
- self.tokenizer,
- prompt=prompt,
- condition_labels=condition_labels,
- image_token_count=image_token_counts or None,
- video_block_token_counts=video_block_counts,
- video_block_timestamps=video_block_timestamps,
- )
- else:
- ids, tags = minimax_h3_ref2va_presentation(
- self.tokenizer,
- prompt=prompt,
- condition_labels=condition_labels,
- image_token_count=image_token_counts or None,
- )
-
- logger.info(
- "MiniMax H3 %s Qwen presentation: %d tokens%s",
- task,
- int(ids.shape[0]),
- (
- f", {len(images)} reference images"
- + (f", {len(prepared_videos)} reference videos" if prepared_videos else "")
- if images
- else (f", {len(prepared_videos)} reference videos" if prepared_videos else "")
- ),
- )
-
- if rank < self.text_encoder_tp_size:
- # Distribute the encode inputs from the DiT main rank to the other
- # encoder TP ranks, then run the distributed encode on all of them.
- ids = self._distribute_encode_inputs(ids, vision_kwargs)
- hidden = self._encode_text_hidden(ids, vision_kwargs)
-
- hidden = _broadcast_tensor(
- hidden,
- dtype=torch.bfloat16,
- device=self.device,
- )
- tags = _broadcast_tensor(
- tags,
- dtype=torch.long,
- device=self.device,
- )
- return hidden, tags
-
def _build_text_encoder_group(self, text_encoder_tp_size: int) -> Any:
"""Create the encoder tensor-parallel process group.
@@ -1565,24 +1223,6 @@ class MiniMaxH3Pipeline(
)
return ids
- def _prepare_reference_videos(
- self,
- values: Any,
- *,
- target_frame_count: int,
- workdir: str,
- start_time_seconds: Any = None,
- ) -> list[dict[str, Any]] | None:
- _, rank, _ = _dit_rank_world()
- if rank != 0:
- return None
- return prepare_reference_videos(
- values,
- target_frame_count=target_frame_count,
- workdir=workdir,
- start_time_seconds=start_time_seconds,
- )
-
def _encode_text_hidden(
self,
input_ids: torch.Tensor,
@@ -1726,197 +1366,6 @@ class MiniMaxH3Pipeline(
return tensor.cpu()
return torch.empty(tensor.shape, dtype=tensor.dtype, device="meta")
- def _encode_visual_conditions(
- self,
- images: list[Image.Image],
- prepared_videos: list[dict[str, Any]] | None,
- *,
- video_count: int,
- ) -> tuple[torch.Tensor | None, list[tuple[int, int, int]]]:
- rows: list[torch.Tensor] = []
- shapes: list[tuple[int, int, int]] = []
- _, rank, _ = _dit_rank_world()
- # Keep image and video references in one residency window when both
- # appear in a request; otherwise the video branch would reload the VAE.
- needs_video_vae = video_count > 0 or (rank == 0 and bool(images))
- video_vae_context = self._component_on_device(self.video_vae) if needs_video_vae else nullcontext()
- with video_vae_context:
- if images:
- image_rows = None
- if rank == 0:
- image_rows = torch.cat([self.video_vae.encode_image(image) for image in images])
- rows.append(
- _broadcast_tensor(
- image_rows,
- dtype=torch.float32,
- device=self.device,
- )
- )
- shapes.extend((1, image.height // 16, image.width // 16) for image in images)
- if video_count:
- video_rows, video_shapes = self._encode_video_conditions_resident(
- prepared_videos,
- count=video_count,
- )
- rows.append(video_rows)
- shapes.extend(video_shapes)
- return (torch.cat(rows) if rows else None), shapes
-
- def _encode_audio_conditions_resident(
- self,
- audios: list[tuple[torch.Tensor, int]],
- *,
- max_duration_seconds: float | None = None,
- ) -> tuple[torch.Tensor | None, list[int]]:
- if not audios:
- return None, []
- if max_duration_seconds is not None:
- max_duration_seconds = float(max_duration_seconds)
- if max_duration_seconds <= 0:
- raise ValueError("max_duration_seconds must be positive")
- _, rank, _ = _dit_rank_world()
- rows = None
- lengths = torch.zeros(len(audios), dtype=torch.long, device=self.device)
- if rank == 0:
- bounded_audios = []
- for waveform, sample_rate in audios:
- if max_duration_seconds is not None:
- max_samples = int(round(max_duration_seconds * int(sample_rate)))
- waveform = waveform[..., :max_samples]
- bounded_audios.append((waveform, sample_rate))
- encoded = [self.audio_vae.encode_waveform(*audio) for audio in bounded_audios]
- rows = torch.cat([item[0] for item in encoded])
- lengths = torch.tensor(
- [int(item[1]) for item in encoded],
- dtype=torch.long,
- device=self.device,
- )
- group, _, world_size = _dit_rank_world()
- if world_size > 1:
- dist.broadcast(lengths, src=0, group=group)
- return (
- _broadcast_tensor(rows, dtype=torch.float32, device=self.device),
- [int(value) for value in lengths.tolist()],
- )
-
- def _encode_video_conditions_resident(
- self,
- prepared_videos: list[dict[str, Any]] | None,
- *,
- count: int,
- ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]:
- group, rank, world_size = _dit_rank_world()
- distributed_encode = self.video_vae.is_distributed_enabled()
- if distributed_encode:
- # Native tiled encode uses collectives, so every VPP rank must
- # enter each reference encode in the same input order.
- prepared_videos_list = [prepared_videos]
- dist.broadcast_object_list(
- prepared_videos_list,
- src=0,
- group=group,
- device=self.device,
- )
- prepared_videos = prepared_videos_list[0]
-
- rows = None
- shapes = torch.zeros((count, 3), dtype=torch.long, device=self.device)
- if rank == 0 or distributed_encode:
- if prepared_videos is None or len(prepared_videos) != count:
- raise ValueError("reference-video preparation is incomplete")
- encoded = [
- self.video_vae.encode_video(load_video_frames(item["prepared_path"])) for item in prepared_videos
- ]
- rows = torch.cat([item[0] for item in encoded])
- shapes = torch.tensor(
- [item[1] for item in encoded],
- dtype=torch.long,
- device=self.device,
- )
- if distributed_encode:
- return (
- rows.to(device=self.device, dtype=torch.float32),
- [tuple(int(value) for value in item) for item in shapes.tolist()],
- )
-
- if world_size > 1:
- dist.broadcast(shapes, src=0, group=group)
- return (
- _broadcast_tensor(rows, dtype=torch.float32, device=self.device),
- [tuple(int(value) for value in item) for item in shapes.tolist()],
- )
-
- def _encode_video_audio_conditions_resident(
- self,
- prepared_videos: list[dict[str, Any]] | None,
- *,
- has_audio: list[bool],
- ) -> tuple[torch.Tensor | None, list[int]]:
- _, rank, _ = _dit_rank_world()
- count = sum(has_audio)
- if count == 0:
- return None, []
- rows = None
- lengths = torch.zeros(count, dtype=torch.long, device=self.device)
- if rank == 0:
- if prepared_videos is None:
- raise ValueError("rank 0 reference-video preparation is incomplete")
- encoded = [
- self.audio_vae.encode_waveform(
- *load_video_audio(
- item["original_path"],
- start_time_seconds=float(item.get("start_time_seconds", 0.0)),
- duration_seconds=item.get(
- "audio_duration_seconds",
- item.get("duration_seconds"),
- ),
- )
- )
- for item in prepared_videos
- if item["input_has_audio"]
- ]
- rows = torch.cat([item[0] for item in encoded])
- lengths = torch.tensor(
- [item[1] for item in encoded],
- dtype=torch.long,
- device=self.device,
- )
- group, _, world_size = _dit_rank_world()
- if world_size > 1:
- dist.broadcast(lengths, src=0, group=group)
- return (
- _broadcast_tensor(rows, dtype=torch.float32, device=self.device),
- [int(value) for value in lengths.tolist()],
- )
-
- def _encode_reference_audio_conditions(
- self,
- prepared_videos: list[dict[str, Any]] | None,
- *,
- has_audio: list[bool],
- standalone_audios: list[tuple[torch.Tensor, int]],
- max_duration_seconds: float,
- ) -> tuple[torch.Tensor | None, list[int], torch.Tensor | None, list[int]]:
- # Embedded and standalone audio are consecutive direct Audio-VAE
- # calls. Keep the component resident across both paths.
- needs_audio_vae = any(has_audio) or bool(standalone_audios)
- audio_vae_context = self._component_on_device(self.audio_vae) if needs_audio_vae else nullcontext()
- with audio_vae_context:
- embedded_condition, embedded_lengths = self._encode_video_audio_conditions_resident(
- prepared_videos,
- has_audio=has_audio,
- )
- external_condition, external_lengths = self._encode_audio_conditions_resident(
- standalone_audios,
- max_duration_seconds=max_duration_seconds,
- )
- return (
- embedded_condition,
- embedded_lengths,
- external_condition,
- external_lengths,
- )
-
def _initial_noise(
self,
*,
@@ -2342,13 +1791,16 @@ class MiniMaxH3Pipeline(
if isinstance(raw_prompt, str):
return None
additional_information = raw_prompt.get("additional_information") or {}
- text_encoder_output = additional_information.get("text_encoder_output")
- if text_encoder_output is None:
+ encoder_output = additional_information.get("encoder_output")
+ if encoder_output is None:
return None
- if not isinstance(text_encoder_output, Mapping):
- raise OmniClientError("text_encoder_output must be a mapping")
+ if not isinstance(encoder_output, Mapping):
+ raise OmniClientError("MiniMax H3 encoder output must be a mapping")
try:
- return MiniMaxH3TextConditioning.from_payload(text_encoder_output)
+ if "hidden_states" in encoder_output and "token_tags" in encoder_output:
+ return MiniMaxH3TextConditioning.from_payload(encoder_output)
+ conditioning = MiniMaxH3EncoderConditioning.from_omni_payload(encoder_output)
+ return MiniMaxH3TextConditioning(conditioning.hidden_states, conditioning.token_tags)
except ValueError as exc:
raise OmniClientError(str(exc)) from exc
@@ -2369,34 +1821,234 @@ class MiniMaxH3Pipeline(
raise OmniClientError(str(exc)) from exc
return videos
- def _prepare_request_inputs(
- self,
- *,
- prompt: str,
- multi_modal_data: dict[str, Any],
- sampling: Any,
- text_conditioning: MiniMaxH3TextConditioning | None = None,
- prepared_reference_videos: list[dict[str, Any]] | None = None,
- ) -> dict[str, Any]:
- """Resolve the task and output shape, then run every request-level encode.
-
- Shared by request-mode :meth:`forward` and step-mode
- :meth:`prepare_encode`; the returned mapping feeds :meth:`diffuse` and
- :meth:`_build_denoise_inputs` unchanged.
- """
- quality = sampling.quality
- logger.debug("MiniMax H3 request quality=%s", quality)
- extra = sampling.extra_args or {}
- preencode_batch_frames = (
- normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17))
- if extra.get("preencode_mp4", False)
- else 17
+ def encode_prompt(self, prepared: PreparedEncoderInputs | None) -> tuple[torch.Tensor, torch.Tensor]:
+ _, rank, _ = _dit_rank_world()
+ ids = tags = hidden = None
+ vision_kwargs: dict[str, torch.Tensor] = {}
+ error = None
+ if rank == 0:
+ try:
+ if prepared is None:
+ raise ValueError("rank 0 must prepare MiniMax H3 text inputs")
+ if prepared.images:
+ vision = self.processor.image_processor(images=prepared.images, return_tensors="pt")
+ vision_kwargs.update(
+ pixel_values=vision["pixel_values"],
+ image_grid_thw=vision["image_grid_thw"],
+ )
+ if prepared.qwen_videos:
+ vision = self.processor.video_processor(
+ videos=[frames for frames, _ in prepared.qwen_videos],
+ do_sample_frames=False,
+ return_tensors="pt",
+ )
+ vision_kwargs.update(
+ pixel_values_videos=vision["pixel_values_videos"],
+ video_grid_thw=vision["video_grid_thw"],
+ )
+ ids, tags = build_minimax_h3_presentation(
+ self.tokenizer,
+ prompt=prepared.prompt,
+ task=prepared.media.task,
+ condition_labels=prepared.condition_labels,
+ image_grid_thw=vision_kwargs.get("image_grid_thw"),
+ video_grid_thw=vision_kwargs.get("video_grid_thw"),
+ video_timestamps=prepared.video_timestamps,
+ merge_size=int(self.processor.image_processor.merge_size),
+ )
+ except Exception as exc:
+ error = exc
+ _broadcast_rank0_exception(error)
+ if rank < self.text_encoder_tp_size:
+ ids = self._distribute_encode_inputs(ids, vision_kwargs)
+ hidden = self._encode_text_hidden(ids, vision_kwargs)
+ return (
+ _broadcast_tensor(hidden, dtype=torch.bfloat16, device=self.device),
+ _broadcast_tensor(tags, dtype=torch.long, device=self.device),
)
+
+ def _distribute_media_inputs(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniMaxH3EncoderMediaInput:
+ group, rank, world_size = _dit_rank_world()
+ if world_size == 1:
+ assert media is not None
+ return media
+ tensors = media.to_mm_tensors() if media is not None else []
+ header = [(media.to_metadata(), [value.dtype for value in tensors])] if media is not None else [None]
+ dist.broadcast_object_list(header, src=0, group=group)
+ metadata, dtypes = header[0]
+ received = []
+ for index, dtype in enumerate(dtypes):
+ value = _broadcast_tensor(tensors[index] if rank == 0 else None, dtype=dtype, device=self.device)
+ if rank != 0:
+ received.append(value.cpu())
+ del value
+ if rank == 0:
+ assert media is not None
+ return media
+ return MiniMaxH3EncoderMediaInput.from_mm_tensors(received, metadata)
+
+ def _broadcast_media_conditioning(
+ self, conditioning: MiniMaxH3EncoderMediaConditioning | None
+ ) -> MiniMaxH3EncoderMediaConditioning:
+ group, rank, world_size = _dit_rank_world()
+ if world_size == 1:
+ assert conditioning is not None
+ return conditioning
+ tensor_names = ("visual_condition", "audio_condition")
+ header = [None]
+ if rank == 0:
+ assert conditioning is not None
+ metadata = {
+ item.name: getattr(conditioning, item.name)
+ for item in fields(conditioning)
+ if item.name not in tensor_names
+ }
+ dtypes = {
+ name: value.dtype if (value := getattr(conditioning, name)) is not None else None
+ for name in tensor_names
+ }
+ header[0] = (metadata, dtypes)
+ dist.broadcast_object_list(header, src=0, group=group)
+ metadata, dtypes = header[0]
+ tensors = {}
+ for name, dtype in dtypes.items():
+ source = getattr(conditioning, name) if rank == 0 else None
+ tensors[name] = _broadcast_tensor(source, dtype=dtype, device=self.device) if dtype is not None else None
+ return MiniMaxH3EncoderMediaConditioning(**metadata, **tensors)
+
+ def _encode_local_media(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniMaxH3EncoderMediaConditioning:
+ group, rank, world_size = _dit_rank_world()
+ distributed_video = self.video_vae.is_distributed_enabled()
+ if distributed_video:
+ media = self._distribute_media_inputs(media)
+ conditioning = None
+ error = None
+ try:
+ if rank == 0 or distributed_video:
+ assert media is not None
+ conditioning = encode_media(
+ media,
+ video_vae=self.video_vae,
+ audio_vae=self.audio_vae,
+ emit_conditioning=rank == 0,
+ component_scope=self._component_on_device,
+ )
+ except ValueError as exc:
+ error = OmniClientError(str(exc))
+ except Exception as exc:
+ error = exc
+ if world_size > 1:
+ # All participants finish the codec phase before any latent broadcast.
+ errors = [None] * world_size
+ info = (str(error), isinstance(error, OmniClientError)) if error is not None else None
+ dist.all_gather_object(errors, info, group=group)
+ for info in errors:
+ if info is not None:
+ if error is not None:
+ raise error
+ message, is_client_error = info
+ raise OmniClientError(message) if is_client_error else RuntimeError(message)
+ elif error is not None:
+ raise error
+ return self._broadcast_media_conditioning(conditioning)
+
+ def _prepare_local_conditioning(
+ self,
+ raw_prompt: Any,
+ sampling: Any,
+ *,
+ require_external_text: bool = False,
+ ) -> MiniMaxH3EncoderConditioning:
+ group, rank, world_size = _dit_rank_world()
+ prepared = text_conditioning = None
+ error = None
+ if rank == 0:
+ try:
+ _, multi_modal_data = self._extract_prompt(raw_prompt)
+ turbo_spec = self._active_turbo_spec(sampling)
+ has_native_lora = self._has_active_native_lora(sampling)
+ task = self._resolve_task(
+ (sampling.extra_args or {}).get("task"),
+ multi_modal_data,
+ turbo_spec=turbo_spec,
+ has_native_lora=has_native_lora,
+ )
+ if turbo_spec is not None:
+ self._validate_turbo_sampling(sampling, turbo_spec)
+ if has_native_lora:
+ self._validate_native_sampling(sampling, task=task)
+ text_conditioning = self._extract_text_conditioning(raw_prompt)
+ if require_external_text and text_conditioning is None:
+ raise OmniClientError(
+ "MiniMax H3 diffusion stage requires text encoder conditioning when text_encoder is not loaded"
+ )
+ prepared = prepare_encoder_inputs(
+ raw_prompt,
+ sampling,
+ task=task,
+ prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt),
+ )
+ except Exception as exc:
+ error = exc
+ _broadcast_rank0_exception(error)
+ reuse_text = [text_conditioning is not None]
+ if world_size > 1:
+ dist.broadcast_object_list(reuse_text, src=0, group=group)
+ if reuse_text[0]:
+ hidden = _broadcast_tensor(
+ text_conditioning.hidden_states if rank == 0 else None,
+ dtype=torch.bfloat16,
+ device=self.device,
+ )
+ tags = _broadcast_tensor(
+ text_conditioning.token_tags if rank == 0 else None, dtype=torch.long, device=self.device
+ )
+ else:
+ hidden, tags = self.encode_prompt(prepared)
+ media = self._encode_local_media(prepared.media if prepared is not None else None)
+ return MiniMaxH3EncoderConditioning.from_components(MiniMaxH3TextConditioning(hidden, tags), media)
+
+ def _prepare_request_inputs(self, raw_prompt: Any, sampling: Any) -> dict[str, Any]:
+ if self.load_text_encoder:
+ conditioning = self._prepare_local_conditioning(raw_prompt, sampling)
+ elif self.load_vae_encoder:
+ conditioning = self._prepare_local_conditioning(
+ raw_prompt,
+ sampling,
+ require_external_text=True,
+ )
+ else:
+ conditioning = self._extract_encoder_conditioning(raw_prompt)
+ return self._prepare_encoder_conditioning_inputs(conditioning, sampling)
+
+ @staticmethod
+ def _extract_encoder_conditioning(prompt: Any) -> MiniMaxH3EncoderConditioning:
+ if isinstance(prompt, list):
+ prompt = prompt[0] if prompt else None
+ additional_information = prompt.get("additional_information") if isinstance(prompt, Mapping) else None
+ payload = additional_information.get("encoder_output") if isinstance(additional_information, Mapping) else None
+ if not isinstance(payload, Mapping) or not payload:
+ raise OmniClientError("MiniMax H3 diffusion stage requires encoder conditioning from the encoder stage")
+ try:
+ return MiniMaxH3EncoderConditioning.from_omni_payload(payload)
+ except ValueError as exc:
+ raise OmniClientError(str(exc)) from exc
+
+ def _prepare_encoder_conditioning_inputs(
+ self,
+ conditioning: MiniMaxH3EncoderConditioning,
+ sampling: Any,
+ ) -> dict[str, Any]:
+ extra = sampling.extra_args or {}
+ requested_task = extra.get("task")
+ if requested_task is not None and str(requested_task).lower() != conditioning.task:
+ raise OmniClientError(
+ f"MiniMax H3 encoder task {conditioning.task!r} does not match diffusion request {requested_task!r}"
+ )
turbo_spec = self._active_turbo_spec(sampling)
has_native_lora = self._has_active_native_lora(sampling)
task = self._resolve_task(
- extra.get("task"),
- multi_modal_data,
+ conditioning.task,
turbo_spec=turbo_spec,
has_native_lora=has_native_lora,
)
@@ -2411,263 +2063,73 @@ class MiniMaxH3Pipeline(
audio_shift=self.default_audio_shift,
)
- raw_image = multi_modal_data.get("image")
- raw_videos = multi_modal_data.get("video")
- raw_audio = multi_modal_data.get("audio")
- images = _load_images(raw_image) if raw_image is not None else []
- video_values = list(raw_videos) if isinstance(raw_videos, (list, tuple)) else raw_videos
- audio_values = list(raw_audio) if isinstance(raw_audio, (list, tuple)) else raw_audio
-
- if task == "t2va" and (images or raw_videos is not None or raw_audio is not None):
- raise OmniClientError("t2va does not accept image, video, or audio conditions")
- if task == "fl2va":
- if not images:
- raise OmniClientError("fl2va requires multi_modal_data.image")
- if len(images) > 2:
- raise OmniClientError("fl2va accepts at most first and last images")
- if raw_videos is not None or raw_audio is not None:
- raise OmniClientError("fl2va accepts image keyframes only")
- if task == "ref2va":
- video_count = (
- len(video_values) if isinstance(video_values, (list, tuple)) else int(video_values is not None)
+ if conditioning.height % 32 or conditioning.width % 32:
+ raise OmniClientError(
+ f"MiniMax H3 encoder canvas must be divisible by 32, got {conditioning.width}x{conditioning.height}"
)
- audio_is_waveform_pair = (
- isinstance(raw_audio, (list, tuple))
- and len(raw_audio) == 2
- and isinstance(raw_audio[1], (int, np.integer))
+ expected_latent_t = MINIMAX_H3_SHAPE_PLANNER.video_latent_t(conditioning.num_frames)
+ expected_audio_t = MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(conditioning.num_frames / MINIMAX_H3_FPS)
+ if (conditioning.latent_t, conditioning.audio_t) != (expected_latent_t, expected_audio_t):
+ raise OmniClientError(
+ "MiniMax H3 encoder latent shape does not match its output frame count: "
+ f"got ({conditioning.latent_t}, {conditioning.audio_t}), expected "
+ f"({expected_latent_t}, {expected_audio_t})"
)
- audio_count = (
- len(audio_values)
- if isinstance(audio_values, (list, tuple)) and not audio_is_waveform_pair
- else int(raw_audio is not None)
- )
- _validate_ref2va_reference_counts(len(images), video_count, audio_count)
- elif raw_videos is not None:
- raise OmniClientError(f"{task} does not accept a video condition")
- image = images[0] if images else None
- height, width, num_frames, latent_t, audio_t = self._resolve_shape(task, sampling, image)
- if task == "fl2va":
- for item in images:
- _validate_reference_image(item)
- prepared_images = [item.resize((width, height), Image.Resampling.LANCZOS) for item in images]
- keyframe_frame_indices = _resolve_fl2va_keyframe_indices(extra, len(images))
- elif task == "ref2va":
- prepared_images = []
- for item in images:
- ref_width, ref_height = _reference_image_shape(item)
- prepared_images.append(item.resize((ref_width, ref_height), Image.Resampling.LANCZOS))
- keyframe_frame_indices = None
- else:
- prepared_images = []
- keyframe_frame_indices = None
+ visual_shapes = list(conditioning.visual_condition_shapes) or None
+ audio_lengths = list(conditioning.audio_condition_lengths) or None
- visual_condition = None
- visual_shape = None
- visual_shapes = None
- audio_condition = None
- ref_audio_t = None
- audio_lengths = None
- ref_blocks = None
- with tempfile.TemporaryDirectory(prefix="minimax_h3_ref2va_") as workdir:
- prepared_videos = None
- has_audio: list[bool] = []
- video_count = 0
- if raw_videos is not None:
- video_count = len(raw_videos) if isinstance(raw_videos, (list, tuple)) else 1
- # File-based reference-video prep runs only on rank 0; other
- # ranks return None without touching disk. If rank 0 raises
- # (e.g. invalid file, unsupported codec) it must not exit
- # ``prepare_encode`` before the downstream broadcasts below --
- # non-zero ranks would then deadlock on them forever. Capture
- # the exception here and let every rank agree on the outcome
- # before starting any subsequent collective.
- prep_error: Exception | None = None
- try:
- _, rank, _ = _dit_rank_world()
- if prepared_reference_videos is not None:
- if rank == 0:
- prepared_videos = _reuse_prepared_reference_videos(
- prepared_reference_videos,
- expected_count=video_count,
- )
- else:
- prepared_videos = self._prepare_reference_videos(
- raw_videos,
- target_frame_count=num_frames,
- workdir=workdir,
- start_time_seconds=extra.get("start_time_seconds"),
- )
- except Exception as exc:
- prep_error = exc
- _broadcast_rank0_exception(prep_error)
- has_audio_tensor = torch.zeros(
- video_count,
- dtype=torch.long,
- device=self.device,
- )
- _, rank, world_size = _dit_rank_world()
- if rank == 0:
- has_audio_tensor = torch.tensor(
- [int(item["input_has_audio"]) for item in prepared_videos or []],
- dtype=torch.long,
- device=self.device,
- )
- if world_size > 1:
- dist.broadcast(
- has_audio_tensor,
- src=0,
- group=get_world_group().device_group,
- )
- has_audio = [bool(value) for value in has_audio_tensor.tolist()]
-
- if raw_audio is not None:
- validate_reference_audio_files(raw_audio)
- standalone_audios = _load_audios(raw_audio) if raw_audio is not None else []
- validate_reference_audio_waveforms(standalone_audios)
- condition_labels: list[tuple[str, int]] = []
- for image_index in range(1, len(prepared_images) + 1):
- condition_labels.append(("image", image_index))
- audio_index = 0
- for video_index, item in enumerate(prepared_videos or (), start=1):
- if item["input_has_audio"]:
- audio_index += 1
- condition_labels.append(("audio", audio_index))
- condition_labels.append(("video", video_index))
- for _ in standalone_audios:
- audio_index += 1
- condition_labels.append(("audio", audio_index))
-
- if text_conditioning is not None:
- text_embeddings = text_conditioning.hidden_states.to(
- device=self.device,
- dtype=torch.bfloat16,
- )
- text_tags = text_conditioning.token_tags.to(
- device=self.device,
- dtype=torch.long,
- )
- elif getattr(self, "text_encoder", None) is not None:
- text_embeddings, text_tags = self.encode_prompt(
- task=task,
- prompt=prompt,
- images=prepared_images,
- prepared_videos=prepared_videos,
- condition_labels=condition_labels if task == "ref2va" else None,
- )
- else:
- raise OmniClientError(
- "MiniMax H3 diffusion stage requires text_encoder_output when text_encoder is not loaded"
- )
-
- # ``prepared_videos`` is intentionally ``None`` on non-zero DiT
- # ranks; the distributed video encoder broadcasts the prepared
- # metadata inside ``_encode_visual_conditions``. Use the global
- # video count here so video + standalone-audio Ref2VA requests do
- # not look like audio-only requests on those ranks.
- if video_count or prepared_images:
- visual_condition, visual_shapes = self._encode_visual_conditions(
- prepared_images,
- prepared_videos,
- video_count=video_count,
- )
- (
- embedded_audio_condition,
- embedded_audio_lengths,
- external_audio_condition,
- external_audio_lengths,
- ) = self._encode_reference_audio_conditions(
- prepared_videos,
- has_audio=has_audio,
- standalone_audios=standalone_audios,
- max_duration_seconds=float(num_frames) / float(sampling.fps or MINIMAX_H3_FPS),
- )
- audio_parts = [
- item for item in (embedded_audio_condition, external_audio_condition) if item is not None
- ]
- audio_condition = torch.cat(audio_parts) if audio_parts else None
- audio_lengths = embedded_audio_lengths + external_audio_lengths
- ref_blocks = []
- image_shapes = visual_shapes[: len(prepared_images)]
- video_shapes = visual_shapes[len(prepared_images) :]
- for shape in image_shapes:
- ref_blocks.append(
- {
- "kind": "image",
- "latent_h": shape[1],
- "latent_w": shape[2],
- }
- )
- audio_iterator = iter(embedded_audio_lengths)
- for shape, contributes_audio in zip(video_shapes, has_audio, strict=True):
- ref_audio = next(audio_iterator) if contributes_audio else 0
- ref_blocks.append(
- {
- "kind": "video_audio" if ref_audio else "video",
- "ref_audio_t": ref_audio,
- "latent_t": shape[0],
- "latent_h": shape[1],
- "latent_w": shape[2],
- }
- )
- for ref_audio_t in external_audio_lengths:
- ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t})
- elif standalone_audios:
- raise OmniClientError("standalone audio references require a Ref2VA visual reference")
-
- if visual_shapes and len(visual_shapes) == 1:
- visual_shape = visual_shapes[0]
- if audio_lengths:
- if any(length < 80 or length > 600 for length in audio_lengths):
- raise OmniClientError("MiniMax H3 audio references must each be between 2 and 15 seconds")
- # Video soundtracks and standalone audio have separate
- # 15-second budgets, validated before encoding. Do not merge
- # those budgets when concatenating their conditioning rows.
- if len(audio_lengths) == 1:
- ref_audio_t = audio_lengths[0]
-
- seed = int(sampling.seed if sampling.seed is not None else 42)
base_schedule, num_steps = self._resolve_sigma_positions(task, sampling)
- video_shift = float(extra.get("flow_shift", self.default_video_shift))
- audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift))
- pad_seq_len = _resolve_pad_seq_len(extra.get("pad_seq_len"))
quality_plan = self._quality_policy.resolve(
- quality=quality,
+ quality=sampling.quality,
num_inference_steps=num_steps,
extra_args=extra,
)
self._cache_dit_runtime.prepare(quality_plan.cache_dit)
- num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt)
return {
"task": task,
- "height": height,
- "width": width,
- "num_frames": num_frames,
- "latent_t": latent_t,
- "latent_h": height // 16,
- "latent_w": width // 16,
- "audio_t": audio_t,
- "text_embeddings": text_embeddings,
- "text_tags": text_tags,
- "visual_condition": visual_condition,
- "visual_condition_shape": visual_shape,
- "audio_condition": audio_condition,
- "ref_audio_t": ref_audio_t,
- "ref_blocks": ref_blocks,
+ "height": conditioning.height,
+ "width": conditioning.width,
+ "num_frames": conditioning.num_frames,
+ "latent_t": conditioning.latent_t,
+ "latent_h": conditioning.height // 16,
+ "latent_w": conditioning.width // 16,
+ "audio_t": conditioning.audio_t,
+ "text_embeddings": conditioning.hidden_states.to(device=self.device, dtype=torch.bfloat16),
+ "text_tags": conditioning.token_tags.to(device=self.device, dtype=torch.long),
+ "visual_condition": (
+ conditioning.visual_condition.to(device=self.device)
+ if conditioning.visual_condition is not None
+ else None
+ ),
+ "visual_condition_shape": visual_shapes[0] if visual_shapes and len(visual_shapes) == 1 else None,
+ "audio_condition": (
+ conditioning.audio_condition.to(device=self.device)
+ if conditioning.audio_condition is not None
+ else None
+ ),
+ "ref_audio_t": audio_lengths[0] if audio_lengths and len(audio_lengths) == 1 else None,
+ "ref_blocks": list(conditioning.ref_blocks) or None,
"visual_condition_shapes": visual_shapes,
"audio_condition_lengths": audio_lengths,
- "keyframe_frame_indices": keyframe_frame_indices,
- "pad_seq_len": pad_seq_len,
- "seed": seed,
+ "keyframe_frame_indices": list(conditioning.keyframe_frame_indices) or None,
+ "pad_seq_len": _resolve_pad_seq_len(extra.get("pad_seq_len")),
+ "seed": int(sampling.seed if sampling.seed is not None else 42),
"num_steps": num_steps,
- "video_shift": video_shift,
- "audio_shift": audio_shift,
+ "video_shift": float(extra.get("flow_shift", self.default_video_shift)),
+ "audio_shift": float(extra.get("audio_flow_shift", self.default_audio_shift)),
"base_schedule": base_schedule,
- "num_outputs": num_outputs,
+ "num_outputs": _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt),
"preencode_mp4": bool(extra.get("preencode_mp4", False)),
- "preencode_batch_frames": preencode_batch_frames,
+ "preencode_batch_frames": (
+ normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17))
+ if extra.get("preencode_mp4", False)
+ else 17
+ ),
"video_codec_options": normalize_video_codec_options(
extra.get("video_codec_options", {"preset": "ultrafast", "threads": "0"})
+ if extra.get("preencode_mp4", False)
+ else None
),
}
@@ -2680,14 +2142,9 @@ class MiniMaxH3Pipeline(
def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
if len(request.prompts) != 1:
raise OmniClientError("MiniMax H3 supports one request at a time")
- raw_prompt = request.prompts[0]
- prompt, multi_modal_data = self._extract_prompt(raw_prompt)
context = self._prepare_request_inputs(
- prompt=prompt,
- multi_modal_data=multi_modal_data,
- sampling=request.sampling_params,
- text_conditioning=self._extract_text_conditioning(raw_prompt),
- prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt),
+ request.prompts[0],
+ request.sampling_params,
)
denoise_kwargs = self._denoise_kwargs(context)
num_outputs = context["num_outputs"]
@@ -2794,13 +2251,9 @@ class MiniMaxH3Pipeline(
"co-batched requests would reuse incompatible cache state. Drop --step-execution "
"or omit quality=high."
)
- prompt, multi_modal_data = self._extract_prompt(state.prompt)
context = self._prepare_request_inputs(
- prompt=prompt,
- multi_modal_data=multi_modal_data,
- sampling=state.sampling,
- text_conditioning=self._extract_text_conditioning(state.prompt),
- prepared_reference_videos=self._extract_prepared_reference_videos(state.prompt),
+ state.prompt,
+ state.sampling,
)
inputs = self._build_denoise_inputs(**self._denoise_kwargs(context))
diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py
index e40251dff..d02649229 100644
--- a/vllm_omni/diffusion/models/minimax_h3/vae.py
+++ b/vllm_omni/diffusion/models/minimax_h3/vae.py
@@ -15,8 +15,10 @@ import torch
import torch.distributed as dist
import torch.nn as nn
from PIL import Image
+from safetensors import safe_open
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from vllm.logger import init_logger
+from vllm.utils.torch_utils import set_default_torch_dtype
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
DistributedVaeMixin,
@@ -106,6 +108,121 @@ def _load_remote_component(
return component_cls.from_pretrained(component_path)
+def _check_trust_remote_code(
+ component_path: str,
+ config: dict[str, Any],
+ trust_remote_code: bool,
+) -> None:
+ """Apply the remote-code gate before either full or encoder-only loading."""
+
+ class_reference = (config.get("auto_map") or {}).get("AutoModel")
+ if isinstance(class_reference, str) and not trust_remote_code:
+ raise ValueError(
+ f"Loading {component_path} executes the modeling code shipped with "
+ f"the checkpoint (auto_map.AutoModel = {class_reference}). Pass "
+ "--trust-remote-code (or trust_remote_code=True) to allow it."
+ )
+
+
+def _load_selected_state_dict(
+ model: nn.Module,
+ weights_path: Path,
+ prefixes: tuple[str, ...],
+) -> None:
+ with safe_open(str(weights_path), framework="pt", device="cpu") as checkpoint:
+ selected = sorted(name for name in checkpoint.keys() if name.startswith(prefixes))
+ if not selected:
+ raise ValueError(f"{weights_path}: no VAE weights match prefixes {prefixes!r}")
+ state_dict = {name: checkpoint.get_tensor(name) for name in selected}
+ model.load_state_dict(state_dict, strict=True)
+
+
+def _remove_modules(model: nn.Module, names: tuple[str, ...]) -> None:
+ for name in names:
+ if not isinstance(getattr(model, name, None), nn.Module):
+ raise ValueError(f"MiniMax H3 VAE model has no module {name!r}")
+ delattr(model, name)
+
+
+def _load_video_vae_encoder(
+ component_path: str,
+ config: dict[str, Any],
+) -> nn.Module:
+ class_reference = (config.get("auto_map") or {}).get("AutoModel")
+ if not isinstance(class_reference, str):
+ raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel")
+ component_cls = get_class_from_dynamic_module(class_reference, component_path)
+ component_module = importlib.import_module(component_cls.__module__)
+ source_classes = getattr(component_module, "_SOURCE_CLASSES", None)
+ source_class_name = config.get("source_class_name")
+ if not isinstance(source_classes, dict) or source_class_name not in source_classes:
+ raise ValueError(f"unsupported MiniMax H3 video VAE source class {source_class_name!r}")
+ if bool(config["vae_parallel_tiling"]):
+ ensure_parallel_state = getattr(component_module, "_ensure_vae_parallel_state", None)
+ if not callable(ensure_parallel_state):
+ raise ValueError("MiniMax H3 video VAE remote module does not expose its parallel-state initializer")
+ ensure_parallel_state()
+
+ source_path = Path(component_path) / str(config["source_path"])
+ weights_path = source_path / str(config["source_safetensors_path"])
+ load_kwargs = {
+ "clip_length": int(config["vae_clip_length"]),
+ "token_drop": int(config["vae_token_drop"]),
+ "encoder_tiling": int(config["vae_encoder_tiling"]),
+ "decoder_tiling": int(config["vae_decoder_tiling"]),
+ "parallel_tiling": int(config["vae_parallel_tiling"]),
+ "tile_size": int(config["vae_tile_size"]),
+ "tile_overlap_min": int(config["vae_tile_overlap_min"]),
+ "encoder_parallel": int(config["vae_encoder_parallel"]),
+ "decoder_parallel": int(config["vae_decoder_parallel"]),
+ "chunk_dim": int(config["vae_chunk_dim"]),
+ }
+ source_cls = source_classes[source_class_name]
+ source_config = source_cls.load_config(str(source_path))
+ with set_default_torch_dtype(torch.float32), torch.device("cpu"):
+ model, _unused = source_cls.from_config(source_config, return_unused_kwargs=True, **load_kwargs)
+ _remove_modules(model, ("decoder", "post_quant_conv"))
+ _load_selected_state_dict(model, weights_path, ("encoder.", "quant_conv."))
+ model.eval()
+ return component_cls(model)
+
+
+def _load_audio_vae_encoder(
+ component_path: str,
+ config: dict[str, Any],
+) -> nn.Module:
+ class_reference = (config.get("auto_map") or {}).get("AutoModel")
+ if not isinstance(class_reference, str):
+ raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel")
+ component_cls = get_class_from_dynamic_module(class_reference, component_path)
+ component_module = importlib.import_module(component_cls.__module__)
+ load_yaml = getattr(component_module, "_load_yaml", None)
+ model_cls = getattr(component_module, "DacAudioVAE", None)
+ if not callable(load_yaml) or not isinstance(model_cls, type):
+ raise ValueError("MiniMax H3 audio VAE remote module does not expose its source loader")
+
+ component_dir = Path(component_path)
+ audio_config = load_yaml(component_dir / str(config["source_config_path"]))
+ metadata_path = component_dir / str(config["source_metadata_path"])
+ metadata_doc = json.loads(metadata_path.read_text(encoding="utf-8"))
+ metadata = metadata_doc["metadata"]["kwargs"]
+ with set_default_torch_dtype(torch.float32), torch.device("cpu"):
+ model = model_cls(
+ encoder_rates=metadata["encoder_rates"],
+ decoder_rates=metadata["decoder_rates"],
+ attn_proj=metadata["attn_proj"],
+ decoder_type=metadata["decoder_type"],
+ decoder_dim=audio_config["model_config"]["decoder_dim"],
+ vae_latent_channels=audio_config["model_config"]["vae_latent_channels"],
+ sample_rate=metadata["sample_rate"],
+ )
+ _remove_modules(model, ("decoder", "dec_in_proj", "logs_proj"))
+ weights_path = component_dir / str(config["source_safetensors_path"])
+ _load_selected_state_dict(model, weights_path, ("encoder.", "pre_block.", "mean_proj."))
+ model.eval()
+ return component_cls(model)
+
+
class _AudioVAEDeterminismContext(AbstractContextManager):
def __enter__(self):
backends = torch.backends
@@ -156,16 +273,24 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
*,
device: torch.device,
load_device: torch.device | None = None,
+ encode_only: bool = False,
+ decode_only: bool = False,
trust_remote_code: bool = False,
) -> None:
super().__init__()
self._device_target = device
+ self.encode_only = bool(encode_only)
+ self.decode_only = bool(decode_only)
+ if self.encode_only and self.decode_only:
+ raise ValueError("MiniMax H3 video VAE cannot be both encode-only and decode-only")
self.config_dict = _load_component_config(component_path)
- self.remote = _load_remote_component(
- component_path,
- self.config_dict,
- trust_remote_code=trust_remote_code,
- )
+ _check_trust_remote_code(component_path, self.config_dict, trust_remote_code)
+ if self.encode_only:
+ self.remote = _load_video_vae_encoder(component_path, self.config_dict)
+ else:
+ self.remote = _load_remote_component(component_path, self.config_dict, trust_remote_code=trust_remote_code)
+ if self.decode_only:
+ _remove_modules(self.remote.model, ("encoder", "quant_conv"))
# Match the reference loader contract before installing inference-only
# decoder fast paths. Keyframe encoding remains FP32; decoder Linear
# weights may be materialized in FP16 because reference decode casts
@@ -222,10 +347,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
self,
parallel_size: int,
mode: str = "tile",
+ process_group: dist.ProcessGroup | None = None,
) -> None:
if mode != "tile":
raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
- group = get_world_group().device_group
+ group = process_group if process_group is not None else get_world_group().device_group
world_size = dist.get_world_size(group)
rank = dist.get_rank(group)
parallel_size = int(parallel_size)
@@ -412,15 +538,40 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
cols, _, _ = self.model.split_tiles(int(latent.shape[-1]) * ratio, True)
return len(rows) * len(cols)
+ def _encoder_tile_count(self, height: int, width: int) -> int:
+ processor = getattr(self.model, "processor", None)
+ align = getattr(processor, "_align_to_total_patch_size", None)
+ if callable(align):
+ height, width = align(int(height), int(width))
+ rows, _, _ = self.model.split_tiles(int(height), False)
+ cols, _, _ = self.model.split_tiles(int(width), False)
+ return len(rows) * len(cols)
+
+ def _encoder_tiling_context(
+ self,
+ height: int,
+ width: int,
+ ) -> AbstractContextManager:
+ parallel_size = int(getattr(self, "parallel_size", 1))
+ if parallel_size <= 1:
+ return nullcontext()
+ num_tiles = self._encoder_tile_count(height, width)
+ if num_tiles >= parallel_size:
+ return nullcontext()
+ logger.warning_once(
+ "MiniMax-H3 VAE encode splits a %dx%d input into %d tile(s) but "
+ "the tile group has %d ranks; encoding rank-locally for this "
+ "shape instead, which avoids ranks without tiles hanging the "
+ "collective.",
+ width,
+ height,
+ num_tiles,
+ parallel_size,
+ )
+ return self._rank_local_tiling()
+
@contextmanager
def _rank_local_tiling(self) -> Iterator[None]:
- """Run one decode with tiling kept on this rank, then restore the group.
-
- Used only when there are fewer tiles than ranks. Every rank then decodes
- every tile, which is slower than sharing the work but is correct and
- involves no collective.
- """
-
state = self._native_parallel_state()
saved_state = dict(state)
saved_tiling = self.model.parallel_tiling
@@ -448,15 +599,24 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
@torch.inference_mode()
def encode_image(self, image: Image.Image) -> torch.Tensor:
+ if getattr(self, "decode_only", False):
+ raise RuntimeError("MiniMax H3 decode-only video VAE cannot encode images")
previous_parallel = self.model.parallel_tiling
- self.model.parallel_tiling = False
+ if int(getattr(self, "parallel_size", 1)) <= 1:
+ self.model.parallel_tiling = False
parameter = next(self.parameters())
previous_dtype = parameter.dtype
if previous_dtype != torch.float32:
self.to(torch.float32)
devices = [parameter.device] if parameter.device.type != "cpu" else []
try:
- with torch.random.fork_rng(devices=devices, device_type=parameter.device.type):
+ with (
+ self._encoder_tiling_context(image.height, image.width),
+ torch.random.fork_rng(
+ devices=devices,
+ device_type=parameter.device.type,
+ ),
+ ):
torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
for device in devices:
with self.device_module.device(device):
@@ -495,13 +655,26 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
self,
frames: Any,
) -> tuple[torch.Tensor, tuple[int, int, int]]:
+ if getattr(self, "decode_only", False):
+ raise RuntimeError("MiniMax H3 decode-only video VAE cannot encode videos")
parameter = next(self.parameters())
previous_dtype = parameter.dtype
if previous_dtype != torch.float32:
self.to(torch.float32)
devices = [parameter.device] if parameter.device.type != "cpu" else []
+ shape = getattr(frames, "shape", None)
+ if int(getattr(self, "parallel_size", 1)) > 1:
+ if shape is None or len(shape) != 4:
+ raise ValueError("parallel MiniMax H3 video encode requires a rank-4 frame array")
+ if int(shape[-1]) in {1, 3, 4}:
+ frame_height, frame_width = int(shape[-3]), int(shape[-2])
+ else:
+ frame_height, frame_width = int(shape[-2]), int(shape[-1])
+ tiling_context = self._encoder_tiling_context(frame_height, frame_width)
+ else:
+ tiling_context = nullcontext()
try:
- with torch.random.fork_rng(devices=devices, device_type=parameter.device.type):
+ with tiling_context, torch.random.fork_rng(devices=devices, device_type=parameter.device.type):
torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
for device in devices:
with self.device_module.device(device):
@@ -561,6 +734,8 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
@torch.inference_mode()
def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
+ if getattr(self, "encode_only", False):
+ raise RuntimeError("MiniMax H3 encode-only video VAE cannot decode latents")
with self._decode_tiling_context(latent):
decoded = self.model.decode_base(self._denormalize_latent(latent))
return self._normalize_decoded_frames(self.model.processor.revert_tensor(decoded))
@@ -598,6 +773,8 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
the complete decode path. After a callback failure, the remaining
chunks are decoded and discarded before the exception is re-raised.
"""
+ if getattr(self, "encode_only", False):
+ raise RuntimeError("MiniMax H3 encode-only video VAE cannot decode latents")
if not callable(on_chunk):
raise TypeError("on_chunk must be callable")
if not callable(getattr(self.model, "_adaptive_decode", None)):
@@ -622,16 +799,24 @@ class MiniMaxH3AudioVAE(nn.Module):
*,
device: torch.device,
load_device: torch.device | None = None,
+ encode_only: bool = False,
+ decode_only: bool = False,
trust_remote_code: bool = False,
) -> None:
super().__init__()
self._device_target = device
+ self.encode_only = bool(encode_only)
+ self.decode_only = bool(decode_only)
+ if self.encode_only and self.decode_only:
+ raise ValueError("MiniMax H3 audio VAE cannot be both encode-only and decode-only")
self.config_dict = _load_component_config(component_path)
- self.remote = _load_remote_component(
- component_path,
- self.config_dict,
- trust_remote_code=trust_remote_code,
- )
+ _check_trust_remote_code(component_path, self.config_dict, trust_remote_code)
+ if self.encode_only:
+ self.remote = _load_audio_vae_encoder(component_path, self.config_dict)
+ else:
+ self.remote = _load_remote_component(component_path, self.config_dict, trust_remote_code=trust_remote_code)
+ if self.decode_only:
+ _remove_modules(self.remote.model, ("encoder", "pre_block", "mean_proj"))
# The checkpoint's audio VAE contract is FP32 for both reference
# encoding and waveform decoding.
initial_device = load_device or device
@@ -674,6 +859,8 @@ class MiniMaxH3AudioVAE(nn.Module):
waveform: torch.Tensor,
sample_rate: int,
) -> tuple[torch.Tensor, int]:
+ if getattr(self, "decode_only", False):
+ raise RuntimeError("MiniMax H3 decode-only audio VAE cannot encode waveforms")
import torchaudio
waveform = waveform.float()
@@ -719,6 +906,8 @@ class MiniMaxH3AudioVAE(nn.Module):
@torch.inference_mode()
def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
+ if getattr(self, "encode_only", False):
+ raise RuntimeError("MiniMax H3 encode-only audio VAE cannot decode latents")
channels = int(self.config_dict["latent_channels"])
mean = torch.tensor(
self.config_dict["latents_mean"],
diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
index 8e4dac3a3..cb1adbb61 100644
--- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
+++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -49,22 +49,17 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
_fused_cuda_supported,
fused_qk_norm_rope,
)
+from vllm_omni.diffusion.layers.qwen_select01_modulation import (
+ can_use_qwen_select01_triton,
+ fused_layernorm_select01,
+ fused_residual_layernorm_select01,
+ select01_modulation_native,
+)
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
logger = init_logger(__name__)
-def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
- """Rotate interleaved pairs before rounding back to the activation dtype.
-
- Qwen-Image's reference uses complex FP32 multiplication. Rounding the
- frequencies to BF16 before rotation loses positional precision; those
- errors accumulate across the denoising steps.
- """
- 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 _qwen_image_qk_norm_rope(
q: torch.Tensor,
k: torch.Tensor,
@@ -97,16 +92,11 @@ def _qwen_image_qk_norm_rope(
fused_k.reshape(batch, seq_len, num_kv_heads, head_dim),
)
+ # Eager path: BF16/activation-dtype RotaryEmbedding on every device.
+ # CUDA used to call FP32 complex multiply here; that matches the Diffusers
+ # helper in unit tests but drops Omni vs Diffusers pipeline PSNR (#7494).
q = norm_q(q)
k = norm_k(k)
- if q.device.type == "cuda":
- return (
- _apply_qwen_image_rotary_emb(q, freqs),
- _apply_qwen_image_rotary_emb(k, freqs),
- )
-
- # Retain the platform-specific kernels on other accelerators, which may
- # not support complex tensors.
cos = torch.real(freqs).to(q.dtype)
sin = torch.imag(freqs).to(q.dtype)
return rope(q, cos, sin), rope(k, cos, sin)
@@ -855,41 +845,11 @@ class QwenImageTransformerBlock(nn.Module):
self.zero_cond_t = zero_cond_t
- def _modulate(self, mod_params, index=None):
+ def _modulate(self, mod_params):
"""Apply modulation to input tensor"""
# shift: b d, scale: b d, gate: b d
shift, scale, gate = mod_params.chunk(3, dim=-1)
-
- if index is not None:
- # Assuming mod_params batch dim is 2*actual_batch (chunked into 2 parts)
- # So shift, scale, gate have shape [2*actual_batch, d]
- actual_batch = shift.size(0) // 2
- shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] # each: [actual_batch, d]
- scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:]
- gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:]
-
- # index: [b, l] where b is actual batch size
- # Expand to [b, l, 1] to match feature dimension
- index_expanded = index.unsqueeze(-1) # [b, l, 1]
-
- # Expand chunks to [b, 1, d] then broadcast to [b, l, d]
- shift_0_exp = shift_0.unsqueeze(1) # [b, 1, d]
- shift_1_exp = shift_1.unsqueeze(1) # [b, 1, d]
- scale_0_exp = scale_0.unsqueeze(1)
- scale_1_exp = scale_1.unsqueeze(1)
- gate_0_exp = gate_0.unsqueeze(1)
- gate_1_exp = gate_1.unsqueeze(1)
-
- # Use torch.where to select based on index
- shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp)
- scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp)
- gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp)
- else:
- shift_result = shift.unsqueeze(1)
- scale_result = scale.unsqueeze(1)
- gate_result = gate.unsqueeze(1)
-
- return scale_result, shift_result, gate_result
+ return scale.unsqueeze(1), shift.unsqueeze(1), gate.unsqueeze(1)
def forward(
self,
@@ -915,8 +875,22 @@ class QwenImageTransformerBlock(nn.Module):
txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim]
# Process image stream - norm1 + modulation
- img_scale1, img_shift1, img_gate1 = self._modulate(img_mod1, modulate_index)
- img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1)
+ use_fused_select01 = modulate_index is not None and can_use_qwen_select01_triton(hidden_states)
+ if use_fused_select01:
+ img_modulated, img_gate1 = fused_layernorm_select01(
+ hidden_states,
+ img_mod1,
+ modulate_index,
+ self.img_norm1.eps,
+ self.img_norm1.layernorm.weight,
+ self.img_norm1.layernorm.bias,
+ )
+ elif modulate_index is not None:
+ img_scale1, img_shift1, img_gate1 = select01_modulation_native(img_mod1, modulate_index)
+ img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1)
+ else:
+ img_scale1, img_shift1, img_gate1 = self._modulate(img_mod1)
+ img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1)
# Process text stream - norm1 + modulation
txt_scale1, txt_shift1, txt_gate1 = self._modulate(txt_mod1)
@@ -941,12 +915,29 @@ class QwenImageTransformerBlock(nn.Module):
img_attn_output, txt_attn_output = attn_output
# Apply attention gates and add residual (like in Megatron)
- hidden_states = hidden_states + img_gate1 * img_attn_output
+ hidden_states_before_attn = hidden_states
+ if not use_fused_select01:
+ hidden_states = hidden_states + img_gate1 * img_attn_output
encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
# Process image stream - norm2 + MLP
- img_scale2, img_shift2, img_gate2 = self._modulate(img_mod2, modulate_index)
- img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2)
+ if use_fused_select01:
+ img_modulated2, hidden_states, img_gate2 = fused_residual_layernorm_select01(
+ img_attn_output,
+ hidden_states_before_attn,
+ img_gate1,
+ img_mod2,
+ modulate_index,
+ self.img_norm2.eps,
+ self.img_norm2.layernorm.weight,
+ self.img_norm2.layernorm.bias,
+ )
+ elif modulate_index is not None:
+ img_scale2, img_shift2, img_gate2 = select01_modulation_native(img_mod2, modulate_index)
+ img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2)
+ else:
+ img_scale2, img_shift2, img_gate2 = self._modulate(img_mod2)
+ img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2)
img_mlp_output = self.img_mlp(img_modulated2)
hidden_states = hidden_states + img_gate2 * img_mlp_output
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
index 15b036864..915fb45f7 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
@@ -179,8 +179,36 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
return {}
+def resolve_wan_transformer_quant_config(
+ config: dict, quant_config: QuantizationConfig | None, component: str
+) -> QuantizationConfig | None:
+ """Resolve the expert before applying its checkpoint's storage contract."""
+ from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+ from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+ # Wan experts are siblings: "transformer_2" must not inherit "transformer"
+ # through the generic layer-prefix resolver. Select the whole expert name.
+ component_quant_config = (
+ quant_config.component_configs.get(component, quant_config.default_config)
+ if isinstance(quant_config, ComponentQuantizationConfig)
+ else quant_config
+ )
+ quantization_disabled = isinstance(quant_config, ComponentQuantizationConfig) and component_quant_config is None
+ resolved_quant_config = resolve_quant_config_from_disk(component_quant_config, config.get("quantization_config"))
+ if quantization_disabled and resolved_quant_config is not None:
+ raise ValueError(
+ f"Quantization is disabled for component {component!r}, but its checkpoint declares quantization. "
+ "Use a BF16 checkpoint for this component or enable its matching quantization method."
+ )
+ return resolved_quant_config
+
+
def create_transformer_from_config(
- config: dict, quant_config: QuantizationConfig | None = None, prefix: str = ""
+ config: dict,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ *,
+ component: str = "transformer",
) -> WanTransformer3DModel:
"""Create WanTransformer3DModel from config dict."""
kwargs: dict = {}
@@ -216,11 +244,7 @@ def create_transformer_from_config(
if "pos_embed_seq_len" in config:
kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"]
- if "quantization_config" in config:
- from vllm_omni.quantization.factory import resolve_quant_config_from_disk
-
- quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"])
-
+ quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
if quant_config is not None:
kwargs["quant_config"] = quant_config
if prefix:
@@ -460,13 +484,13 @@ class Wan22Pipeline(
# Initialize transformers with correct config (weights loaded via load_weights)
if load_transformer:
transformer_config = load_transformer_config(model, "transformer", local_files_only)
- self.transformer = self._create_transformer(transformer_config)
+ self.transformer = self._create_transformer(transformer_config, component="transformer")
else:
self.transformer = None
if load_transformer_2:
transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
- self.transformer_2 = self._create_transformer(transformer_2_config)
+ self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
else:
self.transformer_2 = None
@@ -504,10 +528,17 @@ class Wan22Pipeline(
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
)
- def _create_transformer(self, config: dict) -> WanTransformer3DModel:
+ def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
quant_config = getattr(self.od_config, "quantization_config", None)
- return create_transformer_from_config(config, quant_config=quant_config)
+ # Startup metadata describes the first expert, not a user policy for both.
+ if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+ quant_config = None
+ return create_transformer_from_config(
+ config,
+ quant_config=quant_config,
+ component=component,
+ )
@property
def guidance_scale(self):
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
index 01bd53bb8..1a75c40da 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
@@ -304,25 +304,10 @@ class Wan22I2VPipeline(
# Transformers (weights loaded via load_weights)
# Load config from model directory or HF Hub to get correct in_channels for I2V models
transformer_config = load_transformer_config(model, "transformer", local_files_only)
- self.transformer = create_transformer_from_config(
- transformer_config,
- quant_config=od_config.quantization_config,
- )
+ self.transformer = self._create_transformer(transformer_config, component="transformer")
if self.has_transformer_2:
transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
- t2_quant = transformer_2_config.get("quantization_config")
- if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
- from vllm_omni.quantization.factory import build_quant_config
-
- method = t2_quant["quant_method"]
- kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
- t2_quant = build_quant_config(method, **kwargs)
- else:
- t2_quant = None
- self.transformer_2 = create_transformer_from_config(
- transformer_2_config,
- quant_config=t2_quant,
- )
+ self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
else:
self.transformer_2 = None
@@ -462,10 +447,13 @@ class Wan22I2VPipeline(
image_embeds = self.image_encoder(pixel_values, output_hidden_states=True)
return image_embeds.hidden_states[-2]
- def _create_transformer(self, config: dict) -> WanTransformer3DModel:
+ def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
quant_config = getattr(self.od_config, "quantization_config", None)
- return create_transformer_from_config(config, quant_config=quant_config)
+ # Startup metadata describes the first expert, not a user policy for both.
+ if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+ quant_config = None
+ return create_transformer_from_config(config, quant_config=quant_config, component=component)
def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
sampling_params_list = req.sampling_params_list
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
index aeecff608..0d5f3cd67 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
@@ -28,6 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
Wan22Pipeline,
load_wan_weights_with_optional_gate,
+ resolve_wan_transformer_quant_config,
retrieve_latents,
)
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
@@ -60,6 +61,8 @@ def create_vace_transformer_from_config(
config: dict,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
+ *,
+ component: str = "transformer",
) -> WanVACETransformer3DModel:
"""Create WanVACETransformer3DModel from config dict."""
kwargs = {}
@@ -97,6 +100,7 @@ def create_vace_transformer_from_config(
kwargs["vace_layers"] = config["vace_layers"]
if "vace_in_channels" in config:
kwargs["vace_in_channels"] = config["vace_in_channels"]
+ quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
if quant_config is not None:
kwargs["quant_config"] = quant_config
if prefix:
@@ -205,10 +209,13 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
super().__init__(od_config=od_config, prefix=prefix)
- def _create_transformer(self, config: dict) -> WanVACETransformer3DModel:
+ def _create_transformer(self, config: dict, component: str = "transformer") -> WanVACETransformer3DModel:
"""Build VACE transformer. Respects od_config.quantization_config."""
quant_config = getattr(self.od_config, "quantization_config", None)
- return create_vace_transformer_from_config(config, quant_config=quant_config)
+ # Startup metadata describes the first expert, not a user policy for both.
+ if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+ quant_config = None
+ return create_vace_transformer_from_config(config, quant_config=quant_config, component=component)
def diffuse(
self,
diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
index dbc555604..2808dcee7 100644
--- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
+++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
@@ -18,7 +18,7 @@ from vllm.distributed import (
)
from vllm.logger import init_logger
from vllm.model_executor.layers.conv import Conv3dLayer
-from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.linear import ColumnParallelLinear, LinearBase, QKVParallelLinear, RowParallelLinear
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.utils import (
@@ -45,6 +45,7 @@ from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization.mxfp4_config import NPUMxfp4LinearMethod
logger = init_logger(__name__)
@@ -985,6 +986,68 @@ class WanTransformer3DModel(nn.Module):
# ROPE helper
self._cached_rope_emb = None
self._cached_rope_resolution = None
+ self._validate_w4a8_fallback_layers(quant_config)
+
+ def _validate_w4a8_fallback_layers(self, quant_config: QuantizationConfig | None) -> None:
+ """Reject checkpoint aliases/typos before loading weights or running inference."""
+ requested = tuple(getattr(quant_config, "w4a8_fallback_layers", []))
+ self._w4a8_fallback_layers = requested
+ self._local_w4a8_fallback_layers: tuple[str, ...] = ()
+ if not requested:
+ return
+ modules = dict(self.named_modules())
+ missing_pp_prefixes = [name for name, module in modules.items() if isinstance(module, PPMissingLayer)]
+ local_layers = [
+ name for name in requested if not any(name.startswith(prefix + ".") for prefix in missing_pp_prefixes)
+ ]
+ self._local_w4a8_fallback_layers = tuple(local_layers)
+ invalid = [name for name in local_layers if not isinstance(modules.get(name), LinearBase)]
+ if invalid:
+ raise ValueError(
+ "w4a8_fallback_layers requires exact runtime Linear paths relative to each Wan transformer; "
+ f"unknown/non-Linear paths: {invalid}. Use attn1.to_qkv for fused self-attention, "
+ "not checkpoint attn1.to_q/to_k/to_v names."
+ )
+
+ def _log_w4a8_fallback_load_summary(self, component_name: str) -> None:
+ """Log post-load transforms/cache readiness, not numerical acceptance."""
+ requested = self._w4a8_fallback_layers
+ if not requested:
+ return
+ modules = dict(self.named_modules())
+ local_layers = self._local_w4a8_fallback_layers
+ selected_layers = [
+ name for name in local_layers if getattr(modules[name].quant_method, "is_w4a8_fallback_layer", False)
+ ]
+ bf16_overrides = [name for name in local_layers if name not in selected_layers]
+ processed = [
+ name
+ for name in selected_layers
+ if getattr(modules[name], "_already_called_process_weights_after_loading", False)
+ ]
+ if not selected_layers:
+ processing_state = "not-required"
+ elif len(processed) == len(selected_layers):
+ processing_state = "ready"
+ else:
+ processing_state = "not-ready"
+
+ def qualified(names: Iterable[str]) -> list[str]:
+ return [f"{component_name}.{name}" for name in names]
+
+ logger.info(
+ "Wan W4A8 post-load summary: component=%s; selected=%d %s; BF16 overrides=%d %s; "
+ "weight/A8 processing state=%s (%d/%d selected layers processed); other PP ranks=%d",
+ component_name,
+ len(selected_layers),
+ qualified(selected_layers),
+ len(bf16_overrides),
+ qualified(bf16_overrides),
+ processing_state,
+ len(processed),
+ len(selected_layers),
+ len(requested) - len(local_layers),
+ )
@property
def dtype(self) -> torch.dtype:
@@ -1155,7 +1218,22 @@ class WanTransformer3DModel(nn.Module):
}
params_dict = dict(self.named_parameters())
+ single_scale_prefixes = {
+ name
+ for name, module in self.named_modules()
+ if isinstance(getattr(module, "quant_method", None), NPUMxfp4LinearMethod)
+ }
+
+ def reject_dualscale_tensor(name: str) -> None:
+ if name.endswith(".weight_dual_scale") and name.rsplit(".", 1)[0] in single_scale_prefixes:
+ raise ValueError(
+ f"Single-scale mxfp4 cannot load DualScale tensor {name}; "
+ "changing quant_method does not convert a checkpoint."
+ )
+
loaded_params: set[str] = set()
+ loaded_qkv_shards: dict[str, set[str]] = {}
+ qkv_smooth_scales: dict[str, torch.Tensor] = {}
for name, loaded_weight in weights:
name = weight_name_remapping.get(name, name)
@@ -1166,15 +1244,32 @@ class WanTransformer3DModel(nn.Module):
# Pre-fused to_qkv tensors (from offline MXFP8 merged checkpoint) fall
# through to the else branch and are loaded directly.
for param_name, weight_name, shard_id in stacked_params_mapping:
- if weight_name not in original_name:
+ if f"{weight_name}." not in original_name:
continue
- lookup_name = original_name.replace(weight_name, param_name)
+ lookup_name = original_name.replace(f"{weight_name}.", f"{param_name}.", 1)
+ reject_dualscale_tensor(lookup_name)
# Skip weights that belong to PP stages other than this one
if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict:
break
param = params_dict[lookup_name]
+ if lookup_name.endswith(".mul_scale"):
+ previous_scale = qkv_smooth_scales.get(lookup_name)
+ if previous_scale is not None and not torch.equal(previous_scale, loaded_weight):
+ raise ValueError(f"Fused Q/K/V must share the same Smooth tensor: {lookup_name}")
+ if previous_scale is None:
+ qkv_smooth_scales[lookup_name] = loaded_weight.clone()
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
+ loaded_params.add(original_name)
+ if getattr(param, "output_dim", None) is None and not getattr(param, "needs_scalar_to_array", False):
+ # Shared input-channel Smooth is a complete tensor even
+ # when supplied under a single Q/K/V source name.
+ loaded_params.add(lookup_name)
+ else:
+ shards = loaded_qkv_shards.setdefault(lookup_name, set())
+ shards.add(shard_id)
+ if shards == {"q", "k", "v"}:
+ loaded_params.add(lookup_name)
break
else:
# diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight
@@ -1198,6 +1293,7 @@ class WanTransformer3DModel(nn.Module):
if is_pp_missing_parameter(lookup_name, self):
continue
+ reject_dualscale_tensor(lookup_name)
if lookup_name not in params_dict:
logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
continue
@@ -1222,8 +1318,6 @@ class WanTransformer3DModel(nn.Module):
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
-
- loaded_params.add(original_name)
- loaded_params.add(lookup_name)
+ loaded_params.update((original_name, lookup_name))
return loaded_params
diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py
index 176a8bf69..a646a3a30 100644
--- a/vllm_omni/distributed/omni_connectors/utils/config.py
+++ b/vllm_omni/distributed/omni_connectors/utils/config.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from dataclasses import dataclass, field
from typing import Any
@@ -33,7 +33,7 @@ def get_stage_connector_role(model_config: Any) -> str | None:
def stage_receives_chunks(model_config: Any) -> bool:
"""Whether connector chunks, rather than the orchestrator, feed a stage."""
- return get_stage_connector_role(model_config) != "sender"
+ return bool(getattr(model_config, "async_chunk", True)) and get_stage_connector_role(model_config) != "sender"
def stage_sends_async_output(model_config: Any) -> bool:
diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py
index 52b3e9ff2..d0548207a 100644
--- a/vllm_omni/engine/omni_engine_base.py
+++ b/vllm_omni/engine/omni_engine_base.py
@@ -277,8 +277,9 @@ class OmniEngineBase:
)
self.num_stages = len(self.stage_configs)
- stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None
- self.async_chunk = bool(getattr(stage0_args, "async_chunk", False))
+ self.async_chunk = any(
+ bool(getattr(getattr(stage, "engine_args", None), "async_chunk", False)) for stage in self.stage_configs
+ )
self.stage_pools: list[StagePool] = []
self.stage_clients: list[StageClient] = [] # logical-stage view for external readers
self.input_processor: InputProcessor | None = None
diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index aced2c1b3..aea5c8519 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -30,7 +30,7 @@ from vllm.logger import init_logger
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.pooling_params import PoolingParams
from vllm.sampling_params import RequestOutputKind, SamplingParams
-from vllm.v1.engine import EngineCoreOutputs
+from vllm.v1.engine import EngineCoreOutputs, FinishReason
from vllm.v1.engine.exceptions import EngineDeadError
from vllm.v1.metrics.stats import IterationStats
@@ -777,6 +777,7 @@ class OrchestratorBase:
)
if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state):
raw_terminal_request_ids.add(req_state.request_id)
+ await self._report_duplex_session_request_error(stage_id, replica_id, eco, req_state)
iteration_stats = IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None
processed = await pool.process_llm_raw_outputs(
replica_id,
@@ -1503,6 +1504,49 @@ class OrchestratorBase:
req_state.finished_final_output_stage_ids.add(stage_id)
return True
+ async def _report_duplex_session_request_error(
+ self,
+ stage_id: int,
+ replica_id: int | None,
+ eco: Any,
+ req_state: OrchestratorRequestState,
+ ) -> None:
+ """Turn a scheduler-side error finish of a duplex session request into a request error.
+
+ A duplex session request is resumable, so its terminal outputs are
+ deliberately not turned into processed outputs
+ (``_finish_raw_terminal_requests`` skips duplex sessions), and vLLM's
+ output processor may already have dropped the request state by the
+ time the raw ``FinishReason.ERROR`` output arrives. Without this the
+ session never learns that its request is gone: the next append
+ re-creates the request from scratch and the session silently loses
+ its context. Deliver the error to the request's consumer instead so
+ the serving layer can fail and close the session.
+ """
+ if not req_state.session_owned:
+ return
+ if getattr(eco, "finish_reason", None) != FinishReason.ERROR:
+ return
+ if getattr(eco, "is_segment_finished", False):
+ return
+ stop_reason = getattr(eco, "stop_reason", None)
+ error = stop_reason if isinstance(stop_reason, str) and stop_reason else "duplex session request failed"
+ logger.error(
+ "[Orchestrator] Duplex session request %s failed at stage-%s replica-%s: %s",
+ req_state.request_id,
+ stage_id,
+ replica_id,
+ error,
+ )
+ await self.output_async_queue.put(
+ ErrorMessage(
+ error=error,
+ fatal=False,
+ request_id=req_state.request_id,
+ stage_id=stage_id,
+ )
+ )
+
async def _finish_raw_terminal_requests(
self,
stage_id: int,
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 6114bbea2..1291e68f0 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -236,6 +236,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
req_state = ClientRequestState(
request_id=request_id,
external_request_id=external_request_id,
+ final_stage_id=final_stage_id_for_e2e,
)
req_state.metrics = metrics
req_state.request_arrival_ts = wall_start_ts
diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py
index e7c86de97..e9bf2c77e 100644
--- a/vllm_omni/entrypoints/async_omni_base.py
+++ b/vllm_omni/entrypoints/async_omni_base.py
@@ -479,7 +479,22 @@ class AsyncOmniBase(OmniBase):
if state is not None and rid not in delivered:
queue = getattr(state, "queue", None)
if queue is not None:
- await state.queue.put(self._synthetic_abort_output_message(rid))
+ final_stage_id = getattr(state, "final_stage_id", None)
+ if final_stage_id is None:
+ final_stage_id = 0
+ stage_metadata = getattr(self, "_stage_meta_list", ())
+ final_output_type = (
+ getattr(stage_metadata[final_stage_id], "final_output_type", None)
+ if final_stage_id < len(stage_metadata)
+ else None
+ ) or "text"
+ await state.queue.put(
+ self._synthetic_abort_output_message(
+ rid,
+ stage_id=final_stage_id,
+ final_output_type=final_output_type,
+ )
+ )
delivered.add(rid)
for rid in request_ids:
self._record_request_failure_once(rid, reason="client_abort")
@@ -491,13 +506,18 @@ class AsyncOmniBase(OmniBase):
logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
@staticmethod
- def _synthetic_abort_output_message(request_id: str) -> OutputMessage:
+ def _synthetic_abort_output_message(
+ request_id: str,
+ *,
+ stage_id: int,
+ final_output_type: str,
+ ) -> OutputMessage:
"""Terminal abort OutputMessage used when the engine returned none."""
engine_output = OmniRequestOutput(
request_id=request_id,
finished=True,
- stage_id=0,
- final_output_type="text",
+ stage_id=stage_id,
+ final_output_type=final_output_type,
outputs=[
CompletionOutput(
index=0,
@@ -512,7 +532,7 @@ class AsyncOmniBase(OmniBase):
)
return OutputMessage(
request_id=request_id,
- stage_id=0,
+ stage_id=stage_id,
replica_id=None,
engine_outputs=engine_output,
metrics=None,
diff --git a/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py b/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py
index 748100c5d..c3e55abb7 100644
--- a/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py
+++ b/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py
@@ -28,8 +28,28 @@ from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase
def _common(parser: argparse.ArgumentParser) -> None:
- parser.add_argument("--dataset", default=DEFAULT_DATASET)
- parser.add_argument("--split", default="all")
+ parser.add_argument(
+ "--dataset",
+ default=DEFAULT_DATASET,
+ help=(
+ "Hugging Face dataset id, a JSON/JSONL manifest, or a local Hugging Face "
+ "dataset mirror. A local directory must use a single configuration whose "
+ "data files are named after the RTD_*/PR_* splits "
+ "(data/-00000-of-00001.parquet); a single .parquet file is also "
+ "accepted. --split filters the rows by their preserved split identity, so "
+ "rows without identity fail loudly."
+ ),
+ )
+ parser.add_argument(
+ "--split",
+ default="all",
+ help=(
+ "Restrict to one split (e.g. RTD_OCR) or 'all'. A mistyped split name "
+ "raises a clear error listing the splits actually observed in the data. "
+ "Rows from a manifest or an iterable without split/subset/config identity "
+ "keep the requested split as an override rather than being rejected."
+ ),
+ )
parser.add_argument("--family", choices=("all", "rtd", "pr"), default="all")
parser.add_argument("--media-root")
parser.add_argument("--limit", type=int)
@@ -100,6 +120,8 @@ def run(args: argparse.Namespace) -> int:
limit=args.limit,
ids=args.ids,
)
+ if not samples:
+ raise ValueError("no samples selected; check --dataset/--split/--family/--ids/--limit")
if args.action == "generate":
if args.concurrency < 1:
raise ValueError("--concurrency must be at least 1")
diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py
index 22f1c6d5e..ca9e96bab 100644
--- a/vllm_omni/entrypoints/client_request_state.py
+++ b/vllm_omni/entrypoints/client_request_state.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import asyncio
from vllm_omni.metrics import OrchestratorAggregator
@@ -11,10 +14,12 @@ class ClientRequestState:
request_id: str,
external_request_id: str | None = None,
queue: asyncio.Queue | None = None,
+ final_stage_id: int | None = None,
):
self.request_id = request_id
self.external_request_id = external_request_id
self.stage_id: int | None = None
+ self.final_stage_id: int | None = final_stage_id
self.queue = queue if queue is not None else asyncio.Queue()
self.metrics: OrchestratorAggregator | None = None
# Request-scoped idempotency guard for Prometheus failure counters.
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index e6053f103..94f1ab513 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -22,6 +22,7 @@ from argparse import Namespace
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from http import HTTPStatus
+from pathlib import Path
from typing import Annotated, Any, Literal
import uvloop
@@ -496,21 +497,34 @@ async def build_async_omni(
yield async_omni
-def _is_duplex_model(model: str, kwargs: dict[str, Any]) -> bool:
- """Whether ``model``'s pipeline declares a ``duplex_plugin`` (served by ``DuplexOmni``).
-
- Resolution errors propagate: a duplex model whose pipeline or deploy
- config cannot be resolved must fail startup rather than silently start a
- turn-based server.
- """
+def _should_serve_duplex(model: str, kwargs: dict[str, Any]) -> bool:
+ """Select the serving engine without changing the model's duplex capability."""
from vllm_omni.config.config_factory import StageConfigFactory
+ from vllm_omni.config.stage_config import _DEPLOY_DIR, resolve_deploy_yaml
pipeline_config = StageConfigFactory.get_pipeline_config(
model=model,
trust_remote_code=bool(kwargs.get("trust_remote_code")),
deploy_config_path=kwargs.get("deploy_config"),
)
- return bool(pipeline_config is not None and getattr(pipeline_config, "duplex_plugin", None))
+ if pipeline_config is None or not getattr(pipeline_config, "duplex_plugin", None):
+ return False
+
+ deploy_path = kwargs.get("deploy_config")
+ if deploy_path is None:
+ if pipeline_config.default_deploy_config_name is None:
+ raise ValueError("A duplex-capable model requires a deploy config with session_mode: turn or duplex")
+ deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name
+ else:
+ deploy_path = Path(deploy_path)
+ if not deploy_path.exists() and deploy_path.parent == Path("."):
+ deploy_path = _DEPLOY_DIR / deploy_path
+
+ # Resolve base_config too, so the API and stage workers use the same mode.
+ session_mode = resolve_deploy_yaml(deploy_path).get("session_mode")
+ if session_mode not in ("turn", "duplex"):
+ raise ValueError("A duplex-capable model requires session_mode: turn or duplex in its deploy config")
+ return session_mode == "duplex"
@asynccontextmanager
@@ -577,9 +591,7 @@ async def build_async_omni_from_stage_config(
kwargs.setdefault("log_stats", not args.disable_log_stats)
if client_config is not None:
kwargs["client_config"] = client_config
- if _is_duplex_model(model, kwargs):
- # A duplex model is always served in duplex mode: sessions over
- # /v1/realtime?duplex=1, no turn-based HTTP routes.
+ if _should_serve_duplex(model, kwargs):
async_omni = DuplexOmni(model=model, **kwargs)
else:
async_omni = AsyncOmni(model=model, **kwargs)
@@ -771,9 +783,7 @@ async def omni_init_app_state(
state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
model_name = served_model_names[0] if served_model_names else args.model
- # Duplex mode: a duplex model is served through DuplexOmni only. Sessions
- # run over /v1/realtime?duplex=1; every turn-based HTTP route reports
- # "not available".
+ # Initialize the API surface for the selected engine, not just the model.
if isinstance(engine_client, DuplexOmni):
await _init_duplex_app_state(engine_client, state, args, base_model_paths, vllm_config, request_logger)
return
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index f54ca3e28..80f357651 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -32,6 +32,7 @@ from vllm.entrypoints.launchers.launcher import terminate_if_errored
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
from vllm.logger import init_logger
from vllm.multimodal.media import MediaConnector
+from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.utils import random_uuid
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
@@ -1444,14 +1445,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
tts_params: dict[str, Any] | None = None,
collect: dict | None = None,
target_sample_rate: int | None = None,
+ cumulative_audio: bool = False,
):
"""Generate audio chunks for streaming response.
Handles two audio output modes from the engine:
- Cumulative mode (list): Engine returns growing list of chunks;
we emit only the new tail on each iteration.
- - Per-step mode (tensor): Engine returns single tensor per iteration;
- we emit it directly.
+ - Tensor mode: Emit each tensor directly, or only its new samples when
+ cumulative_audio is set for requests that retain audio for alignment.
Args:
generator: Async generator from the engine
@@ -1462,6 +1464,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
Raw audio bytes for each chunk (with WAV header for first chunk if wav format)
"""
prev_count = 0
+ emitted_samples = 0
sample_rate_val = 24000
first_chunk = True
first_audio_chunk_s: float | None = None
@@ -1563,6 +1566,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
)
if chunk_np.ndim > 1:
chunk_np = chunk_np.squeeze()
+ if cumulative_audio and not isinstance(audio_val, list):
+ total_samples = chunk_np.shape[-1]
+ chunk_np = chunk_np[..., emitted_samples:]
+ emitted_samples = total_samples
+ if chunk_np.size == 0:
+ continue
if resampler is not None:
chunk_np = resampler.process(chunk_np)
if chunk_np.size == 0:
@@ -1957,6 +1966,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY
output_modalities.append(TIMESTAMPS_MODALITY)
+ if is_streaming_request:
+ # The aligner consumes the terminal waveform, so retain earlier
+ # audio chunks instead of draining them after each DELTA output.
+ sampling_params_list = [
+ sp.clone() if isinstance(sp, SamplingParams) else sp for sp in sampling_params_list
+ ]
+ for sp in sampling_params_list:
+ if isinstance(sp, SamplingParams):
+ sp.output_kind = RequestOutputKind.CUMULATIVE
generator = self.engine_client.generate(
prompt=prompt,
@@ -1990,6 +2008,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
tts_params: dict[str, Any] | None = None,
collect: dict | None = None,
target_sample_rate: int | None = None,
+ cumulative_audio: bool = False,
):
"""Yield raw PCM byte chunks from the engine generator.
@@ -2009,6 +2028,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
tts_params=tts_params,
collect=collect,
target_sample_rate=target_sample_rate,
+ cumulative_audio=cumulative_audio,
)
) as chunks:
async for chunk in chunks:
@@ -2024,6 +2044,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
request_id,
tts_params=tts_params,
target_sample_rate=request.sample_rate,
+ cumulative_audio=request.word_timestamps,
)
) as chunks:
async for chunk in chunks:
diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py
index ceeb208c9..45bdd4e90 100644
--- a/vllm_omni/entrypoints/openai/serving_speech_stream.py
+++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py
@@ -483,6 +483,7 @@ class OmniStreamingSpeechHandler:
include_sample_rate=True,
tts_params=tts_params,
collect=collect,
+ cumulative_audio=True,
)
) as stream:
async for chunk, chunk_sample_rate in stream:
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py
index d2851e32c..03f0d3c4b 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py
@@ -4,6 +4,8 @@
"""Breeze-TTS-2 serving adapter for synchronous full-payload generation."""
import asyncio
+import copy
+import math
from typing import TYPE_CHECKING, Any
import numpy as np
@@ -49,6 +51,7 @@ class BreezeTTS2Adapter(ARTTSAdapter):
return set()
def _load_codec_frame_rate(self) -> float | None:
+ assert self.ctx.engine_client is not None
config = self.ctx.engine_client.model_config.hf_config
codec = getattr(config, "codec_config", None)
if isinstance(codec, dict):
@@ -99,6 +102,14 @@ class BreezeTTS2Adapter(ARTTSAdapter):
greedy_error = _validate_greedy_sampling(extra_params)
if greedy_error is not None:
return greedy_error
+ repetition_penalty = extra_params.get("repetition_penalty")
+ if repetition_penalty is not None:
+ try:
+ repetition_penalty = float(repetition_penalty)
+ except (TypeError, ValueError):
+ return "Breeze-TTS-2 repetition_penalty must be finite and positive"
+ if not math.isfinite(repetition_penalty) or repetition_penalty <= 0:
+ return "Breeze-TTS-2 repetition_penalty must be finite and positive"
if request.max_new_tokens is not None:
if request.max_new_tokens < self.max_new_tokens_min:
return f"max_new_tokens must be at least {self.max_new_tokens_min}"
@@ -210,7 +221,16 @@ class BreezeTTS2Adapter(ARTTSAdapter):
request_id: str | None = None,
) -> list:
del prompt, request_id
- return apply_max_new_tokens(sampling_params_list, request)
+ sampling_params_list = apply_max_new_tokens(sampling_params_list, request)
+ repetition_penalty = (request.extra_params or {}).get("repetition_penalty")
+ if repetition_penalty is not None:
+ # The generic speech handler copies extra_params to extra_args;
+ # promote this value so the worker receives sampling metadata.
+ # Only this scalar changes; preserve the caller's params without
+ # copying nested state or the unchanged codec-stage parameters.
+ sampling_params_list = [copy.copy(sampling_params_list[0]), *sampling_params_list[1:]]
+ sampling_params_list[0].repetition_penalty = float(repetition_penalty)
+ return sampling_params_list
def _validate_greedy_sampling(extra_params: dict[str, Any]) -> str | None:
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py
index bed20efc2..555fec42d 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py
@@ -6,8 +6,12 @@ Both variants share the same build/validate flow (``_build_moss_tts_params``
handles each); they are registered under distinct model-type names.
"""
+import asyncio
+import threading
from typing import TYPE_CHECKING, Any, cast
+from transformers import AutoModel, AutoTokenizer
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
from vllm.inputs import tokens_input
from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
@@ -18,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import (
apply_max_new_tokens,
conditioning_cache_salt,
)
+from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt
if TYPE_CHECKING:
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -30,6 +35,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):
super().__init__(ctx)
self._moss_variant = None if self.name == "moss_tts_nano" else self._detect_moss_variant()
self._moss_processor_cache = None
+ self._moss_realtime_components_lock = threading.Lock()
@property
def engine_client(self):
@@ -63,7 +69,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):
# inside the model package; this layer only supplies the processor and
# the process-wide speaker cache.
encoder = build_reference_encoder(
- self._get_moss_processor(),
+ self._get_moss_realtime_components()[2] if self._moss_variant == "realtime" else self._get_moss_processor(),
variant=cast(str, self._moss_variant),
speaker_cache=self._speaker_cache,
)
@@ -157,6 +163,45 @@ class _MossTTSAdapterBase(ARTTSAdapter):
self._moss_processor_cache = proc
return proc
+ def _get_moss_realtime_components(self):
+ with self._moss_realtime_components_lock:
+ cached = getattr(self, "_moss_realtime_components", None)
+ if cached is not None:
+ return cached
+
+ model_id = self.engine_client.model_config.model
+ processor_cls = get_class_from_dynamic_module(
+ "processing_mossttsrealtime.MossTTSRealtimeProcessor",
+ model_id,
+ )
+ tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+ processor = processor_cls(tokenizer=tokenizer)
+ hf_config = self.engine_client.model_config.hf_config
+ codec_path = str(
+ getattr(
+ hf_config,
+ "codec_model_name_or_path",
+ getattr(
+ hf_config,
+ "audio_tokenizer_name_or_path",
+ "OpenMOSS-Team/MOSS-Audio-Tokenizer",
+ ),
+ )
+ )
+ codec = (
+ AutoModel.from_pretrained(
+ codec_path,
+ trust_remote_code=True,
+ )
+ # Reference encoding runs in the API process. Keep the codec on CPU
+ # so it does not reserve accelerator memory outside the model workers.
+ .to("cpu")
+ .eval()
+ )
+ cached = (tokenizer, processor, codec)
+ self._moss_realtime_components = cached
+ return cached
+
async def _build_moss_tts_params(
self,
request: "OpenAICreateSpeechRequest",
@@ -196,24 +241,18 @@ class _MossTTSAdapterBase(ARTTSAdapter):
params["ref_audio_cache_key"] = cache_key
return params
- # ---- MOSS-TTS-Realtime: keep the old prompt_audio_array path ----
- # ``AutoProcessor.from_pretrained`` doesn't auto-discover
- # ``MossTTSRealtimeProcessor`` (no ``processor_config.json`` in the
- # snapshot), and Realtime's prompt format diverges from MossTTSDelay
- # (16-channel grid, separate per-step text feed). The
- # ``prompt_audio_array`` shape lines up well enough with what the
- # talker reads for short prompts; full Realtime support needs a
- # separate processor.from_module path which we don't wire here.
if v == "realtime":
- params = {
- "text": [request.input or ""],
- "mode": ["voice_clone"],
- }
+ tokenizer, processor, _ = await asyncio.to_thread(self._get_moss_realtime_components)
+ references, realtime_resolve_keys = await self._encode_moss_references(
+ request,
+ has_inline_ref_audio=has_inline_ref_audio,
+ two_speaker=False,
+ )
+ params = build_realtime_prompt(tokenizer, processor, request.input or "", references[0])
if request.max_new_tokens is not None:
params["max_new_frames"] = [request.max_new_tokens]
- wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio))
- params["prompt_audio_array"] = [[wav_list, sr]]
- params["ref_audio_cache_key"] = cache_key
+ if 0 in realtime_resolve_keys:
+ params["ref_audio_cache_key"] = realtime_resolve_keys[0]
return params
# ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator)
diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py
index 7addb3934..860d64c44 100644
--- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py
+++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py
@@ -204,6 +204,7 @@ class BreezeTTS2MimiCodec(nn.Module):
f"Breeze codec request id count {len(request_ids)} does not match input request count {len(requests)}"
)
valid: list[tuple[str, torch.Tensor, bool]] = []
+ valid_indices: list[int] = []
for index, request_ids_tensor in enumerate(requests):
runtime_info = (
runtime_additional_information[index]
@@ -238,6 +239,7 @@ class BreezeTTS2MimiCodec(nn.Module):
)
frames = int(flat.numel()) // self._num_codebooks
valid.append((state_id, flat.reshape(self._num_codebooks, frames), finished))
+ valid_indices.append(index)
waveforms = [empty.to(device=ids.device)] * len(requests)
sample_rates = [sr] * len(requests)
@@ -247,7 +249,6 @@ class BreezeTTS2MimiCodec(nn.Module):
multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates},
)
states: list[dict[str, Any]] = []
- valid_indices: list[int] = []
max_frames = max(codes.shape[-1] for _, codes, _ in valid)
batched_codes = valid[0][1].new_zeros((len(valid), self._num_codebooks, max_frames))
for row, (state_id, row_codes, _) in enumerate(valid):
@@ -258,7 +259,6 @@ class BreezeTTS2MimiCodec(nn.Module):
self._decoder_state_cache[state_id] = state
state.setdefault("prefix_frames", 0)
states.append(state)
- valid_indices.append(row)
decoded = decoder.batched_chunked_decode(
batched_codes,
@@ -353,6 +353,27 @@ class BreezeTTS2MimiCodec(nn.Module):
device_map=str(device),
)
self._sample_rate = int(self._audio_tokenizer.model.get_output_sample_rate())
+ if self._async_chunk:
+ connector_config = getattr(self.vllm_config.model_config, "stage_connector_config", None)
+ extra = (
+ connector_config.get("extra", connector_config)
+ if isinstance(connector_config, Mapping)
+ else getattr(connector_config, "extra", None)
+ )
+ value = extra.get("breeze_codec_chunk_frames") if isinstance(extra, Mapping) else None
+ try:
+ chunk_frames = int(value) if value is not None else 8
+ except (TypeError, ValueError) as exc:
+ raise ValueError(f"Invalid Breeze codec_chunk_frames={value!r}") from exc
+ if chunk_frames <= 0:
+ raise ValueError(f"Breeze codec_chunk_frames must be positive, got {chunk_frames}")
+ # Qwen treats chunks shorter than its configured window as
+ # terminal tails and does not advance their rolling cache.
+ # Match the producer's chunk size instead of its default 25.
+ decoder = self._audio_tokenizer.model.decoder
+ decoder._initial_codec_chunk_frames = chunk_frames
+ decoder._incremental_chunk_frames = chunk_frames
+ decoder._incremental_chunk_ramp = []
for _ in weights:
pass
tokenizer_module = getattr(self._audio_tokenizer, "model", None)
diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py
index beaa390c8..765df5fbf 100644
--- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py
+++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py
@@ -13,6 +13,7 @@ generation order without importing HuggingFace GenerationMixin.
from __future__ import annotations
import hashlib
+import math
import os
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
@@ -367,10 +368,12 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
return None
logits = self._main_head_logits(hidden_states)
- # ``make_omni_output`` runs immediately before this hook. It records
- # requests whose previous hidden state selected EOS or reached the
- # frame budget, so the scheduler receives a real stop token instead
- # of depending on an arbitrary deploy-level max_tokens timeout.
+ # ``make_omni_output`` already selected codebook 0 and completed its
+ # frame. Return that same decision to the scheduler, including EOS.
+ # A finite zero survives the sampler's repetition/temperature passes
+ # without applying the repetition penalty twice or changing the frame.
+ # These logits describe deterministic codec selection, not model
+ # confidence; sampled-token logprobs are consequently zero.
states = self._batch_state or []
if states:
for state, row_start, row_end in _iter_request_rows(
@@ -378,10 +381,11 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
getattr(self, "_batch_state_spans", None),
int(logits.shape[0]),
):
- if state.get("breeze_force_eos", False):
+ selected_code0 = state.get("breeze_selected_code0")
+ if selected_code0 is not None:
row = logits[row_start:row_end]
row.fill_(float("-inf"))
- row[:, self.codebook_vocab_size] = 0.0
+ row[:, int(selected_code0)] = 0.0
return logits
def _main_head_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -443,6 +447,13 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
self._batch_state = info_items
self._batch_state_spans = spans
per_request_codes: list[torch.Tensor] = []
+ # The runner passes the same per-request sampling metadata to this
+ # hook and compute_logits, including offline SamplingParams overrides.
+ # When no_penalties is true vLLM does not initialize this GPU buffer.
+ sampling_metadata = kwargs.get("sampling_metadata")
+ repetition_penalties = None
+ if sampling_metadata is not None and not sampling_metadata.no_penalties:
+ repetition_penalties = sampling_metadata.repetition_penalties.tolist()
for index, info in enumerate(info_items):
if spans is not None and index < len(spans):
@@ -458,11 +469,22 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
row_hidden = hidden[row_end - 1]
if self._golden_dump_dir and bool(info.get("_omni_is_prefill", False)):
self._dump_golden_prefill(info, row_hidden)
- main_logits = self._main_head_logits(row_hidden.reshape(1, -1))
+ main_logits = self._main_head_logits(row_hidden.reshape(1, -1)).float()
+ penalty = float(repetition_penalties[index]) if repetition_penalties is not None else 1.0
+ if not math.isfinite(penalty) or penalty <= 0:
+ raise ValueError("Breeze repetition_penalty must be finite and positive")
+ # Match upstream: penalize each previously generated codebook-0
+ # id once, excluding the text prompt and reference audio codes.
+ history = info.setdefault("breeze_generated_code0_ids", [])
+ if penalty != 1.0 and history:
+ history_ids = torch.tensor(history, device=main_logits.device, dtype=torch.long)
+ scores = main_logits[0, history_ids]
+ main_logits[0, history_ids] = torch.where(scores < 0, scores * penalty, scores / penalty)
code0 = int(torch.argmax(main_logits, dim=-1).item())
generated_frames = int(info.get("breeze_generated_frames", 0) or 0)
max_new_frames = int(info.get("breeze_max_new_frames", -1))
terminal = code0 == self.codebook_vocab_size or (max_new_frames > 0 and generated_frames >= max_new_frames)
+ info["breeze_selected_code0"] = self.codebook_vocab_size if terminal else code0
if terminal:
info["breeze_force_eos"] = True
existing = info.get("breeze_audio_codes")
@@ -476,6 +498,7 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
raise RuntimeError(f"Breeze main head produced invalid codec id {code0}")
frame = self._generate_depth_codes(row_hidden, code0)
+ history.append(code0)
info["breeze_current_frame"] = frame[0].detach()
info["breeze_generated_frames"] = generated_frames + 1
info["breeze_force_eos"] = False
diff --git a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py
index d1aaeb13e..2b538df01 100644
--- a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py
+++ b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py
@@ -28,6 +28,12 @@ BREEZE_TTS_2_PIPELINE = PipelineConfig(
custom_process_next_stage_input_func=f"{_PROC}.talker2codec_full_payload",
sampling_constraints={
"detokenize": False,
+ # Text EOS 1 is a valid codec token; only the backbone's
+ # explicit codec EOS below may terminate audio generation.
+ "ignore_eos": True,
+ # The talker has already selected the emitted frame's token;
+ # min_tokens must not mask that token in the sampler.
+ "min_tokens": 0,
"stop_token_ids": [2051],
},
),
diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py
index 7cbd31e9c..2b8f3e61e 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py
@@ -98,6 +98,46 @@ class ResBlock(torch.nn.Module):
remove_weight_norm(self.convs2[idx])
+def _carry_phase_at_boundary(
+ cum_total: torch.Tensor,
+ phase_acc: torch.Tensor | None,
+ next_overlap: int,
+ *,
+ dim: int,
+ trim: int = 0,
+) -> torch.Tensor | None:
+ """Accumulated phase (mod 1) carried into the next streaming window.
+
+ Shared by SineGen and SineGen2 so harmonic phase stays continuous across
+ chunk boundaries when HiFT runs over a bounded mel window. ``cum_total``
+ must already include ``phase_acc``; falls back to it unchanged if
+ ``next_overlap`` (+ ``trim``) covers the whole window.
+ """
+ length = cum_total.shape[dim]
+ if next_overlap >= length + trim:
+ return phase_acc
+ idx = length + trim - next_overlap - 1
+ return cum_total.narrow(dim, idx, 1) % 1
+
+
+def _wrapped_slice(buf: torch.Tensor, offset: int, length: int, dim: int = 1) -> torch.Tensor:
+ """``length`` elements of ``buf`` along ``dim`` starting at ``offset``,
+ wrapping (mod ``buf.shape[dim]``) instead of running past the end and
+ coming back short. Shared by SineGen/SineGen2/SourceModuleHnNSF's
+ position-indexed noise buffers, now that streams can outlast their 300s.
+ """
+ n = buf.shape[dim]
+ start = offset % n
+ parts = []
+ remaining = length
+ while remaining > 0:
+ take = min(remaining, n - start)
+ parts.append(buf.narrow(dim, start, take))
+ remaining -= take
+ start = 0
+ return parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim)
+
+
class SineGen(torch.nn.Module):
"""Definition of sine generator
SineGen(samp_rate, harmonic_num = 0,
@@ -114,13 +154,24 @@ class SineGen(torch.nn.Module):
segment is always sin(np.pi) or cos(0)
"""
- def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0):
+ def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0, causal=False):
super().__init__()
self.sine_amp = sine_amp
self.noise_std = noise_std
self.harmonic_num = harmonic_num
self.sampling_rate = samp_rate
self.voiced_threshold = voiced_threshold
+ self.causal = causal
+ # Fixed per model instance like SineGen2's rand_ini, not redrawn per
+ # call, or each harmonic jumps to a new random phase every window.
+ self.phase_vec = Uniform(low=-np.pi, high=np.pi).sample(sample_shape=(1, harmonic_num + 1, 1))
+ self.phase_vec[:, 0, :] = 0
+ if causal:
+ self.register_buffer(
+ "noise_buf",
+ torch.randn(1, 300 * 24000, harmonic_num + 1),
+ persistent=False,
+ )
def _f02uv(self, f0):
# generate uv signal
@@ -128,22 +179,32 @@ class SineGen(torch.nn.Module):
return uv
@torch.no_grad()
- def forward(self, f0):
+ def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0):
"""sine_tensor, uv = forward(f0)
input F0: tensor(batchsize=1, dim=1, length)
f0 for unvoiced steps should be 0
output sine_tensor: tensor(batchsize=1, length, dim)
output uv: tensor(batchsize=1, length, 1)
+
+ ``next_overlap`` and ``trim`` are sample-indexed here (``f0`` arrives
+ post-``f0_upsamp``), unlike SineGen2 which downsamples to mel rate
+ before its own cumsum -- callers must scale both by samples-per-mel.
"""
f0 = f0.transpose(1, 2)
F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device)
for i in range(self.harmonic_num + 1):
F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate
- theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1)
- u_dist = Uniform(low=-np.pi, high=np.pi)
- phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device)
- phase_vec[:, 0, :] = 0
+ cum = torch.cumsum(F_mat, dim=-1)
+ cum_total = cum if phase_acc is None else phase_acc + cum
+ theta_mat = 2 * np.pi * (cum_total % 1)
+ new_phase_acc = _carry_phase_at_boundary(cum_total, phase_acc, next_overlap, dim=-1, trim=trim)
+ if self.causal:
+ phase_vec = self.phase_vec.to(F_mat.device)
+ else:
+ u_dist = Uniform(low=-np.pi, high=np.pi)
+ phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device)
+ phase_vec[:, 0, :] = 0
# generate sine waveforms
sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec)
@@ -155,11 +216,16 @@ class SineGen(torch.nn.Module):
# std = self.sine_amp/3 -> max value ~ self.sine_amp
# . for voiced regions is self.noise_std
noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
- noise = noise_amp * torch.randn_like(sine_waves)
+ if self.training is False and self.causal is True:
+ noise = noise_amp * _wrapped_slice(self.noise_buf, noise_offset, sine_waves.shape[-1]).transpose(1, 2)
+ else:
+ noise = noise_amp * torch.randn_like(sine_waves)
# first: set the unvoiced part to 0 by uv
# then: additive noise
sine_waves = sine_waves * uv + noise
+ if self.causal:
+ return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise, new_phase_acc
return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise
@@ -215,7 +281,7 @@ class SineGen2(torch.nn.Module):
uv = (f0 > self.voiced_threshold).type(torch.float32)
return uv
- def _f02sine(self, f0_values):
+ def _f02sine(self, f0_values, phase_acc=None, next_overlap=0, trim=0):
"""f0_values: (batchsize, length, dim)
where dim indicates fundamental tone and overtones
"""
@@ -223,9 +289,12 @@ class SineGen2(torch.nn.Module):
# because 2 * np.pi * n doesn't affect phase
rad_values = (f0_values / self.sampling_rate) % 1
- # initial phase noise (no noise for fundamental component)
+ # initial phase noise (no noise for fundamental component). Only applied
+ # on the genuine first chunk of a stream (see rationale above the class);
+ # phase_acc is None exactly when there is no carried-forward phase yet.
if self.training is False and self.causal is True:
- rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device)
+ if phase_acc is None:
+ rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device)
else:
rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device)
rand_ini[:, 0] = 0
@@ -237,7 +306,10 @@ class SineGen2(torch.nn.Module):
rad_values.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear"
).transpose(1, 2)
- phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
+ cum = torch.cumsum(rad_values, dim=1)
+ phase = cum if phase_acc is None else cum + phase_acc
+ new_phase_acc = _carry_phase_at_boundary(phase, phase_acc, next_overlap, dim=1, trim=trim)
+ phase = phase * 2 * np.pi
phase = torch.nn.functional.interpolate(
phase.transpose(1, 2) * self.upsample_scale,
scale_factor=self.upsample_scale,
@@ -272,9 +344,10 @@ class SineGen2(torch.nn.Module):
# get the sines
sines = torch.cos(i_phase * 2 * np.pi)
- return sines
+ new_phase_acc = None
+ return (sines, new_phase_acc) if self.causal else sines
- def forward(self, f0):
+ def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0):
"""sine_tensor, uv = forward(f0)
input F0: tensor(batchsize=1, length, dim=1)
f0 for unvoiced steps should be 0
@@ -284,8 +357,12 @@ class SineGen2(torch.nn.Module):
# fundamental component
fn = torch.multiply(f0, self.harmonic_ids)
- # generate sine waveforms
- sine_waves = self._f02sine(fn) * self.sine_amp
+ if self.causal:
+ sine_waves, new_phase_acc = self._f02sine(fn, phase_acc, next_overlap, trim)
+ else:
+ sine_waves = self._f02sine(fn)
+ new_phase_acc = None
+ sine_waves = sine_waves * self.sine_amp
# generate uv signal
uv = self._f02uv(f0)
@@ -295,14 +372,15 @@ class SineGen2(torch.nn.Module):
# . for voiced regions is self.noise_std
noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
if self.training is False and self.causal is True:
- noise = noise_amp * self.sine_waves[:, : sine_waves.shape[1]].to(sine_waves.device)
+ # indexed by absolute sample position, not window-relative position
+ noise = noise_amp * _wrapped_slice(self.sine_waves, noise_offset, sine_waves.shape[1]).to(sine_waves.device)
else:
noise = noise_amp * torch.randn_like(sine_waves)
# first: set the unvoiced part to 0 by uv
# then: additive noise
sine_waves = sine_waves * uv + noise
- return sine_waves, uv, noise
+ return (sine_waves, uv, noise, new_phase_acc) if self.causal else (sine_waves, uv, noise)
class SourceModuleHnNSF(torch.nn.Module):
@@ -338,10 +416,13 @@ class SourceModuleHnNSF(torch.nn.Module):
self.sine_amp = sine_amp
self.noise_std = add_noise_std
+ self.upsample_scale = upsample_scale
# to produce sine waveforms
if sinegen_type == "1":
- self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold)
+ self.l_sin_gen = SineGen(
+ sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal
+ )
else:
self.l_sin_gen = SineGen2(
sampling_rate, upsample_scale, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal
@@ -354,24 +435,38 @@ class SourceModuleHnNSF(torch.nn.Module):
if causal is True:
self.uv = torch.rand(1, 300 * 24000, 1)
- def forward(self, x):
+ def forward(self, x, phase_acc=None, uv_offset=0, next_overlap=0, trim=0):
"""
Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
F0_sampled (batchsize, length, 1)
Sine_source (batchsize, length, 1)
noise_source (batchsize, length 1)
"""
- # source for harmonic branch
with torch.no_grad():
- sine_wavs, uv, _ = self.l_sin_gen(x)
+ if not self.causal:
+ sine_wavs, uv, _ = self.l_sin_gen(x)
+ new_phase_acc = None
+ elif isinstance(self.l_sin_gen, SineGen):
+ # next_overlap/trim arrive in mel frames; SineGen needs samples.
+ sine_wavs, uv, _, new_phase_acc = self.l_sin_gen(
+ x,
+ phase_acc,
+ next_overlap * self.upsample_scale,
+ trim * self.upsample_scale,
+ noise_offset=uv_offset,
+ )
+ else:
+ sine_wavs, uv, _, new_phase_acc = self.l_sin_gen(
+ x, phase_acc, next_overlap, trim, noise_offset=uv_offset
+ )
sine_merge = self.l_tanh(self.l_linear(sine_wavs))
# source for noise branch, in the same shape as uv
if self.training is False and self.causal is True:
- noise = self.uv[:, : uv.shape[1]] * self.sine_amp / 3
+ noise = _wrapped_slice(self.uv, uv_offset, uv.shape[1]) * self.sine_amp / 3
else:
noise = torch.randn_like(uv) * self.sine_amp / 3
- return sine_merge, noise, uv
+ return sine_merge, noise, uv, new_phase_acc
class HiFTGenerator(nn.Module):
@@ -606,7 +701,7 @@ class HiFTGenerator(nn.Module):
f0 = self.f0_predictor(speech_feat)
# f0->source
s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t
- s, _, _ = self.m_source(s)
+ s, _, _, _ = self.m_source(s)
s = s.transpose(1, 2)
# mel+source->speech
generated_speech = self.decode(x=speech_feat, s=s)
@@ -622,7 +717,7 @@ class HiFTGenerator(nn.Module):
f0 = self.f0_predictor(speech_feat)
# f0->source
s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t
- s, _, _ = self.m_source(s)
+ s, _, _, _ = self.m_source(s)
s = s.transpose(1, 2)
# use cache_source to avoid glitch
if cache_source.shape[2] != 0:
@@ -785,14 +880,28 @@ class CausalHiFTGenerator(HiFTGenerator):
return x
@torch.inference_mode()
- def inference(self, speech_feat: torch.Tensor, finalize: bool = True) -> torch.Tensor:
+ def inference(
+ self,
+ speech_feat: torch.Tensor,
+ finalize: bool = True,
+ phase_acc: torch.Tensor | None = None,
+ uv_offset: int = 0,
+ next_overlap: int = 0,
+ trim: int = 0,
+ f0_margin: int = 0,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
# mel->f0 NOTE f0_predictor precision is crucial for causal inference, move
# self.f0_predictor to cpu if necessary
self.f0_predictor.to("cpu")
f0 = self.f0_predictor(speech_feat.cpu(), finalize=finalize).to(speech_feat)
+ if f0_margin > 0:
+ f0 = f0[:, f0_margin:]
+ speech_feat = speech_feat[:, :, f0_margin:]
# f0->source
s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t
- s, _, _ = self.m_source(s)
+ s, _, _, new_phase_acc = self.m_source(
+ s, phase_acc=phase_acc, uv_offset=uv_offset, next_overlap=next_overlap, trim=trim
+ )
s = s.transpose(1, 2)
if finalize is True:
generated_speech = self.decode(x=speech_feat, s=s, finalize=finalize)
@@ -800,7 +909,7 @@ class CausalHiFTGenerator(HiFTGenerator):
generated_speech = self.decode(
x=speech_feat[:, :, : -self.f0_predictor.condnet[0].causal_padding], s=s, finalize=finalize
)
- return generated_speech, s
+ return generated_speech, s, new_phase_acc
class CausalConv1dUpsample(torch.nn.Conv1d):
@@ -954,6 +1063,11 @@ class CausalConvRNNF0Predictor(nn.Module):
nn.ELU(),
)
self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class)
+ self.left_context_frames = sum(
+ layer.causal_padding
+ for layer in self.condnet
+ if isinstance(layer, CausalConv1d) and layer.causal_type == "left"
+ )
def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor:
if finalize is True:
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
index ef706b6c8..1713467c8 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py
@@ -9,10 +9,8 @@ This module contains the code2wav (token-to-waveform) stage which uses:
3. HiFiGAN vocoder for waveform synthesis
"""
-from __future__ import annotations
-
from collections import Counter
-from typing import cast
+from typing import TypedDict, cast
import numpy as np
import torch
@@ -39,6 +37,21 @@ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
logger = init_logger(__name__)
+class StreamingFlowItem(TypedDict, total=False):
+ """One entry of the batched-streaming item list passed to forward_streaming_batch."""
+
+ index: int
+ req_id: str | None
+ stream_finished: bool
+ token: torch.Tensor
+ prompt_token: torch.Tensor
+ prompt_feat: torch.Tensor
+ embedding: torch.Tensor
+ cache_state: dict[str, torch.Tensor] | None
+ token_offset_tokens: int
+ finalize: bool
+
+
class CosyVoice3Code2Wav(nn.Module):
"""CosyVoice3 Code2Wav stage for token-to-waveform conversion.
@@ -119,6 +132,10 @@ class CosyVoice3Code2Wav(nn.Module):
self.mel_cache_len = 20
self.source_cache_len = int(self.mel_cache_len * 256)
self.speech_window = np.hamming(2 * self.source_cache_len)
+ # Must cover decode()'s own causal receptive field, not just the F0
+ # margin; window_len=48 already passes
+ # test_incremental_hift_bounded_window_is_close, so 64 has headroom.
+ self._hift_window_len = 64
@property
def input_frame_rate(self) -> int:
@@ -226,40 +243,79 @@ class CosyVoice3Code2Wav(nn.Module):
chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
cached_mel = None if not cache_state else cache_state.get("mel")
- speech_offset_obj = None if not cache_state else cache_state.get("speech_offset")
- try:
- speech_offset = int(speech_offset_obj) if speech_offset_obj is not None else 0
- except (TypeError, ValueError):
- speech_offset = 0
-
+ cached_offset = 0 if not cache_state else int(cache_state.get("mel_offset", 0))
+ window_len = self._hift_window_len
+ samples_per_mel = int(np.prod(self.hift.upsample_rates) * self.hift.istft_params["hop_len"])
+ trim = int(self.hift.f0_predictor.condnet[0].causal_padding)
+ f0_margin_frames = int(self.hift.f0_predictor.left_context_frames)
+ f0_margin = 0
if isinstance(cached_mel, torch.Tensor) and cached_mel.numel() > 0:
cached_mel = cached_mel.to(device=chunk_mel.device, dtype=chunk_mel.dtype)
- tts_mel = torch.cat([cached_mel, chunk_mel], dim=-1) if chunk_mel.numel() > 0 else cached_mel
+ total_hist = cached_mel.shape[-1]
+ overlap = min(window_len + trim, total_hist)
+ window_mel = torch.cat([cached_mel[..., -overlap:], chunk_mel], dim=-1)
+ window_offset = cached_offset + total_hist - overlap
+ uv_offset = window_offset * samples_per_mel
+ f0_margin = min(f0_margin_frames, total_hist - overlap)
+ f0_input_mel = (
+ torch.cat(
+ [cached_mel[..., total_hist - overlap - f0_margin : total_hist - overlap], window_mel], dim=-1
+ )
+ if f0_margin > 0
+ else window_mel
+ )
else:
- tts_mel = chunk_mel
+ window_mel = chunk_mel
+ f0_input_mel = chunk_mel
+ window_offset = 0
+ uv_offset = 0
- if tts_mel.shape[-1] == 0:
+ phase_acc = None if not cache_state else cache_state.get("phase_acc")
+ if phase_acc is not None:
+ phase_acc = phase_acc.to(device=chunk_mel.device, dtype=chunk_mel.dtype)
+
+ if window_mel.shape[-1] == 0:
tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype)
+ new_phase_acc = phase_acc
else:
- tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
+ next_overlap = min(window_len + trim, window_mel.shape[-1])
+ # On finalize the F0 predictor consumes the whole window (it does not
+ # hold back `trim` frames), so trim must not inflate the phase-carry
+ # index -- it would point past the end of a short window.
+ carry_trim = 0 if finalize else trim
+ tts_speech, _, new_phase_acc = self.hift.inference(
+ speech_feat=f0_input_mel,
+ finalize=finalize,
+ phase_acc=phase_acc,
+ uv_offset=uv_offset,
+ next_overlap=next_overlap,
+ trim=carry_trim,
+ f0_margin=f0_margin,
+ )
tts_speech = tts_speech.reshape(tts_speech.shape[0], -1)
- speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1])))
- emitted_speech = tts_speech[:, speech_offset:]
+ new_samples = chunk_mel.shape[-1] * samples_per_mel
+ if finalize:
+ frames_withheld_per_chunk = trim + self.hift.conv_pre_look_right + 1
+ released = frames_withheld_per_chunk * samples_per_mel
+ emitted_speech = tts_speech[:, -(new_samples + released) :]
+ else:
+ emitted_speech = tts_speech[:, -new_samples:] if new_samples > 0 else tts_speech[:, :0]
if finalize:
return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None
new_state = {
- "mel": tts_mel.detach().cpu().contiguous(),
- "speech_offset": int(tts_speech.shape[-1]),
+ "mel": f0_input_mel.detach().cpu().contiguous(),
+ "mel_offset": window_offset - f0_margin,
+ "phase_acc": new_phase_acc.detach().cpu().contiguous() if new_phase_acc is not None else None,
}
return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state
@torch.inference_mode()
def forward_streaming_batch(
self,
- items: list[dict[str, object]],
+ items: list[StreamingFlowItem],
*,
n_timesteps: int = 10,
) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]:
@@ -270,20 +326,16 @@ class CosyVoice3Code2Wav(nn.Module):
group and passed to the flow as per-row token lengths.
"""
results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items)
- groups: dict[tuple[int, int, int, bool], list[tuple[int, dict[str, object]]]] = {}
+ groups: dict[tuple[int, int, int, bool], list[tuple[int, StreamingFlowItem]]] = {}
for index, item in enumerate(items):
- token = item["token"]
- prompt_token = item["prompt_token"]
- prompt_feat = item["prompt_feat"]
- embedding = item["embedding"]
- assert isinstance(token, torch.Tensor)
- assert isinstance(prompt_token, torch.Tensor)
- assert isinstance(prompt_feat, torch.Tensor)
- assert isinstance(embedding, torch.Tensor)
+ assert isinstance(item["token"], torch.Tensor)
+ assert isinstance(item["prompt_token"], torch.Tensor)
+ assert isinstance(item["prompt_feat"], torch.Tensor)
+ assert isinstance(item["embedding"], torch.Tensor)
key = (
- int(prompt_token.shape[1]),
- int(prompt_feat.shape[1]),
- int(embedding.shape[1]),
+ int(item["prompt_token"].shape[1]),
+ int(item["prompt_feat"].shape[1]),
+ int(item["embedding"].shape[1]),
bool(item.get("finalize", False)),
)
groups.setdefault(key, []).append((index, item))
@@ -305,11 +357,11 @@ class CosyVoice3Code2Wav(nn.Module):
if len(group) == 1:
index, item = group[0]
result = self.forward_streaming(
- token=item["token"], # type: ignore[arg-type]
- prompt_token=item["prompt_token"], # type: ignore[arg-type]
- prompt_feat=item["prompt_feat"], # type: ignore[arg-type]
- embedding=item["embedding"], # type: ignore[arg-type]
- cache_state=item.get("cache_state"), # type: ignore[arg-type]
+ token=item["token"],
+ prompt_token=item["prompt_token"],
+ prompt_feat=item["prompt_feat"],
+ embedding=item["embedding"],
+ cache_state=item.get("cache_state"),
n_timesteps=n_timesteps,
token_offset_tokens=int(item.get("token_offset_tokens", 0)),
finalize=bool(item.get("finalize", False)),
@@ -318,15 +370,13 @@ class CosyVoice3Code2Wav(nn.Module):
continue
token_tensors = [item["token"] for _, item in group]
- assert all(isinstance(token, torch.Tensor) for token in token_tensors)
token_lens = torch.tensor(
- [int(token.shape[1]) for token in token_tensors], # type: ignore[union-attr]
+ [int(token.shape[1]) for token in token_tensors],
dtype=torch.int32,
)
max_token_len = int(token_lens.max().item())
padded_tokens = []
for token in token_tensors:
- assert isinstance(token, torch.Tensor)
if int(token.shape[1]) == max_token_len:
padded_tokens.append(token)
else:
@@ -337,9 +387,9 @@ class CosyVoice3Code2Wav(nn.Module):
)
padded_tokens.append(torch.cat([token, pad], dim=1))
tokens = torch.cat(padded_tokens, dim=0)
- prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) # type: ignore[list-item]
- prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) # type: ignore[list-item]
- embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) # type: ignore[list-item]
+ prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0)
+ prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0)
+ embeddings = torch.cat([item["embedding"] for _, item in group], dim=0)
prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32)
prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32)
finalize = bool(group[0][1].get("finalize", False))
@@ -370,7 +420,7 @@ class CosyVoice3Code2Wav(nn.Module):
row_feat = row_feat[:, :, trim_mel:]
results[index] = self._stream_hift_from_feat(
row_feat,
- cache_state=item.get("cache_state"), # type: ignore[arg-type]
+ cache_state=item.get("cache_state"),
finalize=finalize,
)
@@ -443,7 +493,7 @@ class CosyVoice3Code2Wav(nn.Module):
dtype=tts_mel.dtype,
)
else:
- tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=True)
+ tts_speech, _, _ = self.hift.inference(speech_feat=tts_mel, finalize=True)
return tts_speech
diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py
index 1317dc044..57f37f3d9 100644
--- a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py
+++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py
@@ -32,6 +32,7 @@ from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
from vllm.sequence import IntermediateTensors
from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.platforms import current_omni_platform
from vllm_omni.utils.speaker_cache import get_speaker_cache
from .configuration_fish_speech import FishSpeechConfig, FishSpeechFastARConfig, FishSpeechSlowARConfig
@@ -195,7 +196,11 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
self.mtp_hidden_size = int(self.text_config.hidden_size)
self.talker_mtp_output_key = ("codes", "audio")
self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")}
- self.talker_mtp_graph_safe = True
+ # NPU (Ascend) ACL graphs cannot capture the Fast-AR random-sampling ops
+ # (torch.nonzero / torch.multinomial used for top-p), so keep the talker
+ # MTP eager on NPU to fall back to native (non-graph) execution. Mirrors
+ # the 310P patch for Qwen3-TTS (talker_mtp_graph_safe=False).
+ self.talker_mtp_graph_safe = not current_omni_platform.is_npu()
# Qwen3 transformer backbone.
self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py
index e8350f41a..0143e9675 100644
--- a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py
+++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py
@@ -1,16 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
-"""MiniMax H3 checkpoint resolution for the AR text-encoder stage."""
-
from pathlib import Path
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
-MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS = {
- "fl2va": ["FL2VA/text_encoder/**"],
- "ref2va": ["Ref2VA/text_encoder/**"],
+MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = {
+ partition: [
+ f"{subdir}/text_encoder/**",
+ f"{subdir}/video_vae/**",
+ f"{subdir}/audio_vae/**",
+ ]
+ for partition, subdir in {"fl2va": "FL2VA", "ref2va": "Ref2VA"}.items()
}
+MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["combined"] = MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["fl2va"]
def resolve_minimax_h3_partition(
@@ -32,7 +35,7 @@ def resolve_minimax_h3_partition(
return "ref2va" if task == "ref2va" else "fl2va"
-def resolve_minimax_h3_model_root(
+def resolve_minimax_h3_encoder_model_root(
model: str,
revision: str | None,
task_type: str | None,
@@ -50,7 +53,7 @@ def resolve_minimax_h3_model_root(
snapshot = download_weights_from_hf_specific(
model_name_or_path=model,
cache_dir=None,
- allow_patterns=MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS[partition],
+ allow_patterns=MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition],
revision=revision,
require_all=True,
)
@@ -59,4 +62,11 @@ def resolve_minimax_h3_model_root(
return str(Path(snapshot) / subdir / "text_encoder")
-__all__ = ["resolve_minimax_h3_model_root", "resolve_minimax_h3_partition"]
+resolve_minimax_h3_model_root = resolve_minimax_h3_encoder_model_root
+
+
+__all__ = [
+ "resolve_minimax_h3_encoder_model_root",
+ "resolve_minimax_h3_model_root",
+ "resolve_minimax_h3_partition",
+]
diff --git a/vllm_omni/model_executor/models/minimax_h3/conditioning.py b/vllm_omni/model_executor/models/minimax_h3/conditioning.py
index efea25771..761292dc3 100644
--- a/vllm_omni/model_executor/models/minimax_h3/conditioning.py
+++ b/vllm_omni/model_executor/models/minimax_h3/conditioning.py
@@ -4,8 +4,8 @@
from __future__ import annotations
-from collections.abc import Mapping
-from dataclasses import dataclass
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, fields
from typing import Any
import torch
@@ -119,3 +119,502 @@ class MiniMaxH3TextConditioning:
"hidden_states": self.hidden_states,
"token_tags": self.token_tags,
}
+
+
+MINIMAX_H3_ENCODER_REQUEST_KEY = "minimax_h3_encoder_request"
+MINIMAX_H3_ENCODER_LAYOUT_KEY = "minimax_h3_encoder_layout"
+STAGE_SCHEMA_VERSION = 1
+_ENCODER_WIRE_SCHEMA_ID = 1
+
+_TASK_TO_CODE = {"t2va": 1, "fl2va": 2, "ref2va": 3}
+_CODE_TO_TASK = {value: key for key, value in _TASK_TO_CODE.items()}
+_REF_KIND_TO_CODE = {"image": 1, "video": 2, "video_audio": 3, "audio": 4}
+_CODE_TO_REF_KIND = {value: key for key, value in _REF_KIND_TO_CODE.items()}
+_INTEGER_DTYPES = frozenset(
+ {
+ torch.int8,
+ torch.int16,
+ torch.int32,
+ torch.int64,
+ torch.uint8,
+ }
+)
+_VIDEO_CONDITION_WIDTH = 96
+_AUDIO_CONDITION_WIDTH = 32
+_AUDIO_CONDITION_CHANNELS = 2
+
+
+def _wire_layout(payload: Mapping[str, Any]) -> torch.Tensor:
+ private_metadata = payload.get("kv_metadata")
+ if not isinstance(private_metadata, Mapping):
+ raise ValueError("MiniMax H3 wire payload requires private layout metadata")
+ layout = private_metadata.get(MINIMAX_H3_ENCODER_LAYOUT_KEY)
+ if isinstance(layout, torch.Tensor):
+ return layout
+ raise ValueError("MiniMax H3 wire payload requires one packed encoder layout tensor")
+
+
+def _validate_condition_tensors(
+ visual: torch.Tensor | None,
+ visual_shapes: Sequence[Sequence[int]],
+ audio: torch.Tensor | None,
+ audio_lengths: Sequence[int],
+) -> None:
+ if (visual is None) != (not visual_shapes):
+ raise ValueError("MiniMax H3 visual condition and shapes must be present together")
+ if (audio is None) != (not audio_lengths):
+ raise ValueError("MiniMax H3 audio condition and lengths must be present together")
+
+ parsed_shapes = [tuple(int(item) for item in shape) for shape in visual_shapes]
+ if any(
+ len(shape) != 3 or any(item <= 0 for item in shape) or shape[1] % 2 or shape[2] % 2 for shape in parsed_shapes
+ ):
+ raise ValueError("MiniMax H3 visual condition shapes must be positive [T, H, W] triplets with even H/W")
+ if visual is not None:
+ expected_rows = sum(t * (h // 2) * (w // 2) for t, h, w in parsed_shapes)
+ if (
+ visual.dtype != torch.float32
+ or visual.ndim != 2
+ or tuple(visual.shape)
+ != (
+ expected_rows,
+ _VIDEO_CONDITION_WIDTH,
+ )
+ ):
+ raise ValueError(
+ "MiniMax H3 visual condition must be FP32 with shape "
+ f"[{expected_rows}, {_VIDEO_CONDITION_WIDTH}], got {visual.dtype} {tuple(visual.shape)}"
+ )
+
+ parsed_lengths = [int(length) for length in audio_lengths]
+ if any(length <= 0 for length in parsed_lengths):
+ raise ValueError("MiniMax H3 audio condition lengths must be positive")
+ if audio is not None:
+ expected_rows = _AUDIO_CONDITION_CHANNELS * sum(parsed_lengths)
+ if (
+ audio.dtype != torch.float32
+ or audio.ndim != 2
+ or tuple(audio.shape)
+ != (
+ expected_rows,
+ _AUDIO_CONDITION_WIDTH,
+ )
+ ):
+ raise ValueError(
+ "MiniMax H3 audio condition must be FP32 with shape "
+ f"[{expected_rows}, {_AUDIO_CONDITION_WIDTH}], got {audio.dtype} {tuple(audio.shape)}"
+ )
+
+
+def _ref_block_rows(ref_blocks: Sequence[Mapping[str, Any]]) -> list[list[int]]:
+ rows: list[list[int]] = []
+ for block in ref_blocks:
+ kind = str(block.get("kind") or "")
+ kind_code = _REF_KIND_TO_CODE.get(kind)
+ if kind_code is None:
+ raise ValueError(f"unsupported MiniMax H3 ref kind {kind!r}")
+ rows.append(
+ [
+ kind_code,
+ int(block.get("ref_audio_t", 0)),
+ int(block.get("latent_t", 0)),
+ int(block.get("latent_h", 0)),
+ int(block.get("latent_w", 0)),
+ ]
+ )
+ return rows
+
+
+def _decode_ref_blocks(value: torch.Tensor) -> tuple[dict[str, Any], ...]:
+ if value.ndim != 2 or value.shape[1] != 5:
+ raise ValueError("MiniMax H3 ref_blocks metadata must have shape [N, 5]")
+ blocks: list[dict[str, Any]] = []
+ for kind_code, ref_audio_t, ref_t, ref_h, ref_w in value.detach().cpu().tolist():
+ kind = _CODE_TO_REF_KIND.get(int(kind_code))
+ if kind is None:
+ raise ValueError(f"unsupported MiniMax H3 ref kind code {kind_code!r}")
+ block: dict[str, Any] = {"kind": kind}
+ for key, item in (
+ ("ref_audio_t", ref_audio_t),
+ ("latent_t", ref_t),
+ ("latent_h", ref_h),
+ ("latent_w", ref_w),
+ ):
+ if item:
+ block[key] = int(item)
+ blocks.append(block)
+ return tuple(blocks)
+
+
+@dataclass(frozen=True)
+class MiniMaxH3EncoderMediaInput:
+ task: str
+ height: int
+ width: int
+ num_frames: int
+ latent_t: int
+ audio_t: int
+ images: tuple[torch.Tensor, ...] = ()
+ videos: tuple[torch.Tensor, ...] = ()
+ video_audios: tuple[tuple[torch.Tensor, int] | None, ...] = ()
+ audios: tuple[tuple[torch.Tensor, int], ...] = ()
+ keyframe_frame_indices: tuple[int, ...] = ()
+
+ @classmethod
+ def from_mm_tensors(
+ cls,
+ values: Sequence[torch.Tensor],
+ metadata: Mapping[str, Any],
+ ) -> MiniMaxH3EncoderMediaInput:
+ if not isinstance(values, Sequence) or isinstance(values, (str, bytes)):
+ raise ValueError("MiniMax H3 encoder media must be a tensor sequence")
+ tensors = list(values)
+ if any(not isinstance(item, torch.Tensor) for item in tensors):
+ raise ValueError("MiniMax H3 encoder media input must contain only tensors")
+ task = str(metadata["task"])
+ if task not in _TASK_TO_CODE:
+ raise ValueError(f"unsupported MiniMax H3 task {task!r}")
+ height = int(metadata["height"])
+ width = int(metadata["width"])
+ num_frames = int(metadata["num_frames"])
+ latent_t = int(metadata["latent_t"])
+ audio_t = int(metadata["audio_t"])
+ image_count = int(metadata.get("image_count", 0))
+ video_count = int(metadata.get("video_count", 0))
+ audio_count = int(metadata.get("audio_count", 0))
+ if min(height, width, num_frames, latent_t, audio_t) <= 0:
+ raise ValueError("MiniMax H3 encoder media dimensions must be positive")
+ if min(image_count, video_count, audio_count) < 0:
+ raise ValueError("MiniMax H3 encoder media counts must be non-negative")
+
+ audio_flags = tuple(bool(value) for value in metadata.get("video_audio_flags", ()))
+ video_audio_sample_rates = tuple(int(value) for value in metadata.get("video_audio_sample_rates", ()))
+ audio_sample_rates = tuple(int(value) for value in metadata.get("audio_sample_rates", ()))
+ if len(audio_flags) != video_count or len(video_audio_sample_rates) != video_count:
+ raise ValueError("MiniMax H3 video audio flags must align with videos")
+ if len(audio_sample_rates) != audio_count:
+ raise ValueError("MiniMax H3 audio sample rates must align with audio inputs")
+ if any((rate > 0) != has_audio for rate, has_audio in zip(video_audio_sample_rates, audio_flags, strict=True)):
+ raise ValueError("MiniMax H3 embedded audio sample rates must align with audio flags")
+ if any(rate <= 0 for rate in audio_sample_rates):
+ raise ValueError("MiniMax H3 audio sample rates must be positive")
+
+ cursor = 0
+ images = tuple(tensors[cursor : cursor + image_count])
+ cursor += image_count
+ videos = tuple(tensors[cursor : cursor + video_count])
+ cursor += video_count
+ if len(images) != image_count or any(image.ndim != 3 or image.shape[-1] != 3 for image in images):
+ raise ValueError("MiniMax H3 encoder images must have shape [H, W, 3]")
+ if len(videos) != video_count or any(video.ndim != 4 or video.shape[-1] != 3 for video in videos):
+ raise ValueError("MiniMax H3 encoder videos must have shape [T, H, W, 3]")
+
+ video_audios: list[tuple[torch.Tensor, int] | None] = []
+ for has_audio, sample_rate in zip(audio_flags, video_audio_sample_rates, strict=True):
+ if not has_audio:
+ video_audios.append(None)
+ continue
+ if cursor >= len(tensors):
+ raise ValueError("MiniMax H3 encoder media input is truncated")
+ waveform = tensors[cursor]
+ cursor += 1
+ if waveform.ndim not in (1, 2):
+ raise ValueError("MiniMax H3 embedded audio must have shape [samples] or [channels, samples]")
+ video_audios.append((waveform, sample_rate))
+
+ audios: list[tuple[torch.Tensor, int]] = []
+ for sample_rate in audio_sample_rates:
+ if cursor >= len(tensors):
+ raise ValueError("MiniMax H3 encoder media input is truncated")
+ waveform = tensors[cursor]
+ cursor += 1
+ if waveform.ndim not in (1, 2):
+ raise ValueError("MiniMax H3 audio must have shape [samples] or [channels, samples]")
+ audios.append((waveform, sample_rate))
+ if cursor != len(tensors):
+ raise ValueError(f"MiniMax H3 encoder media input has {len(tensors) - cursor} trailing tensors")
+ return cls(
+ task=task,
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ latent_t=latent_t,
+ audio_t=audio_t,
+ images=images,
+ videos=videos,
+ video_audios=tuple(video_audios),
+ audios=tuple(audios),
+ keyframe_frame_indices=tuple(int(value) for value in metadata.get("keyframe_frame_indices", ())),
+ )
+
+ def to_mm_tensors(self) -> list[torch.Tensor]:
+ task_code = _TASK_TO_CODE.get(self.task)
+ if task_code is None:
+ raise ValueError(f"unsupported MiniMax H3 task {self.task!r}")
+ if min(self.height, self.width, self.num_frames, self.latent_t, self.audio_t) <= 0:
+ raise ValueError("MiniMax H3 encoder media dimensions must be positive")
+ if len(self.video_audios) != len(self.videos):
+ raise ValueError("MiniMax H3 video audio slots must align with videos")
+ if any(image.ndim != 3 or image.shape[-1] != 3 for image in self.images):
+ raise ValueError("MiniMax H3 encoder images must have shape [H, W, 3]")
+ if any(video.ndim != 4 or video.shape[-1] != 3 for video in self.videos):
+ raise ValueError("MiniMax H3 encoder videos must have shape [T, H, W, 3]")
+ audio_items = [item for item in self.video_audios if item is not None] + list(self.audios)
+ if any(waveform.ndim not in (1, 2) or int(sample_rate) <= 0 for waveform, sample_rate in audio_items):
+ raise ValueError("MiniMax H3 encoder audio must have a waveform and positive sample rate")
+ tensors = [*self.images, *self.videos]
+ for item in self.video_audios:
+ if item is None:
+ continue
+ waveform, _sample_rate = item
+ tensors.append(waveform)
+ tensors.extend(waveform for waveform, _sample_rate in self.audios)
+ return tensors
+
+ def to_metadata(self) -> dict[str, Any]:
+ return {
+ "task": self.task,
+ "height": self.height,
+ "width": self.width,
+ "num_frames": self.num_frames,
+ "latent_t": self.latent_t,
+ "audio_t": self.audio_t,
+ "image_count": len(self.images),
+ "video_count": len(self.videos),
+ "audio_count": len(self.audios),
+ "video_audio_flags": [item is not None for item in self.video_audios],
+ "video_audio_sample_rates": [item[1] if item is not None else 0 for item in self.video_audios],
+ "audio_sample_rates": [sample_rate for _waveform, sample_rate in self.audios],
+ "keyframe_frame_indices": list(self.keyframe_frame_indices),
+ }
+
+
+@dataclass(frozen=True)
+class MiniMaxH3EncoderMediaConditioning:
+ task: str
+ height: int
+ width: int
+ num_frames: int
+ latent_t: int
+ audio_t: int
+ visual_condition: torch.Tensor | None = None
+ visual_condition_shapes: tuple[tuple[int, int, int], ...] = ()
+ audio_condition: torch.Tensor | None = None
+ audio_condition_lengths: tuple[int, ...] = ()
+ ref_blocks: tuple[dict[str, Any], ...] = ()
+ keyframe_frame_indices: tuple[int, ...] = ()
+
+ def to_omni_components(
+ self,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ task_code = _TASK_TO_CODE.get(self.task)
+ if task_code is None:
+ raise ValueError(f"unsupported MiniMax H3 task {self.task!r}")
+ if min(self.height, self.width, self.num_frames, self.latent_t, self.audio_t) <= 0:
+ raise ValueError("MiniMax H3 encoder media dimensions must be positive")
+ _validate_condition_tensors(
+ self.visual_condition,
+ self.visual_condition_shapes,
+ self.audio_condition,
+ self.audio_condition_lengths,
+ )
+ empty = torch.empty((0,), dtype=torch.float32)
+ ref_rows = _ref_block_rows(self.ref_blocks)
+ layout = [
+ _ENCODER_WIRE_SCHEMA_ID,
+ STAGE_SCHEMA_VERSION,
+ task_code,
+ self.height,
+ self.width,
+ self.num_frames,
+ self.latent_t,
+ self.audio_t,
+ len(self.visual_condition_shapes),
+ len(self.audio_condition_lengths),
+ len(ref_rows),
+ len(self.keyframe_frame_indices),
+ ]
+ layout.extend(item for shape in self.visual_condition_shapes for item in shape)
+ layout.extend(self.audio_condition_lengths)
+ layout.extend(item for row in ref_rows for item in row)
+ layout.extend(self.keyframe_frame_indices)
+ return (
+ self.visual_condition if self.visual_condition is not None else empty,
+ self.audio_condition if self.audio_condition is not None else empty,
+ torch.tensor(layout, dtype=torch.int64),
+ )
+
+
+@dataclass(frozen=True)
+class MiniMaxH3EncoderConditioning:
+ hidden_states: torch.Tensor
+ token_tags: torch.Tensor
+ task: str
+ height: int
+ width: int
+ num_frames: int
+ latent_t: int
+ audio_t: int
+ visual_condition: torch.Tensor | None = None
+ visual_condition_shapes: tuple[tuple[int, int, int], ...] = ()
+ audio_condition: torch.Tensor | None = None
+ audio_condition_lengths: tuple[int, ...] = ()
+ ref_blocks: tuple[dict[str, Any], ...] = ()
+ keyframe_frame_indices: tuple[int, ...] = ()
+
+ @classmethod
+ def from_components(
+ cls,
+ text: MiniMaxH3TextConditioning,
+ media: MiniMaxH3EncoderMediaConditioning,
+ ) -> MiniMaxH3EncoderConditioning:
+ """Combine local encoder results without a stage-wire round trip."""
+ text = MiniMaxH3TextConditioning.from_payload(text.to_payload())
+ _validate_condition_tensors(
+ media.visual_condition,
+ media.visual_condition_shapes,
+ media.audio_condition,
+ media.audio_condition_lengths,
+ )
+ return cls(
+ hidden_states=text.hidden_states,
+ token_tags=text.token_tags,
+ **{item.name: getattr(media, item.name) for item in fields(media)},
+ )
+
+ @classmethod
+ def from_omni_payload(cls, payload: Mapping[str, Any]) -> MiniMaxH3EncoderConditioning:
+ hidden_payload = payload.get("hidden_states")
+ embed_payload = payload.get("embed")
+ meta_payload = payload.get("meta")
+ hidden_states = hidden_payload.get("output") if isinstance(hidden_payload, Mapping) else None
+ token_tags = meta_payload.get("token_role_ids") if isinstance(meta_payload, Mapping) else None
+ visual = embed_payload.get("embedding") if isinstance(embed_payload, Mapping) else None
+ audio = embed_payload.get("speech_feat") if isinstance(embed_payload, Mapping) else None
+ if not all(isinstance(item, torch.Tensor) for item in (hidden_states, token_tags, visual, audio)):
+ raise ValueError("MiniMax H3 encoder wire requires text, visual and audio tensors")
+ assert isinstance(hidden_states, torch.Tensor)
+ assert isinstance(token_tags, torch.Tensor)
+ assert isinstance(visual, torch.Tensor)
+ assert isinstance(audio, torch.Tensor)
+ if token_tags.ndim == 2 and token_tags.shape[-1] == 1:
+ token_tags = token_tags.squeeze(-1)
+ text = MiniMaxH3TextConditioning.from_payload({"hidden_states": hidden_states, "token_tags": token_tags})
+ for name, value in (("visual", visual), ("audio", audio)):
+ if value.numel() == 0 and (value.dtype != torch.float32 or tuple(value.shape) != (0,)):
+ raise ValueError(f"MiniMax H3 empty {name} condition slot must be FP32 with shape [0]")
+ if not isinstance(meta_payload, Mapping):
+ raise ValueError("MiniMax H3 wire payload requires a meta mapping")
+ layout = _wire_layout(payload)
+ if layout.dtype not in _INTEGER_DTYPES or layout.ndim != 1:
+ raise ValueError("MiniMax H3 encoder layout must be a one-dimensional integer tensor")
+ values = [int(item) for item in layout.detach().cpu().tolist()]
+ if len(values) < 12:
+ raise ValueError("MiniMax H3 encoder layout header is truncated")
+ (
+ schema_id,
+ version,
+ task_code,
+ height,
+ width,
+ num_frames,
+ latent_t,
+ audio_t,
+ visual_count,
+ audio_count,
+ ref_count,
+ keyframe_count,
+ ) = values[:12]
+ if schema_id != _ENCODER_WIRE_SCHEMA_ID or version != STAGE_SCHEMA_VERSION:
+ raise ValueError(f"unsupported MiniMax H3 encoder wire schema {schema_id}:{version}")
+ task = _CODE_TO_TASK.get(task_code)
+ if task is None:
+ raise ValueError(f"unsupported MiniMax H3 wire task code {task_code!r}")
+ if min(height, width, num_frames, latent_t, audio_t) <= 0:
+ raise ValueError("MiniMax H3 encoder wire dimensions must be positive")
+ if min(visual_count, audio_count, ref_count, keyframe_count) < 0:
+ raise ValueError("MiniMax H3 encoder layout counts must be non-negative")
+ expected_size = 12 + 3 * visual_count + audio_count + 5 * ref_count + keyframe_count
+ if len(values) != expected_size:
+ raise ValueError(f"MiniMax H3 encoder layout has {len(values)} integers, expected {expected_size}")
+ cursor = 12
+ visual_shapes = tuple(
+ tuple(values[cursor + 3 * index : cursor + 3 * (index + 1)]) for index in range(visual_count)
+ )
+ cursor += 3 * visual_count
+ audio_lengths = tuple(values[cursor : cursor + audio_count])
+ cursor += audio_count
+ ref_rows = values[cursor : cursor + 5 * ref_count]
+ cursor += 5 * ref_count
+ ref_blocks_tensor = torch.tensor(ref_rows, dtype=torch.int64).reshape(-1, 5)
+ keyframe_frame_indices = tuple(values[cursor : cursor + keyframe_count])
+ visual_condition = visual if visual.numel() else None
+ audio_condition = audio if audio.numel() else None
+ _validate_condition_tensors(
+ visual_condition,
+ visual_shapes,
+ audio_condition,
+ audio_lengths,
+ )
+ return cls(
+ hidden_states=text.hidden_states,
+ token_tags=text.token_tags,
+ task=task,
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ latent_t=latent_t,
+ audio_t=audio_t,
+ visual_condition=visual_condition,
+ visual_condition_shapes=visual_shapes,
+ audio_condition=audio_condition,
+ audio_condition_lengths=audio_lengths,
+ ref_blocks=_decode_ref_blocks(ref_blocks_tensor),
+ keyframe_frame_indices=keyframe_frame_indices,
+ )
+
+ def to_omni_payload(self) -> dict[str, Any]:
+ text = MiniMaxH3TextConditioning.from_payload(
+ {
+ "hidden_states": self.hidden_states,
+ "token_tags": self.token_tags,
+ }
+ )
+ visual, audio, layout = MiniMaxH3EncoderMediaConditioning(
+ task=self.task,
+ height=self.height,
+ width=self.width,
+ num_frames=self.num_frames,
+ latent_t=self.latent_t,
+ audio_t=self.audio_t,
+ visual_condition=self.visual_condition,
+ visual_condition_shapes=self.visual_condition_shapes,
+ audio_condition=self.audio_condition,
+ audio_condition_lengths=self.audio_condition_lengths,
+ ref_blocks=self.ref_blocks,
+ keyframe_frame_indices=self.keyframe_frame_indices,
+ ).to_omni_components()
+ return {
+ "hidden_states": {"output": text.hidden_states},
+ "embed": {
+ "embedding": visual,
+ "speech_feat": audio,
+ },
+ "meta": {"token_role_ids": text.token_tags},
+ "kv_metadata": {MINIMAX_H3_ENCODER_LAYOUT_KEY: layout},
+ }
+
+
+__all__ = [
+ "MINIMAX_H3_CONDITION_LABELS_KEY",
+ "MINIMAX_H3_ENCODER_LAYOUT_KEY",
+ "MINIMAX_H3_ENCODER_REQUEST_KEY",
+ "MINIMAX_H3_PRESENTATION_TASK_KEY",
+ "MINIMAX_H3_TEXT_CONDITIONING_SCHEMA",
+ "MINIMAX_H3_TEXT_HIDDEN_SIZE",
+ "STAGE_SCHEMA_VERSION",
+ "MiniMaxH3EncoderConditioning",
+ "MiniMaxH3EncoderMediaConditioning",
+ "MiniMaxH3EncoderMediaInput",
+ "MiniMaxH3TextConditioning",
+]
diff --git a/vllm_omni/model_executor/models/minimax_h3/encoder.py b/vllm_omni/model_executor/models/minimax_h3/encoder.py
new file mode 100644
index 000000000..44fefab91
--- /dev/null
+++ b/vllm_omni/model_executor/models/minimax_h3/encoder.py
@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig
+from vllm.distributed.parallel_state import get_tp_group
+
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.minimax_h3.vae import (
+ MiniMaxH3AudioVAE,
+ MiniMaxH3VideoVAE,
+)
+from vllm_omni.model_executor.models.minimax_h3.conditioning import (
+ MINIMAX_H3_ENCODER_LAYOUT_KEY,
+ MINIMAX_H3_ENCODER_REQUEST_KEY,
+ MiniMaxH3EncoderMediaConditioning,
+ MiniMaxH3EncoderMediaInput,
+)
+from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
+from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
+ MiniMaxH3TextEncoderBackbone,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+_MEDIA_CONDITIONING_CACHE_KEY = "_minimax_h3_encoder_media_conditioning"
+_MEDIA_CACHE_KEY = "_minimax_h3_encoder_media"
+_COMPONENT_CONFIG_KEY = "minimax_h3_encoder_components"
+_COMPONENT_ROLES = frozenset({"text_encoder", "video_vae", "audio_vae"})
+
+
+def _partition_root(model_path: str) -> Path:
+ path = Path(model_path)
+ if path.name in {"text_encoder", "video_vae", "audio_vae"}:
+ return path.parent
+ return path
+
+
+@dataclass(frozen=True)
+class MiniMaxH3EncoderComponentConfig:
+ text_parallel_mode: str
+ video_parallel_mode: str
+ audio_parallel_mode: str
+
+ def __post_init__(self) -> None:
+ if self.text_parallel_mode != "tp":
+ raise ValueError("MiniMax H3 text encoder currently supports parallel_mode='tp' only")
+ if self.video_parallel_mode not in {"leader", "patch"}:
+ raise ValueError("MiniMax H3 video VAE parallel_mode must be 'leader' or 'patch'")
+ if self.audio_parallel_mode != "leader":
+ raise ValueError("MiniMax H3 audio VAE currently supports parallel_mode='leader' only")
+
+ @staticmethod
+ def _parallel_mode(raw: Mapping[str, Any], role: str) -> str:
+ value = raw[role]
+ if not isinstance(value, Mapping):
+ raise TypeError(f"MiniMax H3 encoder role {role!r} must be a mapping")
+ if set(value) != {"parallel_mode"}:
+ raise ValueError(f"MiniMax H3 encoder role {role!r} requires only parallel_mode")
+ return str(value["parallel_mode"])
+
+ @classmethod
+ def from_vllm_config(cls, vllm_config: VllmConfig) -> MiniMaxH3EncoderComponentConfig:
+ hf_config = vllm_config.model_config.hf_config
+ raw = getattr(hf_config, _COMPONENT_CONFIG_KEY, None)
+ if not isinstance(raw, Mapping):
+ raise ValueError(f"MiniMax H3 Encoder Stage requires explicit {_COMPONENT_CONFIG_KEY}")
+ roles = set(raw)
+ if roles != _COMPONENT_ROLES:
+ missing = sorted(_COMPONENT_ROLES - roles)
+ unknown = sorted(roles - _COMPONENT_ROLES)
+ raise ValueError(
+ f"{_COMPONENT_CONFIG_KEY} requires exactly {sorted(_COMPONENT_ROLES)}; "
+ f"missing={missing}, unknown={unknown}"
+ )
+ return cls(
+ text_parallel_mode=cls._parallel_mode(raw, "text_encoder"),
+ video_parallel_mode=cls._parallel_mode(raw, "video_vae"),
+ audio_parallel_mode=cls._parallel_mode(raw, "audio_vae"),
+ )
+
+
+class MiniMaxH3Encoder(MiniMaxH3TextEncoderBackbone):
+ requires_request_sample_eligibility = True
+ enable_update_additional_information = True
+ omni_payload_at_request_end = True
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
+ try:
+ tp_group = get_tp_group()
+ tp_rank = int(tp_group.rank_in_group)
+ except (AssertionError, RuntimeError):
+ tp_group = None
+ tp_rank = 0
+ self._component_leader = tp_rank == 0
+ self.component_config = MiniMaxH3EncoderComponentConfig.from_vllm_config(vllm_config)
+ root = _partition_root(vllm_config.model_config.model)
+ missing = [name for name in ("text_encoder", "video_vae", "audio_vae") if not (root / name).is_dir()]
+ if missing:
+ raise RuntimeError(f"MiniMax H3 Encoder Stage requires all three encoders; missing {missing}")
+ device = get_local_device()
+ self.video_vae: MiniMaxH3VideoVAE | None = None
+ self.audio_vae: MiniMaxH3AudioVAE | None = None
+ if self._component_leader or self.component_config.video_parallel_mode == "patch":
+ self.video_vae = MiniMaxH3VideoVAE(
+ str(root / "video_vae"),
+ device=device,
+ encode_only=True,
+ trust_remote_code=vllm_config.model_config.trust_remote_code,
+ )
+ if self.component_config.video_parallel_mode == "patch":
+ if tp_group is None:
+ raise RuntimeError("MiniMax H3 video VAE patch parallelism requires an initialized TP group")
+ self.video_vae.set_parallel_size(
+ tp_group.world_size,
+ process_group=tp_group.device_group,
+ )
+ if self._component_leader:
+ self.audio_vae = MiniMaxH3AudioVAE(
+ str(root / "audio_vae"),
+ device=device,
+ encode_only=True,
+ trust_remote_code=vllm_config.model_config.trust_remote_code,
+ )
+
+ @staticmethod
+ def _media_input(info: dict[str, Any]) -> MiniMaxH3EncoderMediaInput:
+ cached = info.get(_MEDIA_CACHE_KEY)
+ if isinstance(cached, MiniMaxH3EncoderMediaInput):
+ return cached
+ meta = info.get("meta")
+ metadata = meta.get(MINIMAX_H3_ENCODER_REQUEST_KEY) if isinstance(meta, Mapping) else None
+ if not isinstance(metadata, Mapping):
+ raise RuntimeError("MiniMax H3 unified encoder requires request metadata")
+
+ hidden_states = info.get("hidden_states")
+ raw_media = hidden_states.get("layers") if isinstance(hidden_states, Mapping) else None
+ if raw_media is None:
+ media_tensors: list[torch.Tensor] = []
+ elif isinstance(raw_media, Mapping):
+ try:
+ ordered_items = sorted(raw_media.items(), key=lambda item: int(item[0]))
+ except (TypeError, ValueError) as exc:
+ raise RuntimeError("MiniMax H3 encoder media keys must be integer indices") from exc
+ if [int(key) for key, _ in ordered_items] != list(range(len(ordered_items))):
+ raise RuntimeError("MiniMax H3 encoder media indices must be contiguous from zero")
+ media_tensors = [value for _, value in ordered_items]
+ else:
+ raise RuntimeError("MiniMax H3 encoder media must be a namespaced tensor mapping")
+ if any(not isinstance(value, torch.Tensor) for value in media_tensors):
+ raise RuntimeError("MiniMax H3 encoder media must contain only tensors")
+ media = MiniMaxH3EncoderMediaInput.from_mm_tensors(media_tensors, metadata)
+ if media.to_metadata() != dict(metadata):
+ raise RuntimeError("MiniMax H3 encoder media does not match request metadata")
+ info[_MEDIA_CACHE_KEY] = media
+ return media
+
+ def _encode_media(self, media: MiniMaxH3EncoderMediaInput) -> MiniMaxH3EncoderMediaConditioning | None:
+ if not self._component_leader and self.component_config.video_parallel_mode == "leader":
+ return None
+ return encode_media(
+ media,
+ video_vae=self.video_vae,
+ audio_vae=self.audio_vae,
+ emit_conditioning=self._component_leader,
+ )
+
+ @staticmethod
+ def _has_encoder_input(info: object) -> bool:
+ if not isinstance(info, Mapping):
+ return False
+ meta = info.get("meta")
+ return isinstance(meta, Mapping) and isinstance(meta.get(MINIMAX_H3_ENCODER_REQUEST_KEY), Mapping)
+
+ def make_omni_output(
+ self,
+ model_outputs: torch.Tensor | OmniOutput,
+ **kwargs: Any,
+ ) -> OmniOutput:
+ output = super().make_omni_output(model_outputs, **kwargs)
+ infos = kwargs.get("model_intermediate_buffer") or []
+ if not isinstance(infos, Sequence) or not infos:
+ return output
+ sample_eligible = kwargs.get("request_sample_eligible")
+ if sample_eligible is None:
+ sample_eligible = [True] * len(infos)
+ if len(sample_eligible) != len(infos):
+ raise RuntimeError(
+ f"MiniMax H3 unified encoder received {len(sample_eligible)} sampling flags for {len(infos)} requests"
+ )
+ per_request: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None] = []
+ release_workspace = False
+ for eligible, info in zip(sample_eligible, infos, strict=True):
+ if not eligible:
+ per_request.append(None)
+ continue
+ if not isinstance(info, dict) or not self._has_encoder_input(info):
+ raise RuntimeError("MiniMax H3 unified encoder request metadata is missing from the runner")
+ release_workspace = True
+ if _MEDIA_CONDITIONING_CACHE_KEY not in info:
+ info[_MEDIA_CONDITIONING_CACHE_KEY] = self._encode_media(self._media_input(info))
+ if not self._component_leader:
+ continue
+ cached = info.get(_MEDIA_CONDITIONING_CACHE_KEY)
+ if not isinstance(cached, MiniMaxH3EncoderMediaConditioning):
+ raise RuntimeError("MiniMax H3 encoder components did not produce conditioning")
+ per_request.append(cached.to_omni_components())
+
+ if release_workspace:
+ torch.accelerator.empty_cache()
+ if not self._component_leader:
+ return output
+
+ if not any(item is not None for item in per_request):
+ return output
+ empty = torch.empty((0,), dtype=torch.float32)
+ multimodal_outputs = dict(output.multimodal_outputs or {})
+ dotted_key = f"kv_metadata.{MINIMAX_H3_ENCODER_LAYOUT_KEY}"
+ multimodal_outputs[dotted_key] = [
+ item[2] if item is not None else torch.empty((0,), dtype=torch.int64) for item in per_request
+ ]
+
+ embed = dict(multimodal_outputs.get("embed") or {})
+ embed["embedding"] = [item[0] if item is not None else empty for item in per_request]
+ embed["speech_feat"] = [item[1] if item is not None else empty for item in per_request]
+ multimodal_outputs["embed"] = embed
+ return OmniOutput(
+ text_hidden_states=output.text_hidden_states,
+ multimodal_outputs=multimodal_outputs,
+ intermediate_tensors=output.intermediate_tensors,
+ next_token_id=output.next_token_id,
+ )
+
+ def load_weights(
+ self,
+ weights: Iterable[tuple[str, torch.Tensor]],
+ ) -> set[str]:
+ loaded = super().load_weights(weights)
+ for component_name in ("video_vae", "audio_vae"):
+ component = getattr(self, component_name)
+ if component is not None:
+ loaded.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
+ return loaded
+
+
+__all__ = ["MiniMaxH3Encoder"]
diff --git a/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py b/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py
new file mode 100644
index 000000000..2d6998b7a
--- /dev/null
+++ b/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py
@@ -0,0 +1,569 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shared request preparation and condition codec execution for MiniMax H3."""
+
+from __future__ import annotations
+
+import math
+import os
+import tempfile
+from collections.abc import Callable, Mapping, Sequence
+from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+from vllm_omni.errors import OmniClientError
+from vllm_omni.model_executor.models.minimax_h3.conditioning import (
+ MiniMaxH3EncoderMediaConditioning,
+ MiniMaxH3EncoderMediaInput,
+)
+from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
+ MINIMAX_H3_OUTPUT_SHORT_EDGE,
+ load_minimax_h3_images,
+ resolve_minimax_h3_aspect_ratio,
+ resolve_minimax_h3_output_canvas,
+ resolve_minimax_h3_reference_image_shape,
+)
+from vllm_omni.model_executor.models.minimax_h3.reference_video import (
+ MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
+ load_audio_file,
+ load_video_audio,
+ load_video_frames,
+ prepare_reference_videos,
+ sample_reference_video_frames,
+ validate_reference_audio_files,
+ validate_reference_audio_waveforms,
+)
+
+MINIMAX_H3_FPS = 24
+MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0
+MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0
+
+
+def _items(value: Any) -> list[Any]:
+ if value is None:
+ return []
+ if isinstance(value, list):
+ return value
+ if isinstance(value, tuple) and not (len(value) == 2 and isinstance(value[1], Mapping)):
+ return list(value)
+ return [value]
+
+
+def _audio_items(value: Any) -> list[Any]:
+ if value is None:
+ return []
+ if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)):
+ return [value]
+ return list(value) if isinstance(value, (list, tuple)) else [value]
+
+
+def _load_audio(value: Any) -> tuple[torch.Tensor, int]:
+ if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))):
+ audios = _load_audios(value)
+ if len(audios) != 1:
+ raise OmniClientError(f"MiniMax H3 expected one audio, got {len(audios)}")
+ return audios[0]
+ if isinstance(value, (str, os.PathLike)):
+ return load_audio_file(str(value))
+ if isinstance(value, (list, tuple)) and len(value) == 2:
+ waveform, sample_rate = value
+ return torch.as_tensor(waveform).float(), int(sample_rate)
+ if isinstance(value, dict):
+ waveform = value.get("waveform", value.get("array"))
+ sample_rate = value.get("sample_rate", value.get("sampling_rate"))
+ if waveform is not None and sample_rate is not None:
+ return torch.as_tensor(waveform).float(), int(sample_rate)
+ raise OmniClientError("MiniMax H3 audio input must be a path, (waveform, sample_rate), or a waveform mapping")
+
+
+def _load_audios(value: Any) -> list[tuple[torch.Tensor, int]]:
+ if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))):
+ if not value:
+ raise OmniClientError("MiniMax H3 audio input must not be empty")
+ return [_load_audio(item) for item in value]
+ return [_load_audio(value)]
+
+
+def _as_int_list(value: Any, *, name: str) -> list[int]:
+ if isinstance(value, bool):
+ raise OmniClientError(f"{name} must be an integer or a list of integers")
+ if isinstance(value, (int, np.integer)):
+ return [int(value)]
+ if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
+ result = list(value)
+ if not result:
+ raise OmniClientError(f"{name} must not be empty")
+ if any(isinstance(item, bool) or not isinstance(item, (int, np.integer)) for item in result):
+ raise OmniClientError(f"{name} must contain only integers")
+ return [int(item) for item in result]
+ raise OmniClientError(f"{name} must be an integer or a list of integers")
+
+
+def _resolve_fl2va_keyframe_indices(extra: Mapping[str, Any], image_count: int) -> list[int]:
+ target = extra.get("target")
+ target = target if isinstance(target, Mapping) else {}
+ raw = extra.get("frame_indices", extra.get("frame_index"))
+ if raw is None:
+ raw = target.get("frame_indices", target.get("frame_index"))
+ raw_indices = ([0] if image_count == 1 else [0, -1]) if raw is None else _as_int_list(raw, name="frame_indices")
+ if len(raw_indices) != image_count:
+ raise OmniClientError(
+ f"MiniMax H3 FL2VA requires one frame index per image: got {raw_indices!r} for {image_count} image(s)"
+ )
+ if tuple(raw_indices) not in ((0,), (-1,), (0, -1)):
+ raise OmniClientError("MiniMax H3 FL2VA frame_indices must be [0], [-1], or [0, -1]")
+ return raw_indices
+
+
+def _validate_ref2va_reference_counts(image_count: int, video_count: int, audio_count: int) -> None:
+ if min(image_count, video_count, audio_count) < 0:
+ raise OmniClientError("MiniMax H3 reference counts must be non-negative")
+ if image_count + video_count == 0:
+ raise OmniClientError("ref2va requires at least one image or video reference")
+ if image_count > 9:
+ raise OmniClientError("ref2va accepts at most 9 image references")
+ if video_count > 3:
+ raise OmniClientError("ref2va accepts at most 3 video references")
+ if audio_count > 3:
+ raise OmniClientError("ref2va accepts at most 3 standalone audio references")
+ if image_count + video_count + audio_count > 12:
+ raise OmniClientError("ref2va accepts at most 12 total references")
+
+
+def _validate_reference_image(image: Image.Image) -> None:
+ resolve_minimax_h3_reference_image_shape(image)
+
+
+def resolve_minimax_h3_shape(
+ task: str,
+ sampling: Any,
+ image: Image.Image | None,
+) -> tuple[int, int, int, int, int]:
+ # The diffusion package exports a pipeline that imports this shared module.
+ from vllm_omni.diffusion.models.minimax_h3.time_request import (
+ MINIMAX_H3_SHAPE_PLANNER,
+ minimax_h3_align_frame_count,
+ )
+
+ fps = int(getattr(sampling, "fps", None) or MINIMAX_H3_FPS)
+ if fps != MINIMAX_H3_FPS:
+ raise OmniClientError(f"MiniMax H3 output fps is fixed at {MINIMAX_H3_FPS}")
+ extra = sampling.extra_args or {}
+ target = extra.get("target")
+ if target is not None and not isinstance(target, Mapping):
+ raise OmniClientError("MiniMax H3 extra_args['target'] must be an object")
+ target = target if isinstance(target, Mapping) else {}
+ duration = target.get("duration_seconds", extra.get("duration_seconds", extra.get("duration")))
+ if duration is not None:
+ if isinstance(duration, bool):
+ raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}")
+ try:
+ duration = float(duration)
+ except (TypeError, ValueError) as exc:
+ raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}") from exc
+ if (
+ not math.isfinite(duration)
+ or not MINIMAX_H3_MIN_OUTPUT_SECONDS <= duration <= MINIMAX_H3_MAX_OUTPUT_SECONDS
+ ):
+ raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration}")
+ requested_frames = int(round(duration * fps))
+ elif int(getattr(sampling, "num_frames", None) or 1) > 1:
+ requested_frames = int(sampling.num_frames)
+ else:
+ requested_frames = 124 if task == "ref2va" else 209
+ if not MINIMAX_H3_MIN_OUTPUT_SECONDS <= requested_frames / fps <= MINIMAX_H3_MAX_OUTPUT_SECONDS:
+ raise OmniClientError(
+ f"MiniMax H3 output duration must be in [4, 15] seconds, got {requested_frames / fps:.3f}"
+ )
+ num_frames = minimax_h3_align_frame_count(requested_frames)
+
+ height = getattr(sampling, "height", None)
+ width = getattr(sampling, "width", None)
+ aspect_ratio = resolve_minimax_h3_aspect_ratio(
+ task,
+ target.get("aspect_ratio", extra.get("aspect_ratio")),
+ image,
+ )
+ raw_short_edge = target.get("short_edge", extra.get("short_edge", MINIMAX_H3_OUTPUT_SHORT_EDGE))
+ if isinstance(raw_short_edge, bool) or not isinstance(raw_short_edge, (int, np.integer)):
+ raise OmniClientError(
+ f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {raw_short_edge!r}"
+ )
+ if height is None or width is None:
+ height, width = resolve_minimax_h3_output_canvas(aspect_ratio, int(raw_short_edge))
+ height = int(height) // 32 * 32
+ width = int(width) // 32 * 32
+ if min(height, width) <= 0:
+ raise OmniClientError(f"invalid MiniMax H3 canvas {width}x{height}")
+ if width > 4 * height or height > 4 * width:
+ raise OmniClientError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]")
+ return (
+ height,
+ width,
+ num_frames,
+ MINIMAX_H3_SHAPE_PLANNER.video_latent_t(num_frames),
+ MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(num_frames / fps),
+ )
+
+
+def _resolve_task(
+ extra_args: Mapping[str, Any],
+ multi_modal_data: Mapping[str, Any],
+) -> str:
+ requested = extra_args.get("task")
+ if requested is not None:
+ return str(requested).lower()
+ if multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None:
+ return "ref2va"
+ if multi_modal_data.get("image") is not None:
+ return "fl2va"
+ return "t2va"
+
+
+def _prepare_encoder_images(
+ task: str,
+ images: list[Image.Image],
+ *,
+ height: int,
+ width: int,
+) -> list[Any]:
+ if not images:
+ return []
+ if task == "ref2va":
+ return [
+ image.resize(
+ resolve_minimax_h3_reference_image_shape(image),
+ Image.Resampling.LANCZOS,
+ )
+ for image in images
+ ]
+ if task != "fl2va":
+ return images
+ return [image.resize((width, height), Image.Resampling.LANCZOS) for image in images]
+
+
+def _image_to_tensor(image: Image.Image) -> torch.Tensor:
+ array = np.asarray(image.convert("RGB"), dtype=np.uint8)
+ return torch.from_numpy(np.array(array, copy=True)).contiguous()
+
+
+def _frames_to_tensor(frames: Sequence[Any]) -> torch.Tensor:
+ if len(frames) == 0:
+ raise ValueError("MiniMax H3 reference video must contain frames")
+ return torch.stack(
+ [
+ _image_to_tensor(frame if isinstance(frame, Image.Image) else Image.fromarray(np.asarray(frame)))
+ for frame in frames
+ ]
+ )
+
+
+def _reuse_prepared_reference_videos(
+ prepared: list[dict[str, Any]] | None,
+ *,
+ expected_count: int,
+) -> list[dict[str, Any]] | None:
+ if prepared is None:
+ return None
+ if len(prepared) != expected_count:
+ raise OmniClientError("MiniMax H3 prepared-reference-video count does not match the request")
+ for item in prepared:
+ if not os.path.isfile(item["prepared_path"]):
+ raise OmniClientError(f"MiniMax H3 prepared reference video is unavailable: {item['prepared_path']}")
+ return prepared
+
+
+@dataclass(frozen=True)
+class PreparedEncoderInputs:
+ """CPU-owned media and Qwen presentation inputs, independent of a runner."""
+
+ prompt: str
+ media: MiniMaxH3EncoderMediaInput
+ images: list[Image.Image]
+ qwen_videos: list[tuple[np.ndarray, dict[str, Any]]]
+ video_timestamps: list[list[float]]
+ condition_labels: list[tuple[str, int]]
+
+
+def _effective_audio_inputs(
+ video_audios: Sequence[tuple[torch.Tensor, int] | None],
+ standalone_audios: Sequence[tuple[torch.Tensor, int]],
+ *,
+ max_standalone_seconds: float,
+) -> list[tuple[torch.Tensor, int]]:
+ """Return the waveforms exactly as the audio WVAE will consume them."""
+ audio_inputs = [item for item in video_audios if item is not None]
+ audio_inputs.extend(
+ (waveform[..., : int(round(max_standalone_seconds * sample_rate))], sample_rate)
+ for waveform, sample_rate in standalone_audios
+ )
+ return audio_inputs
+
+
+def prepare_encoder_inputs(
+ prompt: Any,
+ sampling: Any,
+ *,
+ task: str | None = None,
+ prepared_reference_videos: list[dict[str, Any]] | None = None,
+) -> PreparedEncoderInputs:
+ """Decode references once and prepare the shared text/codec presentation."""
+ if isinstance(prompt, str):
+ prompt = {"prompt": prompt}
+ if not isinstance(prompt, dict):
+ raise TypeError(f"MiniMax H3 expects a string or dict prompt, got {type(prompt)!r}")
+
+ text = str(prompt.get("prompt") or "")
+ if not text:
+ raise OmniClientError("MiniMax H3 requires a non-empty prompt")
+ multi_modal_data = prompt.get("multi_modal_data") or {}
+ if not isinstance(multi_modal_data, Mapping):
+ raise TypeError("multi_modal_data must be a mapping")
+
+ image_values = _items(multi_modal_data.get("image"))
+ videos = _items(multi_modal_data.get("video"))
+ raw_audio = multi_modal_data.get("audio")
+ audio_values = _audio_items(raw_audio)
+ diffusion_sampling_params = sampling
+ extra_args = getattr(diffusion_sampling_params, "extra_args", None) or {}
+ task = task if task is not None else _resolve_task(extra_args, multi_modal_data)
+ raw_images = load_minimax_h3_images(image_values) if image_values else []
+
+ if task == "t2va":
+ if raw_images or videos or audio_values:
+ raise OmniClientError("t2va does not accept image, video, or audio conditions")
+ elif task == "fl2va":
+ if not raw_images or videos or audio_values:
+ raise OmniClientError("fl2va requires image conditions only")
+ if len(raw_images) > 2:
+ raise OmniClientError("fl2va accepts at most first and last images")
+ for image in raw_images:
+ _validate_reference_image(image)
+ elif task == "ref2va":
+ _validate_ref2va_reference_counts(len(raw_images), len(videos), len(audio_values))
+ if not raw_images and not videos:
+ raise OmniClientError("ref2va requires an image or video condition")
+ else:
+ raise OmniClientError(f"unsupported MiniMax H3 task {task!r}")
+
+ height, width, num_frames, latent_t, audio_t = resolve_minimax_h3_shape(
+ task,
+ diffusion_sampling_params,
+ raw_images[0] if raw_images else None,
+ )
+ images = _prepare_encoder_images(
+ task,
+ raw_images,
+ height=height,
+ width=width,
+ )
+ keyframe_indices = _resolve_fl2va_keyframe_indices(extra_args, len(images)) if task == "fl2va" else []
+ video_timestamps: list[list[float]] = []
+ qwen_video_inputs: list[tuple[np.ndarray, dict[str, Any]]] = []
+ condition_labels: list[tuple[str, int]] = []
+ encoded_video_inputs: list[torch.Tensor] = []
+ video_audio_inputs: list[tuple[torch.Tensor, int] | None] = []
+
+ if task == "fl2va":
+ condition_labels.extend(("image", index) for index in range(1, len(images) + 1))
+ elif task == "ref2va":
+ condition_labels.extend(("image", index) for index in range(1, len(images) + 1))
+ prepared_videos: list[dict[str, Any]] = []
+ if videos:
+ with tempfile.TemporaryDirectory(prefix="minimax_h3_encoder_") as workdir:
+ prepared_videos = _reuse_prepared_reference_videos(
+ prepared_reference_videos, expected_count=len(videos)
+ )
+ if prepared_videos is None:
+ prepared_videos = prepare_reference_videos(
+ videos,
+ target_frame_count=num_frames,
+ workdir=workdir,
+ start_time_seconds=extra_args.get("start_time_seconds"),
+ )
+ for item in prepared_videos:
+ full_frames = load_video_frames(item["prepared_path"])
+ encoded_video_inputs.append(_frames_to_tensor(full_frames))
+ sampled = sample_reference_video_frames(
+ item["prepared_path"],
+ decoded_frames=full_frames,
+ )
+ video_timestamps.append(sampled["block_timestamps"])
+ frames = np.stack(sampled["frames"])
+ frame_count = int(frames.shape[0])
+ qwen_video_inputs.append(
+ (
+ frames,
+ {
+ "total_num_frames": frame_count,
+ "fps": MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
+ "duration": frame_count / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
+ "video_backend": "minimax_h3",
+ "frames_indices": list(range(frame_count)),
+ "do_sample_frames": False,
+ },
+ )
+ )
+ if item["input_has_audio"]:
+ waveform, sample_rate = load_video_audio(
+ item["original_path"],
+ start_time_seconds=float(item.get("start_time_seconds", 0.0)),
+ duration_seconds=item.get(
+ "audio_duration_seconds",
+ item.get("duration_seconds"),
+ ),
+ )
+ video_audio_inputs.append((waveform.float().contiguous(), int(sample_rate)))
+ else:
+ video_audio_inputs.append(None)
+ audio_index = 0
+ for video_index, item in enumerate(prepared_videos, start=1):
+ if item["input_has_audio"]:
+ audio_index += 1
+ condition_labels.append(("audio", audio_index))
+ condition_labels.append(("video", video_index))
+ for _ in audio_values:
+ audio_index += 1
+ condition_labels.append(("audio", audio_index))
+
+ if raw_audio is not None:
+ validate_reference_audio_files(raw_audio)
+ standalone_audios = _load_audios(raw_audio) if raw_audio is not None else []
+ validate_reference_audio_waveforms(standalone_audios)
+ media_input = MiniMaxH3EncoderMediaInput(
+ task=task,
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ latent_t=latent_t,
+ audio_t=audio_t,
+ images=tuple(_image_to_tensor(image) for image in images),
+ videos=tuple(encoded_video_inputs),
+ video_audios=tuple(video_audio_inputs),
+ audios=tuple((waveform.float().contiguous(), int(sample_rate)) for waveform, sample_rate in standalone_audios),
+ keyframe_frame_indices=tuple(keyframe_indices),
+ )
+ audio_inputs = _effective_audio_inputs(
+ media_input.video_audios,
+ media_input.audios,
+ max_standalone_seconds=float(media_input.num_frames) / MINIMAX_H3_FPS,
+ )
+ embedded_audio_count = sum(item is not None for item in media_input.video_audios)
+ # Video soundtracks and standalone references have separate 15-second budgets.
+ validate_reference_audio_waveforms(audio_inputs[:embedded_audio_count])
+ validate_reference_audio_waveforms(audio_inputs[embedded_audio_count:])
+
+ return PreparedEncoderInputs(
+ prompt=text,
+ media=media_input,
+ images=images,
+ qwen_videos=qwen_video_inputs,
+ video_timestamps=video_timestamps,
+ condition_labels=condition_labels,
+ )
+
+
+def _image_from_tensor(value: torch.Tensor) -> Image.Image:
+ if value.ndim != 3 or value.shape[-1] != 3:
+ raise ValueError(f"MiniMax H3 image tensor must have shape [H, W, 3], got {tuple(value.shape)}")
+ array = value.detach().cpu().to(torch.uint8).contiguous().numpy()
+ return Image.fromarray(array, mode="RGB")
+
+
+def encode_media(
+ media: MiniMaxH3EncoderMediaInput,
+ *,
+ video_vae: Any,
+ audio_vae: Any,
+ emit_conditioning: bool,
+ component_scope: Callable[[Any], AbstractContextManager[Any]] = nullcontext,
+) -> MiniMaxH3EncoderMediaConditioning | None:
+ """Run codecs on participating ranks; only the output rank runs audio.
+
+ The caller selects participating ranks and configures the VAE process group.
+ Every visual participant must call this function with the same references.
+ Component scopes allow the diffusion runner to stage one codec at a time.
+ """
+ visual_rows: list[torch.Tensor] = []
+ visual_shapes: list[tuple[int, int, int]] = []
+ if media.images or media.videos:
+ if video_vae is None:
+ raise RuntimeError("MiniMax H3 video VAE is not resident on this rank")
+ with component_scope(video_vae):
+ for value in media.images:
+ image = _image_from_tensor(value)
+ visual_rows.append(video_vae.encode_image(image))
+ visual_shapes.append((1, image.height // 16, image.width // 16))
+ for value in media.videos:
+ frames = np.asarray(value.detach().cpu().to(torch.uint8).contiguous().numpy())
+ rows, shape = video_vae.encode_video(frames)
+ visual_rows.append(rows)
+ visual_shapes.append(tuple(int(item) for item in shape))
+
+ if not emit_conditioning:
+ return None
+
+ audio_rows: list[torch.Tensor] = []
+ audio_lengths: list[int] = []
+ embedded_audio_count = sum(item is not None for item in media.video_audios)
+ audio_inputs = _effective_audio_inputs(
+ media.video_audios,
+ media.audios,
+ max_standalone_seconds=float(media.num_frames) / MINIMAX_H3_FPS,
+ )
+ if audio_inputs:
+ if audio_vae is None:
+ raise RuntimeError("MiniMax H3 audio WVAE is not resident on the encoder leader")
+ with component_scope(audio_vae):
+ for waveform, sample_rate in audio_inputs:
+ rows, length = audio_vae.encode_waveform(waveform, sample_rate)
+ audio_rows.append(rows)
+ audio_lengths.append(int(length))
+ if audio_lengths:
+ if any(length < 80 or length > 600 for length in audio_lengths):
+ raise ValueError("MiniMax H3 audio references must each be between 2 and 15 seconds")
+ if sum(audio_lengths[:embedded_audio_count]) > 600 or sum(audio_lengths[embedded_audio_count:]) > 600:
+ raise ValueError("MiniMax H3 audio references must be at most 15 seconds in total")
+
+ ref_blocks: list[dict[str, Any]] = []
+ image_shapes = visual_shapes[: len(media.images)]
+ video_shapes = visual_shapes[len(media.images) :]
+ ref_blocks.extend({"kind": "image", "latent_h": shape[1], "latent_w": shape[2]} for shape in image_shapes)
+ embedded_lengths = iter(audio_lengths[:embedded_audio_count])
+ for shape, embedded in zip(video_shapes, media.video_audios, strict=True):
+ ref_audio_t = int(next(embedded_lengths)) if embedded is not None else 0
+ ref_blocks.append(
+ {
+ "kind": "video_audio" if ref_audio_t else "video",
+ "ref_audio_t": ref_audio_t,
+ "latent_t": shape[0],
+ "latent_h": shape[1],
+ "latent_w": shape[2],
+ }
+ )
+ ref_blocks.extend(
+ {
+ "kind": "audio",
+ "ref_audio_t": int(length),
+ }
+ for length in audio_lengths[embedded_audio_count:]
+ )
+ return MiniMaxH3EncoderMediaConditioning(
+ task=media.task,
+ height=media.height,
+ width=media.width,
+ num_frames=media.num_frames,
+ latent_t=media.latent_t,
+ audio_t=media.audio_t,
+ visual_condition=torch.cat(visual_rows) if visual_rows else None,
+ visual_condition_shapes=tuple(visual_shapes),
+ audio_condition=torch.cat(audio_rows) if audio_rows else None,
+ audio_condition_lengths=tuple(audio_lengths),
+ ref_blocks=tuple(ref_blocks),
+ keyframe_frame_indices=media.keyframe_frame_indices,
+ )
diff --git a/vllm_omni/model_executor/models/minimax_h3/pipeline.py b/vllm_omni/model_executor/models/minimax_h3/pipeline.py
index 5ac94b5ee..085e2508a 100644
--- a/vllm_omni/model_executor/models/minimax_h3/pipeline.py
+++ b/vllm_omni/model_executor/models/minimax_h3/pipeline.py
@@ -16,18 +16,18 @@ MINIMAX_H3_PIPELINE = PipelineConfig(
model_type="minimax_h3_disaggregated",
default_deploy_config_name="minimax_h3_disaggregated.yaml",
stage_cli_aliases={"text_encoder_tp_size": (0, "tensor_parallel_size")},
- model_arch="MiniMaxH3TextEncoder",
+ model_arch="MiniMaxH3Encoder",
stages=(
StagePipelineConfig(
stage_id=0,
- model_stage="text_encoder",
+ model_stage="encoder",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
owns_tokenizer=True,
requires_multimodal_data=True,
- model_arch="MiniMaxH3TextEncoder",
+ model_arch="MiniMaxH3Encoder",
engine_output_type="latent",
- prompt_transform_func=f"{_PROCESSOR}.prepare_text_encoder_prompt",
+ prompt_transform_func=f"{_PROCESSOR}.prepare_encoder_prompt",
sampling_constraints={
"max_tokens": 1,
"temperature": 0.0,
@@ -42,9 +42,9 @@ MINIMAX_H3_PIPELINE = PipelineConfig(
input_sources=(0,),
final_output=True,
final_output_type="video",
- requires_multimodal_data=True,
+ requires_multimodal_data=False,
model_arch="MiniMaxH3Pipeline",
- custom_process_input_func=f"{_PROCESSOR}.text_encoder2diffusion",
+ custom_process_input_func=f"{_PROCESSOR}.encoder2diffusion",
omni_kv_config={"need_recv_cache": False},
model_path_resolver=f"{_DIFFUSION_PIPELINE}.resolve_minimax_h3_diffusion_model_path",
inline_diffusion=True,
diff --git a/vllm_omni/model_executor/models/minimax_h3/reference_video.py b/vllm_omni/model_executor/models/minimax_h3/reference_video.py
index 4f98d9e13..4b361b26c 100644
--- a/vllm_omni/model_executor/models/minimax_h3/reference_video.py
+++ b/vllm_omni/model_executor/models/minimax_h3/reference_video.py
@@ -531,14 +531,27 @@ def _decode_video_frames_ffmpeg(
)
-def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]:
- meta = _probe_video(prepared_path)
+def sample_reference_video_frames(
+ prepared_path: str,
+ *,
+ decoded_frames: np.ndarray | None = None,
+) -> dict[str, Any]:
+ if decoded_frames is None:
+ meta = _probe_video(prepared_path)
+ frame_count = int(meta["frame_count"])
+ else:
+ decoded_frames = np.asarray(decoded_frames)
+ if decoded_frames.ndim != 4 or decoded_frames.shape[-1] != 3 or len(decoded_frames) <= 0:
+ raise OmniClientError(
+ f"decoded reference video frames must have shape [T, H, W, 3], got {decoded_frames.shape}"
+ )
+ frame_count = len(decoded_frames)
ratio = MINIMAX_H3_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS
indices: list[int] = []
cursor = 0.0
while True:
frame_index = int(round(cursor))
- if frame_index >= meta["frame_count"]:
+ if frame_index >= frame_count:
break
if not indices or frame_index > indices[-1]:
indices.append(frame_index)
@@ -546,16 +559,16 @@ def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]:
if not indices:
raise OmniClientError(f"no frames sampled from {prepared_path}")
- # The preparation step emits a lossless RGB stream. Decode only the sampled
- # frames in one ffmpeg process so Qwen3VL sees the exact prepared pixels
- # without decoding the entire high-bitrate stream into host memory.
- decoded_frames = _decode_video_frames_ffmpeg(
- prepared_path,
- frame_count=int(meta["frame_count"]),
- indices=indices,
- width=meta["width"],
- height=meta["height"],
- )
+ if decoded_frames is None:
+ decoded_frames = _decode_video_frames_ffmpeg(
+ prepared_path,
+ frame_count=frame_count,
+ indices=indices,
+ width=meta["width"],
+ height=meta["height"],
+ )
+ else:
+ decoded_frames = decoded_frames[indices]
frames = [np.asarray(frame) for frame in decoded_frames]
timestamps = [index / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))]
diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py
index b1933af95..9ae74036e 100644
--- a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py
+++ b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py
@@ -176,7 +176,7 @@ class _ResidualMerge(nn.Module):
info=Qwen3VLProcessingInfo,
dummy_inputs=Qwen3VLDummyInputsBuilder,
)
-class MiniMaxH3TextEncoder(Qwen3VLForConditionalGeneration):
+class MiniMaxH3TextEncoderBackbone(Qwen3VLForConditionalGeneration):
"""Qwen3-VL encoder used by MiniMax H3.
MiniMax H3 consumes the residual stream after decoder layer 50. The
@@ -305,3 +305,6 @@ class MiniMaxH3TextEncoder(Qwen3VLForConditionalGeneration):
encoder_weights(),
mapper=self.hf_to_vllm_mapper,
)
+
+
+MiniMaxH3TextEncoder = MiniMaxH3TextEncoderBackbone
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
index 9d2d2d171..8d36086ef 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
+++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
@@ -44,10 +44,10 @@ def _maybe_prefix(prefix: str, name: str) -> str:
def _iter_state_row_spans(
- states: Sequence[dict[str, Any]],
+ states: Sequence[dict[str, Any] | None],
spans: Sequence[tuple[int, int]] | None,
num_rows: int,
-) -> Iterable[tuple[dict[str, Any], int, int]]:
+) -> Iterable[tuple[dict[str, Any] | None, int, int]]:
"""Yield state/logit-row ranges without assuming equal request row counts."""
if num_rows <= 0 or not states:
return
@@ -798,6 +798,9 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
have_multimodal_outputs: bool = True
has_preprocess: bool = True
has_postprocess: bool = True
+ # Intermediate prefill chunks reach make_omni_output but must not produce
+ # audio until the scheduler has completed the prompt.
+ requires_request_sample_eligibility: bool = True
AUDIO_BOS = 1025
AUDIO_EOS = 1026
@@ -858,7 +861,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
# We still need a logits_processor so vLLM downstream stays happy with
# something callable; a minimal pass-through suffices.
self.logits_processor = LogitsProcessor(self.text_vocab_size)
- self._batch_state: list[dict[str, Any]] | None = None
+ self._batch_state: list[dict[str, Any] | None] | None = None
# Stacked weight cache built after load_weights().
# shape: (n_vq, audio_vocab_size, hidden_size)
@@ -972,13 +975,14 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
span_len = int(input_ids.shape[0])
audio_state = info_dict.get("audio_state")
is_first_call = not isinstance(audio_state, dict)
+ is_prefill = bool(info_dict.get("_omni_is_prefill", span_len > 1))
- if span_len > 1 or is_first_call:
+ if is_prefill or is_first_call:
# Prefill: read the full reference-audio code grid (channels x T)
# from ``info_dict["codes"]["ref"]``; slice the chunk that aligns
# with this prefill window.
ref_codes = (info_dict.get("codes", {}) or {}).get("ref")
- ref_offset = int(info_dict.get("ref_offset", 0))
+ ref_offset = int(info_dict.get("_omni_num_computed_tokens", info_dict.get("ref_offset", 0)))
chunk_audio = None
if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0:
if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0:
@@ -994,22 +998,24 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
if chunk_audio is not None
else torch.full((self.n_vq,), self.audio_pad_token, dtype=torch.long, device=device)
)
- # Capture the streaming-text list from the request (set by the
- # realtime end2end builder; empty for the delay variants which
- # don't take this code path anyway).
- remaining = (info_dict.get("ids", {}) or {}).get("all")
- if not isinstance(remaining, list):
- remaining = []
info_update: dict[str, Any] = {
- "audio_state": {
+ "audio_codes": {"current": current_codes},
+ "ref_offset": ref_offset + span_len,
+ }
+ if is_first_call:
+ remaining = (info_dict.get("ids", {}) or {}).get("all")
+ if not isinstance(remaining, list):
+ remaining = []
+ max_new_frames = info_dict.get("max_new_frames")
+ if isinstance(max_new_frames, (list, tuple)):
+ max_new_frames = max_new_frames[0] if max_new_frames else None
+ info_update["audio_state"] = {
"is_stopping": False,
"step": 0,
"text_cursor": 0,
"remaining_text": list(remaining),
- },
- "audio_codes": {"current": current_codes},
- "ref_offset": ref_offset + span_len,
- }
+ "max_new_frames": int(max_new_frames) if max_new_frames is not None else -1,
+ }
return input_ids, embeds, info_update
# Decode step: text_token from the just-sampled vLLM logit, plus the
@@ -1046,13 +1052,21 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
info_dicts: list[dict[str, Any]] = (
kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
)
+ sample_eligible = kwargs.get("request_sample_eligible")
+ eligible = [
+ sample_eligible is None or i >= len(sample_eligible) or bool(sample_eligible[i])
+ for i in range(len(info_dicts))
+ ]
# Defensive state init.
for info in info_dicts:
if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict):
info["audio_state"] = {"is_stopping": False, "step": 0}
- self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
+ self._batch_state = [
+ (info["audio_state"] if isinstance(info, dict) and eligible[i] else None)
+ for i, info in enumerate(info_dicts)
+ ]
self._batch_state_spans = kwargs.get("request_token_spans")
# See delay talker: real per-request row spans from the runner are
@@ -1074,6 +1088,8 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
for i, info in enumerate(info_dicts):
if not isinstance(info, dict):
continue
+ if not eligible[i]:
+ continue
row_start, row_end = spans[i]
row_end = min(int(row_end), num_rows)
if row_start >= row_end:
@@ -1115,10 +1131,13 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
)
ch0 = int(new_codes[0].item())
+ state["step"] = int(state.get("step", 0)) + 1
+ max_new_frames = int(state.get("max_new_frames", -1))
+ if max_new_frames > 0 and state["step"] >= max_new_frames:
+ state["is_stopping"] = True
# Stop condition mirrors upstream: codebook 0 == eos_audio_id.
if ch0 == self.AUDIO_EOS:
state["is_stopping"] = True
- state["step"] = int(state.get("step", 0)) + 1
info["audio_codes"] = {
"current": new_codes,
"accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"),
@@ -1127,7 +1146,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
if ch0 in (self.AUDIO_BOS, self.audio_pad_token):
# Skip the bos / pad frames — they don't decode to real audio.
- state["step"] = int(state.get("step", 0)) + 1
info["audio_codes"] = {
"current": new_codes,
"accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"),
@@ -1141,7 +1159,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
updated_acc = new_codes.unsqueeze(0)
info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
- state["step"] = int(state.get("step", 0)) + 1
per_req_codes[i] = new_codes.unsqueeze(0)
have_codes = True
diff --git a/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py b/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py
new file mode 100644
index 000000000..985f76205
--- /dev/null
+++ b/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Prompt construction for MOSS-TTS-Realtime."""
+
+from typing import Any
+
+import numpy as np
+import torch
+
+REALTIME_AUDIO_CHANNELS = 16
+REALTIME_PREFILL_TEXT_TOKENS = 12
+
+
+def build_realtime_prompt(
+ tokenizer: Any,
+ processor: Any,
+ text: str,
+ reference_codes: torch.Tensor,
+) -> dict[str, Any]:
+ """Build the mixed text/audio prompt consumed by the Realtime talker."""
+ audio_tokens = reference_codes.detach().cpu().numpy()
+ system_grid = np.asarray(processor.make_ensemble(prompt_audio_tokens=audio_tokens), dtype=np.int64)
+ expected_width = REALTIME_AUDIO_CHANNELS + 1
+ if system_grid.ndim != 2 or system_grid.shape[1] != expected_width:
+ raise ValueError(
+ f"MOSS-TTS-Realtime processor returned shape {system_grid.shape}; "
+ f"expected (sequence_length, {expected_width})"
+ )
+
+ assistant_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
+ assistant_grid = np.full(
+ (len(assistant_ids), expected_width),
+ processor.audio_channel_pad,
+ dtype=np.int64,
+ )
+ assistant_grid[:, 0] = assistant_ids
+
+ text_ids = tokenizer.encode(text, add_special_tokens=False)
+ if not text_ids:
+ raise ValueError("MOSS-TTS-Realtime input is empty after tokenization")
+
+ prefill_len = min(len(text_ids), REALTIME_PREFILL_TEXT_TOKENS)
+ text_grid = np.full(
+ (prefill_len, expected_width),
+ processor.audio_channel_pad,
+ dtype=np.int64,
+ )
+ text_grid[:, 0] = text_ids[:prefill_len]
+ text_grid[-1, 1] = processor.audio_bos_token
+
+ prompt_grid = np.concatenate((system_grid, assistant_grid, text_grid), axis=0)
+ params: dict[str, Any] = {
+ "prompt_token_ids": prompt_grid[:, 0].tolist(),
+ "codes": {"ref": torch.from_numpy(prompt_grid[:, 1:].copy())},
+ }
+ remaining_text_ids = text_ids[prefill_len:]
+ if remaining_text_ids:
+ params["ids"] = {"all": remaining_text_ids}
+ return params
diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py
index 36b2616e7..a828d59e0 100644
--- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py
+++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Reference-audio encoding + speaker cache for the MOSS-TTS-family talker.
This lives in the model package (not the shared serving layer) so all
@@ -360,6 +363,18 @@ class MossReferenceEncoder:
await self._batcher.aclose()
+class MossRealtimeReferenceEncoder(MossReferenceEncoder):
+ """Use the realtime codec with the shared reference cache and batcher."""
+
+ def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]:
+ with torch.no_grad():
+ encoded = self._processor.batch_encode([wav.squeeze(0) for wav in prepared], num_quantizers=self._n_vq)
+ return [
+ encoded.audio_codes[:, i, : int(length.item())].transpose(0, 1).contiguous()
+ for i, length in enumerate(encoded.audio_codes_lengths)
+ ]
+
+
def build_reference_encoder(
processor: Any,
*,
@@ -369,9 +384,17 @@ def build_reference_encoder(
"""Build the per-server encoder for a MOSS-TTS ``variant``.
Derives the encode geometry (``n_vq`` and the working sample rate) from the
- upstream processor's ``model_config`` so the variant knowledge stays in the
- model package rather than in ``serving_speech.py``.
+ upstream processor's ``model_config``; Realtime takes its codec directly
+ and uses 16 codebooks at 24 kHz. Variant knowledge stays in the model package.
"""
+ if variant == "realtime":
+ return MossRealtimeReferenceEncoder(
+ processor,
+ variant=variant,
+ n_vq=16,
+ sr_target=24000,
+ speaker_cache=speaker_cache,
+ )
n_vq = int(getattr(processor.model_config, "n_vq", 32))
# Local-v1.5 encodes reference audio at a fixed 24 kHz working rate
# regardless of its 48 kHz stereo *output* codec -- mirrors the offline
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
index 263ffa2ab..b7b7d074c 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
@@ -758,6 +758,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
tokenizer = self.info.get_tokenizer()
processor = self.info.get_hf_processor()
audio_token_id = tokenizer.get_vocab()[processor.audio_token]
+ video_token_id = tokenizer.get_vocab()[processor.video_token]
result_placeholders = dict(placeholders)
audio_placeholders = []
@@ -765,7 +766,10 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
audio_idx = 0
for video_idx, video_placeholder in enumerate(placeholders["video"]):
- audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id
+ placeholder_tokens = torch.tensor(video_placeholder.tokens)
+ audio_is_embed = placeholder_tokens == audio_token_id
+ # Audio boundary tokens keep their text embeddings.
+ video_is_embed = placeholder_tokens == video_token_id
if video_use_audio_in_video[video_idx]:
audio_placeholder = PlaceholderFeaturesInfo(
@@ -783,7 +787,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
item_idx=video_idx,
start_idx=video_placeholder.start_idx,
tokens=video_placeholder.tokens,
- is_embed=~audio_is_embed,
+ is_embed=video_is_embed,
)
video_placeholders.append(video_placeholder_with_mask)
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
index 8d0389d63..c78679c55 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
@@ -588,7 +588,10 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
if layer_idx in capture_set:
hs = captured_hidden_states.setdefault("hidden_states", {})
layers = hs.setdefault("layers", {})
- layers[layer_idx] = hidden_states.clone().view(-1, hidden_states.shape[-1])
+ # vLLM defers the residual addition until the next RMSNorm.
+ # Reconstruct the logical decoder state before capturing it.
+ captured = hidden_states.clone() if residual is None else hidden_states + residual
+ layers[layer_idx] = captured.view(-1, captured.shape[-1])
hidden_states, residual = layer(
positions,
diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py
index e5d7b4115..1a519917c 100644
--- a/vllm_omni/model_executor/models/registry.py
+++ b/vllm_omni/model_executor/models/registry.py
@@ -476,10 +476,10 @@ _OMNI_MODELS = {
"qwen3_vl",
"AuraQwen3VLForConditionalGeneration",
),
- "MiniMaxH3TextEncoder": (
+ "MiniMaxH3Encoder": (
"minimax_h3",
- "text_encoder",
- "MiniMaxH3TextEncoder",
+ "encoder",
+ "MiniMaxH3Encoder",
),
## MiniMax-Music3 (text-to-music; AR talker -> flow-matching acoustic decoder)
"MiniMaxMusic3TalkerForConditionalGeneration": (
diff --git a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py
index 87808adc3..78fd0225f 100644
--- a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py
+++ b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py
@@ -1,73 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""MiniMax H3 text-encoder stage input and output adapters."""
+"""Adapters between MiniMax H3 conditioning and the Omni stage runner."""
from __future__ import annotations
import copy
-import shutil
-import tempfile
from collections.abc import Mapping, Sequence
from typing import Any
-import numpy as np
-from PIL import Image
-
-from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY
-from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_align_frame_count
-from vllm_omni.errors import OmniClientError
+from vllm_omni.data_entry_keys import unflatten_payload
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
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,
- MiniMaxH3TextConditioning,
+ MiniMaxH3EncoderConditioning,
)
-from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
- MINIMAX_H3_OUTPUT_SHORT_EDGE,
- load_minimax_h3_images,
- resolve_minimax_h3_aspect_ratio,
- resolve_minimax_h3_output_canvas,
- resolve_minimax_h3_reference_image_shape,
-)
-from vllm_omni.model_executor.models.minimax_h3.reference_video import (
- MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
- MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
- prepare_reference_videos,
- sample_reference_video_frames,
- serialize_prepared_reference_videos,
-)
-
-
-def _items(value: Any) -> list[Any]:
- if value is None:
- return []
- if isinstance(value, list):
- return value
- if isinstance(value, tuple) and not (len(value) == 2 and isinstance(value[1], Mapping)):
- return list(value)
- return [value]
-
-
-def _audio_items(value: Any) -> list[Any]:
- if value is None:
- return []
- if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)):
- return [value]
- return list(value) if isinstance(value, (list, tuple)) else [value]
-
-
-def _resolve_task(
- extra_args: Mapping[str, Any],
- multi_modal_data: Mapping[str, Any],
-) -> str:
- requested = extra_args.get("task")
- if requested is not None:
- return str(requested).lower()
- if multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None:
- return "ref2va"
- if multi_modal_data.get("image") is not None:
- return "fl2va"
- return "t2va"
+from vllm_omni.model_executor.models.minimax_h3.encoder_processing import prepare_encoder_inputs
def _diffusion_sampling_params(sampling_params_list: Sequence[Any]) -> Any:
@@ -78,190 +27,29 @@ def _diffusion_sampling_params(sampling_params_list: Sequence[Any]) -> Any:
]
if len(diffusion_params) != 1:
raise RuntimeError(
- "MiniMax H3 text encoding requires exactly one OmniDiffusionSamplingParams stage parameter, "
+ "MiniMax H3 encoding requires exactly one OmniDiffusionSamplingParams stage parameter, "
f"got {len(diffusion_params)}"
)
return diffusion_params[0]
-def _ref2va_target_frame_count(sampling_params_list: Sequence[Any]) -> int:
- sampling = _diffusion_sampling_params(sampling_params_list)
- extra_args = sampling.extra_args or {}
- target = extra_args.get("target")
- target = target if isinstance(target, Mapping) else {}
- duration = target.get("duration_seconds", extra_args.get("duration_seconds", extra_args.get("duration")))
- if duration is not None:
- requested = int(round(float(duration) * 24))
- elif int(getattr(sampling, "num_frames", None) or 1) > 1:
- requested = int(sampling.num_frames)
- else:
- requested = 124
- return minimax_h3_align_frame_count(requested)
-
-
-def _prepare_qwen_images(
- task: str,
- values: list[Any],
- sampling_params_list: Sequence[Any],
-) -> list[Any]:
- if not values:
- return []
- images = load_minimax_h3_images(values)
- if task == "ref2va":
- return [
- image.resize(
- resolve_minimax_h3_reference_image_shape(image),
- Image.Resampling.LANCZOS,
- )
- for image in images
- ]
- if task != "fl2va":
- return images
-
- sampling = _diffusion_sampling_params(sampling_params_list)
- extra_args = sampling.extra_args or {}
- target = extra_args.get("target")
- if target is not None and not isinstance(target, Mapping):
- raise OmniClientError("MiniMax H3 extra_args['target'] must be an object")
- target = target if isinstance(target, Mapping) else {}
- aspect_ratio = resolve_minimax_h3_aspect_ratio(
- task,
- target.get("aspect_ratio", extra_args.get("aspect_ratio")),
- images[0],
- )
- if not 0.25 <= aspect_ratio <= 4.0:
- raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be in [1:4, 4:1], got {aspect_ratio}")
- height = sampling.height
- width = sampling.width
- if height is None or width is None:
- short_edge = target.get(
- "short_edge",
- extra_args.get("short_edge", MINIMAX_H3_OUTPUT_SHORT_EDGE),
- )
- if isinstance(short_edge, bool) or not isinstance(short_edge, (int, np.integer)):
- raise OmniClientError(
- f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {short_edge!r}"
- )
- height, width = resolve_minimax_h3_output_canvas(aspect_ratio, int(short_edge))
- height = int(height) // 32 * 32
- width = int(width) // 32 * 32
- if min(height, width) <= 0:
- raise OmniClientError(f"invalid MiniMax H3 canvas {width}x{height}")
- if width > 4 * height or height > 4 * width:
- raise OmniClientError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]")
- return [image.resize((width, height), Image.Resampling.LANCZOS) for image in images]
-
-
-def prepare_text_encoder_prompt(
+def prepare_encoder_prompt(
prompt: Any,
sampling_params_list: Sequence[Any],
) -> Any:
- """Build H3's labeled Qwen3-VL presentation for Stage 0.
-
- The upstream Qwen3-VL multimodal processor expands each image/video
- placeholder into the exact number of vision tokens and adds timestamped
- video blocks. Audio is represented only by its H3 text label and is not
- sent to Qwen3-VL.
- """
+ prepared = prepare_encoder_inputs(prompt, _diffusion_sampling_params(sampling_params_list))
if isinstance(prompt, str):
- return prompt
- if not isinstance(prompt, dict):
- raise TypeError(f"MiniMax H3 expects a string or dict prompt, got {type(prompt)!r}")
-
- prompt = copy.copy(prompt)
- additional_information = dict(prompt.get("additional_information") or {})
- meta = dict(additional_information.get("meta") or {})
- meta.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, None)
- additional_information["meta"] = meta
- prompt["additional_information"] = additional_information
-
- text = str(prompt.get("prompt") or "")
- if not text:
- raise OmniClientError("MiniMax H3 requires a non-empty prompt")
- multi_modal_data = prompt.get("multi_modal_data") or {}
- if not isinstance(multi_modal_data, Mapping):
- raise TypeError("multi_modal_data must be a mapping")
-
- image_values = _items(multi_modal_data.get("image"))
- videos = _items(multi_modal_data.get("video"))
- audios = _audio_items(multi_modal_data.get("audio"))
- diffusion_sampling = _diffusion_sampling_params(sampling_params_list)
- extra_args = diffusion_sampling.extra_args or {}
- task = _resolve_task(extra_args, multi_modal_data)
- images = _prepare_qwen_images(task, image_values, sampling_params_list)
- qwen_video_inputs: list[tuple[np.ndarray, dict[str, Any]]] = []
- condition_labels: list[tuple[str, int]] = []
-
- if task == "t2va":
- if images or videos or audios:
- raise OmniClientError("t2va does not accept image, video, or audio conditions")
- elif task == "fl2va":
- if not images or videos or audios:
- raise OmniClientError("fl2va requires image conditions only")
- condition_labels.extend(("image", index) for index in range(1, len(images) + 1))
- elif task == "ref2va":
- if not images and not videos:
- raise OmniClientError("ref2va requires an image or video condition")
- condition_labels.extend(("image", index) for index in range(1, len(images) + 1))
- prepared_videos: list[dict[str, Any]] = []
- artifact_dir: str | None = None
- if videos:
- artifact_dir = tempfile.mkdtemp(prefix="minimax_h3_ref2va_")
- try:
- prepared_videos = prepare_reference_videos(
- videos,
- target_frame_count=_ref2va_target_frame_count(sampling_params_list),
- workdir=artifact_dir,
- start_time_seconds=extra_args.get("start_time_seconds"),
- )
- for item in prepared_videos:
- sampled = sample_reference_video_frames(item["prepared_path"])
- frames = np.stack(sampled["frames"])
- frame_count = int(frames.shape[0])
- qwen_video_inputs.append(
- (
- frames,
- {
- "total_num_frames": frame_count,
- "fps": MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
- "duration": frame_count / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
- "video_backend": "minimax_h3",
- "frames_indices": list(range(frame_count)),
- "do_sample_frames": False,
- },
- )
- )
- except BaseException:
- shutil.rmtree(artifact_dir, ignore_errors=True)
- raise
- audio_index = 0
- for video_index, item in enumerate(prepared_videos, start=1):
- if item["input_has_audio"]:
- audio_index += 1
- condition_labels.append(("audio", audio_index))
- condition_labels.append(("video", video_index))
- for _ in audios:
- audio_index += 1
- condition_labels.append(("audio", audio_index))
- else:
- raise OmniClientError(f"unsupported MiniMax H3 task {task!r}")
-
+ prompt = {"prompt": prompt}
+ text = prepared.prompt
+ images = prepared.images
+ qwen_video_inputs = prepared.qwen_videos
+ condition_labels = prepared.condition_labels
+ media_input = prepared.media
+ task = media_input.task
transformed = copy.copy(prompt)
- if isinstance(prompt.get("additional_information"), Mapping):
- transformed["additional_information"] = dict(prompt["additional_information"])
+ additional_information = dict(prompt.get("additional_information") or {})
transformed["prompt"] = text
- if task == "ref2va" and prepared_videos and artifact_dir is not None:
- additional_information = dict(transformed.get("additional_information") or {})
- meta = dict(additional_information.get("meta") or {})
- meta[MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY] = serialize_prepared_reference_videos(
- prepared_videos,
- artifact_dir,
- )
- additional_information["meta"] = meta
- transformed["additional_information"] = additional_information
- transformed[REQUEST_ARTIFACT_DIRS_KEY] = [artifact_dir]
- qwen_mm_data = dict(multi_modal_data)
- qwen_mm_data.pop("audio", None)
+ qwen_mm_data: dict[str, Any] = {}
if images:
qwen_mm_data["image"] = images
if qwen_video_inputs:
@@ -271,7 +59,16 @@ def prepare_text_encoder_prompt(
mm_processor_kwargs = dict(prompt.get("mm_processor_kwargs") or {})
mm_processor_kwargs[MINIMAX_H3_PRESENTATION_TASK_KEY] = task
mm_processor_kwargs[MINIMAX_H3_CONDITION_LABELS_KEY] = condition_labels
+ media_tensors = media_input.to_mm_tensors()
transformed["mm_processor_kwargs"] = mm_processor_kwargs
+
+ hidden_states = dict(additional_information.get("hidden_states") or {})
+ hidden_states["layers"] = dict(enumerate(media_tensors))
+ additional_information["hidden_states"] = hidden_states
+ meta = dict(additional_information.get("meta") or {})
+ meta[MINIMAX_H3_ENCODER_REQUEST_KEY] = media_input.to_metadata()
+ additional_information["meta"] = meta
+ transformed["additional_information"] = additional_information
return transformed
@@ -295,18 +92,20 @@ def _global_request_id(prompt: Mapping[str, Any]) -> str | None:
return str(value) if value is not None else None
-def text_encoder2diffusion(
+def encoder2diffusion(
source_outputs: list[Any],
prompt: Any = None,
requires_multimodal_data: bool = False,
streaming_context: Any | None = None,
) -> dict[str, Any] | None:
- """Attach Stage 0 hidden states and token tags to the original request."""
+ """Reuse the encoder handoff for all three H3 encoder components."""
del requires_multimodal_data, streaming_context
if not source_outputs:
return None
if len(source_outputs) != 1:
- raise RuntimeError(f"MiniMax H3 diffusion requires exactly one text-encoder source, got {len(source_outputs)}")
+ raise RuntimeError(f"MiniMax H3 DiT requires exactly one encoder source, got {len(source_outputs)}")
+ if not getattr(source_outputs[0], "finished", True):
+ return None
diffusion_prompt = _original_prompt(prompt)
source_output = source_outputs[0]
@@ -318,25 +117,32 @@ def text_encoder2diffusion(
and str(source_request_id) != expected_request_id
):
raise RuntimeError(
- "MiniMax H3 text-encoder request ID does not match the diffusion request: "
+ "MiniMax H3 encoder request ID does not match the diffusion request: "
f"source={source_request_id!r}, expected={expected_request_id!r}"
)
outputs = getattr(source_output, "outputs", None)
if not isinstance(outputs, list) or len(outputs) != 1:
output_count = len(outputs) if isinstance(outputs, list) else 0
- raise RuntimeError(f"MiniMax H3 text encoder must return exactly one completion, got {output_count}")
-
- completion = outputs[0]
- payload = completion.multimodal_output
+ raise RuntimeError(f"MiniMax H3 encoder must return exactly one completion, got {output_count}")
+ payload = getattr(outputs[0], "multimodal_output", None)
if not isinstance(payload, Mapping):
- raise RuntimeError("MiniMax H3 text encoder returned no conditioning payload")
+ raise RuntimeError("MiniMax H3 encoder returned no conditioning payload")
try:
- conditioning = MiniMaxH3TextConditioning.from_omni_payload(payload)
+ conditioning = MiniMaxH3EncoderConditioning.from_omni_payload(unflatten_payload(dict(payload)))
except ValueError as exc:
raise RuntimeError(str(exc)) from exc
additional_information = dict(diffusion_prompt.get("additional_information") or {})
- additional_information["text_encoder_output"] = conditioning.to_payload()
+ additional_information.pop("hidden_states", None)
+ meta = dict(additional_information.get("meta") or {})
+ meta.pop(MINIMAX_H3_ENCODER_REQUEST_KEY, None)
+ if meta:
+ additional_information["meta"] = meta
+ else:
+ additional_information.pop("meta", None)
+ additional_information["encoder_output"] = conditioning.to_omni_payload()
diffusion_prompt["additional_information"] = additional_information
+ diffusion_prompt["multi_modal_data"] = None
+ diffusion_prompt.pop("model_intermediate_buffer", None)
return diffusion_prompt
diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py
index 8f15aaf50..c0975f651 100644
--- a/vllm_omni/quantization/factory.py
+++ b/vllm_omni/quantization/factory.py
@@ -446,6 +446,33 @@ def _disk_marks_serialized(qc_kwargs: dict[str, Any], quant_config: object) -> b
return False
+_MXFP4_SERIALIZED_FLAGS = {
+ "mxfp4": "is_checkpoint_mxfp4_serialized",
+ "mxfp4_dualscale": "is_checkpoint_serialized",
+}
+
+
+def _normalize_serialized_mxfp4_layer_policy(qc_method: str, qc_kwargs: dict[str, Any]) -> None:
+ """Canonicalize the BF16 layer policy owned by serialized MXFP4 metadata.
+
+ A serialized checkpoint config is a complete storage contract: an omitted
+ layer list means no explicit BF16 layers. Method-only online metadata does
+ not carry that contract and must leave caller-provided layer policy intact.
+ """
+ method = _normalize_quant_method_alias(qc_method)
+ if method is None:
+ return
+ serialized_flag = _MXFP4_SERIALIZED_FLAGS.get(method)
+ if serialized_flag is None or not qc_kwargs.get(serialized_flag):
+ return
+
+ ignored_layers = qc_kwargs.get("ignored_layers")
+ if not ignored_layers:
+ ignored_layers = qc_kwargs.get("modules_to_not_convert")
+ qc_kwargs["ignored_layers"] = ignored_layers or []
+ qc_kwargs.pop("modules_to_not_convert", None)
+
+
def resolve_quant_config_from_disk(
quant_config: QuantizationConfig | None,
disk_qc: dict[str, Any] | str | None,
@@ -460,22 +487,24 @@ def resolve_quant_config_from_disk(
- quant_config is None: auto-detect from disk_qc (full build).
- Methods mismatch: raise ValueError — prevents silent weight corruption.
- Disk marks serialized but quant_config is online: rebuild from disk.
- - ignored_layers differ: rebuild from disk (per-transformer BF16 routing).
+ - Serialized MXFP4 metadata owns the complete BF16 layer policy; omitted
+ ignored_layers means the checkpoint default (no explicit BF16 layers).
+ - Method-only metadata carries no storage policy and preserves caller
+ ignored_layers and runtime fallback policies.
"""
if disk_qc is None:
return quant_config
if isinstance(disk_qc, str):
- if quant_config is None:
- logger.info("Auto-detected quantization from config.json: method=%s", disk_qc)
- return build_quant_config(disk_qc)
- return quant_config
+ qc_method = disk_qc
+ qc_kwargs: dict[str, Any] = {}
+ else:
+ if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc:
+ return quant_config
+ qc_method = disk_qc["quant_method"]
+ qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
- if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc:
- return quant_config
-
- qc_method: str = disk_qc["quant_method"]
- qc_kwargs: dict[str, Any] = {k: v for k, v in disk_qc.items() if k != "quant_method"}
+ _normalize_serialized_mxfp4_layer_policy(qc_method, qc_kwargs)
if quant_config is None:
logger.info(
@@ -494,6 +523,26 @@ def resolve_quant_config_from_disk(
"Pass a matching --quantization flag or omit it for auto-detection."
)
+ if isinstance(disk_qc, str):
+ return quant_config
+
+ # Serialized checkpoint metadata owns storage and BF16 layer routing. The
+ # active config owns runtime step/layer policies, including explicit empty
+ # lists. Preserve them when each cascade transformer rebuilds from disk.
+ for policy in ("w4a8_fallback_steps", "w4a8_fallback_layers"):
+ if hasattr(quant_config, policy):
+ qc_kwargs[policy] = list(getattr(quant_config, policy))
+ if hasattr(quant_config, "mxfp4_scale_alg"):
+ # Runtime activation policy is independent of how offline W4 was made.
+ qc_kwargs["mxfp4_scale_alg"] = quant_config.mxfp4_scale_alg
+ if hasattr(quant_config, "require_smooth_scale"):
+ disk_requires_smooth = qc_kwargs.get("require_smooth_scale", False)
+ if not isinstance(disk_requires_smooth, bool):
+ raise ValueError("require_smooth_scale must be a boolean")
+ # Either a calibrated checkpoint or an explicit caller can require
+ # Smooth. Per-expert rebuilds must never weaken that requirement.
+ qc_kwargs["require_smooth_scale"] = quant_config.require_smooth_scale or disk_requires_smooth
+
if _disk_marks_serialized(qc_kwargs, quant_config):
logger.info(
"config.json marks checkpoint as serialized; switching to offline %s mode.",
@@ -501,6 +550,14 @@ def resolve_quant_config_from_disk(
)
return build_quant_config(qc_method, **qc_kwargs)
+ if (
+ "require_smooth_scale" in qc_kwargs
+ and hasattr(quant_config, "require_smooth_scale")
+ and qc_kwargs["require_smooth_scale"] != quant_config.require_smooth_scale
+ ):
+ logger.info("config.json Smooth requirement differs from active config; rebuilding quant_config.")
+ return build_quant_config(qc_method, **qc_kwargs)
+
# AutoRound MXFP checkpoints use data_type="mx_fp" instead of
# is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is
# selected according to the checkpoint's bit width.
diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py
index ac93eaa00..43f3efa9d 100644
--- a/vllm_omni/quantization/mxfp4_config.py
+++ b/vllm_omni/quantization/mxfp4_config.py
@@ -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
"""W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers.
Architecture mirrors mxfp8_config.py:
@@ -41,7 +41,7 @@ Dual-scale (W4A4_MXFP4_DUALSCALE):
and npu_dual_level_quant_matmul. Weight stored in NZ hardware format via npu_format_cast.
Checkpoint tensor shapes for dual-scale:
- weight : (N, K) float8_e4m3fn – FP4 packed
+ weight : (N, K) float8_e4m3fn – numeric FP4 values (packed after loading)
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits)
weight_dual_scale : (N, K//512, 1) float32 – coarse scale (extra dim avoids shape assert)
mul_scale : (K,) float32 – per-input-channel activation pre-scale
@@ -51,9 +51,9 @@ Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/q
from __future__ import annotations
-import re
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
+import regex as re
import torch
from torch.nn import Module
from vllm.logger import init_logger
@@ -83,9 +83,54 @@ from vllm_omni.quantization.mxfp8_config import (
if TYPE_CHECKING:
from vllm.model_executor.models.utils import WeightsMapper
+
logger = init_logger(__name__)
+def _validate_w4a8_fallback_steps(steps: list[int] | None) -> list[int]:
+ if steps is None:
+ return []
+ if not isinstance(steps, list) or any(type(step) is not int or step < 0 for step in steps):
+ raise ValueError("w4a8_fallback_steps must be a list of non-negative integer denoise step indices.")
+ return sorted(set(steps))
+
+
+def _validate_w4a8_fallback_layers(layers: list[str] | None) -> list[str]:
+ """Names are exact runtime Linear paths, not checkpoint aliases or patterns."""
+ if layers is None:
+ return []
+ if not isinstance(layers, list) or any(
+ not isinstance(name, str) or re.fullmatch(r"[A-Za-z_][A-Za-z_0-9]*(?:\.[A-Za-z_0-9]+)*", name) is None
+ for name in layers
+ ):
+ raise ValueError("w4a8_fallback_layers must be a list of non-empty exact runtime Linear paths.")
+ return sorted(set(layers))
+
+
+def _is_w4a8_fallback_step(steps: list[int]) -> bool:
+ """Read the request's zero-based denoise index; never keep state on a Linear."""
+ if not steps:
+ return False
+
+ from vllm_omni.diffusion.forward_context import (
+ get_forward_context,
+ is_forward_context_available,
+ )
+
+ step = get_forward_context().denoise_step_idx if is_forward_context_available() else None
+ if step is None:
+ raise RuntimeError("W4A8 step fallback requires a pipeline that publishes its denoise step in ForwardContext.")
+ return step in steps
+
+
+def _validate_smooth_scale(scale: torch.Tensor) -> None:
+ """Reject malformed calibration before loading or activation dtype conversion."""
+ if scale.ndim != 1 or not scale.is_floating_point():
+ raise ValueError("MXFP4 mul_scale must be a one-dimensional floating-point Smooth tensor.")
+ if not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()):
+ raise ValueError("MXFP4 mul_scale must contain finite, strictly positive Smooth values.")
+
+
# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
@@ -105,10 +150,24 @@ class DiffusionMXFP4Config(QuantizationConfig):
self,
is_checkpoint_mxfp4_serialized: bool = False,
ignored_layers: list[str] | None = None,
+ w4a8_fallback_steps: list[int] | None = None,
+ require_smooth_scale: bool = False,
+ w4a8_fallback_layers: list[str] | None = None,
+ mxfp4_scale_alg: int = 0,
) -> None:
super().__init__()
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
self.ignored_layers = ignored_layers or []
+ self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps)
+ self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers)
+ if not isinstance(require_smooth_scale, bool):
+ raise ValueError("require_smooth_scale must be a boolean")
+ if require_smooth_scale and not self.is_checkpoint_mxfp4_serialized:
+ raise ValueError("require_smooth_scale requires an offline MXFP4 checkpoint")
+ self.require_smooth_scale = require_smooth_scale
+ if type(mxfp4_scale_alg) is not int or mxfp4_scale_alg not in (0, 2):
+ raise ValueError("mxfp4_scale_alg must be 0 (OCP MX) or 2 (UOS with dst_type_max=7.25).")
+ self.mxfp4_scale_alg = mxfp4_scale_alg
@classmethod
def get_name(cls) -> QuantizationMethods:
@@ -139,6 +198,10 @@ class DiffusionMXFP4Config(QuantizationConfig):
return cls(
is_checkpoint_mxfp4_serialized=is_serialized,
ignored_layers=ignored_layers,
+ w4a8_fallback_steps=config.get("w4a8_fallback_steps"),
+ w4a8_fallback_layers=config.get("w4a8_fallback_layers"),
+ require_smooth_scale=config.get("require_smooth_scale", False),
+ mxfp4_scale_alg=config.get("mxfp4_scale_alg", 0),
)
def get_quant_method(
@@ -155,9 +218,15 @@ class DiffusionMXFP4Config(QuantizationConfig):
return UnquantizedLinearMethod()
if current_omni_platform.is_npu():
if self.is_checkpoint_mxfp4_serialized:
- return NPUMxfp4LinearMethod(self)
- return NPUMxfp4OnlineLinearMethod(self)
+ return NPUMxfp4LinearMethod(self, prefix=prefix)
+ return NPUMxfp4OnlineLinearMethod(self, prefix=prefix)
if current_omni_platform.is_rocm():
+ if self.mxfp4_scale_alg != 0:
+ raise NotImplementedError("MXFP4 UOS quantization is currently only supported on NPU (Ascend).")
+ if self.w4a8_fallback_steps or self.w4a8_fallback_layers:
+ raise NotImplementedError(
+ "MXFP4 W4A8 step/layer fallback is currently only supported on NPU (Ascend)."
+ )
gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName
if "gfx950" not in gcn_arch:
raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}")
@@ -176,6 +245,56 @@ class DiffusionMXFP4Config(QuantizationConfig):
# ---------------------------------------------------------------------------
+def _npu_quantize_mxfp4(x: torch.Tensor, scale_alg: int) -> tuple[torch.Tensor, torch.Tensor]:
+ """Quantize weight or activation using the same explicit FP4 algorithm."""
+ import torch_npu
+
+ # Universal Optimal Scaling (Qmax=7.25), MXAttention Section 4.1:
+ # https://arxiv.org/abs/2607.24377
+ scale_kwargs = {"dst_type_max": 7.25} if scale_alg == 2 else {}
+ return torch_npu.npu_dynamic_mx_quant(
+ x,
+ dst_type=torch_npu.float4_e2m1fn_x2,
+ axis=-1,
+ block_size=32,
+ round_mode="rint",
+ scale_alg=scale_alg,
+ **scale_kwargs,
+ )
+
+
+def _npu_w4a8_matmul(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ weight_scale: torch.Tensor,
+ bias: torch.Tensor | None,
+ output_dtype: torch.dtype,
+) -> torch.Tensor:
+ """MXFP8 activations with prepared single-scale MXFP4 weights (N, K)."""
+ import torch_npu
+
+ x_q, x_scale = torch_npu.npu_dynamic_mx_quant(
+ x, dst_type=torch_npu.float8_e4m3fn, axis=-1, block_size=32, round_mode="rint", scale_alg=0
+ )
+ # Use the BF16 output/bias pair supported by the mixed FP8/FP4 kernel
+ # in the target A5 runtime. Restore the caller's dtype afterwards.
+ if bias is not None:
+ bias = bias.to(torch.bfloat16).reshape(1, -1)
+ output = torch_npu.npu_quant_matmul(
+ x_q,
+ weight.transpose(0, 1),
+ weight_scale.transpose(0, 1),
+ scale_dtype=torch_npu.float8_e8m0fnu,
+ x2_dtype=torch_npu.float4_e2m1fn_x2,
+ pertoken_scale=x_scale,
+ pertoken_scale_dtype=torch_npu.float8_e8m0fnu,
+ bias=bias,
+ output_dtype=torch.bfloat16,
+ group_sizes=[0, 0, 32],
+ )
+ return output.to(output_dtype)
+
+
class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
"""NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints.
@@ -187,7 +306,8 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared.
"""
- def __init__(self, quant_config: DiffusionMXFP4Config) -> None:
+ def __init__(self, quant_config: DiffusionMXFP4Config, *, prefix: str = "") -> None:
+ self.is_w4a8_fallback_layer = prefix in quant_config.w4a8_fallback_layers
self.quant_config = quant_config
self.out_dtype = torch.get_default_dtype()
@@ -210,11 +330,14 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
layer.orig_dtype = params_dtype
layer.weight_block_size = None
- # BF16 placeholder; cast to float4_e2m1fn_x2 in process_weights.
layer.register_parameter(
"weight",
ModelWeightParameter(
- data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype),
+ data=torch.empty(
+ output_size_per_partition,
+ input_size_per_partition,
+ dtype=params_dtype,
+ ),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
@@ -234,12 +357,42 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
),
)
+ # Older single-scale checkpoints have no Smooth tensor. Keep their
+ # identity transform while allowing calibrated checkpoints to load it.
+ def smooth_weight_loader(
+ param: ModelWeightParameter, loaded_weight: torch.Tensor, *args: Any, **kwargs: Any
+ ) -> None:
+ _validate_smooth_scale(loaded_weight)
+ if not callable(weight_loader):
+ raise ValueError("Loading MXFP4 Smooth requires a callable weight_loader")
+ weight_loader(param, loaded_weight, *args, **kwargs)
+
+ layer.register_parameter(
+ "mul_scale",
+ ModelWeightParameter(
+ data=torch.ones(input_size_per_partition, dtype=torch.float32),
+ input_dim=0,
+ output_dim=None,
+ weight_loader=smooth_weight_loader,
+ ),
+ )
+ setattr(layer.mul_scale, "ignore_warning", True)
+ setattr(layer.mul_scale, "is_checkpoint_optional", not self.quant_config.require_smooth_scale)
+ if self.quant_config.is_checkpoint_mxfp4_serialized:
+ setattr(layer.weight_scale, "is_checkpoint_required", True)
+
def process_weights_after_loading(self, layer: Module) -> None:
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
import torch_npu
+ smooth_param = cast(torch.Tensor, layer.mul_scale)
+ _validate_smooth_scale(smooth_param)
+ smooth = smooth_param.data.to(dtype=cast(torch.dtype, layer.orig_dtype)).contiguous()
+ # Finite FP32 values can overflow or underflow in the activation dtype.
+ _validate_smooth_scale(smooth)
+
# NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose.
w = layer.weight
if w.dtype != torch_npu.float4_e2m1fn_x2:
@@ -258,15 +411,31 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
replace_parameter(layer, "weight", w)
replace_parameter(layer, "weight_scale", s)
+ # Keep loader metadata while converting the calibration tensor once.
+ layer.mul_scale.data = smooth.to(device=cast(torch.Tensor, w).device)
layer._already_called_process_weights_after_loading = True
# --- NPU MXFP4 ops — shared with online path via inheritance ---
- def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
- import torch_npu
+ def _apply_inner(
+ self,
+ layer: torch.nn.Module,
+ x: torch.Tensor,
+ bias: torch.Tensor | None,
+ ori_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ # Online quantization has no calibrated Smooth tensor. For offline
+ # weights the same pre-scale must precede both A4 and A8 quantization.
+ if (mul_scale := getattr(layer, "mul_scale", None)) is not None:
+ x = x * mul_scale.to(x.dtype)
+ if self.is_w4a8_fallback_layer or _is_w4a8_fallback_step(self.quant_config.w4a8_fallback_steps):
+ return _npu_w4a8_matmul(
+ x, cast(torch.Tensor, layer.weight), cast(torch.Tensor, layer.weight_scale), bias, ori_dtype
+ )
+ return super()._apply_inner(layer, x, bias, ori_dtype)
- # No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2.
- return torch_npu.npu_dynamic_mx_quant(x)
+ def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ return _npu_quantize_mxfp4(x, self.quant_config.mxfp4_scale_alg)
def _quant_matmul(
self,
@@ -318,8 +487,6 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
- import torch_npu
-
if layer.weight.device == torch.device("meta"):
weight = ModelWeightParameter(
data=torch.empty_like(layer.weight, device=layer._load_device),
@@ -331,8 +498,10 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
layer.register_parameter("weight", weight)
initialize_single_dummy_weight(layer.weight)
- # NPU: quantize BF16/FP16 (N, K) → FP4. No dst_type → float4_e2m1fn_x2.
- weight_fp4, weight_scale_raw = torch_npu.npu_dynamic_mx_quant(layer.weight)
+ # Prepare once; offline checkpoints keep their existing weight algorithm.
+ weight_fp4, weight_scale_raw = _npu_quantize_mxfp4(
+ cast(torch.Tensor, layer.weight), self.quant_config.mxfp4_scale_alg
+ )
# Weight stays (N, K) — no pre-transpose for FP4 packed format.
# Scale: (N, S) → (N, S/2, 2). Not pre-transposed; done inline in _quant_matmul.
@@ -497,7 +666,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
"""NPU W4A4 MXFP4 dual-scale offline method for pre-quantized checkpoints.
Checkpoint tensors and their canonical post-load shapes:
- weight : (N, K) float8_e4m3fn – FP4 packed (2 values per byte)
+ weight : (N, K) float8_e4m3fn – numeric FP4 values; packed into NZ after loading
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits); reshaped to (N, K//64, 2)
weight_dual_scale : (N, K//512, 1) float32 – coarse scale; transposed to (K//512, N)
mul_scale : (K,) float32 – per-input-channel activation pre-scale (from calibration)
@@ -509,7 +678,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
Reference: MindIE-SD W4A4MXFP4DualQuantLinear.
"""
- def __init__(self, quant_config: Any) -> None:
+ def __init__(self, quant_config: DiffusionMXFP4DualScaleMixedConfig) -> None:
self.quant_config = quant_config
self.out_dtype = torch.get_default_dtype()
@@ -532,7 +701,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
layer.orig_dtype = params_dtype
layer.weight_block_size = None
- # FP4 packed: 2 values per float8_e4m3fn byte → checkpoint stores as float8_e4m3fn; register as BF16
+ # Checkpoint FP4 values are stored numerically in FP8; BF16 preserves them exactly.
layer.register_parameter(
"weight",
ModelWeightParameter(
@@ -587,6 +756,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
),
)
setattr(layer.mul_scale, "ignore_warning", True)
+ setattr(layer.weight_scale, "is_checkpoint_required", True)
def process_weights_after_loading(self, layer: Module) -> None:
if getattr(layer, "_already_called_process_weights_after_loading", False):
@@ -594,7 +764,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
import torch_npu
- # float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29).
+ # Numeric FP4 values → packed float4_e2m1fn_x2 → NZ hardware format (format ID 29).
# NZ layout is required by npu_dual_level_quant_matmul for FP4 weight matrices.
w = torch_npu.npu_dtype_cast(layer.weight.data.npu(), torch_npu.float4_e2m1fn_x2)
w = torch_npu.npu_format_cast(w.view(torch.int8), 29, customize_dtype=torch.int8)
@@ -631,7 +801,10 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
"""
if ori_dtype not in (torch.bfloat16, torch.float16):
x = x.to(torch.bfloat16)
- x_q, l0_scale, l1_scale = self._quantize_activation(x, layer.mul_scale)
+ # The dual-level activation kernel requires x and Smooth to have the
+ # same dtype. Prepared calibration scales are BF16, including online.
+ smooth_scale = cast(torch.Tensor, layer.mul_scale).to(x.dtype)
+ x_q, l0_scale, l1_scale = self._quantize_activation(x, smooth_scale)
return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype)
def _quantize_activation( # type: ignore[override]
@@ -785,11 +958,20 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig):
is_checkpoint_serialized: bool = False,
ignored_layers: list[str] | None = None,
num_bf16_fallback_layers: int = 5,
+ w4a8_fallback_steps: list[int] | None = None,
+ w4a8_fallback_layers: list[str] | None = None,
) -> None:
super().__init__()
self.is_checkpoint_serialized = is_checkpoint_serialized
self.ignored_layers = ignored_layers or []
self.num_bf16_fallback_layers = num_bf16_fallback_layers
+ self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps)
+ self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers)
+ if self.w4a8_fallback_steps or self.w4a8_fallback_layers:
+ raise ValueError(
+ "mxfp4_dualscale does not support W4A8 fallback. Use mxfp4 with single-scale weights; "
+ "changing a DualScale checkpoint's method name does not convert its storage format."
+ )
@classmethod
def get_name(cls) -> QuantizationMethods:
@@ -822,6 +1004,8 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig):
is_checkpoint_serialized=is_serialized,
ignored_layers=ignored_layers,
num_bf16_fallback_layers=num_bf16_fallback_layers,
+ w4a8_fallback_steps=config.get("w4a8_fallback_steps"),
+ w4a8_fallback_layers=config.get("w4a8_fallback_layers"),
)
def get_quant_method(
diff --git a/vllm_omni/utils/forced_aligner.py b/vllm_omni/utils/forced_aligner.py
index 3b578ee51..e6dcecb57 100644
--- a/vllm_omni/utils/forced_aligner.py
+++ b/vllm_omni/utils/forced_aligner.py
@@ -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
"""Forced-aligner config + word-timestamp decoding for TTS.
The aligner runs as a pooling stage appended to the pipeline by
@@ -153,6 +153,7 @@ def inject_forced_aligner_stage(
# Pooling is one-shot -> sync AR scheduler (the async one assumes
# token-by-token generation).
async_scheduling=False,
+ async_chunk=False,
)
deploy = replace(deploy, stages=list(deploy.stages) + [aligner_ds])
logger.info("[stage_init] Injected forced-aligner stage %d (model=%s)", new_id, fa.model)
+