Merge main and preserve event-driven defaults in shared engine bases

Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
Sy03 committed 2026-09-16 20:14:33 +08:00
commit d2bb944863
371 files changed
+43235 -38379

No files matched your search

+65
View File
@@ -205,6 +205,71 @@ steps:
pytest -s -v tests/engine/test_async_omni_engine_abort.py
'
- label: "ROCm · Entrypoints GPU Coverage (R2-02)"
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: NonBlocking
timeout_in_minutes: 30
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
artifact_paths:
- "artifacts/rocm-r2-02/**/*"
commands:
# Native AMD tests run from /app/vllm-omni, while Buildkite uploads
# artifacts relative to its checkout directory.
- export R2_02_ARTIFACT_DIR="$$BUILDKITE_BUILD_CHECKOUT_PATH/artifacts/rocm-r2-02"
- mkdir -p "$$R2_02_ARTIFACT_DIR"
- >-
ps -eo pid=,ppid=,lstart=,stat=,comm= >
"$$R2_02_ARTIFACT_DIR/processes-before.txt"
# Verbose pytest output and JUnit retain selected node IDs without a
# second collection-only pass over the same GPU-backed test tree.
# The inner timeout leaves five minutes for teardown verification and
# artifact generation before the Buildkite step timeout.
- >-
timeout --signal=TERM --kill-after=1m 25m
pytest -s -v -ra
tests/entrypoints/
-m 'core_model and cuda and not
(cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)'
--run-level "core_model"
--durations=0
--junitxml="$$R2_02_ARTIFACT_DIR/pytest.xml"
2>&1 | tee "$$R2_02_ARTIFACT_DIR/pytest.log"
- >-
tail -n 100 "$$R2_02_ARTIFACT_DIR/pytest.log" >
"$$R2_02_ARTIFACT_DIR/pytest-summary.txt"
- |
sleep 5
ps -eo pid=,ppid=,lstart=,stat=,comm= > "$$R2_02_ARTIFACT_DIR/processes-after.txt"
# Compare PID plus process start time so PID reuse cannot hide a leak.
# Double dollars preserve awk field references through Buildkite interpolation.
awk '
NR == FNR {
if ($$9 != "ps") {
baseline[$$1 "|" $$3 "|" $$4 "|" $$5 "|" $$6 "|" $$7] = 1
}
next
}
$$9 != "ps" {
identity = $$1 "|" $$3 "|" $$4 "|" $$5 "|" $$6 "|" $$7
if (!(identity in baseline)) {
print
leaked = 1
}
}
END {
if (leaked) {
print "status=FAIL leaked_processes_detected=1"
exit 1
}
print "status=PASS leaked_processes=0"
}
' "$$R2_02_ARTIFACT_DIR/processes-before.txt" \
"$$R2_02_ARTIFACT_DIR/processes-after.txt" \
2>&1 | tee "$$R2_02_ARTIFACT_DIR/process-cleanup.txt"
# - label: "Omni Model Test Qwen3-Omni"
# agent_pool: mi300_2
# depends_on: amd-build
@@ -38,9 +38,13 @@ x-minicpmo-4-5: &minicpmo_4_5
x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/engine/duplex/
- vllm_omni/engine/duplex_omni_engine.py
- vllm_omni/engine/duplex_orchestrator.py
- vllm_omni/engine/omni_engine_base.py
- vllm_omni/entrypoints/duplex/
- vllm_omni/entrypoints/duplex_omni.py
- vllm_omni/entrypoints/async_omni_base.py
- vllm_omni/clients/
- vllm_omni/entrypoints/duplex_request_client.py
- vllm_omni/outputs/duplex.py
- vllm_omni/model_executor/duplex_sampling.py
- vllm_omni/experimental/fullduplex/
@@ -249,7 +253,11 @@ source_file_dependencies:
- tests/assets/minicpmo_4_5/
- tests/clients/
- tests/engine/duplex/
- tests/engine/test_duplex_omni_engine.py
- tests/engine/test_duplex_orchestrator.py
- tests/entrypoints/duplex/
- tests/entrypoints/test_duplex_omni.py
- tests/entrypoints/openai_api/test_duplex_api_server.py
- tests/model_executor/models/minicpmo_4_5/duplex/
omni_minicpmo_4_5_duplex_perf:
+3 -1
View File
@@ -66,7 +66,9 @@ steps:
# 0.26 startup regression is understood and fixed.
timeout_in_minutes: 90
commands:
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
# Turn-based server VAD over /v1/realtime is deferred to RFC #7181 PR 5
# (see the P0 thread on #7413); the test stays in the tree and runs again there.
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" --deselect tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_server_vad_multi_turn_without_client_commit
mirror_hardwares: h100_2
- label: "Entrypoint Test with L4"
+2 -1
View File
@@ -191,13 +191,14 @@ steps:
# from downgrading transformers; see
# examples/offline_inference/text_to_speech/README.md.
- label: "TTS · Gepard-1.0 · L4"
timeout_in_minutes: 60
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt
- uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt
- >-
pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py
tests/e2e/online_serving/test_gepard_tts.py
-m "tts and cuda"
--run-level "full_model"
mirror_hardwares: l4_1
+83
View File
@@ -139,6 +139,89 @@ You can configure per-stage sampling parameters for multi-stage models.
>
> Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect.
For MiniMax-H3 Ref2VA, **Video References** accepts up to 9 images (`image_1`–`image_9`),
3 videos (`video_1`–`video_3`), and 3 audio clips (`audio_1`–`audio_3`), with at most
12 connected inputs in total. Any mixture containing at least one image or video
is supported; empty and audio-only references are rejected. For example, you can
combine 6 images, 3 videos, and 3 audio clips in one request.
Within each media type, references follow slot-number order, skipping unconnected
slots. For example, connecting `image_2` and `image_9` sends `image_2` as the first
image and `image_9` as the second. Each image slot uses the first image in its batch.
Existing connections to `image_1`, `image_2`, `audio_1`, `audio_2`, `video_1`, and
`video_2` remain valid in saved workflows.
#### MiniMax-H3 Reference to Video
Load [MiniMax-H3 Reference to Video](example_workflows/vLLM-Omni%20MiniMax-H3%20Reference%20to%20Video.json)
from **Templates → ComfyUI-vLLM-Omni**, or drag the JSON onto the canvas.
It adapts the [official ComfyUI R2V workflow](https://github.com/Comfy-Org/workflow_templates/blob/main/templates/video_minimax_h3_r2v.json)
to remote vLLM-Omni execution. Model loading, sampling, and VAE decoding run on
the server; ComfyUI loads references and saves the returned video.
Use an extension version containing the audio-preserving output fix from
[#7456](https://github.com/vllm-project/vllm-omni/pull/7456). Older versions can
drop the server's audio while decoding the response.
1. Start a Ref2VA-capable service using the [MiniMax-H3 recipe](../../recipes/MiniMaxAI/MiniMax-H3.md).
Set **Generate Video** to its `/v1` URL and served model name.
2. Select your own image in **Load Image**. The template's filenames are placeholders
relative to ComfyUI's input directory; no sample assets or model weights are bundled.
3. Connect **Load Video** and **Load Audio** when needed. Duplicate the loaders to
fill more reference slots, or disconnect the image when using video-only input.
Keep the references connected to **Generate Video** and leave `frame` disconnected.
4. Match prompt tags such as `<Picture 1>`, `<Video 1>`, and `<Audio 1>` to the
connected references, counting each media type separately and skipping empty slots.
5. Run the workflow. **Save Video** writes an MP4 under `output/video/` and preserves
the generated audio when the audio-output prerequisite is installed.
The default is **1344×768, 24 FPS, 124 frames** (about 5.17 seconds), 50 sampling
points, video flow shift 12, audio flow shift 3, and seed 42. The duration widget
is set to 5.167 seconds, which converts to 124 frames at 24 FPS. Other H3 canvas presets
are 1024×768, 768×768, 768×1024, and 768×1344. Keep frame counts at `17k+5` within
the 4–15 second output range; examples are 107, 124, 209, and 345 frames.
Reference video/audio clips must each be 2–15 seconds, with at most 15 seconds of
reference video and at most 15 seconds of standalone reference audio per request.
For **Ref2VA Turbo**, use the same template. Configure a Ref2VA-only server with
`--lora-backend peft --lora-path "$TURBO_LORA" --task-type ref2va`, following the
recipe's hardware setup. For `minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors`,
set **LoRA** `local_path` to that file's path on the server and connect it to
**Generate Video**. Use `num_inference_steps=5`, `flow_shift=12`,
`audio_flow_shift=3`, and LoRA `scale=1`. The API counts sigma points, so five
points produce four denoiser evaluations. Use the Diffusers artifact, not its
`_comfyui_` export. To return to base mode, disconnect LoRA and restore 50 steps.
The 8-step v1.0 Ref2VA adapter instead requires nine sampling points and video
flow shift 6. FastH3 is a separate startup-fused T2VA deployment.
Validate the template's wiring and defaults locally with:
```bash
python -m pytest tests/e2e/features/comfyui/test_h3_reference_workflow.py -q
```
For a real-model check, run the workflow against H3 and inspect the saved file:
```bash
ffprobe -v error -show_entries stream=codec_type,width,height,r_frame_rate,sample_rate,channels \
-show_entries format=duration -of json "$OUTPUT_MP4"
```
Check for 24 FPS video, stereo audio, the requested canvas, and aligned duration;
also play the result to assess reference conditioning and audio/video synchronization.
Record the server commit, model/adapter, input assets, prompt, settings, and output
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).
@@ -10,6 +10,9 @@ from .utils.api_client import VLLMOmniClient
from .utils.logger import get_logger
from .utils.models import lookup_model_spec
from .utils.types import (
MAX_REFERENCE_AUDIOS,
MAX_REFERENCE_IMAGES,
MAX_REFERENCE_VIDEOS,
AudioFormat,
AutoregressionSamplingParams,
DiffusionSamplingParams,
@@ -918,6 +921,10 @@ class VLLMOmniVideoReferences:
"audio_2": ("AUDIO",),
"video_1": ("VIDEO",),
"video_2": ("VIDEO",),
# Append ports to preserve connections in saved workflows.
**{f"image_{i}": ("IMAGE",) for i in range(3, MAX_REFERENCE_IMAGES + 1)},
**{f"audio_{i}": ("AUDIO",) for i in range(3, MAX_REFERENCE_AUDIOS + 1)},
**{f"video_{i}": ("VIDEO",) for i in range(3, MAX_REFERENCE_VIDEOS + 1)},
},
}
@@ -934,21 +941,26 @@ class VLLMOmniVideoReferences:
audio_2: AudioInput | None = None,
video_1: VideoInput | None = None,
video_2: VideoInput | None = None,
image_3: torch.Tensor | None = None,
image_4: torch.Tensor | None = None,
image_5: torch.Tensor | None = None,
image_6: torch.Tensor | None = None,
image_7: torch.Tensor | None = None,
image_8: torch.Tensor | None = None,
image_9: torch.Tensor | None = None,
audio_3: AudioInput | None = None,
video_3: VideoInput | None = None,
**kwargs,
):
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
refs = VideoReferences()
if image_1 is not None:
refs["image_1"] = image_1
if image_2 is not None:
refs["image_2"] = image_2
if audio_1 is not None:
refs["audio_1"] = audio_1
if audio_2 is not None:
refs["audio_2"] = audio_2
if video_1 is not None:
refs["video_1"] = video_1
if video_2 is not None:
refs["video_2"] = video_2
for kind, values in (
("image", (image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9)),
("video", (video_1, video_2, video_3)),
("audio", (audio_1, audio_2, audio_3)),
):
for index, value in enumerate(values, start=1):
if value is not None:
refs[f"{kind}_{index}"] = value
return (refs,)
@@ -20,6 +20,7 @@ from comfy_api.input import AudioInput, VideoInput
from .format import (
audio_to_base64,
audio_to_bytes,
base64_to_audio,
base64_to_image_tensor,
bytes_to_audio,
@@ -31,7 +32,13 @@ from .format import (
)
from .logger import get_logger, pretty_printer
from .models import lookup_model_spec
from .types import AudioFormat
from .types import (
MAX_REFERENCE_AUDIOS,
MAX_REFERENCE_IMAGES,
MAX_REFERENCE_VIDEOS,
MAX_TOTAL_REFERENCES,
AudioFormat,
)
logger = get_logger(__name__)
@@ -304,38 +311,44 @@ class VLLMOmniClient:
# === multimodal inputs (first-last-frames, references, etc.) ===
input_reference_image: torch.Tensor | None = None
audio_reference: AudioInput | None = None
reference_videos: list[VideoInput] = []
video_task: str | None = None
if frame is not None:
input_reference_image = frame
video_task = "fl2va"
elif references is not None:
images = [references[k] for k in ("image_1", "image_2") if k in references and references[k] is not None]
audios = [references[k] for k in ("audio_1", "audio_2") if k in references and references[k] is not None]
videos = [references[k] for k in ("video_1", "video_2") if k in references and references[k] is not None]
if not images and not audios and not videos:
raise ValueError("references is empty; connect at least one image/audio/video.")
if videos:
if images or audios:
raise ValueError(
"Video references cannot be combined with image or audio references. "
"Connect only video_1/video_2 for multi-video Ref2VA."
)
reference_videos = videos
video_task = "ref2va"
elif len(images) == 1 and len(audios) == 1:
input_reference_image = images[0]
audio_reference = audios[0]
video_task = "ref2va"
else:
reference_formats = (
("image", MAX_REFERENCE_IMAGES, "png", "image/png", image_tensor_to_png_bytes),
("video", MAX_REFERENCE_VIDEOS, "mp4", "video/mp4", video_to_bytes),
("audio", MAX_REFERENCE_AUDIOS, "mp3", "audio/mpeg", audio_to_bytes),
)
supported_inputs = {f"{kind}_{i}" for kind, limit, *_ in reference_formats for i in range(1, limit + 1)}
connected = {name: value for name, value in references.items() if value is not None}
unsupported = connected.keys() - supported_inputs
if unsupported:
raise ValueError(f"Unsupported reference input(s): {', '.join(sorted(unsupported))}.")
if not any(name.startswith(("image_", "video_")) for name in connected):
raise ValueError(
"Invalid references combination. Supported modes: "
"(1) one or more videos only, or (2) exactly one image and one audio."
"references requires at least one image or video; audio-only inputs are not supported."
)
if len(connected) > MAX_TOTAL_REFERENCES:
raise ValueError(
f"references supports at most {MAX_TOTAL_REFERENCES} inputs in total "
f"(up to {MAX_REFERENCE_IMAGES} images, {MAX_REFERENCE_VIDEOS} videos, "
f"and {MAX_REFERENCE_AUDIOS} audios)."
)
for kind, limit, extension, content_type, encode in reference_formats:
for index in range(1, limit + 1):
name = f"{kind}_{index}"
if name in connected:
filename = f"{name}.{extension}"
form.add_field(
"input_references",
encode(connected[name], filename),
filename=filename,
content_type=content_type,
)
video_task = "ref2va"
else:
video_task = "t2va"
@@ -348,21 +361,6 @@ class VLLMOmniClient:
content_type="image/png",
)
if audio_reference is not None:
form.add_field(
"audio_reference",
json.dumps({"audio_url": audio_to_base64(audio_reference)}, ensure_ascii=False),
)
for idx, video in enumerate(reference_videos, start=1):
video_filename = f"reference_{idx}.mp4"
form.add_field(
"input_references",
video_to_bytes(video, video_filename),
filename=video_filename,
content_type="video/mp4",
)
# === model specific params. Either use a specialized builder, or add flattened fields as-is ===
if model_params is not None:
model_params = dict(model_params)
@@ -30,6 +30,12 @@ class MiniMaxH3ModelSpecificParams(dict):
pass
MAX_REFERENCE_IMAGES = 9
MAX_REFERENCE_VIDEOS = 3
MAX_REFERENCE_AUDIOS = 3
MAX_TOTAL_REFERENCES = 12
class VideoReferences(dict):
pass
@@ -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.
@@ -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
}
@@ -0,0 +1,786 @@
{
"id": "e7a4e606-6510-4c4e-b127-9a8b6bf3ea71",
"revision": 0,
"last_node_id": 11,
"last_link_id": 5,
"nodes": [
{
"id": 1,
"type": "LoadImage",
"pos": [
40,
100
],
"size": [
320,
300
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "COMBO",
"widget": {
"name": "image"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": [
1
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"h3_reference.png",
"image"
],
"title": "Reference image \u2014 select a local file"
},
{
"type": "LoadVideo",
"size": [
320,
200
],
"flags": {},
"inputs": [],
"outputs": [
{
"name": "VIDEO",
"type": "VIDEO",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadVideo"
},
"widgets_values": [
"h3_reference.mp4"
],
"id": 2,
"pos": [
40,
470
],
"order": 1,
"mode": 0,
"title": "Optional reference video"
},
{
"id": 3,
"type": "LoadAudio",
"pos": [
40,
1100
],
"size": [
320,
160
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "audio",
"name": "audio",
"type": "COMBO",
"widget": {
"name": "audio"
},
"link": null
},
{
"localized_name": "audioUI",
"name": "audioUI",
"type": "AUDIO_UI",
"widget": {
"name": "audioUI"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "AUDIOUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "AUDIO",
"name": "AUDIO",
"type": "AUDIO",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadAudio"
},
"widgets_values": [
"h3_reference.wav",
null,
""
],
"title": "Optional reference audio"
},
{
"id": 4,
"type": "VLLMOmniVideoReferences",
"pos": [
420,
100
],
"size": [
300,
365
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "image_1",
"type": "IMAGE",
"shape": 7,
"link": 1
},
{
"name": "image_2",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "audio_1",
"type": "AUDIO",
"shape": 7,
"link": null
},
{
"name": "audio_2",
"type": "AUDIO",
"shape": 7,
"link": null
},
{
"name": "video_1",
"type": "VIDEO",
"shape": 7,
"link": null
},
{
"name": "video_2",
"type": "VIDEO",
"shape": 7,
"link": null
},
{
"name": "image_3",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "image_4",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "image_5",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "image_6",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "image_7",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "image_8",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "image_9",
"type": "IMAGE",
"shape": 7,
"link": null
},
{
"name": "audio_3",
"type": "AUDIO",
"shape": 7,
"link": null
},
{
"name": "video_3",
"type": "VIDEO",
"shape": 7,
"link": null
}
],
"outputs": [
{
"localized_name": "references",
"name": "references",
"type": "VIDEO_REFERENCES",
"links": [
2
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniVideoReferences"
},
"widgets_values": [],
"title": "References \u2014 up to 12 total"
},
{
"id": 5,
"type": "VLLMOmniDiffusionSampling",
"pos": [
420,
530
],
"size": [
310,
275
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
3
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1.0,
1.0,
false,
false,
42,
"fixed"
],
"title": "Sampling \u2014 base 50 / Turbo 5"
},
{
"id": 6,
"type": "VLLMOmniMiniMaxH3Params",
"pos": [
790,
750
],
"size": [
310,
105
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "audio_flow_shift",
"name": "audio_flow_shift",
"type": "FLOAT",
"widget": {
"name": "audio_flow_shift"
},
"link": null
},
{
"localized_name": "flow_shift",
"name": "flow_shift",
"type": "FLOAT",
"widget": {
"name": "flow_shift"
},
"link": null
}
],
"outputs": [
{
"localized_name": "MiniMax-H3 video params",
"name": "MiniMax-H3 video params",
"type": "VIDEO_PARAMS",
"links": [
4
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniMiniMaxH3Params"
},
"widgets_values": [
3.0,
12.0
],
"title": "H3 video/audio flow shifts"
},
{
"type": "VLLMOmniRemoteLoRA",
"size": [
310,
180
],
"flags": {},
"inputs": [],
"outputs": [
{
"name": "lora",
"type": "REMOTE_LORA",
"links": null
}
],
"properties": {
"Node name for S&R": "VLLMOmniRemoteLoRA"
},
"widgets_values": [
"",
"h3-ref2va-turbo",
1.0,
0
],
"id": 7,
"pos": [
420,
940
],
"order": 6,
"mode": 0,
"title": "Optional Turbo LoRA \u2014 connect to enable"
},
{
"id": 8,
"type": "VLLMOmniGenerateVideo",
"pos": [
800,
100
],
"size": [
470,
560
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"localized_name": "frame",
"name": "frame",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"localized_name": "references",
"name": "references",
"shape": 7,
"type": "VIDEO_REFERENCES",
"link": 2
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 3
},
{
"localized_name": "lora",
"name": "lora",
"shape": 7,
"type": "REMOTE_LORA",
"link": null
},
{
"localized_name": "model_params",
"name": "model_params",
"shape": 7,
"type": "VIDEO_PARAMS",
"link": 4
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
},
{
"localized_name": "fps",
"name": "fps",
"type": "INT",
"widget": {
"name": "fps"
},
"link": null
},
{
"localized_name": "duration",
"name": "duration",
"type": "FLOAT",
"widget": {
"name": "duration"
},
"link": null
}
],
"outputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"links": [
5
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateVideo"
},
"widgets_values": [
"http://localhost:8000/v1",
"MiniMaxAI/MiniMax-H3",
"Use the subject and visual style from <Picture 1>. Create a cinematic shot with a gentle camera move and natural ambient sound.",
"",
1344,
768,
24,
5.167
],
"title": "MiniMax-H3 Reference to Video"
},
{
"id": 9,
"type": "SaveVideo",
"pos": [
1360,
100
],
"size": [
410,
420
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"link": 5
},
{
"localized_name": "filename_prefix",
"name": "filename_prefix",
"type": "STRING",
"widget": {
"name": "filename_prefix"
},
"link": null
},
{
"localized_name": "format",
"name": "format",
"type": "COMBO",
"widget": {
"name": "format"
},
"link": null
},
{
"localized_name": "codec",
"name": "codec",
"type": "COMFY_DYNAMICCOMBO_V3",
"widget": {
"name": "codec"
},
"link": null
}
],
"outputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"links": null
}
],
"properties": {},
"widgets_values": [
"video/MiniMax-H3-Ref2VA",
"mp4",
"h264"
],
"title": "Save generated video with audio"
},
{
"id": 10,
"type": "MarkdownNote",
"pos": [
1360,
590
],
"size": [
440,
595
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Reference inputs and output settings",
"properties": {},
"widgets_values": [
"## Reference to Video\nSelect an image, set the server URL and served model name, then run.\n\nConnect the optional video/audio loaders to Video References, or duplicate loaders for more inputs: up to 9 images, 3 videos, 3 audio clips, 12 total. At least one image or video is required.\n\nSlots are numbered separately per media type. Empty slots are skipped. Update prompt tags such as <Picture 1>, <Video 1>, <Audio 1> to match the connected references. Videos keep their own soundtrack in the reference upload.\n\n**Output:** 24 FPS; 124 frames (~5.17 s). Valid frame counts are 107, 124, 141, 158, 175, 192, 209, 226, 243, 260, 277, 294, 311, 328, 345.\n\n**Width x height:** 1344x768 (16:9), 1024x768 (4:3), 768x768 (1:1), 768x1024 (3:4), 768x1344 (9:16).\n\nEach reference video/audio clip must be 2\u201315 s. Total reference-video duration <=15 s; total standalone-audio duration <=15 s.\n\nRequires the expanded Video References node and the audio-preserving video output fix. See the extension README for setup and validation status."
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 11,
"type": "MarkdownNote",
"pos": [
790,
940
],
"size": [
480,
455
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Optional Turbo \u2014 same workflow",
"properties": {},
"widgets_values": [
"## Turbo Ref2VA\nBase mode: leave LoRA disconnected; use 50 steps, video flow shift 12, audio flow shift 3.\n\nFor the 4-step v0.1 Ref2VA Turbo adapter:\n1. Start a Ref2VA-only server with the Diffusers adapter loaded (see README).\n2. Set LoRA local_path to its path on the server, then connect LoRA to Generate Video.\n3. Set Sampling num_inference_steps to **5**; keep shifts at **12 / 3**, scale **1**.\n\nvLLM-Omni counts sigma points: 5 points give 4 denoiser evaluations. Use minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors, not the _comfyui_ export.\n\nDisconnect LoRA and restore 50 steps to return to base mode. FastH3 is a separate T2VA deployment, not this option.\n\nAdapted from the official ComfyUI MiniMax-H3 R2V workflow: https://github.com/Comfy-Org/workflow_templates/blob/main/templates/video_minimax_h3_r2v.json"
],
"color": "#432",
"bgcolor": "#000"
}
],
"links": [
[
1,
1,
0,
4,
0,
"IMAGE"
],
[
2,
4,
0,
8,
1,
"VIDEO_REFERENCES"
],
[
3,
5,
0,
8,
2,
"SAMPLING_PARAMS"
],
[
4,
6,
0,
8,
4,
"VIDEO_PARAMS"
],
[
5,
8,
0,
9,
0,
"VIDEO"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 0.85,
"offset": [
100,
100
]
}
},
"version": 0.4
}
@@ -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())
+1
View File
@@ -44,6 +44,7 @@ RUN python3 -m pip install -U pip && \
"onnx>=1.16.0" \
"librosa>=0.10.0" \
"datasets>=3.0.0" \
"torchmetrics==1.9.0" \
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
--default-timeout=60
+6 -1
View File
@@ -5,6 +5,9 @@
Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
- [vllm_omni.entrypoints.async_omni_base.AsyncOmniBase][]
- [vllm_omni.entrypoints.duplex_omni.DuplexOmni][]
- [vllm_omni.entrypoints.duplex_omni.DuplexSessionHandle][]
- [vllm_omni.engine.cfg_companion_tracker.CfgCompanionTracker][]
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
@@ -17,12 +20,13 @@ Main entry points for vLLM-Omni inference and serving.
## Clients
Client libraries that talk to a running vLLM-Omni server (the full-duplex Realtime client and its per-model session presets).
Client libraries for full-duplex sessions: the Realtime WebSocket client, the in-process client over `DuplexOmni`, and the per-model session presets.
- [vllm_omni.clients.duplex.AudioDelta][]
- [vllm_omni.clients.duplex.AudioFormat][]
- [vllm_omni.clients.duplex.ConnectionResumed][]
- [vllm_omni.clients.duplex.DuplexClient][]
- [vllm_omni.clients.duplex.DuplexClientBase][]
- [vllm_omni.clients.duplex.DuplexClientError][]
- [vllm_omni.clients.duplex.DuplexConnectionError][]
- [vllm_omni.clients.duplex.DuplexEvent][]
@@ -44,6 +48,7 @@ Client libraries that talk to a running vLLM-Omni server (the full-duplex Realti
- [vllm_omni.clients.duplex.TextDelta][]
- [vllm_omni.clients.duplex.TranscriptDelta][]
- [vllm_omni.clients.duplex.WebSocketTransport][]
- [vllm_omni.clients.inline_duplex.InlineDuplexClient][]
## Inputs
@@ -82,6 +82,7 @@ depends on the installed kernels and model path.
| --- | --- | --- | --- | --- |
| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
| `VLLM_OMNI_ABORT_TIMEOUT` | Float seconds; default `2` | Engine abort wait for Videos DELETE and `generate()` cancel/error cleanup; read when `async_omni` imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off), except Qwen3-TTS defaults to on | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. An explicit value always wins. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental |
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
+2 -1
View File
@@ -17,6 +17,7 @@ Common `PipelineConfig` fields include:
| ------- | ------------- |
| `model_type` | Pipeline identifier used during model and config resolution. |
| `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. |
| `duplex_plugin` | Dotted path of the model's `DuplexModelPlugin`. Set only for full-duplex models; it makes `vllm-omni serve` run the model through `DuplexOmni` (a server whose every surface runs on a duplex session) and the engine host a `DuplexOrchestrator` with the plugin loaded. |
| `model_arch` | Default Hugging Face architecture for the pipeline. |
| `hf_architectures` | Architecture names used to identify checkpoints whose `model_type` is shared. |
| `hf_config_predicate` | Optional predicate used to select between pipelines with otherwise identical HF metadata. |
@@ -52,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. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. |
| `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. |
| `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). |
+11 -12
View File
@@ -131,13 +131,13 @@ for the validated batching, attention, and quantization combinations.
### MiniMax-H3: conditioning, denoising, and decode
MiniMax-H3 illustrates logical component and phase boundaries between text
encoding, task-specific denoising, and media decoding:
`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
these components in one diffusion stage. These boundaries are useful for
profiling and offload; independent placement would require a future
multi-stage `PipelineConfig`. The stage mapping follows the implementation and
serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
MiniMax-H3 combines text and reference-media encoding, task-specific denoising,
and media decoding. Standard serving co-locates these components in one diffusion
stage. The opt-in
[disaggregated topology](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md)
runs Qwen and the video/audio VAE encoders in Stage 0, then sends their conditioning
to a Stage 1 containing the DiTs and VAE decoders. The stage mapping follows the
implementation and serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
notes that a detailed technical report is forthcoming.
#### Model architecture by component
@@ -149,16 +149,15 @@ notes that a detailed technical report is forthcoming.
| **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
Both DiT partitions remain behind one task-selection point; `task` chooses
which partition runs. In the current implementation, the three logical
components execute inside one diffusion stage. This is different from
Qwen3-Omni's three independently scheduled stages, even though both models
produce synchronized audio and video-like outputs.
which partition runs.
#### Representative component execution
The following table describes the standard single-stage deployment.
| Logical component | Batching | Attention and execution | Parallelism | Quantization |
| --- | --- | --- | --- | --- |
| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
| **Text Encoder** | Follows the diffusion request scheduler | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
| **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
| **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
@@ -76,6 +76,34 @@ preserve common `self` or `cross` policies. A model also decides whether its
path is compatible with automatic TRTLLM selection, because only the model
knows whether masking and packing satisfy the kernel contract.
### Model-owned implementation specializations
Use `Attention(..., impl_overrides={backend_name: implementation_class})`
when a model needs different computation within an existing backend contract.
The model defines the implementation in its own directory and passes it when
constructing the layer:
```python
# MyModelVSAImpl is a model-owned subclass of FastVideoVSAImpl.
self.attention = Attention(
num_heads=num_heads,
head_size=head_size,
softmax_scale=head_size**-0.5,
causal=False,
role="self",
impl_overrides={"FASTVIDEO_VSA": MyModelVSAImpl},
)
```
The override is applied after role/platform selection, keyed by the selected
backend's `get_name()`. It must subclass the **selected implementation**;
incompatible platform implementations raise `TypeError`. Unselected overrides
have no effect. Backend capabilities, constructor options, and shared
parallel dispatch remain unchanged and must be honored by the specialization.
Use a new backend for a new provider/capability contract. `custom_attention`
instead owns communication and requires `skip_sequence_parallel=True`.
## Adding or changing a backend
An implementation change is complete only when it:
@@ -91,8 +119,9 @@ An implementation change is complete only when it:
## FastVideo VSA model metadata
FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in
the original latent video grid. Wan integrations therefore attach the
FastWan's text-to-video and image-to-video modes both use `Wan22Pipeline`.
It operates on a flattened DiT sequence but partitions tokens in the original
latent video grid. Wan integrations therefore attach the
post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The
backend validates that `T * H * W` equals the sequence length, derives the
runtime block count, and selects the configured top-k key/value blocks for
@@ -104,3 +133,8 @@ the checkpoint contains those weights. Checkpoints without the projection do
not allocate or execute it. When top-k selects every block, native checkpoints
route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path
to retain checkpoint semantics.
MiniMax-H3 supplies `MiniMaxH3VSAImpl` through `impl_overrides`.
`models/minimax_h3/attention/fastvideo_h3.py` owns its prefix layout, video tiling, and
learned compression gate. Pooling, block-map construction, and tile64 provider
calls are shared through `attention/ops/block_sparse.py`.
+359 -1056
View File
File diff suppressed because it is too large. Load diff
+9 -2
View File
@@ -16,6 +16,7 @@ required_reviewers:
- "@yenuo26"
- "@NickCao"
primary_code_paths:
- vllm_omni/engine/omni_engine_base.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/engine/async_engine_utils.py
- vllm_omni/engine/orchestrator.py
@@ -28,8 +29,12 @@ related_code_paths:
- vllm_omni/engine/stage_pool.py
- vllm_omni/outputs/output_processor.py
- vllm_omni/distributed/omni_coordinator/**
# Duplex kernel design is owned by ../fullduplex.md; listed here because the
# orchestrator hosts the duplex control plane.
# The duplex siblings (DuplexOmniEngine / DuplexOrchestrator) and the
# engine-resident session framework are owned by ../fullduplex.md; listed
# here because they fill the template seams of OmniEngineBase and
# OrchestratorBase.
- vllm_omni/engine/duplex_omni_engine.py
- vllm_omni/engine/duplex_orchestrator.py
- vllm_omni/engine/duplex/**
depends_on:
- input_output_modality_contracts.md
@@ -48,6 +53,8 @@ validation_paths:
- tests/engine/test_cfg_companion_lifecycle.py
- tests/engine/test_rpc_result_router.py
- tests/engine/duplex/**
- tests/engine/test_duplex_orchestrator.py
- tests/engine/test_duplex_omni_engine.py
- tests/engine/test_duplex_import_boundary.py
upstream_refs:
- vllm.v1.engine.EngineCoreRequest
+8 -3
View File
@@ -19,7 +19,9 @@ required_reviewers:
primary_code_paths:
- vllm_omni/entrypoints/omni.py
- vllm_omni/entrypoints/async_omni.py
- vllm_omni/entrypoints/async_omni_base.py
- vllm_omni/entrypoints/omni_base.py
- vllm_omni/entrypoints/duplex_omni.py
- vllm_omni/entrypoints/cli/**
- vllm_omni/entrypoints/openai/**
- vllm_omni/entrypoints/openpi/**
@@ -27,14 +29,13 @@ primary_code_paths:
- vllm_omni/entrypoints/stage_utils.py
- vllm_omni/entrypoints/utils.py
- vllm_omni/entrypoints/duplex/**
- vllm_omni/entrypoints/duplex_request_client.py
- vllm_omni/clients/**
primary_path_exceptions:
- path: vllm_omni/entrypoints/openai/errors.py
owner: error_contracts.md
- path: vllm_omni/entrypoints/duplex/**
owner: ../fullduplex.md
- path: vllm_omni/entrypoints/duplex_request_client.py
- path: vllm_omni/entrypoints/duplex_omni.py
owner: ../fullduplex.md
- path: vllm_omni/clients/**
owner: ../fullduplex.md
@@ -42,7 +43,9 @@ related_code_paths:
- vllm_omni/errors.py
- vllm_omni/inputs/**
- vllm_omni/outputs/**
- vllm_omni/engine/omni_engine_base.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/engine/duplex_omni_engine.py
- vllm_omni/config/**
- vllm_omni/deploy/**
depends_on:
@@ -53,7 +56,9 @@ validation_paths:
- tests/entrypoints/test_omni_entrypoints.py
- tests/entrypoints/test_async_omni.py
- tests/entrypoints/test_async_omni_pause_sleep_routing.py
- tests/entrypoints/test_async_omni_duplex.py
- tests/entrypoints/test_duplex_omni.py
- tests/entrypoints/duplex/**
- tests/entrypoints/openai_api/test_duplex_api_server.py
- tests/entrypoints/test_serve.py
- tests/entrypoints/test_stream_finish_reason.py
- tests/entrypoints/openai/**
+12
View File
@@ -81,6 +81,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.
+2 -2
View File
@@ -25,8 +25,8 @@ CATEGORIES = {
"clients": {
"name": "Clients",
"description": (
"Client libraries that talk to a running vLLM-Omni server "
"(the full-duplex Realtime client and its per-model session presets)."
"Client libraries for full-duplex sessions: the Realtime WebSocket client, "
"the in-process client over `DuplexOmni`, and the per-model session presets."
),
},
"inputs": {
+1 -1
View File
@@ -107,7 +107,7 @@ th {
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
| `DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenBMB/MiniCPM-o-4_5.md) |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
| `HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
+12 -5
View File
@@ -145,15 +145,22 @@ interchangeably.
| Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) |
| Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) |
| Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) |
| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` (alias `WS /v1/duplex`) | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
| Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) |
| Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) |
All six routes are model- or configuration-dependent. In particular,
`/v1/realtime` is not full duplex unless the client sets `duplex=1` and the
deployment explicitly enables duplex sessions. Clients should also verify the
duplex capability payload because the query-parameter form falls back to the
ordinary realtime handler when duplex is unavailable.
`/v1/realtime` is full duplex when the model is a duplex model (its pipeline
declares a `duplex_plugin` and the deploy configuration sets
`session_mode: duplex`); a stock Realtime client needs no vendor query
parameter, and `duplex=0` is the explicit opt-out. Such a server serves the
websocket route plus `POST /v1/chat/completions`, and no other turn-based
HTTP route. On a duplex server that chat route is not the turn-based path:
each request runs on a short-lived duplex session, so it holds one of
`duplex_session.max_sessions` for its lifetime and answers at the model's
real-time pace -- see [Full Duplex](full_duplex_api.md). Clients should
also verify the duplex capability payload because the query-parameter form
falls back to the ordinary realtime handler when duplex is unavailable.
## Related Endpoints
+47 -34
View File
@@ -15,40 +15,52 @@ capability gates, see the [Realtime Duplex API guide](realtime_duplex_api.md).
| Endpoint | Protocol | Recommended use |
| --- | --- | --- |
| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients |
| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing |
Both endpoints use the same duplex engine and require the same model-side
adapter. Prefer the Realtime projection unless the native lifecycle events are
specifically required.
| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style events (the normative contract) | Applications and browser clients |
| `WS /v1/duplex` | Alias of `/v1/realtime?duplex=1` (same protocol, same handler) | Clients that prefer a dedicated path |
| Python: `DuplexOmni` / `InlineDuplexClient` | Typed commands and events in-process | Embedding the model without a server |
## Enable Full Duplex
The route becomes usable only when the deployment configuration explicitly
sets:
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
`/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`).
`/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
alongside the live websocket sessions, with the request options the chat
service supports. It is wired only when the model's plugin declares
`DuplexCapabilities.supports_chat_completions`; on any other duplex model the
route reports "not available". A model that should not serve the route at all
lists it in the deploy configuration's `endpoint_restrictions`.
The deploy configuration of such a model must agree:
```yaml
session_mode: duplex
```
The selected model pipeline must also provide a duplex serving adapter.
Model-native deployments configure the corresponding engine runtime extension
and control plane as part of their registered pipeline.
A duplex model started with a deploy configuration that does not set it fails
at startup rather than falling back to turn-based serving.
!!! warning
`WS /v1/duplex` fails with `Duplex API is not available` when duplex is not
enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary
turn-based realtime handler when no duplex handler exists. Confirm that
On a deployment that is *not* duplex, `WS /v1/duplex` fails with
`Duplex API is not available` and `/v1/realtime?duplex=1` falls back to the
ordinary turn-based realtime handler. Confirm that
`session.created.session.capabilities` is present before treating the
connection as full duplex.
The current unified-runtime integrations are:
- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5.yaml`
(`session_mode: duplex`);
- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex;
- Nemotron VoiceChat, via its registered duplex plugin package.
**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`) is the only model
served over this endpoint today. PersonaPlex and Nemotron VoiceChat still carry
their pre-framework duplex code: their pipelines declare no `duplex_plugin`, so
they run turn-based until the follow-up PRs port them to the plugin contract
(RFC [vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)).
JoyVL is a separate HTTP interaction orchestrator and does not use these
WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
@@ -101,22 +113,23 @@ deferred, treated as a short acknowledgement, or used to interrupt output.
The `session.created` payload includes capability fields such as
`supports_barge_in`, `supports_playback_ack`, `supports_multi_session`,
`supports_session_resume`, and `chunk_period_ms`. Treat this payload as the
runtime contract.
runtime contract and branch on the flags, never on the model name: a model
that supports native overlapping speech may still advertise
`supports_barge_in=false` when destructive output interruption and model-state
rewind are not validated for it. Capacity and session-resume behavior also
depend on the selected deployment configuration. The per-model table lives in
the [Realtime Duplex API guide](realtime_duplex_api.md#capability-negotiation-by-model).
For example, PersonaPlex supports native overlapping speech but currently
advertises `supports_barge_in=false`; destructive output interruption and
model-state rewind have not been validated for that integration. Capacity and
session-resume behavior also depend on the selected deployment configuration.
## Python API
## Native Protocol
`vllm_omni.entrypoints.duplex_omni.DuplexOmni` runs the same engine-resident
sessions in-process: `open_session()` returns a `DuplexSessionHandle` whose
`submit()` takes typed `DuplexCommand` objects and whose `events()` yields
typed `DuplexEvent` objects (each with a `to_realtime()` wire rendering).
`vllm_omni.clients.inline_duplex.InlineDuplexClient` exposes that handle
behind the `DuplexClient` API. See the
[Realtime Duplex API guide](realtime_duplex_api.md#using-the-python-api).
`WS /v1/duplex` exposes lower-level lifecycle names including
`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`,
and `session.close`. It returns native session, input, response, overlap, and
error events. This protocol is experimental and may evolve with the runtime;
applications should use the provided Realtime client where possible.
See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo),
[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex),
See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo)
and the [full-duplex runtime design](../design/fullduplex.md)
for model-specific validation and architecture details.
+214 -135
View File
@@ -1,17 +1,18 @@
# Realtime Duplex API
vLLM-Omni serves full-duplex speech models over a WebSocket endpoint,
`/v1/realtime`, that speaks the OpenAI Realtime vocabulary plus a
`/v1/realtime?duplex=1`, that speaks the OpenAI Realtime vocabulary plus a
set of duplex extensions: the model decides when to listen and when to speak,
the user can talk over the assistant, playback progress is reported back so
history stays honest, and a dropped connection can resume the same session.
This page covers how to run a duplex deployment, how to drive it from Python
with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract.
The handler is enabled by either a model-native `session_mode: duplex`
deployment or an explicit `duplex_session` configuration for turn-based
Server VAD. The runtime architecture is described in
[Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
The endpoint is served for models that ship a duplex plugin — currently
MiniCPM-o 4.5 only — and just for deploy configurations that declare
`session_mode: duplex`. PersonaPlex and Nemotron VoiceChat are ported to the
plugin contract in follow-up PRs and are not served over this endpoint yet. The runtime architecture
is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
## Quick Start
@@ -26,11 +27,38 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
```
`vllm_omni/deploy/minicpmo_4_5.yaml` declares `session_mode: duplex` and
`duplex_session.max_sessions: 4`, which mounts the WebSocket routes
`ws://<host>:8099/v1/realtime` (this page) and `ws://<host>:8099/v1/duplex`
(the native dialect for raw-protocol clients) next to the usual HTTP API.
PersonaPlex and Nemotron VoiceChat use `vllm_omni/deploy/personaplex.yaml` and
`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml` in the same way.
`duplex_session.max_sessions: 4`. Because the MiniCPM-o 4.5 pipeline declares a
`duplex_plugin`, `vllm-omni serve` runs it through `DuplexOmni`: the server
mounts `ws://<host>:8099/v1/realtime?duplex=1` (this page; `ws://<host>:8099/v1/duplex`
is an alias of the same route), `POST /v1/chat/completions`, `/v1/models` and
`/health`, and nothing else. The remaining turn-based HTTP routes (speech,
batch, embeddings, video, ...) are not served by a duplex model.
`/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:
- **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.
### Run the Example Client
@@ -81,25 +109,24 @@ ships a preset:
| Preset | Input / output audio | What it sets |
| --- | --- | --- |
| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `extra_body.native_duplex=True` (opt in to the model-native lane), `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip |
| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip |
| `vllm_omni.clients.personaplex.create_duplex_session_config(voice="NATF2.pt", persona="")` | 24 kHz `pcm_f32le` / 24 kHz `pcm16` | bundled `.pt` voice prompt and the persona as `instructions` |
| `SessionConfig(...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | model-neutral defaults; pass `extra_body`, `turn_detection`, `overlap_policy`, `playback_commit_policy`, `instructions`, `voice`, `temperature` yourself |
Keyword overrides on a preset replace the corresponding `SessionConfig`
field, e.g. `create_duplex_session_config(ref_audio=..., temperature=0.6)`.
Constructor options: `session_id` (names the session the client creates; it
does not resume or take over an existing session — entering the client always
performs the `session.update` handshake), `reconnect` (a `ReconnectPolicy`,
default five attempts with 0.25–4 s jittered backoff; `None` disables
auto-resume), `heartbeat_interval_s` (default 30 s; `None` disables the lease
heartbeat), `handshake_timeout_s`, and `connect` (a custom transport factory,
used by tests and benchmarks). Resume is supported only as automatic reconnect
within the same `DuplexClient` instance (see
[Reconnect and resume](#reconnect-and-resume)); re-attaching from a new client
or process requires the wire-level `session.resume` handshake
(`resume_token`, `incarnation`, `last_received_server_event_seq`), which this
client does not expose.
Constructor options: `reconnect` (a `ReconnectPolicy`, default five
attempts with 0.25–4 s jittered backoff; `None` disables auto-resume),
`heartbeat_interval_s` (default 30 s; `None` disables the lease heartbeat),
`handshake_timeout_s`, and `connect` (a custom transport factory, used by
tests and benchmarks). The session id is allocated by the server and becomes
`client.session_id` once `session.created` arrives; a client cannot choose
it. Resume is supported only as automatic reconnect within the same
`DuplexClient` instance (see [Reconnect and resume](#reconnect-and-resume));
re-attaching from a new client or process requires the wire-level
`session.resume` handshake (`resume_token`,
`last_received_server_event_seq`), which this client does not expose.
### Stream audio and commit a turn
@@ -195,7 +222,7 @@ the stamped `event_id`.
On a transport drop the client reconnects with the `ReconnectPolicy` backoff
and sends `session.resume` from the last acknowledged `server_event_seq`.
Replayed events are deduplicated, consumers see one `ConnectionResumed`
event, and `client.incarnation` / `client.resume_token` are refreshed. The
event, and `client.resume_token` is refreshed. The
server keeps a detached session alive for `disconnect_grace_s` (30 s by
default); after that a resume fails with `session_resume_expired`. Resume is
gated by `capabilities.supports_session_resume` (MiniCPM-o 4.5 and Nemotron
@@ -238,9 +265,67 @@ definitions out of the dependency-free client. `audio_bytes(response_id)`, `coun
`openai-realtime-duplex` benchmark backend (`vllm-omni bench serve`) is built
on the same helpers.
## Using the Python API
`vllm_omni.entrypoints.duplex_omni.DuplexOmni` is the in-process Python API
of a duplex model: sessions run inside the engine, and the class only opens,
resumes and closes them and pipes typed commands in and typed events out
through a `DuplexSessionHandle`.
```python
from vllm_omni.engine.duplex.commands import AppendAudio, Commit
from vllm_omni.engine.duplex.events import AudioDelta, SessionClosed
from vllm_omni.entrypoints.duplex_omni import DuplexOmni
omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True)
async with await omni.open_session({"modalities": ["audio", "text"], "ref_audio": ref_audio_data_url}) as session:
async def consume():
async for event in session.events(): # typed DuplexEvent objects
if isinstance(event, AudioDelta):
play(event.audio)
elif isinstance(event, ResponseDone):
break # one turn is enough here
task = asyncio.create_task(consume())
await session.submit(AppendAudio(audio=pcm16_bytes, format="pcm16", sample_rate_hz=16000))
await session.submit(Commit())
await task
# leaving the block closes the session (the last event is SessionClosed)
```
`open_session` takes a `DuplexSessionConfig` or the same Realtime `session`
object the WebSocket handshake takes; the session id is always allocated by
`DuplexOmni`. `DuplexSessionHandle` offers `submit(command)` plus one
convenience wrapper per command (`append_audio`, `commit`, `create_response`,
`cancel_response`, `barge_in`, `ack_playback`, `update`, `create_item`,
`delete_item`, `truncate_item`, `heartbeat`, ...). `events()` is a
single-consumer async iterator that ends after `session.closed` /
`session.expired`; rejected commands come back as `ErrorEvent`s on it.
Every event renders the wire JSON with `to_realtime()`, so anything written
against the WebSocket protocol works unchanged on the typed stream.
`vllm_omni.clients.inline_duplex.InlineDuplexClient` wraps a `DuplexOmni`
behind the `DuplexClient` API, so the same application code runs in-process
(no server) or over the network:
```python
from vllm_omni.clients.inline_duplex import InlineDuplexClient
from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config
config = create_duplex_session_config(ref_audio=audio_data_url(Path("reference_voice.wav")))
async with InlineDuplexClient(omni, model="openbmb/MiniCPM-o-4_5", config=config) as client:
await client.stream_pcm(question)
await client.commit()
async for response in client.responses():
...
```
`examples/online_serving/barge_in_client.py --inline` drives the barge-in
scenario this way. The inline client has no transport, so it has no
reconnect, resume, heartbeat or event-acknowledgement machinery.
## API Reference
This section is the complete wire contract of `/v1/realtime`: how
This section is the complete wire contract of `/v1/realtime?duplex=1`: how
the dialect relates to the OpenAI Realtime protocol, which surfaces each
model's capability flags gate, a catalogue of every client-to-server and
server-to-client event, and a JSON example for each. `DuplexClient` speaks
@@ -253,21 +338,26 @@ its protocol tests.
### Endpoint and transport
The endpoint is `ws(s)://<host>/v1/realtime`. Bare connections select the
session handler when the deployment declares either `session_mode: duplex`
or `duplex_session`. `?duplex=1` remains a compatibility alias, while
`?duplex=0` selects the legacy Realtime handler. Optional query parameters
are `model`, `session_id`,
`autostart` (`0` means resume-only), `resume`, and `native_duplex`
(`minicpmo45_native_duplex` is accepted as a deprecated alias and folded
into the canonical name). Every message is one JSON object per WebSocket
The endpoint is `ws(s)://<host>/v1/realtime`. On a deployment whose pipeline
declares a `duplex_plugin` and whose deploy configuration declares
`session_mode: duplex`, a bare connection is already a duplex session, so a
stock Realtime client needs no vendor query parameter. `?duplex=1` (`1`,
`true` or `on`) is the spelling this page uses and stays valid;
`?duplex=0` (`0`, `false`, `off`) explicitly opts out and selects the
turn-based Realtime handler, which a duplex server does not mount --
that connection is answered with `Realtime API is not available`.
`ws(s)://<host>/v1/duplex` is an alias that always selects the duplex
handler. Optional query parameters are `model`, `autostart` (`0` means
resume-only) and `resume`. Every message is one JSON object per WebSocket
text frame, discriminated by `type`. Inbound events are applied in arrival
order through one per-session mailbox; outbound events preserve that order
and carry a monotonically increasing `server_event_seq`, which is the replay
cursor for `session.resume`. Client events may carry an `event_id`, echoed in
`error.error.event_id` when the event is rejected.
Session identity is `(session_id, incarnation)`; `attachment_generation`
Session identity is the server-allocated `session_id` (`duplex-<uuid>`,
announced in `session.created`; a `session_id` the client puts in the session
object is ignored, and ids are never reused); `attachment_generation`
names the current socket; `epoch` is the interruption fence that advances on
every cancel, clear, or barge-in; `turn_id` counts committed model turns;
`response_id` (item id `item_<response_id>`) names one assistant utterance.
@@ -281,12 +371,13 @@ with no OpenAI counterpart.
The event vocabulary is uniform, but several surfaces are gated by the
`capabilities` object the server returns in `session.created`; a client must
branch on those flags rather than on the model name. The current model
plugins advertise:
branch on those flags rather than on the model name. MiniCPM-o 4.5 is the
only model on the plugin contract today; the other two columns record what
their integrations advertise once the follow-up PRs port them:
| Capability | MiniCPM-o 4.5 | PersonaPlex | Nemotron VoiceChat | Gated surface |
| --- | --- | --- | --- | --- |
| `implementation_level` | `model_native_duplex` (when `extra_body.native_duplex=true`) | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak`; the chat-fallback lane otherwise |
| `implementation_level` | `model_native_duplex` | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak` (constant: every duplex model is model-native) |
| `chunk_period_ms` | 1000 | 80 | 80 | the model unit that `response.listen` decisions and camera frames align to |
| `supports_session_resume` | yes | no | yes | `session.resume`, `session.resumed`, `session.replaced`, replay after a transport drop |
| `supports_barge_in` | yes | no | no | `barge_in`, `turn.signal{event:"barge_in"}`, `overlap_policy=barge_in_on_speech`, `turn_detection.server_vad` |
@@ -316,14 +407,8 @@ work unmodified. The overlap falls into the three tiers used by the event catalo
| Object shapes | `realtime.session`, `realtime.item`, `realtime.response`; item content parts `input_text` / `input_audio` / `output_text` / `output_audio`; `function_call` and `function_call_output` items; `previous_item_id` chaining; `output_index` / `content_index` addressing |
| Sequencing | `response.created → output_item.added → content_part.added → deltas → *.done → content_part.done → output_item.done → response.done → rate_limits.updated`; `speech_started → speech_stopped → committed` |
Audio/transcript output uses only the current (non-beta) OpenAI event names
(`response.output_audio.delta`, not the deprecated `response.audio.delta`),
matching `openai.types.realtime` in the official `openai` Python SDK and what
clients built against it (e.g. `livekit-plugins-openai`'s non-Azure code path)
expect; a client still on the old beta SDK types would not recognize these
events. Session fields and `conversation.item.*` events, in contrast, are
deliberately dual: the server emits and accepts **both** the beta and the GA
spellings at once (`output_modalities` alongside `modalities`,
The server deliberately emits and accepts **both** the beta and the GA
spellings at once (`response.output_audio.delta` alongside `output_modalities`,
`conversation.item.added` alongside `conversation.item.created`) so either
generation of OpenAI client parses the stream.
@@ -333,7 +418,7 @@ A stock client ignores the extra keys; the extensions are additive.
| Message | Extension |
| --- | --- |
| `session.created` / `session.updated` | top-level `incarnation`, `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` |
| `session.created` / `session.updated` | top-level `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` |
| `response.created` / `response.done` / `response.listen` | `response_id` at top level; the raw duplex event under `response.metadata` (`duplex_event`); `status_details.reason` uses vLLM reasons (`barge_in`, `client_cancelled`, `new_response`, …) |
| `response.output_audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` |
| `response.speak` (inserted before the first delta) | not an OpenAI event, but rides the OpenAI response envelope (`response_id`, `item_id`, `output_index`, `content_index`) |
@@ -349,14 +434,9 @@ Semantic divergences hidden behind shared names:
before any commit; a commit may end in `response.listen` and no
response, and the model may open a response with no commit at all.
OpenAI: commit ⇒ item, `response.create` ⇒ exactly one response.
- With `turn_detection={"type":"server_vad"}`, omitted `interrupt_response`
defaults to `false` on turn-based models such as Qwen3-Omni, enabling
endpointing and automatic responses without interruption. Explicit `true`
is rejected there. Model-native duplex keeps the `true` default and rejects
`false`; its VAD implies `overlap_policy=barge_in_on_speech`.
This capability-dependent default differs from OpenAI's `true` default;
the effective configuration is returned in `session.created` / `session.updated`.
`semantic_vad` is unsupported.
- `turn_detection.interrupt_response=false` is rejected;
`create_response` is ignored; `semantic_vad` is unsupported; the VAD
runs per session (Silero) and implies `overlap_policy=barge_in_on_speech`.
- `rate_limits.updated` is always an empty list (compatibility only).
- Cancellation never reuses a `response_id`; truncation is driven by
`playback.ack` as well as `conversation.item.truncate`.
@@ -365,11 +445,11 @@ Semantic divergences hidden behind shared names:
| Area | Messages / fields |
| --- | --- |
| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `incarnation`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume`, `native_duplex` |
| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume` |
| Turn-taking | `response.speak`, `response.listen`, `overlap.decision`, `overlap_policy`, `barge_in`, `turn.signal`, `input.cancel`, `input.text.append`, `epoch` / `turn_id` fencing, `force_listen` |
| Playback truth | `playback.ack` / `playback.acknowledged`, `playback_commit_policy` (`ack_only` \| `commit_all_on_done`), `playback{generated_ms, sent_ms, played_ms, committed_ms}`, `history_committed`, `audio_text_marks` |
| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `native_duplex` — deprecated alias `minicpmo45_native_duplex` is folded into it, `force_listen_count`, `duplex_initial_user_text`) |
| Diagnostics | `runtime.control`, `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/entrypoints/duplex/realtime_state.py`) |
| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `force_listen_count`, `duplex_initial_user_text`) |
| Diagnostics | `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/engine/duplex/events.py`) |
#### Consequences for clients
@@ -433,11 +513,15 @@ remains an approximation, not word-level audio alignment.
### Event catalogue
Every event type the Realtime route accepts or emits, with a worked example in the next section: 21
client→server and 42 server→client. Aliases (`push_chunk`,
`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`,
`audio.playback_ack`, `close_session`, `close`, `session_close`) share the
payload of their canonical event and are not listed separately; the alias
table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
client→server and 42 server→client. Aliases (`input_text.append`,
`push_text`, `signal_turn`, `audio.playback_ack`, `close_session`, `close`)
share the payload of their canonical event and are not listed
separately; the alias table is `translate_realtime_command` in
`vllm_omni/engine/duplex/realtime_commands.py`. The pre-Realtime WAV-append
aliases `push_chunk` and `input.audio.append` are **not** accepted on this
route: send `input_audio_buffer.append` with one of the supported input
formats (`pcm16`, `pcm_s16le`, `s16le`, `pcm_f32le`, `g711_ulaw`,
`g711_alaw`).
#### Client to server
@@ -447,18 +531,18 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o
| `session.heartbeat` | 3 | Keepalive; refreshes the engine session lease. |
| `session.event_ack` | 3 | Acknowledges received server events by `server_event_seq`; trims the resume replay journal. |
| `session.close` | 3 | Graceful close; answered with `session.closed`. |
| `session.resume` | 3 | Re-attach a live session on a new socket (`resume_token`, `incarnation`, `last_received_server_event_seq`). |
| `session.resume` | 3 | Re-attach a live session on a new socket (`session_id`, `resume_token`, `last_received_server_event_seq`). |
| `input_audio_buffer.append` | 2 | Append one audio chunk. OpenAI name; extended with `is_speech`, `video_frames`, `duration_ms`, `audio_end_ms`, per-event `format`/`sample_rate_hz`. |
| `input_audio_buffer.commit` | 2 | Seal the buffered utterance into a user item. OpenAI name; extended with `final`, `response_create`, `is_speech:false`. |
| `input.commit` | 3 | Alias of `input_audio_buffer.commit` accepted by the runner. |
| `input_audio_buffer.clear` | 1 | Drop un-committed input audio. |
| `input.cancel` | 3 | Cancel pending input; advances `epoch`. |
| `input.text.append` | 3 | Append text to the open input item (chat-fallback lane only). |
| `input.text.append` | 3 | Append text input; model-native sessions reject it with `native_text_append_unsupported`. |
| `conversation.item.create` | 1 | Add a completed user item (text, audio, or `function_call_output`) to history. |
| `conversation.item.retrieve` | 1 | Fetch a stored item by id. |
| `conversation.item.delete` | 1 | Delete a stored item by id. |
| `conversation.item.truncate` | 1 | Truncate an assistant item's audio/transcript at `audio_end_ms`. |
| `response.create` | 1 | Explicitly request a response with per-response overrides. |
| `response.create` | 1 | Explicitly request a response with per-response overrides. Requires committed audio input; otherwise `response_create_without_input`. |
| `response.cancel` | 1 | Cancel the active (or named) response; advances `epoch`. |
| `output_audio_buffer.clear` | 1 | Discard queued output audio; advances `epoch` (OpenAI: WebRTC-only, here also WebSocket). |
| `barge_in` | 3 | Explicit hard interrupt; requires `capabilities.supports_barge_in`. |
@@ -469,10 +553,10 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o
| Event type | Tier | Description |
| --- | --- | --- |
| `session.created` | 2 | Session opened. OpenAI name; adds `incarnation`, `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`epoch`, `turn_id`, `playback`, `capabilities`, ...). |
| `session.created` | 2 | Session opened. OpenAI name; adds `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`id` is the server-allocated session id, `epoch`, `turn_id`, `playback`, `capabilities`, ...). |
| `session.updated` | 2 | Echo of the effective session config after every `session.update` (same extended `session` object). |
| `session.heartbeat_ack` | 3 | Reply to `session.heartbeat`. |
| `session.closed` | 3 | Last event on the socket; `reason` ∈ `session_close`, `timeout`, `disconnect`, ... |
| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, ... |
| `session.resumed` | 3 | Resume accepted; carries the new `attachment_generation` and rotated `resume_token`; journaled events are replayed after it. |
| `session.replaced` | 3 | Sent to the superseded socket when another socket resumes the session. |
| `session.expired` | 3 | Engine lease reaped (`disconnect_grace_expired`, `idle_ttl_expired`); socket closes after it. |
@@ -508,8 +592,8 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o
| `duplex.function_call.done` | 3 | Raw fallback when a model tool call is malformed. |
| `output_audio_buffer.cleared` | 1 | Reply to `output_audio_buffer.clear`. |
| `playback.acknowledged` | 3 | Reply to `playback.ack`; returns the playback ledger and `history_committed`. |
| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`). |
| `runtime.control` | 3 | Diagnostic: an engine control signal was partially unsupported. |
| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`); details under `metadata`. |
| `turn.event` | 3 | A turn-taking signal the session accepted (`event`, resulting `turn_state`). |
| `error` | 2 | Error envelope. OpenAI shape and classes; vLLM-Omni error codes. |
### Message examples
@@ -528,7 +612,6 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
"event_id": "evt_0001",
"session": {
"model": "openbmb/MiniCPM-o-4_5",
"session_id": "sess_demo_01",
"modalities": ["audio", "text"],
"instructions": "You are a concise voice assistant.",
"input_audio_format": "pcm16",
@@ -541,7 +624,6 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
"idle_timeout_s": 300,
"extra_body": {
"auto_response": true,
"native_duplex": true,
"force_listen_count": 0
}
}
@@ -553,13 +635,12 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
```json
{
"type": "session.created",
"incarnation": 0,
"attachment_generation": 1,
"resume_token": "rt_3f9c…",
"session": {
"object": "realtime.session",
"type": "realtime",
"id": "sess_demo_01",
"id": "duplex-5f2a…",
"model": "openbmb/MiniCPM-o-4_5",
"state": "open",
"turn_state": "idle",
@@ -612,7 +693,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
`session.created` and every later `session.update`)
```json
{"type": "session.updated", "session": {"object": "realtime.session", "id": "sess_demo_01", "…": "…"}}
{"type": "session.updated", "session": {"object": "realtime.session", "id": "duplex-5f2a…", "…": "…"}}
```
`C→S session.heartbeat` / `S→C session.heartbeat_ack`
@@ -622,7 +703,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
```
```json
{"type": "session.heartbeat_ack", "session_id": "sess_demo_01"}
{"type": "session.heartbeat_ack", "session_id": "duplex-5f2a…"}
```
`C→S session.event_ack`
@@ -638,7 +719,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
```
```json
{"type": "session.closed", "session_id": "sess_demo_01", "reason": "session_close"}
{"type": "session.closed", "session_id": "duplex-5f2a…", "epoch": 3, "reason": "client_close", "event": {"type": "session.closed", "reason": "client_close"}}
```
`C→S session.resume` (new socket, `?duplex=1&autostart=0`)
@@ -646,8 +727,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
```json
{
"type": "session.resume",
"session_id": "sess_demo_01",
"incarnation": 0,
"session_id": "duplex-5f2a…",
"resume_token": "rt_3f9c…",
"last_received_server_event_seq": 40
}
@@ -659,8 +739,8 @@ events with `server_event_seq > 40` are replayed in order)
```json
{
"type": "session.resumed",
"session_id": "sess_demo_01",
"incarnation": 0,
"session_id": "duplex-5f2a…",
"session": {"id": "duplex-5f2a…", "state": "open", "epoch": 3, "turn_id": 2, "...": "the engine's current session object"},
"attachment_generation": 2,
"resume_token": "rt_a71e…"
}
@@ -669,19 +749,19 @@ events with `server_event_seq > 40` are replayed in order)
`S→C session.replaced` (sent to the superseded socket, which is then closed)
```json
{"type": "session.replaced", "session_id": "sess_demo_01", "attachment_generation": 1}
{"type": "session.replaced", "session_id": "duplex-5f2a…", "attachment_generation": 1}
```
`S→C session.expired`
```json
{"type": "session.expired", "session_id": "sess_demo_01", "incarnation": 0, "reason": "disconnect_grace_expired"}
{"type": "session.expired", "session_id": "duplex-5f2a…", "reason": "disconnect_grace_expired"}
```
`S→C session.resync_required`
```json
{"type": "session.resync_required", "session_id": "sess_demo_01", "reason": "journal_gap"}
{"type": "session.resync_required", "session_id": "duplex-5f2a…", "reason": "journal_gap"}
```
#### Input
@@ -772,7 +852,7 @@ silence/noise declaration (no runtime append; answered with `response.listen`)
"item_id": "item_9b2d…",
"event": {
"type": "input.committed",
"session_id": "sess_demo_01",
"session_id": "duplex-5f2a…",
"turn_id": 1,
"epoch": 0,
"history_len": 1,
@@ -802,7 +882,7 @@ deferred commit during an active response (`event.response_create_deferred`)
"type": "input_audio_buffer.committed",
"previous_item_id": "item_9b2d…",
"item_id": "item_c044…",
"event": {"type": "input.committed", "session_id": "sess_demo_01", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}}
"event": {"type": "input.committed", "session_id": "duplex-5f2a…", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}}
}
```
@@ -911,7 +991,7 @@ deferred commit during an active response (`event.response_create_deferred`)
"status_details": null,
"output": [],
"modalities": ["audio", "text"],
"metadata": {"duplex_event": {"type": "response.created", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "turn_id": 1}}
"metadata": {"duplex_event": {"type": "response.created", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "turn_id": 1}}
}
}
```
@@ -947,7 +1027,7 @@ deferred commit during an active response (`event.response_create_deferred`)
"item_id": "item_resp_01",
"output_index": 0,
"content_index": 0,
"metadata": {"session_id": "sess_demo_01", "epoch": 0, "model_speak": true}
"metadata": {"session_id": "duplex-5f2a…", "epoch": 0, "model_speak": true}
}
```
@@ -964,7 +1044,7 @@ deferred commit during an active response (`event.response_create_deferred`)
"format": "pcm16",
"sample_rate_hz": 24000,
"metadata": {
"session_id": "sess_demo_01",
"session_id": "duplex-5f2a…",
"epoch": 0,
"model_speak": true,
"end_of_turn": false,
@@ -1029,7 +1109,7 @@ deferred commit during an active response (`event.response_create_deferred`)
"status": "completed",
"status_details": null,
"output": [{"id": "item_resp_01", "object": "realtime.item", "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_audio", "transcript": "It is sunny, 24 degrees."}]}],
"metadata": {"type": "response.done", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}}
"metadata": {"type": "response.done", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}}
}
}
```
@@ -1055,9 +1135,9 @@ deferred commit during an active response (`event.response_create_deferred`)
```json
{
"type": "response.listen",
"session_id": "sess_demo_01",
"session_id": "duplex-5f2a…",
"epoch": 0,
"response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "sess_demo_01", "epoch": 0, "reason": "model_listen"}}
"response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "duplex-5f2a…", "epoch": 0, "reason": "model_listen"}}
}
```
@@ -1104,7 +1184,7 @@ function-call response (Nemotron VoiceChat tools)
raw fallback when a model tool call is malformed (`call_id`/`name` missing)
```json
{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "sess_demo_01", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}}
{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "duplex-5f2a…", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}}
```
#### Playback
@@ -1120,7 +1200,7 @@ raw fallback when a model tool call is malformed (`call_id`/`name` missing)
"type": "playback.acknowledged",
"event": {
"type": "playback.acknowledged",
"session_id": "sess_demo_01",
"session_id": "duplex-5f2a…",
"epoch": 0,
"item_id": "item_resp_01",
"played_ms": 1000,
@@ -1145,20 +1225,17 @@ truncating ack after an interruption (what was actually heard)
```json
{
"type": "overlap.decision",
"session_id": "sess_demo_01",
"session_id": "duplex-5f2a…",
"epoch": 0,
"policy": "listen_only",
"action": "listen",
"reason": "long_overlap_speech",
"overlap_ms": 1400,
"buffer_audio": true,
"defer_runtime_append": true,
"force_listen": true
"metadata": {"overlap_ms": 1400, "buffer_audio": true, "defer_runtime_append": true, "force_listen": true}
}
```
```json
{"type": "overlap.decision", "session_id": "sess_demo_01", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"}
{"type": "overlap.decision", "session_id": "duplex-5f2a…", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"}
```
server-VAD session (`session.update` fragment) and the annotated append the
@@ -1172,12 +1249,6 @@ translator forwards internally (shown for reference; clients never see it)
{"type": "input_audio_buffer.append", "audio": "…", "format": "pcm16", "sample_rate_hz": 16000, "is_speech": true, "force_listen": true, "vad": {"backend": "silero", "is_speech": true, "speech_active": true, "speech_started": true, "speech_stopped": false, "speech_probability": 0.93}}
```
`S→C runtime.control` (only when an engine signal was partially unsupported)
```json
{"type": "runtime.control", "session_id": "sess_demo_01", "epoch": 1, "result": {"…": "redacted engine control result"}}
```
#### Errors
`S→C error` — Realtime shape
@@ -1203,71 +1274,79 @@ translator forwards internally (shown for reference; clients never see it)
```
```json
{"type": "error", "error": {"type": "rate_limit_error", "code": "duplex_session_busy", "message": "duplex session capacity reached (max_sessions=4)"}}
{"type": "error", "error": {"type": "rate_limit_error", "code": "resource_exhausted", "message": "duplex_session_capacity_exhausted: limit=4"}}
```
Handshake-stage errors (before the Realtime projector is bound) use the
flat native shape:
```json
{"type": "error", "error": "Unknown or expired duplex session: sess_demo_01", "code": "session_resume_expired"}
{"type": "error", "error": "Unknown or expired duplex session: duplex-5f2a…", "code": "session_resume_expired"}
```
#### Coverage
The examples above cover every event type the Realtime route emits or
accepts (verified by diffing the `type` literals in
`vllm_omni/entrypoints/duplex/` against this document). Not shown, because
they never reach a `/v1/realtime?duplex=1` client unprojected: the native
`/v1/duplex` dialect (`session.create` / `open_session` / `session.config`,
`input.committed`, `input.cancelled`, `audio.cancelled`,
`response.output_audio.delta`, `response.text.delta`, `response.message`,
`function_call.done`) — their Realtime projections are the examples above
(see the name map at the end of this section) — and the runner-internal markers `__timeout__`,
`__disconnect__`, `__replaced_attachment__`. Input aliases (`push_chunk`,
`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`,
`audio.playback_ack`, `close_session`, `close`, `session_close`) share the
payload of their canonical event (see the alias table in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md)).
`vllm_omni/engine/duplex/events.py` and `vllm_omni/engine/duplex/commands.py`
against this document). Not shown, because they never reach a client
unprojected: the session-internal events the runner produces
(`input.committed`, `input.cancelled`, `audio.cancelled`,
`response.output_audio.delta`, `response.text.delta`, `function_call.done`)
— their Realtime projections are the examples above (see the name map at the
end of this section). Every public event is a typed `DuplexEvent` whose
`to_realtime()` is the wire object; every client event is parsed into a
typed `DuplexCommand` by `command_from_realtime()`.
#### Native to Realtime name map
#### Internal to Realtime name map
The native `/v1/duplex` dialect is what the session runner produces; the
Realtime projector renames and fans out these events before they reach a
`/v1/realtime?duplex=1` client.
The session runner emits these session-internal events; the Realtime
projector (`vllm_omni/engine/duplex/realtime_events.py`) renames and fans
them out into typed events before they reach a client.
| Native / internal event | Realtime projection |
| Internal event | Realtime projection |
| --- | --- |
| `session.create` (`open_session`, `session.config`) | first `session.update` |
| `input.committed` | `conversation.item.added` / `conversation.item.created`, `input_audio_buffer.committed`, `conversation.item.input_audio_transcription.completed`, `conversation.item.done` |
| `input.cancelled` | `input_audio_buffer.cleared` |
| `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.output_audio.delta` and `response.output_audio_transcript.delta` |
| `response.output_audio.done` / `response.output_text.done` | `response.output_audio.done` / `response.output_text.done` |
| `response.text.delta` | `response.output_text.delta` |
| `response.message` (chat-fallback raw chunk) | passed through |
| `audio.cancelled` | optional `output_audio_buffer.cleared`, then the cancelled terminal events and `response.done` with `status: "cancelled"` |
| `function_call.done` | `response.created`, `response.output_item.added` (`function_call` item), `response.function_call_arguments.delta` / `.done`, `response.output_item.done`, `response.done` |
| `playback.acknowledged` | `playback.acknowledged` (wrapped verbatim under `event`) |
| `runtime.control`, `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged |
| `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged |
## Known Limitations
- Only MiniCPM-o 4.5 is served over this endpoint today; PersonaPlex and
Nemotron VoiceChat arrive with the follow-up PRs that port them to the
plugin contract.
- Several surfaces are capability-gated per model (see *Capability
negotiation by model* above): PersonaPlex does not support session resume,
barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in
or audio truncation; camera frames are consumed only by MiniCPM-o 4.5; tool
calls are produced only by Nemotron VoiceChat.
- `turn_detection` supports `server_vad` or `null`, not `semantic_vad`.
Turn-based models support endpointing without interruption, including
`create_response=false` for explicit response creation. Model-native duplex
requires `interrupt_response=true` and `create_response=true`.
- A response is generated from committed **audio**, so the OpenAI text-prompt
shape does not drive one: `conversation.item.create` with an `input_text`
part adds the item to history, but a following `response.create` with no
committed audio is rejected with `response_create_without_input`, and
`input.text.append` is rejected with `native_text_append_unsupported`.
Text-to-speech has its own shape here — put the text in
`extra_body.duplex_initial_user_text` on `session.update`, then stream audio
units (silence is enough). The seeded turn is what the model answers.
- `turn_detection` supports only `server_vad` with `interrupt_response=true`;
`semantic_vad`, `interrupt_response=false`, and `create_response` are not
supported.
- `input_audio_transcription` and `input_audio_noise_reduction` are accepted
and echoed but no separate transcription or noise-reduction stage runs;
transcripts come from the model.
- `rate_limits.updated` is emitted for compatibility only and always carries
an empty list.
- Session capacity is bounded by `duplex_session.max_sessions` in the deploy
configuration; admission beyond it fails with `duplex_session_busy` or
`resource_exhausted`.
- The native `/v1/duplex` route speaks the internal dialect and does not
support `session.resume`; use `/v1/realtime?duplex=1` for the contract on
this page.
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`.
+47 -2
View File
@@ -6,6 +6,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
- **Gepard-1.0** (`nineninesix/gepard-1.0`) -- single-stage native-AR TTS with a 22.05 kHz NanoCodec. Zero-shot default voice only.
See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, MOSS-TTS-Nano, and Breeze-TTS-2.
@@ -40,11 +41,16 @@ vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
# CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request)
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \
--omni --port 8091 --trust-remote-code
# Gepard-1.0 (zero-shot default voice; packaged gepard.yaml)
vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \
--stage-init-timeout 900 \
--deploy-config vllm_omni/deploy/gepard.yaml
```
### Generate Speech
**Using curl:**
**Qwen3-TTS CustomVoice, using curl:**
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
@@ -56,7 +62,24 @@ curl -X POST http://localhost:8091/v1/audio/speech \
}' --output output.wav
```
**Using Python:**
**Gepard-1.0, using curl:**
Gepard accepts `input` (required), `voice` (`"default"`), `response_format`,
`stream` / `stream_format`, `max_new_tokens`, and `seed`. Unsupported fields
include `speed`, `language`, `instructions`, `task_type`, `ref_audio`,
`ref_text`, `extra_params`, and `word_timestamps`.
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, this is Gepard speaking.",
"voice": "default",
"seed": 7
}' --output output.wav
```
**Qwen3-TTS CustomVoice, using Python:**
```python
import httpx
@@ -761,6 +784,28 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
| ------- | ------------- |
| `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. |
### Gepard-1.0
| Model | Description |
| ----- | ----------- |
| `nineninesix/gepard-1.0` | Zero-shot native-AR TTS. 22.05 kHz mono. `voice` must be omitted or `"default"`. |
Gepard request fields:
| Field | Behavior |
| ----- | -------- |
| `input` | Required. Empty/whitespace-only returns 400. |
| `voice` | Omitted or `"default"` only. Other values 400. |
| `response_format` | `wav` default. Non-streaming: `wav`/`pcm`/`flac`/`mp3`. `opus` returns 400 (22.05 kHz is not an Opus sample rate). Streaming: `pcm`/`wav` only. |
| `stream` / `stream_format` | SSE (`speech.audio.*`) and raw `audio` byte streaming. |
| `speed` | Must be `1.0`. Gepard has no native speed control. |
| `max_new_tokens` | Frame budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms). Default 1000 from deploy YAML; adapter bounds 1..4096. |
| `seed` | Optional. Reaches the in-model 32-head sampler. The packaged YAML currently pins `seed: 42`, so serving is deterministic by default until that pin is removed. |
| `extra_params` | Any key returns 400, including `temperature`/`top_p`/`top_k`. |
| Cloning / style fields | `ref_audio`, `ref_text`, `speaker_embedding`, `instructions`, `language`, `task_type`, `word_timestamps`, and similar declared-but-unsupported fields return 400. |
See the [Gepard section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#gepard-10) for launch commands. Native-AR recompute preemption is a known limitation under concurrency.
### OmniVoice
| Model | Description |
+8
View File
@@ -50,6 +50,14 @@ WebSocket /v1/video/chat/stream
| Server -> Client | `session.done` | none | Session closed. |
| Server -> Client | `error` | `message` | Recoverable protocol or generation error. |
### Frame Consumption Reporting
When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt.
Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path.
**Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`.
### `session.config` Fields
| Field | Type | Default | Description |
+10
View File
@@ -109,6 +109,16 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
The final content is available from `/v1/videos/{video_id}/content` after the
job status becomes `completed`.
`queued` means the request is still waiting for diffusion scheduler admission.
The status changes to `in_progress` when the scheduler first selects the
request for execution.
`DELETE /v1/videos/{video_id}` issues a bounded engine abort
(`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend
task. Cancellation cleanup is also bounded and best-effort: it confirms
the abort was queued, and the current request batch may still drain.
The job is then re-read so a completed save is not orphaned.
### Synchronous Response
`POST /v1/videos/sync` blocks until generation finishes and returns raw video
@@ -1,24 +1,33 @@
# FastVideo VSA
FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's
self-attention by partitioning the post-patch latent grid into spatiotemporal
FastVideo Variable Sparse Attention (VSA) partitions the post-patch latent grid into spatiotemporal
blocks. For every query block, VSA scores the key/value blocks and computes
attention only against the selected top-k blocks.
The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
## Supported models
MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route
through the same backend. See the
[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry,
supported topologies, and commands.
| Model / checkpoint | Required adapter | Tasks | Sequence parallelism |
| --- | --- | --- | --- |
| `FastVideo/FastWan2.2-TI2V-5B-Diffusers` | None | T2V, I2V through `Wan22Pipeline` | Disabled |
| `MiniMaxAI/MiniMax-H3` | [FastH3 VSA (recipe)](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter) | T2VA | Disabled or pure Ulysses |
| Wan I2V-14B, S2V, VACE | — | Unsupported | — |
VSA is a CUDA-only, explicitly selected backend. It requires the
`fastvideo-kernel` package and currently supports non-causal self-attention
with equal query and key/value sequence lengths. Unsupported shapes, masks,
dtypes, sequence-parallel execution, or kernel failures fall back to
`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
sequence-parallel context is one of those fallbacks; the H3 route supports pure
Ulysses and rejects ring or all-gather sequence parallelism at startup.
## Installation
Install vLLM-Omni from this checkout with the `vsa` extra:
```bash
uv pip install -e '.[vsa]'
```
The extra installs the tested kernel dependency automatically. Prebuilt kernels require
Linux, Python 3.12, and glibc 2.34 or newer (x86-64 or aarch64). The full
FastVideo framework and provider environment variables are not required.
`FASTVIDEO_VSA` selects the attention algorithm. On SM120, the H3 integration
currently executes FastVideo's 64-token Triton block-sparse kernel; it does not
dispatch to FlashInfer. The `vsa` extra installs `fastvideo-kernel==0.3.4` for
this execution path.
## Enable the backend
@@ -27,24 +36,18 @@ flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained
for every query block:
```bash
vllm-omni serve <model> \
vllm serve <model> --omni \
--diffusion-attention-backend FASTVIDEO_VSA \
--fastvideo-vsa-topk 64
```
The backwards-compatible environment variable selects the backend with the
default `topk=64`:
<details markdown="1">
<summary>Alternative configuration formats</summary>
Equivalent structured configuration:
```bash
export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA
vllm-omni serve <model>
```
To tune top-k, pass the CLI backend and top-k flags together as shown above,
or use the equivalent structured configuration:
```bash
vllm-omni serve <model> \
vllm serve <model> --omni \
--diffusion-attention-config \
'{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}'
```
@@ -73,8 +76,15 @@ stages:
fastvideo_vsa_topk: 64
```
</details>
## Choose top-k
H3 uses 64-token video blocks and keeps
all prefix blocks; see the [FastH3 VSA recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter).
### Wan top-k behavior
At runtime the backend logs the sequence shape and derived block count:
```text
@@ -117,8 +127,16 @@ selecting VSA does not turn a native checkpoint into a distilled model.
## Verify routing and fallback
Check the startup and first-forward logs instead of assuming that selecting
the backend guarantees sparse execution:
VSA is a CUDA-only, explicitly selected backend. It requires the
`fastvideo-kernel` package and currently supports non-causal self-attention
with equal query and key/value sequence lengths. Unsupported shapes, masks,
dtypes, sequence-parallel execution, or kernel failures fall back to
`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
sequence-parallel context is one of those fallbacks; the H3 route supports pure
Ulysses and rejects ring or all-gather sequence parallelism at startup.
H3 accelerator faults propagate instead of attempting dense recovery.
On the Wan route, check the startup and first-forward logs:
- `route=VSA` means top-k block selection is active.
- `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all
@@ -126,6 +144,10 @@ the backend guarantees sparse execution:
- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
executed; the warning includes the reason.
For H3, check `FastH3 adapter active` at startup and
`FASTVIDEO_VSA H3 routing` during DiT execution, as described in the
[model recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter).
The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks,
standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
and no active sequence-parallel context. The MiniMax-H3 route uses 64-token
@@ -18,6 +18,7 @@ For the full list of supported architectures across all modalities, see
| --- | --- | --- | --- | --- | --- |
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
| Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — |
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
@@ -92,7 +93,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
```
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others).
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for Gepard, 24 kHz for the others).
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md).
@@ -212,6 +213,61 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
---
## Gepard-1.0
Single-stage native AR TTS at 22.05 kHz mono. Zero-shot only: omit `voice` or pass `"default"`. Voice cloning from reference audio is not available yet.
### Prerequisites
Same NeMo NanoCodec install as the [offline Gepard section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#gepard-10). On a host whose CUDA toolkit cannot build kernels — no `nvcc`/`ninja`, or a consumer Blackwell (`sm_120`) card — also `export VLLM_USE_FLASHINFER_SAMPLER=0` before launch. That env is not a deploy-YAML field.
### Launch
```bash
vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \
--stage-init-timeout 900 \
--deploy-config vllm_omni/deploy/gepard.yaml
# or:
./gepard/run_server.sh
```
`--stage-init-timeout 900` matches the online e2e fixture and the offline example. Serve defaults to 300s, which is often too short for a cold download of the talker plus NanoCodec.
The packaged `vllm_omni/deploy/gepard.yaml` must be passed with `--deploy-config`. The checkpoint self-identifies as `qwen3_5_text`, so omitting the YAML launches a diffusion fallback instead of the Gepard pipeline. The YAML sets `async_chunk: false`, `max_num_seqs: 4`, and currently pins `seed: 42`, so serving is deterministic by default until that YAML seed is removed. Pass an explicit per-request `seed` in tests and clients rather than depending on either default.
### Sending requests
```bash
python examples/online_serving/text_to_speech/gepard/speech_client.py \
--text "Hello, this is Gepard speaking."
python examples/online_serving/text_to_speech/gepard/speech_client.py \
--text "Hello, this is Gepard speaking." --seed 7 --stream --output output.pcm
```
Gepard accepts `input` (required), `voice` (`"default"`), `response_format`, `stream` / `stream_format`, `max_new_tokens`, and `seed`. Unsupported fields include `speed`, `language`, `instructions`, `task_type`, `ref_audio`, `ref_text`, `extra_params`, and `word_timestamps`.
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, this is Gepard speaking.",
"voice": "default",
"seed": 7
}' --output output.wav
```
### Notes
- Output: 22.05 kHz mono. `max_new_tokens` is a **frame** budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms at 21.5 fps; adapter bounds 1..4096).
- Supported request fields: `input` (required), `voice` (`default` only), `response_format` (`wav` default; `wav/pcm/flac/mp3` non-streaming; `opus` 400 because 22.05 kHz is not an Opus sample rate; streaming `pcm/wav` only), `stream` / `stream_format`, `max_new_tokens`, `seed`.
- Unsupported: `speed`, `extra_params` (including `temperature`/`top_p`/`top_k`), `ref_audio`, `ref_text`, `speaker_embedding`, `task_type`, `instructions`, `language`, and `word_timestamps`.
- Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands.
- Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI.
- For the request-field table, see the [Speech API Gepard section](https://github.com/vllm-project/vllm-omni/blob/main/docs/serving/speech_api.md#gepard-10).
---
## higgs-audio v2
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
@@ -19,7 +19,7 @@ list of supported architectures across all modalities, see
| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | 2 (talker + codec) | ✓ | ✓ (async chunk) | voice design / voice direction (`--instruction`) | 24 kHz |
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning later) | ✓ (online) | zero-shot | 22.05 kHz |
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
@@ -167,7 +167,7 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
```
### Architecture
```
```text
Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
(Llama-based) (32k vocab) (Flow Matching)
```
@@ -491,8 +491,8 @@ Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Ref
### Notes
- Output: 22.05 kHz mono WAV.
- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
- Offline and online: `examples/offline_inference/text_to_speech/gepard/end2end.py` and `/v1/audio/speech` (see the [online Gepard section](../../online_serving/text_to_speech/README.md#gepard-10)).
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it). Online serving **must** pass `--deploy-config vllm_omni/deploy/gepard.yaml` because the checkpoint self-identifies as `qwen3_5_text`.
- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
+143 -85
View File
@@ -14,7 +14,10 @@ Scenario (see ``barge_in_client_flow.md`` for the diagram)
The demo is model-neutral: it drives any duplex model through the public
:class:`vllm_omni.clients.duplex.DuplexClient` over
``/v1/realtime?duplex=1``. Model-specific session shape comes from the
``/v1/realtime?duplex=1``, or — with ``--inline`` — through
:class:`vllm_omni.clients.inline_duplex.InlineDuplexClient` on an in-process
:class:`vllm_omni.entrypoints.duplex_omni.DuplexOmni` (no server needed; the
model loads in this process). Model-specific session shape comes from the
per-model presets under ``vllm_omni.clients.<model>`` selected with
``--preset`` (default ``minicpmo_4_5``, whose model decides listen/speak on
its own — no VAD).
@@ -31,6 +34,15 @@ then:
--interrupt-wav follow_up_16k.wav \
--output-dir ./duplex_out
or, without a server, in-process (``--deploy-config`` is optional and defaults
to the model's bundled duplex profile):
python barge_in_client.py --inline \
--model /path/to/MiniCPM-o-4_5 \
--ref-audio /path/to/reference_voice.wav \
--question-wav question_16k.wav \
--interrupt-wav follow_up_16k.wav
Input WAVs must be mono 16 kHz PCM16; when a preset's session takes a
different capture format (the PersonaPlex preset streams 24 kHz float32),
the audio is converted before streaming. ``--ref-audio`` is required by the
@@ -55,6 +67,7 @@ if str(REPO_ROOT) not in sys.path:
from vllm_omni.clients.duplex import ( # noqa: E402
AudioFormat,
DuplexClient,
DuplexClientBase,
EventCollector,
SessionConfig,
acknowledge_collected_playback,
@@ -139,6 +152,39 @@ def _fold_responses(collector: EventCollector) -> dict[str, dict[str, object]]:
return out
def _make_client(args: argparse.Namespace, config: SessionConfig) -> DuplexClientBase:
"""Websocket client by default; ``--inline`` drives an in-process DuplexOmni."""
if not args.inline:
return DuplexClient(
args.url,
model=args.model,
config=config,
reconnect=None,
heartbeat_interval_s=None,
handshake_timeout_s=args.timeout_s,
)
# Imported lazily: the in-process path loads the model (and its engine
# dependencies) in this process, which the websocket path never needs.
from vllm_omni.clients.inline_duplex import InlineDuplexClient
from vllm_omni.entrypoints.duplex_omni import DuplexOmni
# Models shipped with custom HF code (MiniCPM-o) need trust_remote_code to
# resolve their pipeline; harmless for the others.
omni_kwargs: dict[str, object] = {"model": args.model, "trust_remote_code": True}
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
omni = DuplexOmni(**omni_kwargs)
client = InlineDuplexClient(
omni,
model=args.model,
config=config,
handshake_timeout_s=args.timeout_s,
)
# The client owns no engine; stop the in-process one when the run ends.
client.owned_omni = omni # type: ignore[attr-defined]
return client
async def run(args: argparse.Namespace) -> int:
config = _session_config(args.preset, ref_audio_path=args.ref_audio)
question = _convert_to_session_format(read_pcm16_wav(Path(args.question_wav)), config.input_audio)
@@ -148,105 +194,117 @@ async def run(args: argparse.Namespace) -> int:
output_dir.mkdir(parents=True, exist_ok=True)
collector = EventCollector()
client = DuplexClient(
args.url,
model=args.model,
config=config,
reconnect=None,
heartbeat_interval_s=None,
handshake_timeout_s=args.timeout_s,
)
client = _make_client(args, config)
# GREET: open the voice session (handshake happens on enter).
async with client:
consume_task = asyncio.create_task(collector.consume(client))
try:
async with client:
consume_task = asyncio.create_task(collector.consume(client))
# ASK: stream the question at realtime pace, then commit the turn.
# The model decides on its own to answer (no VAD).
print("[1] streaming question ...", file=sys.stderr)
await client.stream_pcm(question, chunk_ms=chunk_ms)
await client.commit(create_response=False)
await wait_for_condition(
lambda: collector.count("response.created") > 0,
timeout_s=args.timeout_s,
label="first response.created",
)
print("[1] assistant is answering", file=sys.stderr)
# ASK: stream the question at realtime pace, then commit the turn.
# The model decides on its own to answer (no VAD).
print("[1] streaming question ...", file=sys.stderr)
await client.stream_pcm(question, chunk_ms=chunk_ms)
await client.commit(create_response=False)
await wait_for_condition(
lambda: collector.count("response.created") > 0,
timeout_s=args.timeout_s,
label="first response.created",
)
print("[1] assistant is answering", file=sys.stderr)
# Let the answer play for a moment, as a listening user would.
await asyncio.sleep(args.listen_s)
# Let the answer play for a moment, as a listening user would.
await asyncio.sleep(args.listen_s)
# INTERRUPT: talk over the assistant while its audio still streams.
# The serving overlap policy decides: enough overlapped speech barges
# in (the first response is cancelled) or a short remark is deferred
# until the model yields the turn.
print("[2] interrupting mid-answer ...", file=sys.stderr)
ids_before_interrupt = set(collector.response_ids)
await client.stream_pcm(interrupt, chunk_ms=chunk_ms)
await client.commit(create_response=False)
# INTERRUPT: talk over the assistant while its audio still streams.
# The serving overlap policy decides: enough overlapped speech barges
# in (the first response is cancelled) or a short remark is deferred
# until the model yields the turn.
print("[2] interrupting mid-answer ...", file=sys.stderr)
ids_before_interrupt = set(collector.response_ids)
await client.stream_pcm(interrupt, chunk_ms=chunk_ms)
await client.commit(create_response=False)
# ANSWER AGAIN: wait for a response created *after* the interruption
# to finish. The first answer may still be streaming here; matching on
# any response.done would let its completion satisfy the wait and cut
# off the actual follow-up answer at close.
def follow_up_answer_done() -> bool:
return any(
event.get("type") == "response.done"
and collector.response_id(event) is not None
and collector.response_id(event) not in ids_before_interrupt
for event in collector.events
# ANSWER AGAIN: wait for a response created *after* the interruption
# to finish. The first answer may still be streaming here; matching on
# any response.done would let its completion satisfy the wait and cut
# off the actual follow-up answer at close.
def follow_up_answer_done() -> bool:
return any(
event.get("type") == "response.done"
and collector.response_id(event) is not None
and collector.response_id(event) not in ids_before_interrupt
for event in collector.events
)
await wait_for_condition(
follow_up_answer_done,
timeout_s=args.timeout_s,
label="response.done for the response answering the interruption",
)
await wait_for_condition(
follow_up_answer_done,
timeout_s=args.timeout_s,
label="response.done for the response answering the interruption",
)
# HANG UP: report playback, close the session cleanly.
await acknowledge_collected_playback(client, collector)
await client.close(timeout_s=args.timeout_s)
try:
await asyncio.wait_for(consume_task, timeout=5.0)
except (asyncio.TimeoutError, asyncio.CancelledError):
consume_task.cancel()
# HANG UP: report playback, close the session cleanly.
await acknowledge_collected_playback(client, collector)
await client.close(timeout_s=args.timeout_s)
try:
await asyncio.wait_for(consume_task, timeout=5.0)
except (asyncio.TimeoutError, asyncio.CancelledError):
consume_task.cancel()
# Report: one WAV per response; statuses show what the overlap did.
responses = _fold_responses(collector)
summary: list[dict[str, object]] = []
for index, (response_id, item) in enumerate(responses.items(), start=1):
audio = item["audio"]
assert isinstance(audio, bytes)
sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type]
wav_path = output_dir / f"response_{index}_{item['status']}.wav"
if audio:
write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz)
summary.append(
{
"response_id": response_id,
"status": item["status"],
"audio_s": round(len(audio) / (sample_rate_hz * 2), 2),
"text": item["text"],
"wav": wav_path.name if audio else None,
}
# Report: one WAV per response; statuses show what the overlap did.
responses = _fold_responses(collector)
summary: list[dict[str, object]] = []
for index, (response_id, item) in enumerate(responses.items(), start=1):
audio = item["audio"]
assert isinstance(audio, bytes)
sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type]
wav_path = output_dir / f"response_{index}_{item['status']}.wav"
if audio:
write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz)
summary.append(
{
"response_id": response_id,
"status": item["status"],
"audio_s": round(len(audio) / (sample_rate_hz * 2), 2),
"text": item["text"],
"wav": wav_path.name if audio else None,
}
)
(output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False))
(output_dir / "events.jsonl").write_text(
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events),
encoding="utf-8",
)
(output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False))
(output_dir / "events.jsonl").write_text(
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events),
encoding="utf-8",
)
print(json.dumps(summary, indent=2, ensure_ascii=False))
cancelled = sum(1 for item in summary if item["status"] == "cancelled")
print(
f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}",
file=sys.stderr,
)
print(json.dumps(summary, indent=2, ensure_ascii=False))
cancelled = sum(1 for item in summary if item["status"] == "cancelled")
print(
f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}",
file=sys.stderr,
)
finally:
owned_omni = getattr(client, "owned_omni", None)
if owned_omni is not None:
owned_omni.shutdown()
return 0
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime")
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
parser.add_argument(
"--model",
default="openbmb/MiniCPM-o-4_5",
help="served model name, or the model path with --inline",
)
parser.add_argument(
"--inline",
action="store_true",
help="run the model in this process through DuplexOmni + InlineDuplexClient instead of connecting to --url",
)
parser.add_argument(
"--deploy-config",
default=None,
help="deploy YAML for --inline (default: the model's bundled profile)",
)
parser.add_argument(
"--preset",
choices=["minicpmo_4_5", "personaplex", "none"],
@@ -1,8 +1,10 @@
# `barge_in_client.py` — what happens, at a glance
The demo plays one everyday conversation pattern against the full-duplex
server: ask something, get talked-over mid-answer, and see the assistant
cope. High-level flow below — for the exact wire events, read the
The demo plays one everyday conversation pattern against a full-duplex
model: ask something, get talked-over mid-answer, and see the assistant
cope. It runs either against a server over `/v1/realtime?duplex=1` or, with
`--inline`, against an in-process `DuplexOmni` — the conversation below is the
same in both modes. High-level flow — for the exact wire events, read the
docstring in `barge_in_client.py`.
```mermaid
@@ -48,7 +50,8 @@ answer is saved as `response_1_cancelled.wav`, a completed one as
duration per answer.
Companion: the runtime architecture lives in
[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client
API is `vllm_omni.clients.duplex.DuplexClient`, documented with the wire
protocol in
[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client API is
`vllm_omni.clients.duplex.DuplexClient` over the websocket and
`vllm_omni.clients.inline_duplex.InlineDuplexClient` in-process — both are
documented with the wire protocol in
[`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md).
@@ -5,10 +5,10 @@
# Server side for barge_in_client.py: start MiniCPM-o 4.5 in native
# full-duplex mode.
#
# The duplex deploy profile registers the WS /v1/duplex endpoint (with
# /v1/realtime?duplex=1 projecting onto the same handler), enables the
# engine duplex control plane, and bounds the Thinker/Talker/Code2Wav
# stages to four live duplex sessions on one GPU. See
# MiniCPM-o 4.5 declares a duplex plugin, so vllm-omni serve runs it through
# DuplexOmni: the server is duplex-only (WS /v1/realtime?duplex=1, alias
# /v1/duplex, plus /v1/models and /health) and bounds the
# Thinker/Talker/Code2Wav stages to four live duplex sessions on one GPU. See
# vllm_omni/deploy/minicpmo_4_5.yaml (session_mode: duplex) for the session
# limits; every other duplex_session knob keeps its runtime default
# (idle TTL 300s, disconnect grace 30s, 16MiB pending input per session).
@@ -22,4 +22,4 @@ exec vllm-omni serve "$MODEL" \
--deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \
--trust-remote-code \
--host 0.0.0.0 --port "$PORT"
# duplex endpoints: ws://<host>:$PORT/v1/duplex and /v1/realtime?duplex=1
# duplex endpoint: ws://<host>:$PORT/v1/realtime?duplex=1 (alias /v1/duplex)
@@ -445,7 +445,6 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
ref_audio=_ref_audio_data_url(args.ref_audio),
temperature=args.temperature,
),
session_id=args.session_id,
reconnect=None,
heartbeat_interval_s=None,
handshake_timeout_s=args.timeout_s,
@@ -638,7 +637,6 @@ def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
parser.add_argument("--session-id")
parser.add_argument(
"--input-wav",
help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.",
@@ -86,7 +86,6 @@
url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
url.searchParams.set('duplex', '1');
url.searchParams.set('model', config.model || 'openbmb/MiniCPM-o-4_5');
url.searchParams.set('native_duplex', '1');
url.searchParams.set('autostart', '0');
return url.toString();
}
@@ -458,7 +457,6 @@
settled = true;
const extraBody = {
auto_response: true,
native_duplex: true,
};
const session = {
modalities: ['audio', 'text'],
@@ -17,6 +17,7 @@ For the full list of supported architectures across all modalities, see
| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — |
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
| Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — |
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
| IndexTTS-2.5 | native `checkpoints/` bundle | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
@@ -96,12 +97,52 @@ curl -X POST http://localhost:8091/v1/audio/speech \
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
```
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others).
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for Gepard and IndexTTS-2, and 24 kHz for many others).
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md).
---
## Gepard-1.0
Single-stage native AR TTS at 22.05 kHz mono. Zero-shot only: omit `voice` or pass `"default"`. Voice cloning from reference audio is not available yet.
### Prerequisites
Same NeMo NanoCodec install as the [offline Gepard section](../../offline_inference/text_to_speech/README.md#gepard-10). On a host whose CUDA toolkit cannot build kernels — no `nvcc`/`ninja`, or a consumer Blackwell (`sm_120`) card — also `export VLLM_USE_FLASHINFER_SAMPLER=0` before launch (same note as offline). That env is not a deploy-YAML field.
### Launch
```bash
vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \
--stage-init-timeout 900 \
--deploy-config vllm_omni/deploy/gepard.yaml
# or:
./gepard/run_server.sh
```
`--stage-init-timeout 900` matches the online e2e fixture and the offline example. Serve defaults to 300s, which is often too short for a cold download of the talker plus NanoCodec.
The packaged `vllm_omni/deploy/gepard.yaml` must be passed with `--deploy-config`. The checkpoint self-identifies as `qwen3_5_text`, so omitting the YAML launches a diffusion fallback instead of the Gepard pipeline. The YAML sets `async_chunk: false`, `max_num_seqs: 4`, and currently pins `seed: 42`, so serving is deterministic by default until that YAML seed is removed. Pass an explicit per-request `seed` in tests and clients rather than depending on either default.
### Sending requests
```bash
python examples/online_serving/text_to_speech/gepard/speech_client.py \
--text "Hello, this is Gepard speaking."
python examples/online_serving/text_to_speech/gepard/speech_client.py \
--text "Hello, this is Gepard speaking." --seed 7 --stream --output output.pcm
```
### Notes
- Output: 22.05 kHz mono. `max_new_tokens` is a **frame** budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms at 21.5 fps; adapter bounds 1..4096).
- Supported request fields: `input` (required), `voice` (`default` only), `response_format` (`wav` default; `wav/pcm/flac/mp3` non-streaming; `opus` 400 because 22.05 kHz is not an Opus sample rate; streaming `pcm/wav` only), `stream` / `stream_format`, `max_new_tokens`, `seed`.
- Unsupported: `speed`, `extra_params` (including `temperature`/`top_p`/`top_k`), `ref_audio`, `ref_text`, `speaker_embedding`, `task_type`, `instructions`, `language`, and `word_timestamps`.
- Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands.
- Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI.
## dots.tts
Single-stage, text-only TTS at 48 kHz. The current integration supports
@@ -0,0 +1,33 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Launch vLLM-Omni server for Gepard-1.0 TTS (zero-shot default voice).
#
# Usage:
# ./run_server.sh
# CUDA_VISIBLE_DEVICES=0 ./run_server.sh
#
# Packaged deploy config: vllm_omni/deploy/gepard.yaml (async_chunk=false,
# 22.05 kHz mono, default seed 42 until the YAML seed is removed).
# Cold start downloads the talker and NeMo NanoCodec; the serve default of
# --stage-init-timeout 300 is too tight (offline end2end.py uses 900).
#
# On a host whose CUDA toolkit cannot JIT FlashInfer (no nvcc/ninja, or a
# consumer Blackwell sm_120 card), also:
# export VLLM_USE_FLASHINFER_SAMPLER=0
set -e
MODEL="${MODEL:-nineninesix/gepard-1.0}"
PORT="${PORT:-8091}"
echo "Starting Gepard-1.0 server with model: $MODEL"
vllm-omni serve "$MODEL" \
--host 0.0.0.0 \
--port "$PORT" \
--trust-remote-code \
--omni \
--stage-init-timeout 900 \
--deploy-config vllm_omni/deploy/gepard.yaml
@@ -0,0 +1,116 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Client for Gepard-1.0 TTS via /v1/audio/speech.
Gepard is zero-shot: omit ``voice`` or pass ``"default"``. Output is 22.05 kHz
mono. ``seed`` is optional; the packaged deploy YAML currently pins ``seed: 42``
until that default is removed, so requests without an explicit seed are still
deterministic.
Examples:
python speech_client.py --text "Hello, this is Gepard speaking."
python speech_client.py --text "Hello, this is Gepard speaking." --seed 7
python speech_client.py --text "Hello, this is Gepard speaking." --stream --output output.pcm
"""
from __future__ import annotations
import argparse
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
DEFAULT_MODEL = "nineninesix/gepard-1.0"
def run_tts(args) -> None:
payload = {
"model": args.model,
"input": args.text,
"voice": args.voice,
"response_format": args.response_format,
}
if args.seed is not None:
payload["seed"] = args.seed
if args.max_new_tokens is not None:
payload["max_new_tokens"] = args.max_new_tokens
if args.stream:
payload["stream"] = True
payload["stream_format"] = "audio"
if args.response_format not in ("pcm", "wav"):
print(f"Note: streaming requires pcm/wav; overriding --response-format {args.response_format} to 'pcm'")
payload["response_format"] = "pcm"
print(f"Model: {args.model}")
print(f"Text: {args.text}")
print(f"Voice: {args.voice}")
if args.seed is not None:
print(f"Seed: {args.seed}")
print("Generating audio...")
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
if args.stream:
output_path = args.output or ("output.wav" if payload["response_format"] == "wav" else "output.pcm")
with httpx.Client(timeout=300.0) as client:
with client.stream("POST", api_url, json=payload, headers=headers) as resp:
if resp.status_code != 200:
print(f"Error: {resp.status_code}")
print(resp.read().decode())
return
total_bytes = 0
with open(output_path, "wb") as f:
for chunk in resp.iter_bytes():
f.write(chunk)
total_bytes += len(chunk)
print(f"Streamed {total_bytes} bytes to: {output_path}")
return
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return
output_path = args.output or "output.wav"
with open(output_path, "wb") as f:
f.write(response.content)
print(f"Audio saved to: {output_path}")
def main() -> None:
parser = argparse.ArgumentParser(description="Gepard-1.0 TTS client")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument("--voice", default="default", help="Voice name (only 'default' is supported)")
parser.add_argument("--seed", type=int, default=None, help="Sampling seed")
parser.add_argument(
"--max-new-tokens",
type=int,
default=None,
help="Frame budget (1 token = 1 frame ≈ 46.4 ms at 21.5 fps)",
)
parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "mp3", "flac", "pcm"],
help="Audio format (default: wav). Streaming is pcm/wav only. Opus is not supported at 22.05 kHz.",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Batch speech client for Qwen3-TTS via /v1/audio/speech/batch endpoint.
This script demonstrates how to synthesize multiple texts in a single request.
@@ -7,7 +10,7 @@ the reference for each item.
Start the server (with batch-optimized stage settings for best throughput):
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--omni \
--trust-remote-code \
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Gradio demo for Qwen3-TTS with gapless streaming audio playback.
Uses a custom AudioWorklet-based player for gap-free streaming,
@@ -351,7 +354,7 @@ def generate_speech(
raise gr.Error(f"Failed to decode audio: {e}")
def create_app(api_base: str):
def create_app(api_base: str, task_type: str = "CustomVoice"):
"""Create the FastAPI app with streaming proxy + Gradio UI."""
fastapi_app = FastAPI()
@@ -463,10 +466,11 @@ def create_app(api_base: str):
lines=4,
)
with gr.Row():
task_type = gr.Radio(
task_type_input = gr.Radio(
choices=TASK_TYPES,
value="CustomVoice",
value=task_type,
label="Task Type",
interactive=False,
scale=2,
)
language = gr.Dropdown(
@@ -479,17 +483,21 @@ def create_app(api_base: str):
choices=voices,
value=voices[0] if voices else None,
label="speaker",
visible=True,
visible=task_type == "CustomVoice",
allow_custom_value=True,
)
instructions = gr.Textbox(
label="Instructions",
placeholder="e.g., Speak with excitement / A warm, friendly female voice",
lines=2,
visible=True,
info="Optional style/emotion instructions",
visible=task_type != "Base",
info=(
"Required: describe the voice style"
if task_type == "VoiceDesign"
else "Optional style/emotion instructions"
),
)
with gr.Column(visible=False) as ref_group:
with gr.Column(visible=task_type == "Base") as ref_group:
ref_audio = gr.Audio(
label="Reference Audio (upload for voice cloning)",
type="numpy",
@@ -563,7 +571,7 @@ def create_app(api_base: str):
autoplay=True,
visible=False,
)
with gr.Column(visible=True) as examples_cv:
with gr.Column(visible=task_type == "CustomVoice") as examples_cv:
gr.Examples(
examples=[
[
@@ -588,7 +596,7 @@ def create_app(api_base: str):
inputs=[text_input, voice, language, instructions],
label="examples",
)
with gr.Column(visible=False) as examples_vd:
with gr.Column(visible=task_type == "VoiceDesign") as examples_vd:
gr.Examples(
examples=[
[
@@ -605,7 +613,7 @@ def create_app(api_base: str):
inputs=[text_input, language, instructions],
label="examples",
)
with gr.Column(visible=False) as examples_base:
with gr.Column(visible=task_type == "Base") as examples_base:
gr.Examples(
examples=[
[
@@ -653,9 +661,9 @@ def create_app(api_base: str):
gr.update(visible=is_base), # examples_base
)
task_type.change(
task_type_input.change(
fn=on_task_change,
inputs=[task_type],
inputs=[task_type_input],
outputs=[
voice,
instructions,
@@ -703,7 +711,7 @@ def create_app(api_base: str):
all_inputs = [
text_input,
task_type,
task_type_input,
voice,
language,
instructions,
@@ -788,6 +796,9 @@ def parse_args():
description="Gradio demo for Qwen3-TTS with gapless AudioWorklet streaming.",
)
add_common_args(parser)
parser.add_argument(
"--task-type", choices=TASK_TYPES, default="CustomVoice", help="Task matching the served checkpoint"
)
return parser.parse_args()
@@ -796,7 +807,7 @@ def main():
logging.basicConfig(level=logging.INFO)
print(f"Connecting to vLLM server at: {args.api_base}")
app = create_app(args.api_base)
app = create_app(args.api_base, task_type=args.task_type)
import uvicorn
@@ -8,25 +8,28 @@ to generate audio from text using Qwen3-TTS models.
Examples:
# CustomVoice task (predefined speaker)
python openai_speech_client.py --text "Hello, how are you?" --voice vivian
python openai_speech_client.py --text "Hello, how are you?" --speaker vivian
# CustomVoice with emotion instruction
python openai_speech_client.py --text "I'm so happy!" --voice vivian \
python openai_speech_client.py --text "I'm so happy!" --speaker vivian \
--instructions "Speak with excitement"
# VoiceDesign task (voice from description)
python openai_speech_client.py --text "Hello world" \
--model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
--task-type VoiceDesign \
--instructions "A warm, friendly female voice"
# Base task (voice cloning)
python openai_speech_client.py --text "Hello world" \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--task-type Base \
--ref-audio "https://example.com/reference.wav" \
--ref-text "This is the reference transcript"
# Base task with pre-computed speaker embedding
python openai_speech_client.py --text "Hello world" \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--task-type Base \
--speaker-embedding embedding.json
"""
@@ -1,4 +1,7 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Launch both vLLM server and Gradio demo for Qwen3-TTS
#
# Usage:
@@ -179,7 +182,7 @@ fi
echo ""
echo "Starting Gradio demo..."
cd "$SCRIPT_DIR"
GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--host" "$GRADIO_IP" "--port" "$GRADIO_PORT")
GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--host" "$GRADIO_IP" "--port" "$GRADIO_PORT" "--task-type" "$TASK_TYPE")
if [ "$GRADIO_SHARE" = true ]; then
GRADIO_CMD+=("--share")
fi
+5
View File
@@ -123,6 +123,11 @@ fa4 = [
"flash-attn-4[cu13]==4.0.0b18",
]
# FastVideo VSA kernels for the supported Linux/Python 3.12 wheel environment.
vsa = [
"fastvideo-kernel==0.3.4",
]
docs = [
"mkdocs>=1.5.0",
"mkdocs-api-autonav",
+20 -7
View File
@@ -1,9 +1,10 @@
# MiniMax H3 disaggregated text encoder
# MiniMax H3 disaggregated encoder stage
This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
its hidden states and token-role metadata to an encoder-free diffusion stage.
The standard MiniMax H3 recipes remain single-stage and continue to load the
text encoder inside the diffusion pipeline.
This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and
audio WVAE encoder together in a vLLM stage. It sends their conditioning
payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes
remain single-stage and continue to run all components inside the diffusion
pipeline.
## Start the server
@@ -23,6 +24,14 @@ deployment override for the available hardware. Diffusion quantization,
layerwise offload, distributed layerwise offload, VAE parallelism, and USP
settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
Stage 0's `tensor_parallel_size` defines one shared rank group for all three
encoder roles; it does not choose each role's execution strategy. The required
`hf_overrides.minimax_h3_encoder_components` mapping declares those strategies
explicitly: `text_encoder` uses tensor parallelism, `video_vae` uses patch
parallelism, and `audio_vae` runs on the group leader. All three role entries
are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit;
the configuration does not create independent role schedulers or world sizes.
For example, this five-GPU topology assigns one GPU to the encoder and four
to the diffusion stage. `--stage-overrides` keeps placement and parallelism
scoped to the owning stage rather than broadcasting an override to both:
@@ -57,8 +66,12 @@ The Stage 1 VAE patch-parallel options remain independent of offload and
quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
hardware-qualified profiles before combining these options.
Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
text-encoder weights. This H3 topology explicitly keeps its single-replica
Stage 1 sets `model_loaded.text_encoder: false` and
`model_loaded.vae_encoder: false`. It loads only the VAE decoders and skips
tokenizer, processor, and text-encoder downloads. It requires complete text and
media conditioning from Stage 0, with no fallback to local encoding.
This H3 topology explicitly keeps its single-replica
diffusion stage inline, avoiding serialization of decoded video through a
subprocess. It expects both stages to run in one deployment; cross-node payload
transport is outside this configuration.
+10 -9
View File
@@ -23,6 +23,12 @@ One vLLM-Omni diffusion stage can load both DiTs while instantiating the
tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only
once. Requests select the DiT with `extra_params.task`.
This single-stage recipe uses the default `model_loaded.text_encoder: true` and
`model_loaded.vae_encoder: true`. Precomputed text conditioning can replace the
local Qwen call. Setting `model_loaded.text_encoder: false` while leaving
`model_loaded.vae_encoder: true` requires precomputed text conditioning;
reference media is still encoded locally.
The generated MP4 contains H.264 video and synchronized stereo audio.
## Prerequisites
@@ -35,7 +41,7 @@ hf auth login
export MODEL=MiniMaxAI/MiniMax-H3
```
The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
By default, the pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
`Ref2VA/transformer/**`. It does not download or load the diffusers-format
`transformer`, `transformer_ref`, or `vae` weights at the repository root, nor
duplicate Ref2VA copies of shared components.
@@ -1230,11 +1236,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
mode does not support `cache_backend`.
- The first regional-compile request is a warmup and should not be included in
steady-state performance measurements.
- The serving path accepts fewer references than the model supports. H3 documents up
to 9 images, 3 video clips, and 3 audio clips (12 files) per Omni Reference
request; the current vLLM-Omni path takes exactly one image plus one audio
reference, or one or more videos with no separate `audio_reference` (it uses the
source soundtracks).
- Ref2VA accepts up to 9 images, 3 video clips, and 3 audio clips, with at most
12 references in total. At least one image or video is required; audio-only
requests are rejected.
- The 768 px short-edge mode is available for T2VA and FL2VA; 1344x768 is the
documented 16:9 request shape.
- `--cfg-parallel-size > 1` is rejected by design (CFG-distilled, no negative branch).
@@ -1247,9 +1251,6 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
in host memory on every rank during startup before retaining only each rank's
shard. Size startup host memory for that transient peak.
- TeaCache and Cache-DiT cannot be enabled on the same server.
- Image+audio Ref2VA accepts exactly one image and one audio reference.
- Video Ref2VA accepts one or more video files, but not an additional standalone
audio reference.
- Pure Ulysses still replicates the full DiT on every rank, so smaller-memory GPUs
cannot use `--usp N --tp 1` as a resident capacity path. Use DiT tensor parallelism
or model-level CPU offload; text-encoder TP alone is not sufficient.
+60 -15
View File
@@ -1,6 +1,6 @@
# MiniCPM-o 4.5
> Online serving and offline inference for omni multimodal chat
> Online full-duplex and chat-completions serving, and offline inference, for omni multimodal chat
> (text / image / audio / video → text + 24 kHz speech)
## Summary
@@ -9,8 +9,12 @@
- Model: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
- Task: Omni multimodal chat — accepts text / image / audio / video input;
emits text and 24 kHz mono speech in the same response
- Mode: Online serving via the OpenAI-compatible `/v1/chat/completions`
API (plus Gradio demo), and offline inference via `Omni.generate`
- Mode: Online full-duplex serving over the OpenAI Realtime protocol
(`/v1/realtime?duplex=1`, alias `/v1/duplex`), the OpenAI-compatible
`/v1/chat/completions` API (plus Gradio demo) on the same server, and
offline turn-based inference via `Omni.generate`. Besides those the server
exposes only `/v1/models` and `/health`; the other turn-based HTTP routes
are not available on a duplex deployment.
- Maintainer: [`@tc-mb`](https://github.com/tc-mb) (MiniCPM-V / MiniCPM-o team)
## When to use this recipe
@@ -18,9 +22,10 @@
Use this recipe as a known-good starting point for serving
`openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member
of the MiniCPM-o family — it runs a multimodal thinker, a streaming
MiniCPMTTS codec talker, and a separate batched Code2Wav stage so a single
`/v1/chat/completions` call can return text and 24 kHz speech in one
shot. The recommended batching deploy isolates the Thinker on GPU 0 and
MiniCPMTTS codec talker, and a separate batched Code2Wav stage, so a live
duplex session streams text and 24 kHz speech while it keeps listening, and a
single `/v1/chat/completions` call returns both in one shot. The
recommended batching deploy isolates the Thinker on GPU 0 and
co-locates Talker and Code2Wav on GPU 1; 1-GPU, 3-GPU, and 8x4090 layouts are
also provided.
@@ -36,8 +41,11 @@ also provided.
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
- 8x RTX 4090 layout:
[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
- Online example + Gradio demo:
- Online examples (Realtime CLI demo, browser client, barge-in client,
chat-completions client, curl script, Gradio demo):
[`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/)
- Duplex API and Python client docs:
[`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md)
- Offline end-to-end example:
[`examples/offline_inference/minicpmo/`](../../examples/offline_inference/minicpmo/)
- Pipeline / talker source:
@@ -179,6 +187,13 @@ as one aggregated chunk, so the table reports Stage 0 metrics.
#### Verification
**Health and model listing**:
```bash
curl http://localhost:8099/health
curl http://localhost:8099/v1/models
```
**Quick smoke test (text-only output)**:
```bash
@@ -191,8 +206,31 @@ curl http://localhost:8099/v1/chat/completions \
}'
```
**Text + speech in one response** (the headline 4.5 feature). The model
bridge conditions the Talker from the generated assistant span, so the
**One duplex turn over the Realtime WebSocket** (text + speech stream back
while the session keeps listening):
```bash
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--url ws://localhost:8099/v1/realtime?duplex=1 \
--model openbmb/MiniCPM-o-4_5 \
--input-wav /path/to/input_16k_mono_pcm16.wav \
--ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \
--output-dir /tmp/minicpmo_realtime_duplex_demo
```
The demo prints transcript deltas as they arrive and saves the streamed
24 kHz audio to WAV files. The same session shape is available from Python
through `vllm_omni.clients.duplex.DuplexClient` (over the server) or
`vllm_omni.clients.inline_duplex.InlineDuplexClient` (in-process, no server;
`examples/online_serving/barge_in_client.py --inline`).
**Barge-in**: `examples/online_serving/barge_in_client.py` plays a question,
interrupts the answer with a second WAV and checks that the first response is
cancelled and the second one is answered.
**Text + speech in one response over `/v1/chat/completions`** (the headline
4.5 feature, served by the ordinary chat service on the duplex engine). The
model bridge conditions the Talker from the generated assistant span, so the
generic serving layer does not inject MiniCPM-specific template defaults.
`use_tts_template=true` remains supported when explicitly requested:
@@ -243,6 +281,12 @@ python examples/online_serving/minicpmo/gradio_demo.py \
Open `http://<host>:7862` and try a text prompt with the **"Generate
speech output (TTS)"** checkbox on / off.
**Turn-based generation** is also available offline
(`examples/offline_inference/minicpmo/`, `Omni.generate`).
The Daily-Omni numbers above were measured on the `/v1/chat/completions`
route.
#### Notes
- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and
@@ -326,9 +370,13 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
and defaults to deterministic seed 42. Stage-1 deploy sampling parameters
control only vLLM's binary continue/stop token.
- **Output audio**: 24 kHz mono WAV inside the OpenAI-style
`message.audio.data` (base64). The Gradio demo's WAV player decodes
this automatically.
- **Output audio**: 24 kHz mono. On `/v1/chat/completions` it is base64 WAV
inside the OpenAI-style `message.audio.data` (the Gradio demo's WAV player
decodes this automatically); on a duplex session it streams as base64
`response.audio.delta` events; offline `Omni.generate` returns WAV.
- **Response choices**: on `/v1/chat/completions` text and audio are separate
choices. SDK clients should select the choice whose `message.audio.data` is
populated rather than assuming `choices[0]` contains audio.
- **Routing**: MiniCPM-o 4.5 and 2.6 both ship `architectures=
["MiniCPMO"]` in HF config; routing is disambiguated by
@@ -340,6 +388,3 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
- **Async chunking**: enabled in all deploy configs. Talker sends
25-code chunks with three-code left context to Code2Wav through
`SharedMemoryConnector`; terminal chunks flush held lookahead state.
- **Response choices**: text and audio are separate choices. SDK clients
should select the choice whose `message.audio.data` is populated rather
than assuming `choices[0]` contains audio.
+1 -1
View File
@@ -71,7 +71,7 @@ recipes/
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
| [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online full-duplex serving and offline inference for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`OpenMOSS/MOSS-TTS-Nano-NPU.md`](./OpenMOSS/MOSS-TTS-Nano-NPU.md) | Online multilingual voice-cloned TTS | 1x Ascend 910B3 64GB |
+40 -21
View File
@@ -32,13 +32,13 @@ mode is selected per request:
T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
- `policy` — given a first frame or video and a language instruction,
- `policy` — given a first frame or video and a language instruction,
**predict** the action trajectory (and a rollout video). Use the async
`POST /v1/videos` endpoint and read the predicted action from the top-level
`action` field.
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
the async `POST /v1/videos` endpoint and read the recovered action from
the top-level `action` field
(`{data, shape, dtype, raw_action_dim, domain_id}`).
@@ -106,6 +106,25 @@ vllm serve nvidia/Cosmos3-Nano \
--init-timeout 1800
```
For Cosmos3, SeaCache is the recommended default choice when opting into
diffusion caching. Caching remains opt-in; add `--cache-backend sea_cache` to
the command above. Its default threshold and maximum cached-step streak are
tuned for Cosmos3, so no `--cache-config` is required.
Override individual defaults with a JSON cache configuration; for example:
```bash
vllm serve nvidia/Cosmos3-Nano \
--omni \
--cache-backend sea_cache \
--cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}'
```
Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are
more conservative. Higher values allow more cached steps and may improve
speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
`0` removes the streak cap.
To run **without** guardrails (you are responsible for license compliance),
add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
`--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
@@ -382,12 +401,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's
#### Notes
- **Measured latency (1x B300, bf16, guardrails off):**
- T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
- T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
- I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
- Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
- Guardrails-on overhead: ~8% on T2I, negligible on video.
- **Measured latency (1x B300, bf16, guardrails off, diffusion cache off):**
- T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
- T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
- I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
- Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
- Guardrails-on overhead: ~8% on T2I, negligible on video.
- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
~33 GB on disk.
@@ -443,12 +462,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's
`conditioning_fps`, `action_chunk_size`, `raw_action_dim`, `deterministic_seed`,
and `session_id`.
- **Known limitations:**
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
`nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
the server fails at pipeline build with a gated-repo / safety-checker error.
- A guardrail-blocked prompt currently returns HTTP 500
- A guardrail-blocked prompt currently returns HTTP 500
(`"Guardrail blocked prompt"`).
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
online. Use async `POST /v1/videos` when you need the predicted/recovered
action payload under the top-level `action` field; sync `/v1/videos/sync`
returns raw MP4 bytes and does not expose action metadata in the response body.
@@ -579,14 +598,14 @@ so an md5 comparison across two runs also works as a smoke check.
a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images
does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is
**540 s cold vs 161 s warm, with byte-identical output**.
- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps
- **Measured on 1x MI350X (bf16, guardrails off, diffusion cache off, warm):** T2I 1024² @ 50 steps
**~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is
the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the
denoise loop. The optional flags act on different terms of the memory bill and
are therefore complementary rather than redundant:
| Flag | Latency | Peak reserved | Peak allocated | Acts on |
|---|---|---|---|---|
| --- | --- | --- | --- | --- |
| *(none)* | 161 s | 120 GiB | 95 GiB | — |
| `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations |
| `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights |
@@ -768,11 +787,11 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
#### Notes
- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):**
- T2I 1024² — 10 steps → ~8 s
- T2V 1280×720 @ 20 steps — 49 frames → ~55 s
- I2V 1280×720 @ 10 steps — 25 frames → ~25 s
- V2V 480×320 @ 10 steps — 17 frames → ~12 s
- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off, diffusion cache off):**
- T2I 1024² — 10 steps → ~8 s
- T2V 1280×720 @ 20 steps — 49 frames → ~55 s
- I2V 1280×720 @ 10 steps — 25 frames → ~25 s
- V2V 480×320 @ 10 steps — 17 frames → ~12 s
- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU;
full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
~33 GB on disk.
@@ -785,9 +804,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
(for model loading), `--tensor-parallel-size 8` for multi-NPU, and
`--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class.
- **Known limitations:**
- Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
- Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
resolution-parsing bug; basic V2V without transfer hints works.
- FP8 online quantization and layerwise offload are not supported on NPU.
- FP8 online quantization and layerwise offload are not supported on NPU.
### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation
+21 -6
View File
@@ -62,6 +62,21 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V
composition can shift at the same seed).
SeaCache is the recommended default choice when opting into diffusion caching,
which remains opt-in. Add `--cache-backend sea_cache` to either command above;
its defaults are tuned for Cosmos3.
Override individual defaults with `--cache-config`; for example:
```bash
--cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}'
```
Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are
more conservative. Higher values allow more cached steps and may improve
speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
`0` removes the streak cap.
#### Verification
Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
@@ -109,12 +124,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
#### Notes
- **Measured (2x B300, bf16, guardrails off, official 2-GPU config above):**
- T2I 1024², 50 steps → **~6 s**
- T2V 1280×720, 189 frames, 35 steps → **~197 s**
- I2V 1280×720, 189 frames, 35 steps → **~200 s**
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
- **Measured (2x B300, bf16, guardrails off, diffusion cache off, official 2-GPU config above):**
- T2I 1024², 50 steps → **~6 s**
- T2V 1280×720, 189 frames, 35 steps → **~197 s**
- I2V 1280×720, 189 frames, 35 steps → **~200 s**
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
- Same generation defaults, supported sizes, V2V reference-video controls
(`condition_frame_indexes_vision`, `condition_video_keep`), and
+1 -1
View File
@@ -178,7 +178,7 @@ mistralai/Voxtral-4B-TTS-2603 checkpoint. running the
voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../examples/offline_inference/text_to_speech/README.md#voice-cloning-capability-gated-upstream)
) or adding a `ref_audio` request parameter will fail
because the public checkpoint does not provide the encoder weights needed to
convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`
convert reference audio into conditioning features. The expected failure is a 400 BadRequest `Voice cloning with 'ref_audio' requires the full Voxtral checkpoint with encoder weights. The open-source variant only supports preset voices via the 'voice' parameter.`.
### 1x NVIDIA RTX A6000 48GB
@@ -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
@@ -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")
+41 -27
View File
@@ -135,6 +135,7 @@ class MockResponse:
async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch):
class FakeRealtimeClient:
last_instance = None
instances: list = []
def __init__(self, url):
assert url == "ws://localhost:8000/v1/realtime?duplex=1"
@@ -143,6 +144,9 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
self.sent = []
self.response_count = 0
self.ack_count = 0
self.silence_seconds = []
self.closed = 0
type(self).instances.append(self)
type(self).last_instance = self
async def __aenter__(self):
@@ -155,10 +159,22 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
assert model == "openbmb/MiniCPM-o-4_5"
self.configure_kwargs = kwargs
async def stream_silence(self, *, seconds, chunk_ms=200, until=None):
# A model-native session speaks off its seeded context once audio
# units arrive; the silence itself carries no content. The probe
# stops the silence as soon as the turn settles.
self.silence_seconds.append(seconds)
self._emit_response()
assert until is not None and until()
return 1.0
async def send(self, event):
self.sent.append(event)
if event["type"] != "response.create":
return
self._emit_response()
def _emit_response(self):
self.response_count += 1
response_id = f"resp-{self.response_count}"
now = time.monotonic()
@@ -210,7 +226,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
self.ack_count += 1
async def close_session(self, **_kwargs):
return None
self.closed += 1
monkeypatch.setattr(
"vllm_omni.benchmarks.patch.patch._RealtimeTTSProbe",
@@ -239,30 +255,27 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
session=None,
)
client = FakeRealtimeClient.last_instance
assert client.configure_kwargs["native_duplex"] is False
assert client.configure_kwargs["extra_body"] == {
"ref_audio": "data:audio/wav;base64,AAAA",
"return_stage_metrics": True,
}
assert client.sent == [
event
for index in range(4)
for event in (
{
"type": "conversation.item.create",
"item": {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": f"text {index}"}],
},
},
{"type": "response.create"},
)
]
assert client.ack_count == 4
# One session per utterance: a model-native duplex session takes its text
# once, in the session context, so it cannot be re-seeded for a second one.
sessions = FakeRealtimeClient.instances
assert len(sessions) == 4
for index, client in enumerate(sessions):
assert "native_duplex" not in client.configure_kwargs
assert client.configure_kwargs["extra_body"] == {
"ref_audio": "data:audio/wav;base64,AAAA",
"return_stage_metrics": True,
"duplex_initial_user_text": f"text {index}",
"force_listen_count": 0,
}
# The target text rides the session context, never a conversation item:
# a text-only response.create is rejected by a model-native session.
assert client.sent == []
assert client.silence_seconds and all(seconds > 0 for seconds in client.silence_seconds)
assert client.ack_count == 1
assert client.closed == 1
assert output.success is True
assert output.generated_text == "turn 1 turn 2 turn 3 turn 4"
assert output.generated_text == "turn 1 turn 1 turn 1 turn 1"
assert output.audio_duration == pytest.approx(0.4)
assert output.ttft > 0
assert output.audio_ttfp > output.ttft
@@ -281,11 +294,12 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
"session_id": session_id,
"request_index": index,
"utterance_id": f"utt-{index}",
"response_id": f"resp-{index + 1}",
# Response ids are per session, and each utterance now gets its own.
"response_id": "resp-1",
"source": "client_monotonic_receive",
"measurement_origin": {
"ttft": "conversation.item.create client send to first non-empty text delta",
"ttfp": "conversation.item.create client send to first audio packet",
"ttft": "first silence append client send to first non-empty text delta",
"ttfp": "first silence append client send to first audio packet",
"rtf": "request-start-to-last-audio receive time divided by emitted audio duration",
"tpot": "Stage-0 engine mean time per output token",
},
+24 -3
View File
@@ -467,6 +467,27 @@ def test_tolerated_playback_ack_rejection_is_a_warning_not_a_failure():
oi._raise_if_session_terminated(collector, 0, warnings=warnings)
def test_our_own_close_is_expected_even_though_the_server_stamps_a_reason():
"""``session.close`` is answered with ``session.closed`` carrying ``client_close``.
The guard exists to catch a session that ended for a reason we did not ask
for. A close we requested is not that, whether or not the server names it.
"""
collector = _collector(({"type": "session.closed", "reason": "client_close"}, 1.0))
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
collector = _collector(({"type": "session.closed", "event": {"reason": "client_close"}}, 1.0))
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
collector = _collector(({"type": "session.closed", "reason": "disconnect"}, 1.0))
with pytest.raises(RuntimeError, match="Unexpected session.closed: disconnect"):
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
collector = _collector(({"type": "session.expired", "reason": "timeout"}, 1.0))
with pytest.raises(RuntimeError, match="session.expired: timeout"):
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
@pytest.mark.parametrize(
("event", "match"),
[
@@ -561,8 +582,8 @@ class _RealtimeClient(oi._RealtimeSession):
instances: list[_RealtimeClient] = []
def __init__(self, config: oi.OmniInteractBenchmarkConfig, session_id: str, reference_audio: str):
super().__init__(config, session_id, reference_audio)
def __init__(self, config: oi.OmniInteractBenchmarkConfig, reference_audio: str):
super().__init__(config, reference_audio)
self.acks: list[tuple[str, int]] = []
self.instances.append(self)
@@ -621,7 +642,7 @@ async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monke
assert "autostart=0" in _RealtimeClient.instances[-1].url
session_config = _RealtimeClient.instances[-1].session_config
assert session_config.extra_body["custom"] == "value"
assert session_config.extra_body["native_duplex"] is True
assert "native_duplex" not in session_config.extra_body
assert session_config.ref_audio == "data:audio/wav;base64,ref"
acks = _RealtimeClient.instances[-1].acks
# Cumulative acks for the one response: an optional 0 ms checkpoint the
+179
View File
@@ -0,0 +1,179 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The Seed-TTS Realtime probe queues for a duplex session slot instead of failing."""
from __future__ import annotations
import pytest
from vllm_omni.benchmarks.patch import patch as bench_patch
from vllm_omni.clients import duplex as duplex_client
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _AdmissionGate:
"""Stands in for ``DuplexClient``: refuses admission ``refusals`` times, then admits."""
instances: list[_AdmissionGate] = []
refusals = 0
refusal_code = "resource_exhausted"
def __init__(self, url: str, **kwargs: object) -> None:
del url, kwargs
self.entered = False
type(self).instances.append(self)
async def __aenter__(self) -> _AdmissionGate:
cls = type(self)
if cls.refusals > 0:
cls.refusals -= 1
raise duplex_client.DuplexProtocolError("duplex_session_capacity_exhausted: limit=4", code=cls.refusal_code)
self.entered = True
return self
async def __aexit__(self, exc_type, exc, tb) -> None:
del exc_type, exc, tb
async def events(self): # the collector's consume loop ends at once
return
yield # pragma: no cover
@pytest.fixture
def gate(monkeypatch: pytest.MonkeyPatch) -> type[_AdmissionGate]:
_AdmissionGate.instances = []
_AdmissionGate.refusals = 0
_AdmissionGate.refusal_code = "resource_exhausted"
monkeypatch.setattr(duplex_client, "DuplexClient", _AdmissionGate)
return _AdmissionGate
async def _configure(probe: bench_patch._RealtimeTTSProbe) -> None:
try:
await probe.configure("test-model", auto_response=True)
finally:
await probe.__aexit__(None, None, None)
@pytest.mark.asyncio
async def test_configure_waits_for_a_free_session_slot(gate) -> None:
"""Two refusals, then a slot: the request is served, not counted as failed."""
gate.refusals = 2
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
await _configure(probe)
assert len(gate.instances) == 3
assert gate.instances[-1].entered is True
assert probe._client is gate.instances[-1]
@pytest.mark.asyncio
async def test_configure_raises_other_rejections_at_once(gate) -> None:
gate.refusals = 1
gate.refusal_code = "unsupported_turn_detection"
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
with pytest.raises(duplex_client.DuplexProtocolError):
await _configure(probe)
assert len(gate.instances) == 1
@pytest.mark.asyncio
async def test_configure_gives_up_when_the_slot_wait_budget_is_spent(gate, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(bench_patch._RealtimeTTSProbe, "_SESSION_SLOT_WAIT_S", 0.0)
gate.refusals = 5
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
with pytest.raises(duplex_client.DuplexProtocolError) as excinfo:
await _configure(probe)
assert excinfo.value.code == "resource_exhausted"
assert len(gate.instances) == 1
class _SilenceSink:
"""Stands in for a configured ``DuplexClient``: records every appended chunk."""
def __init__(self) -> None:
self.config = duplex_client.SessionConfig()
self.chunks: list[bytes] = []
async def append_audio(self, pcm: bytes, *, is_speech: bool | None = None, video_frames=None) -> None:
del video_frames
assert is_speech is False
self.chunks.append(pcm)
@pytest.mark.asyncio
async def test_stream_silence_stops_at_the_turn_end(monkeypatch: pytest.MonkeyPatch) -> None:
"""Silence is fed until the turn settles, not for the whole cap."""
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
sink = _SilenceSink()
probe._client = sink
async def no_wait(_: float) -> None:
return
monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait)
streamed = await probe.stream_silence(seconds=12.0, chunk_ms=200, until=lambda: len(sink.chunks) >= 3)
assert len(sink.chunks) == 3
assert streamed == pytest.approx(0.6)
assert all(len(chunk) == sink.config.input_audio.byte_count(200) for chunk in sink.chunks)
@pytest.mark.asyncio
async def test_stream_silence_runs_to_the_cap_when_the_turn_never_settles(monkeypatch: pytest.MonkeyPatch) -> None:
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
sink = _SilenceSink()
probe._client = sink
async def no_wait(_: float) -> None:
return
monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait)
streamed = await probe.stream_silence(seconds=1.0, chunk_ms=200, until=lambda: False)
assert len(sink.chunks) == 5
assert streamed == pytest.approx(1.0)
class _TurnEvents:
def __init__(self, response_ids: list[str], audio: dict[str, bytes], text: dict[str, str] | None = None) -> None:
self.response_ids = response_ids
self._audio = audio
self._text = text or {}
def audio_bytes(self, response_id: str) -> bytes:
return self._audio.get(response_id, b"")
def response_text(self, response_id: str) -> str:
return self._text.get(response_id, "")
def test_turn_response_is_the_first_one_with_audio_even_if_the_model_spoke_again(caplog) -> None:
events = _TurnEvents(["r0", "r1", "r2"], {"r1": b"\x01\x00", "r2": b"\x02\x00"})
with caplog.at_level("WARNING", logger=bench_patch.logger.name):
assert bench_patch._seed_tts_turn_response_id(events, 0, 0) == "r1"
assert "spoke again" in caplog.text
def test_turn_response_ignores_responses_from_before_the_turn() -> None:
events = _TurnEvents(["r0", "r1"], {"r0": b"\x01\x00", "r1": b"\x02\x00"})
assert bench_patch._seed_tts_turn_response_id(events, 1, 0) == "r1"
def test_turn_without_audio_is_an_error() -> None:
events = _TurnEvents(["r0"], {})
with pytest.raises(RuntimeError, match="no audio response"):
bench_patch._seed_tts_turn_response_id(events, 0, 0)
def test_stall_report_says_the_model_never_answered() -> None:
events = _TurnEvents(["r0"], {"r0": b"\x01\x00"})
report = bench_patch._seed_tts_turn_stall_report(events, 1, 3, 30.0)
assert "turn 3 never started a response" in report
assert "30.0s of silence" in report
def test_stall_report_lists_a_response_that_never_finished() -> None:
events = _TurnEvents(["r1"], {"r1": b"\x01\x00\x02\x00"}, {"r1": "hello"})
report = bench_patch._seed_tts_turn_stall_report(events, 0, 0, 4.2)
assert "started 1 response(s) after 4.2s of silence" in report
assert "r1 (4 audio bytes, text 'hello')" in report
+425 -114
View File
@@ -1,6 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the public duplex client (fake transport, no server)."""
"""Unit tests for the public duplex clients (fake transports, no server).
``DuplexClient`` runs over a fake websocket; ``InlineDuplexClient`` over a
fake ``DuplexOmni`` whose handle records typed commands and replays typed
events. The demux / handshake tests that only need "feed a server event" are
parameterized over both clients.
"""
from __future__ import annotations
@@ -8,6 +14,8 @@ import asyncio
import base64
import json
import wave
from collections.abc import Callable
from typing import Any
from urllib.parse import parse_qs, urlsplit
import pytest
@@ -16,8 +24,11 @@ from vllm_omni.clients.duplex import (
AudioFormat,
ConnectionResumed,
DuplexClient,
DuplexClientBase,
DuplexConnectionError,
DuplexProtocolError,
DuplexSessionClosedError,
ErrorEvent,
EventCollector,
ReconnectPolicy,
SessionConfig,
@@ -31,17 +42,26 @@ from vllm_omni.clients.duplex import (
summarize_session_request_metrics,
write_pcm16_wav,
)
from vllm_omni.clients.inline_duplex import InlineDuplexClient
from vllm_omni.engine.duplex import commands as duplex_commands
from vllm_omni.engine.duplex import events as duplex_events
from vllm_omni.engine.duplex.messages import DuplexSessionError
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SESSION_ID = "duplex-0123456789abcdef"
# session.created is not journaled by the server, so it carries no server_event_seq.
SESSION_CREATED = {
"type": "session.created",
"session": {"session_id": "sess-1"},
"incarnation": 0,
"session": {"id": SESSION_ID, "model": "test-model"},
"attachment_generation": 1,
"resume_token": "tok-1",
"server_event_seq": 1,
}
SESSION_CLOSED = {"type": "session.closed", "session_id": "sess-1", "server_event_seq": 99}
SESSION_CLOSED = {"type": "session.closed", "session_id": SESSION_ID, "reason": "client_close", "server_event_seq": 99}
# ---------------------------------------------------------------------------
# Fake websocket transport
class FakeSocket:
@@ -88,14 +108,139 @@ def make_client(*sockets: FakeSocket, **kwargs) -> tuple[DuplexClient, list[str]
return client, calls
# ---------------------------------------------------------------------------
# Fake DuplexOmni for the inline client
class _WireEvent:
"""A typed-event stand-in that renders a fixed wire payload."""
def __init__(self, raw: dict[str, object]) -> None:
self.raw = raw
@property
def is_terminal(self) -> bool:
return self.raw.get("type") in {"session.closed", "session.expired"}
def to_realtime(self) -> dict[str, object]:
return dict(self.raw)
class FakeInlineHandle:
def __init__(self, session_id: str) -> None:
self.session_id = session_id
self.closed = False
self.close_calls = 0
self.commands: list[duplex_commands.DuplexCommand] = []
self.consumer_active = False
self._outbox: asyncio.Queue = asyncio.Queue()
def deliver(self, event: Any) -> None:
self._outbox.put_nowait(event)
def feed(self, raw: dict[str, object]) -> None:
self.deliver(_WireEvent(raw))
async def submit(self, command: duplex_commands.DuplexCommand) -> None:
if self.closed:
raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id)
self.commands.append(command)
async def events(self):
self.consumer_active = True
try:
while True:
event = await self._outbox.get()
yield event
if event.is_terminal:
return
finally:
self.consumer_active = False
async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None:
self.close_calls += 1
if self.closed:
return
self.closed = True
self.deliver(duplex_events.SessionClosed(session_id=self.session_id, reason=reason))
class FakeInlineOmni:
def __init__(self) -> None:
self.opened: list[dict[str, object]] = []
self.handles: list[FakeInlineHandle] = []
async def open_session(self, config: dict[str, object]) -> FakeInlineHandle:
self.opened.append(dict(config))
handle = FakeInlineHandle(SESSION_ID)
self.handles.append(handle)
handle.deliver(
duplex_events.SessionCreated(
session_id=SESSION_ID,
session={"id": SESSION_ID, "model": config.get("model")},
)
)
return handle
def make_inline_client(**kwargs) -> tuple[InlineDuplexClient, FakeInlineOmni]:
omni = FakeInlineOmni()
kwargs.setdefault("handshake_timeout_s", 5.0)
return InlineDuplexClient(omni, model="test-model", **kwargs), omni
# ---------------------------------------------------------------------------
# One rig over both clients: ``feed`` a server event, drive ``client``
class Rig:
def __init__(self, client: DuplexClientBase, feed: Callable[[dict[str, object]], None]) -> None:
self.client = client
self.feed = feed
def _websocket_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig:
sock = FakeSocket()
if handshake is not None:
sock.feed(handshake)
client, _ = make_client(sock)
return Rig(client, sock.feed)
def _inline_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig:
client, omni = make_inline_client()
original_open = omni.open_session
async def open_session(config: dict[str, object]) -> FakeInlineHandle:
handle = await original_open(config)
if handshake is not SESSION_CREATED:
# Replace the default session.created with the requested first event.
handle._outbox = asyncio.Queue()
if handshake is not None:
handle.feed(handshake)
return handle
omni.open_session = open_session # type: ignore[method-assign]
def feed(raw: dict[str, object]) -> None:
omni.handles[-1].feed(raw)
return Rig(client, feed)
@pytest.fixture(params=["websocket", "inline"])
def make_rig(request: pytest.FixtureRequest) -> Callable[..., Rig]:
return _websocket_rig if request.param == "websocket" else _inline_rig
# ---------------------------------------------------------------------------
# SessionConfig
def test_session_config_payload_defaults():
payload = SessionConfig().to_session_payload(model="m", session_id="s")
payload = SessionConfig().to_session_payload(model="m")
assert payload["model"] == "m"
assert payload["session_id"] == "s"
assert "session_id" not in payload
assert payload["modalities"] == ["audio", "text"]
assert payload["input_audio_format"] == "pcm16"
assert payload["output_audio_format"] == "pcm16"
@@ -137,19 +282,59 @@ def test_audio_format_math():
# ---------------------------------------------------------------------------
# Handshake and lifecycle
# Handshake and lifecycle (both clients)
async def test_handshake_adopts_session_and_acks():
@pytest.mark.asyncio
async def test_handshake_adopts_server_allocated_session_id(make_rig):
rig = make_rig()
async with rig.client as client:
assert client.session_id == SESSION_ID
assert client.session_info == {"id": SESSION_ID, "model": "test-model"}
rig.feed(SESSION_CLOSED)
@pytest.mark.asyncio
async def test_handshake_error_raises_protocol_error(make_rig):
rig = make_rig(handshake={"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}})
with pytest.raises(DuplexProtocolError) as excinfo:
async with rig.client:
pass
assert excinfo.value.code == "unsupported_audio_format"
@pytest.mark.asyncio
async def test_session_expired_raises_from_event_stream(make_rig):
rig = make_rig()
async with rig.client as client:
rig.feed(
{"type": "session.expired", "session_id": SESSION_ID, "reason": "lease_expired", "server_event_seq": 2}
)
with pytest.raises(DuplexSessionClosedError) as excinfo:
async for _ in client.events():
pass
assert "lease_expired" in excinfo.value.reason
# ---------------------------------------------------------------------------
# Websocket transport specifics
@pytest.mark.asyncio
async def test_websocket_handshake_sends_session_update_and_acks_sequenced_events():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
client, calls = make_client(sock)
async with client:
assert client.session_id == "sess-1"
assert client.session_id == SESSION_ID
assert client.resume_token == "tok-1"
assert calls == ["ws://test-host:8099/v1/realtime?duplex=1&model=test-model&autostart=0"]
assert sock.sent[0]["type"] == "session.update"
assert sock.sent[0]["session"]["model"] == "test-model"
assert "session_id" not in sock.sent[0]["session"]
# session.created is unjournaled (no seq): nothing to ack yet.
assert _acks(sock) == []
sock.feed({"type": "response.listen", "server_event_seq": 1})
await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 1} in _acks(sock))
close_task = asyncio.create_task(client.close())
await _drain(lambda: "session.close" in sock.sent_types())
@@ -158,28 +343,7 @@ async def test_handshake_adopts_session_and_acks():
assert "session.close" in sock.sent_types()
async def test_handshake_error_raises_protocol_error():
sock = FakeSocket()
sock.feed({"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}})
client, _ = make_client(sock)
with pytest.raises(DuplexProtocolError) as excinfo:
async with client:
pass
assert excinfo.value.code == "unsupported_audio_format"
async def test_session_expired_raises_from_event_stream():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
client, _ = make_client(sock)
async with client:
sock.feed({"type": "session.expired", "reason": "lease_expired", "server_event_seq": 2})
with pytest.raises(DuplexSessionClosedError) as excinfo:
async for _ in client.events():
pass
assert "lease_expired" in excinfo.value.reason
@pytest.mark.asyncio
async def test_resync_required_drops_resume_credential():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -187,7 +351,7 @@ async def test_resync_required_drops_resume_credential():
async with client:
assert client.resume_token == "tok-1"
stream = client.events()
sock.feed({"type": "session.resync_required", "session_id": "sess-1", "server_event_seq": 2})
sock.feed({"type": "session.resync_required", "session_id": SESSION_ID, "server_event_seq": 2})
async for event in stream:
if event.type == "session.resync_required":
break
@@ -202,9 +366,10 @@ async def test_resync_required_drops_resume_credential():
# ---------------------------------------------------------------------------
# Input events
# Input events (websocket wire shape)
@pytest.mark.asyncio
async def test_append_audio_tracks_cumulative_end_ms():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -224,6 +389,7 @@ async def test_append_audio_tracks_cumulative_end_ms():
sock.feed(SESSION_CLOSED)
@pytest.mark.asyncio
async def test_append_audio_strips_video_frame_data_url_prefix():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -238,6 +404,7 @@ async def test_append_audio_strips_video_frame_data_url_prefix():
sock.feed(SESSION_CLOSED)
@pytest.mark.asyncio
async def test_stream_pcm_chunking():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -249,6 +416,7 @@ async def test_stream_pcm_chunking():
sock.feed(SESSION_CLOSED)
@pytest.mark.asyncio
async def test_interruption_primitives_send_documented_events():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -262,22 +430,21 @@ async def test_interruption_primitives_send_documented_events():
# ---------------------------------------------------------------------------
# Response demultiplexing
# Response demultiplexing (both clients)
def _b64(data: bytes) -> str:
return base64.b64encode(data).decode("ascii")
async def test_response_handle_flow():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@pytest.mark.asyncio
async def test_response_handle_flow(make_rig):
rig = make_rig()
chunk = b"\x00\x01" * 2400 # 100 ms of pcm16 @ 24 kHz
client, _ = make_client(sock)
async with client:
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
sock.feed(
async with rig.client as client:
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
rig.feed(
{
"type": "response.output_audio.delta",
"response_id": "resp-1",
@@ -286,7 +453,7 @@ async def test_response_handle_flow():
"server_event_seq": 4,
}
)
sock.feed(
rig.feed(
{
"type": "response.output_audio_transcript.delta",
"response_id": "resp-1",
@@ -294,7 +461,7 @@ async def test_response_handle_flow():
"server_event_seq": 5,
}
)
sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6})
rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6})
async for response in client.responses():
chunks = [piece async for piece in response.audio()]
@@ -304,49 +471,46 @@ async def test_response_handle_flow():
assert response.played_ms == pytest.approx(100.0)
assert response.finished
break
sock.feed(SESSION_CLOSED)
rig.feed(SESSION_CLOSED)
async def test_listen_decision_yields_finished_silent_handle():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
client, _ = make_client(sock)
async with client:
sock.feed({"type": "response.listen", "server_event_seq": 2})
@pytest.mark.asyncio
async def test_listen_decision_yields_finished_silent_handle(make_rig):
rig = make_rig()
async with rig.client as client:
rig.feed({"type": "response.listen", "server_event_seq": 2})
async for response in client.responses():
assert response.decision == "listen"
assert response.finished
assert [piece async for piece in response.audio()] == []
break
sock.feed(SESSION_CLOSED)
rig.feed(SESSION_CLOSED)
async def test_listen_terminates_active_response():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
client, _ = make_client(sock)
async with client:
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3})
@pytest.mark.asyncio
async def test_listen_terminates_active_response(make_rig):
rig = make_rig()
async with rig.client as client:
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3})
async for response in client.responses():
await response.wait(timeout_s=5.0)
assert response.decision == "listen"
break
sock.feed(SESSION_CLOSED)
rig.feed(SESSION_CLOSED)
async def test_id_less_listen_never_closes_an_in_flight_response():
@pytest.mark.asyncio
async def test_id_less_listen_never_closes_an_in_flight_response(make_rig):
# An id-less listen is a standalone decision beat (e.g. a silence-skip
# while a response streams); it must surface as its own finished handle
# and leave the in-flight response to its real terminal.
sock = FakeSocket()
sock.feed(SESSION_CREATED)
client, _ = make_client(sock)
async with client:
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
sock.feed({"type": "response.listen", "server_event_seq": 4})
sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5})
rig = make_rig()
async with rig.client as client:
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
rig.feed({"type": "response.listen", "server_event_seq": 4})
rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5})
seen: list[tuple[str | None, str | None, bool]] = []
async for response in client.responses():
await response.wait(timeout_s=5.0)
@@ -355,20 +519,19 @@ async def test_id_less_listen_never_closes_an_in_flight_response():
break
assert ("resp-1", "speak", True) in seen
assert (None, "listen", True) in seen
sock.feed(SESSION_CLOSED)
rig.feed(SESSION_CLOSED)
async def test_terminal_listen_keeps_spoken_decision():
@pytest.mark.asyncio
async def test_terminal_listen_keeps_spoken_decision(make_rig):
# A terminal listen after the response spoke is the model yielding the
# turn; the handle must stay decision="speak" with its audio intact.
sock = FakeSocket()
sock.feed(SESSION_CREATED)
rig = make_rig()
chunk = b"\x00\x01" * 2400
client, _ = make_client(sock)
async with client:
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
sock.feed(
async with rig.client as client:
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
rig.feed(
{
"type": "response.output_audio.delta",
"response_id": "resp-1",
@@ -377,24 +540,23 @@ async def test_terminal_listen_keeps_spoken_decision():
"server_event_seq": 4,
}
)
sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5})
rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5})
async for response in client.responses():
chunks = [piece async for piece in response.audio()]
assert chunks == [chunk]
assert response.decision == "speak"
assert response.finished
break
sock.feed(SESSION_CLOSED)
rig.feed(SESSION_CLOSED)
async def test_error_event_surfaces_on_response_path():
@pytest.mark.asyncio
async def test_error_event_surfaces_on_response_path(make_rig):
# A rejected send produces no response; a consumer waiting in responses()
# must see the rejection instead of waiting forever.
sock = FakeSocket()
sock.feed(SESSION_CREATED)
client, _ = make_client(sock)
async with client:
sock.feed(
rig = make_rig()
async with rig.client as client:
rig.feed(
{
"type": "error",
"error": {"code": "input_audio_buffer_empty", "message": "empty commit"},
@@ -405,9 +567,37 @@ async def test_error_event_surfaces_on_response_path():
async for _ in client.responses():
pass
assert excinfo.value.code == "input_audio_buffer_empty"
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
@@ -425,24 +615,24 @@ async def test_slow_audio_consumer_drops_oldest_instead_of_stalling():
# ---------------------------------------------------------------------------
# Resume
# Resume (websocket only)
@pytest.mark.asyncio
async def test_resume_after_transport_drop():
first = FakeSocket()
first.feed(SESSION_CREATED)
second = FakeSocket()
# Mirror the real resume activation payload
# (entrypoints/duplex/serving.py, activation_payload_factory): it carries
# no nested "session" object, only the identity fields.
# the identity fields plus the (empty) public session object.
second.feed(
{
"type": "session.resumed",
"session_id": "sess-1",
"incarnation": 0,
"attachment_generation": 1,
"session_id": SESSION_ID,
"session": {},
"attachment_generation": 2,
"resume_token": "tok-2",
"server_event_seq": 5,
}
)
client, calls = make_client(
@@ -453,34 +643,39 @@ async def test_resume_after_transport_drop():
async with client:
stream = client.events()
async def take_two():
async def collect_until_resumed():
seen = []
async for event in stream:
seen.append(event)
if len(seen) == 2:
if isinstance(event, SessionResumed):
return seen
consumer = asyncio.create_task(take_two())
consumer = asyncio.create_task(collect_until_resumed())
await asyncio.sleep(0) # let the consumer subscribe before the drop
await asyncio.sleep(0)
first.feed({"type": "response.listen", "server_event_seq": 7})
await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 7} in _acks(first))
first.feed(RuntimeError("transport dropped"))
seen = await asyncio.wait_for(consumer, timeout=5.0)
assert isinstance(seen[0], ConnectionResumed)
assert isinstance(seen[1], SessionResumed)
assert [type(event) for event in seen[-2:]] == [ConnectionResumed, SessionResumed]
assert client.resume_token == "tok-2"
# The superseded transport is closed, not leaked.
assert first.closed
# The activation payload has no session object; the info captured at
# the original handshake must survive the resume.
assert client.session_info == {"session_id": "sess-1"}
assert client.session_info == {"id": SESSION_ID, "model": "test-model"}
assert len(calls) == 2
resume = second.sent[0]
assert resume["type"] == "session.resume"
assert resume["session_id"] == "sess-1"
assert resume["incarnation"] == 0
assert resume["resume_token"] == "tok-1"
assert resume["last_received_server_event_seq"] == 1
assert resume == {
"type": "session.resume",
"session_id": SESSION_ID,
"resume_token": "tok-1",
"last_received_server_event_seq": 7,
}
second.feed(SESSION_CLOSED)
@pytest.mark.asyncio
async def test_no_reconnect_policy_surfaces_closed():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -493,6 +688,7 @@ async def test_no_reconnect_policy_surfaces_closed():
pass
@pytest.mark.asyncio
async def test_heartbeat_survives_transient_send_failure():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -514,6 +710,7 @@ async def test_heartbeat_survives_transient_send_failure():
sock.feed(SESSION_CLOSED)
@pytest.mark.asyncio
async def test_aexit_sends_session_close_on_error_path():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -534,6 +731,7 @@ def test_target_url_forces_autostart_off():
assert query["autostart"] == ["0"]
@pytest.mark.asyncio
async def test_fatal_resume_error_gives_up():
first = FakeSocket()
first.feed(SESSION_CREATED)
@@ -553,6 +751,121 @@ async def test_fatal_resume_error_gives_up():
assert len(calls) == 2 # no retry after a fatal resume error
# ---------------------------------------------------------------------------
# Inline client over a fake DuplexOmni
@pytest.mark.asyncio
async def test_inline_handshake_opens_session_with_payload_and_adopts_id():
client, omni = make_inline_client(config=SessionConfig(ref_audio="data:audio/wav;base64,AAA="))
async with client:
assert omni.opened == [client.config.to_session_payload(model="test-model")]
assert omni.opened[0]["ref_audio"] == "data:audio/wav;base64,AAA="
assert client.session_id == SESSION_ID
handle = omni.handles[0]
# The session was opened with the full session object already; the
# base handshake's session.update must not be re-submitted as a patch
# (a ref_audio patch would be rejected by the runner).
assert handle.commands == []
await client.send({"type": "session.update", "session": {"instructions": "later"}})
assert [type(command) for command in handle.commands] == [duplex_commands.UpdateSession]
assert handle.commands[0].patch == {"instructions": "later"}
@pytest.mark.asyncio
async def test_inline_client_submits_typed_commands_in_order():
client, omni = make_inline_client()
async with client:
handle = omni.handles[0]
pcm = b"\x01\x00" * 1600
await client.append_audio(pcm, is_speech=True)
await client.commit(create_response=True)
await client.cancel_response("resp-1")
await client.ack_playback(120.9, response_id="resp-1")
assert [type(command) for command in handle.commands] == [
duplex_commands.AppendAudio,
duplex_commands.Commit,
duplex_commands.CancelResponse,
duplex_commands.AckPlayback,
]
append, commit, cancel, ack = handle.commands
assert append.is_speech is True
assert append.audio_end_ms == 100
assert append.event_id
assert commit.final is True and commit.create_response is True
assert cancel.response_id == "resp-1"
assert ack.played_ms == 120 and ack.response_id == "resp-1"
@pytest.mark.asyncio
async def test_inline_session_close_maps_to_handle_close():
client, omni = make_inline_client()
async with client:
pass
handle = omni.handles[0]
assert handle.closed
assert handle.close_calls >= 1
assert not any(isinstance(command, duplex_commands.CloseSession) for command in handle.commands)
# The typed session.closed reached the client as the expected end.
with pytest.raises(DuplexSessionClosedError):
await client.send({"type": "session.heartbeat"})
@pytest.mark.asyncio
async def test_inline_client_forwards_heartbeats_and_drops_event_acks():
client, omni = make_inline_client()
async with client:
handle = omni.handles[0]
await client.send({"type": "session.heartbeat"})
await client.send({"type": "session.event_ack", "server_event_seq": 3})
# A heartbeat is a lease touch the engine must see; the ack is a journal concern.
assert [type(command).__name__ for command in handle.commands] == ["Heartbeat"]
@pytest.mark.asyncio
async def test_inline_malformed_payload_surfaces_as_error_event():
client, omni = make_inline_client()
async with client:
handle = omni.handles[0]
# Subscribe before sending: the translation error is dispatched inline.
first_event = asyncio.create_task(client.wait_for("error", timeout_s=5.0))
await asyncio.sleep(0)
event_id = await client.send({"type": "playback.ack"}) # played_ms missing
event = await first_event
assert isinstance(event, ErrorEvent)
assert event.code == "bad_event"
assert event.related_event_id == event_id
assert handle.commands == []
with pytest.raises(DuplexProtocolError) as excinfo:
async for _ in client.responses():
pass
assert excinfo.value.code == "bad_event"
@pytest.mark.asyncio
async def test_inline_teardown_cancels_pump_and_closes_handle():
client, omni = make_inline_client()
async with client:
handle = omni.handles[0]
await _drain(lambda: handle.consumer_active)
pump = client._pump_task
assert pump is not None and not pump.done()
assert pump.done()
assert not handle.consumer_active
assert handle.closed
@pytest.mark.asyncio
async def test_inline_send_after_handle_closed_raises_connection_error():
client, omni = make_inline_client()
async with client:
handle = omni.handles[0]
handle.closed = True # the engine ended the session without an event yet
with pytest.raises(DuplexConnectionError, match="not open"):
await client.send({"type": "input_audio_buffer.clear"})
handle.closed = False
# ---------------------------------------------------------------------------
# Collector
@@ -600,8 +913,7 @@ def test_build_realtime_url_with_model_extra_query():
url = build_realtime_url(
"ws://localhost:8099/v1/realtime?custom=1",
"openbmb/MiniCPM-o-4_5",
session_id="session-a",
extra_query={"native_duplex": "1"},
extra_query={"probe": "1"},
)
query = parse_qs(urlsplit(url).query)
@@ -609,8 +921,7 @@ def test_build_realtime_url_with_model_extra_query():
"custom": ["1"],
"duplex": ["1"],
"model": ["openbmb/MiniCPM-o-4_5"],
"native_duplex": ["1"],
"session_id": ["session-a"],
"probe": ["1"],
}
@@ -619,12 +930,11 @@ def test_build_realtime_url_resume_only_when_autostart_disabled():
"ws://localhost:8099/v1/realtime?duplex=1",
"openbmb/MiniCPM-o-4_5",
autostart=False,
extra_query={"native_duplex": "1"},
)
query = parse_qs(urlsplit(url).query)
assert query["autostart"] == ["0"]
assert query["native_duplex"] == ["1"]
assert "session_id" not in query
def test_event_collector_partitions_audio_by_response():
@@ -989,6 +1299,7 @@ def test_duplex_unit_boundary_and_residual_math():
assert chunk_period_ms([{"type": "session.created", "session": {}}]) == 1000
@pytest.mark.asyncio
async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
sock = FakeSocket()
sock.feed(SESSION_CREATED)
@@ -1011,16 +1322,16 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
sock.feed(SESSION_CLOSED)
def test_build_realtime_url_native_duplex_flag_and_http_scheme():
url = build_realtime_url("http://localhost:8099/v1/realtime", None, native_duplex=None)
def test_build_realtime_url_rewrites_http_scheme_and_rejects_others():
url = build_realtime_url("http://localhost:8099/v1/realtime", None)
parts = urlsplit(url)
assert parts.scheme == "ws"
assert parse_qs(parts.query) == {"duplex": ["1"]}
url = build_realtime_url("https://host/v1/realtime", "m", native_duplex=False)
url = build_realtime_url("https://host/v1/realtime", "m", autostart=True)
parts = urlsplit(url)
assert parts.scheme == "wss"
assert parse_qs(parts.query)["native_duplex"] == ["0"]
assert parse_qs(parts.query) == {"duplex": ["1"], "model": ["m"], "autostart": ["1"]}
with pytest.raises(ValueError):
build_realtime_url("ftp://host/v1/realtime", "m")
+4 -3
View File
@@ -18,13 +18,13 @@ from vllm_omni.clients.personaplex import (
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_minicpmo45_session_config_matches_native_duplex_deployment():
def test_minicpmo45_session_config_matches_duplex_deployment():
config = create_minicpmo45_session_config(ref_audio="data:audio/wav;base64,AAA=", temperature=0.0)
payload = config.to_session_payload(model="openbmb/MiniCPM-o-4_5")
assert payload["ref_audio"] == "data:audio/wav;base64,AAA="
assert payload["overlap_policy"] == "listen_only"
assert payload["playback_commit_policy"] == "ack_only"
assert payload["extra_body"]["native_duplex"] is True
assert "native_duplex" not in payload["extra_body"]
assert payload["extra_body"]["auto_response"] is True
assert payload["temperature"] == 0.0
@@ -41,8 +41,9 @@ def test_personaplex_session_config_matches_deployment():
@pytest.mark.parametrize(
"duplex_package",
[
# PersonaPlex and Nemotron VoiceChat rejoin this list with the follow-up
# PRs that port them to the duplex plugin framework.
"vllm_omni.model_executor.models.minicpmo_4_5.duplex",
"vllm_omni.model_executor.models.personaplex.duplex",
],
)
def test_model_duplex_packages_stay_clear_of_client_library(duplex_package: str):
+42 -4
View File
@@ -1328,7 +1328,11 @@ class TestDeployConfigLoading:
stages = merge_pipeline_deploy(pipeline, deploy)
assert deploy.session_mode == "duplex"
assert deploy.active_stream_window == max_sessions
# ``0`` disables the limiter: the configured session/stage capacity is
# what bounds the pipeline. Any positive window serializes concurrent
# audio first-packet generation (measured: TTFP -23% at 0 vs 4 on
# L20X, 128 requests / concurrency 8).
assert deploy.active_stream_window == 0
assert deploy.duplex_session.max_sessions == max_sessions
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
@@ -1558,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]
@@ -1583,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"]
@@ -1600,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}
@@ -2587,6 +2605,26 @@ class TestPlatformOverrides:
assert rocm.stages[0].enforce_eager is None
assert rocm.stages[1].enforce_eager is True
@pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"])
@pytest.mark.parametrize("platform", ["cuda", "npu", "rocm"])
def test_qwen3_tts_default_code2wav_dtype_is_bf16(self, deploy_name, platform):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
deploy = _apply_platform_overrides(deploy, platform=platform)
stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy)
assert stages[1].yaml_engine_args["dtype"] == "bfloat16"
@pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"])
@pytest.mark.parametrize("platform", ["cuda", "npu"])
@pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"])
def test_qwen3_tts_propagates_explicit_code2wav_dtype(self, deploy_name, platform, dtype):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
deploy.stages[1].engine_extras["dtype"] = dtype
deploy = _apply_platform_overrides(deploy, platform=platform)
stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy)
assert stages[1].yaml_engine_args["dtype"] == dtype
def test_higgs_audio_v3_rocm_uses_triton_attention(self):
deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
+1 -1
View File
@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
"""Make an inventory expansion an explicit review decision."""
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
assert category_counts == {
EnvironmentVariableCategory.PUBLIC_OMNI: 27,
EnvironmentVariableCategory.PUBLIC_OMNI: 28,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 64,
@@ -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
@@ -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
+5 -81
View File
@@ -26,8 +26,8 @@
{
"dataset_name": "seed-tts",
"dataset_path": "zhaochenyang20/seed-tts-eval",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"backend": "openai-realtime-tts",
"endpoint": "/v1/realtime",
"num_prompts": [
32,
64,
@@ -42,73 +42,7 @@
"disable_shuffle": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
"modalities": [
"text",
"audio"
],
"chat_template_kwargs": {
"enable_thinking": false,
"use_tts_template": true
}
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
0.9731,
1.1986,
1.686
],
"mean_ttft_ms": [
152.5665,
270.2559,
301.455
],
"mean_e2el_ms": [
1027.2953,
3306.3735,
4686.5366
],
"mean_audio_ttfp_ms": [
458.1553,
2030.2335,
2068.6233
],
"mean_audio_rtf": [
0.2357,
0.791,
1.0324
]
},
"A3": {
"request_throughput": [
0.5383,
0.6042,
0.7547
],
"mean_ttft_ms": [
333.2633,
533.8656,
547.3388
],
"mean_e2el_ms": [
1857.2154,
6600.5101,
10450.6811
],
"mean_audio_ttfp_ms": [
986.4666,
3411.0754,
3352.7515
],
"mean_audio_rtf": [
0.4423,
1.5734,
2.3024
]
}
}
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
}
]
},
@@ -138,8 +72,8 @@
{
"dataset_name": "seed-tts",
"dataset_path": "zhaochenyang20/seed-tts-eval",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"backend": "openai-realtime-tts",
"endpoint": "/v1/realtime",
"num_prompts": [
32,
64,
@@ -154,16 +88,6 @@
"disable_shuffle": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
"modalities": [
"text",
"audio"
],
"chat_template_kwargs": {
"enable_thinking": false,
"use_tts_template": true
}
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
}
]
@@ -43,43 +43,7 @@
"extra_body": {
"save_duplex_request_metrics": true
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
0.3038
],
"mean_ttft_ms": [
139.997
],
"mean_e2el_ms": [
3269.8807
],
"mean_audio_ttfp_ms": [
456.073
],
"mean_audio_rtf": [
0.2254
]
},
"A3": {
"request_throughput": [
0.1594
],
"mean_ttft_ms": [
274.164
],
"mean_e2el_ms": [
6250.7696
],
"mean_audio_ttfp_ms": [
994.422
],
"mean_audio_rtf": [
0.46163
]
}
}
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
}
]
}
@@ -26,8 +26,8 @@
{
"dataset_name": "seed-tts",
"dataset_path": "zhaochenyang20/seed-tts-eval",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"backend": "openai-realtime-tts",
"endpoint": "/v1/realtime",
"num_prompts": [
10,
40
@@ -40,41 +40,7 @@
"disable_shuffle": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
"modalities": [
"text",
"audio"
],
"chat_template_kwargs": {
"enable_thinking": false,
"use_tts_template": true
}
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
0.9731,
1.1986
],
"mean_ttft_ms": [
152.5665,
270.2559
],
"mean_e2el_ms": [
1027.2953,
3306.3735
],
"mean_audio_ttfp_ms": [
458.1553,
2030.2335
],
"mean_audio_rtf": [
0.2357,
0.791
]
}
}
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
}
]
}
@@ -157,14 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
"ref_audio": None,
"ref_text": None,
},
# NOTE: This is delegating down the TTS adapter, so the message is both specific
# to the Qwen3TTS adapter and dependent on whether or not the model has any speakers
# (as opposed to having valida speakers, but getting an invalid speaker). it would
# be more ideal to validate that this is landing on `.validate()` on the adapter
# and ensuring that the returned message matches what we would get from calling the
# adapter validate directly.
("no speakers configured",),
id="customvoice_invalid_voice",
("Base checkpoint does not support task_type='CustomVoice'",),
id="base_checkpoint_customvoice_task_mismatch",
),
pytest.param(
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import sys
import types
import pytest
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
from vllm_omni.diffusion.attention.ops.block_sparse import (
build_prefix_dense_block_map,
fastvideo_block_sparse_attn_bshd,
mean_pool_tiles,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.mark.parametrize("prefix_blocks,topk", [(0, 1), (2, 0), (2, 1), (2, 3), (2, 5)])
def test_block_map_keeps_prefix_dense_and_selects_non_prefix_keys(prefix_blocks, topk):
sparse_blocks = 3
blocks = prefix_blocks + sparse_blocks
scores = torch.arange(2 * blocks * blocks, dtype=torch.float32).reshape(1, 2, blocks, blocks)
block_map = build_prefix_dense_block_map(scores, prefix_blocks, sparse_blocks, topk)
expected = torch.zeros_like(scores, dtype=torch.bool)
# Scores increase with key index, so top-k always selects the last keys.
for query_block in range(blocks):
for key_block in range(blocks):
if (
query_block < prefix_blocks
or key_block < prefix_blocks
or key_block >= blocks - min(topk, sparse_blocks)
):
expected[..., query_block, key_block] = True
assert torch.equal(block_map, expected)
@pytest.mark.parametrize("block_size", [4, 32, 64])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_pooling_uses_valid_tile_sizes_and_fp32_accumulation(block_size, dtype):
sizes = torch.tensor([block_size, block_size - 1, 0], dtype=torch.int32)
tiles = torch.arange(2 * 3 * block_size * 2 * 3).reshape(2, 3, block_size, 2, 3).to(dtype)
tiles[:, 1, -1] = 0
tiles[:, 2] = 0
expected = torch.stack(
[tiles[:, 0].float().mean(dim=1), tiles[:, 1, :-1].float().mean(dim=1), torch.zeros(2, 2, 3)], dim=2
)
result = mean_pool_tiles(tiles.flatten(1, 2), sizes, block_size)
assert result.dtype == torch.float32
torch.testing.assert_close(result, expected, rtol=0, atol=0)
def test_provider_wrapper_restores_transport_rows_and_bshd_layout(monkeypatch):
monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False)
calls = []
def block_sparse_attn(q, k, v, block_map, sizes):
calls.append((q.shape, block_map.shape, sizes.tolist()))
assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
assert sizes.dtype == torch.int32
return q + k + v, None
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
setattr(module, "block_sparse_attn", block_sparse_attn)
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
query = torch.randn(1, 4 * 64, 2, 8)
block_map = torch.ones(1, 2, 4, 4, dtype=torch.bool)
sizes = torch.tensor([64, 17, 64, 0])
output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
expected = torch.zeros_like(query)
expected[:, : 3 * 64] = query[:, : 3 * 64] * 3
torch.testing.assert_close(output, expected)
assert calls == [(torch.Size([1, 2, 3 * 64, 8]), torch.Size([1, 2, 3, 3]), [64, 17, 64])]
def test_custom_op_fake_preserves_metadata_without_loading_provider(monkeypatch):
monkeypatch.setitem(sys.modules, "fastvideo_kernel", None)
with FakeTensorMode():
query = torch.empty(1, 4 * 64, 2, 8, dtype=torch.bfloat16)
block_map = torch.empty(1, 2, 4, 4, dtype=torch.bool)
sizes = torch.empty(4, dtype=torch.int32)
output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
assert output.shape == query.shape
assert output.dtype == query.dtype
assert output.device == query.device
@pytest.mark.parametrize("error_type", [ImportError, RuntimeError, torch.AcceleratorError])
def test_native_provider_failure_only_recovers_without_accelerator_fault(monkeypatch, mocker, error_type):
monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1")
error = error_type("native provider failure")
native = mocker.Mock()
native.is_supported.return_value = True
native.block_sparse_attn_sm100a.side_effect = error
provider = types.ModuleType("fastvideo_kernel")
setattr(provider, "block_sparse_attn_sm100a", native)
index = types.ModuleType("fastvideo_kernel.triton_kernels.index")
setattr(
index,
"map_to_index",
mocker.Mock(return_value=(torch.zeros(1, dtype=torch.int32), torch.ones(1, dtype=torch.int32))),
)
triton = types.ModuleType("fastvideo_kernel.block_sparse_attn")
fallback = mocker.Mock(side_effect=lambda q, *args: (q.clone(), None))
setattr(triton, "block_sparse_attn", fallback)
monkeypatch.setitem(sys.modules, "fastvideo_kernel", provider)
monkeypatch.setitem(
sys.modules, "fastvideo_kernel.triton_kernels", types.ModuleType("fastvideo_kernel.triton_kernels")
)
monkeypatch.setitem(sys.modules, "fastvideo_kernel.triton_kernels.index", index)
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", triton)
query = torch.arange(128 * 2 * 8, dtype=torch.float32).reshape(1, 128, 2, 8)
block_map = torch.ones(1, 2, 2, 2, dtype=torch.bool)
sizes = torch.tensor([64, 64], dtype=torch.int32)
if error_type is torch.AcceleratorError:
with pytest.raises(torch.AcceleratorError, match="native provider failure") as raised:
fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
assert raised.value is error
fallback.assert_not_called()
else:
output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
torch.testing.assert_close(output, query, rtol=0, atol=0)
fallback.assert_called_once()
+3 -161
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import math
import sys
import types
from typing import Any
@@ -12,14 +11,10 @@ import torch
from vllm_omni.diffusion.attention.backends.abstract import (
AttentionMetadata,
PackedPaddingMetadata,
VideoTokenLayout,
VideoTokenSpan,
)
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
FastVideoVSABackend,
FastVideoVSAImpl,
_build_h3_block_map,
_get_h3_tile_metadata,
)
from vllm_omni.diffusion.attention.backends.registry import (
DiffusionAttentionBackendEnum,
@@ -218,144 +213,8 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
extra: dict[str, Any] = {
"vsa_h3_prefix_segments": prefix_segments,
}
if gate is not None:
extra["gate_compress"] = gate
return AttentionMetadata(
packed_padding=packed_padding,
video_layout=VideoTokenLayout(
used_len=sum(prefix_segments) + math.prod(video_shape),
video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
),
extra=extra,
)
def _fake_block_sparse_kernel(monkeypatch, calls):
"""Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
calls["block_map"] = block_map.clone()
calls["variable_block_sizes"] = variable_block_sizes.tolist()
calls["q_shape"] = tuple(q.shape)
blocks = block_map.shape[-1]
assert blocks * 64 == q.shape[2]
# Expand the per-block map to a per-token mask so the stand-in honours
# the sparsity contract the real kernel implements, and drop the rows
# each partial block pads with, which the real kernel skips through
# variable_block_sizes.
token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
token_mask = token_mask & within.reshape(1, 1, 1, -1)
out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
return out.to(q.dtype), None
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
setattr(module, "block_sparse_attn", block_sparse_attn)
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
def _h3_impl(**backend_kwargs):
kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
kwargs.update(backend_kwargs)
return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
calls: dict[str, Any] = {}
_fake_block_sparse_kernel(monkeypatch, calls)
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5, 70), (2, 4, 4)
seq_len = sum(prefix_segments) + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
assert output.shape == query.shape
# topk covers every video block, so the block map is dense and the tiled
# kernel must reproduce plain attention over the same rows.
reference = torch.nn.functional.scaled_dot_product_attention(
*(x.float().transpose(1, 2) for x in (query, query, query))
).transpose(1, 2)
torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
# 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
assert calls["variable_block_sizes"] == [5, 64, 6, 32]
def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
calls: dict[str, Any] = {}
_fake_block_sparse_kernel(monkeypatch, calls)
impl = _h3_impl(topk=1)
prefix_segments, video_shape = (5,), (8, 4, 4)
seq_len = 5 + 8 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
# 1 prefix block + 2 video blocks. The odd count is padded to an even one
# for the paired-CTA kernel contract, and the Triton route slices that
# transport-only partner off again before launching.
block_map = calls["block_map"]
assert block_map.shape[-1] == 3
assert calls["variable_block_sizes"] == [5, 64, 64]
assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
assert block_map[..., :1].all(), "prefix keys are exempt from selection"
assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
# The same odd geometry must still untile back onto the original rows.
dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
reference = torch.nn.functional.scaled_dot_product_attention(
*(x.float().transpose(1, 2) for x in (query, query, query))
).transpose(1, 2)
torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
_fake_block_sparse_kernel(monkeypatch, {})
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5,), (2, 4, 4)
seq_len = 5 + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
zero_gate = torch.zeros_like(query)
with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
live_gate = torch.full_like(query, 0.5)
with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
assert not torch.allclose(with_live_gate.float(), without_gate.float())
def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
_fake_block_sparse_kernel(monkeypatch, {})
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5,), (2, 4, 4)
valid = 5 + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
padding = PackedPaddingMetadata(
q_length=valid,
kv_length=valid,
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
)
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
assert output.shape == query.shape
assert torch.count_nonzero(output[:, valid:]) == 0
def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
# The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
# The backend advertises supports_packed_mask_free, so a model may skip
# building the padding mask. Every fallback out of the VSA path must honour
# that contract itself; SDPA reads attn_mask and nothing else.
impl = FastVideoVSAImpl(
@@ -374,6 +233,8 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
)
monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
perturbed_input = query.clone()
perturbed_input[:, valid:] += 100.0
@@ -384,22 +245,3 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
assert baseline.shape == query.shape
torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
assert torch.count_nonzero(baseline[:, valid:]) == 0
def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
(5, 70, 9), (5, 6, 6), torch.device("cpu")
)
assert sizes[:4].tolist() == [5, 64, 6, 9]
assert prefix_blocks == 4
assert video_blocks == 8
assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
assert block_map[:, :, :2].all()
assert block_map[..., :2].all()
assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()
@@ -0,0 +1,136 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
from vllm_omni.diffusion.attention import layer
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSABackend, FastVideoVSAImpl
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
from vllm_omni.diffusion.data import AttentionSpec
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
class _SpecializedImpl(FastVideoVSAImpl):
pass
@pytest.fixture
def local_attention(monkeypatch):
monkeypatch.setattr(layer, "get_current_diffusion_config_or_none", lambda: None)
monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention())
@pytest.mark.parametrize("explicit", [False, True])
@pytest.mark.parametrize("override", [False, True])
def test_specialization_preserves_selected_backend_and_options(mocker, local_attention, explicit, override):
spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=7) if explicit else None
select = mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, spec))
attention = layer.Attention(
num_heads=2,
head_size=8,
softmax_scale=8**-0.5,
causal=False,
role="video.self",
role_category="self",
qkv_layout="BSND",
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl} if override else None,
)
assert attention.attn_backend is FastVideoVSABackend
assert type(attention.attention) is (_SpecializedImpl if override else FastVideoVSAImpl)
assert attention.attention.topk == (7 if explicit else 64)
assert attention.attention.qkv_layout == "BSND"
assert attention.backend_pref == "FASTVIDEO_VSA"
assert attention.backend_explicit is explicit
assert attention.attn_spec is spec
assert isinstance(attention.parallel_strategy, NoParallelAttention)
select.assert_called_once_with(
role="video.self", head_size=8, attention_config=None, role_category="self", allow_trtllm_default=False
)
def test_unselected_specialization_does_not_replace_dense_backend(mocker, local_attention):
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(SDPABackend, None))
attention = layer.Attention(
num_heads=2,
head_size=8,
softmax_scale=8**-0.5,
causal=False,
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
)
assert attention.attn_backend is SDPABackend
assert type(attention.attention) is SDPABackend.get_impl_cls()
def test_specialization_does_not_bypass_selection_errors(mocker, local_attention):
mocker.patch.object(layer, "get_attn_backend_for_role", side_effect=ImportError("optional kernel unavailable"))
with pytest.raises(ImportError, match="optional kernel unavailable"):
layer.Attention(
num_heads=2,
head_size=8,
softmax_scale=8**-0.5,
causal=False,
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
)
def test_specialization_must_extend_selected_implementation(mocker, local_attention):
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, None))
with pytest.raises(TypeError, match="must subclass the selected implementation"):
layer.Attention(
num_heads=2,
head_size=8,
softmax_scale=8**-0.5,
causal=False,
impl_overrides={"FASTVIDEO_VSA": SDPABackend.get_impl_cls()},
)
def test_platform_backend_capabilities_survive_model_specialization(mocker, local_attention):
from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import MiniMaxH3VSAImpl
class PlatformBackend(FastVideoVSABackend):
@classmethod
def supports_packed_mask_free(cls):
return False
@staticmethod
def get_supported_head_sizes():
return [128]
# A registry/platform may refine capabilities without changing the
# implementation. The model must not replace those capability decisions.
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
attention = layer.Attention(
num_heads=2,
head_size=128,
softmax_scale=128**-0.5,
causal=False,
impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl},
)
assert attention.attn_backend is PlatformBackend
assert not attention.attn_backend.supports_packed_mask_free()
assert attention.attn_backend.get_supported_head_sizes() == [128]
assert type(attention.attention) is MiniMaxH3VSAImpl
def test_incompatible_platform_implementation_is_not_discarded(mocker, local_attention):
class PlatformImpl(FastVideoVSAImpl):
pass
class PlatformBackend(FastVideoVSABackend):
@staticmethod
def get_impl_cls():
return PlatformImpl
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
with pytest.raises(TypeError, match="selected implementation .*PlatformImpl"):
layer.Attention(
num_heads=2,
head_size=8,
softmax_scale=8**-0.5,
causal=False,
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
)
+635
View File
@@ -0,0 +1,635 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import datetime
from contextlib import nullcontext
from types import SimpleNamespace
import pytest
import torch
from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.cache.seacache import (
SeaCacheBackend,
SeaCacheConfig,
SeaCacheRootHook,
apply_sea_cache_hook,
)
from vllm_omni.diffusion.cache.seacache.sea_filter import (
apply_sea_filter,
extrapolate_residual,
indicator_distance,
)
from vllm_omni.diffusion.cache.selector import get_cache_backend
from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext, get_extractor
from vllm_omni.diffusion.data import DiffusionCacheConfig
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TinyCosmos3Transformer(torch.nn.Module):
"""Small model that exposes a cache-neutral execution boundary."""
def _run_gen_layers(self, hidden_gen: torch.Tensor) -> torch.Tensor:
return hidden_gen * 2
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
text_ids: torch.Tensor | None = None,
text_mask: torch.Tensor | None = None,
video_shape: tuple[int, int, int] | None = None,
noisy_frame_mask: torch.Tensor | None = None,
control_latents: list[torch.Tensor] | torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
del timestep, text_ids, text_mask, video_shape, noisy_frame_mask
if kwargs:
raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
controls = (
[]
if control_latents is None
else [control_latents]
if isinstance(control_latents, torch.Tensor)
else list(control_latents)
)
inputs = [*controls, hidden_states]
gen_input = torch.cat(
[value.movedim(1, -1).flatten(1, 3) for value in inputs],
dim=1,
)
return self._run_gen_layers(gen_input)
def _extract_tiny_cosmos3_context(
module: TinyCosmos3Transformer,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
text_ids: torch.Tensor | None = None,
text_mask: torch.Tensor | None = None,
video_shape: tuple[int, int, int] | None = None,
noisy_frame_mask: torch.Tensor | None = None,
control_latents: list[torch.Tensor] | torch.Tensor | None = None,
**kwargs,
) -> CacheContext:
del timestep, text_ids, text_mask, video_shape
if kwargs:
raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
controls = (
[]
if control_latents is None
else [control_latents]
if isinstance(control_latents, torch.Tensor)
else list(control_latents)
)
vision_items = [*controls, hidden_states]
gen_input = torch.cat(
[value.movedim(1, -1).flatten(1, 3) for value in vision_items],
dim=1,
)
return CacheContext(
modulated_input=gen_input,
hidden_states=gen_input,
encoder_hidden_states=None,
temb=torch.zeros_like(gen_input[:, 0]),
run_transformer_blocks=lambda: (module._run_gen_layers(gen_input),),
postprocess=lambda output: output,
extra_states={
"sea_cache_latents": vision_items,
"sea_cache_noisy_frame_mask": noisy_frame_mask,
},
)
class Cosmos3OmniDiffusersPipeline:
def __init__(self) -> None:
self.transformer = TinyCosmos3Transformer()
self._current_step_index: int | None = None
self._current_sigma: float | None = None
self._num_timesteps: int | None = None
@property
def current_step_index(self) -> int | None:
return self._current_step_index
@property
def current_sigma(self) -> float | None:
return self._current_sigma
@property
def num_timesteps(self) -> int | None:
return self._num_timesteps
def _latent(value: float) -> torch.Tensor:
return torch.full((1, 2, 2, 2, 2), value)
def _run_step(
transformer: TinyCosmos3Transformer,
timestep: int,
value: float,
*,
hook: SeaCacheRootHook | None = None,
context: str = "cond",
control: float | None = None,
noisy_frame_mask: torch.Tensor | None = None,
) -> torch.Tensor:
controls = None if control is None else [_latent(control)]
cache_context = hook.cache_context(context) if hook is not None else nullcontext()
with torch.inference_mode(), cache_context:
return transformer(
hidden_states=_latent(value),
timestep=torch.tensor([timestep]),
noisy_frame_mask=noisy_frame_mask,
control_latents=controls,
)
def _apply_test_hook(
transformer: TinyCosmos3Transformer,
metadata: SimpleNamespace,
config: SeaCacheConfig | None = None,
) -> SeaCacheRootHook:
return apply_sea_cache_hook(
transformer,
config or SeaCacheConfig(threshold=100.0),
current_step_callback=lambda: metadata.step,
current_sigma_callback=lambda: metadata.sigma,
num_inference_steps_callback=lambda: metadata.num_steps,
extractor_fn=_extract_tiny_cosmos3_context,
)
def test_config_validation() -> None:
assert SeaCacheConfig().threshold == 0.25
with pytest.raises(ValueError, match="residual_order"):
SeaCacheConfig(residual_order=-1)
with pytest.raises(ValueError, match="max_consecutive_cached"):
SeaCacheConfig(max_consecutive_cached=-1)
def test_sea_filter_matches_reference_equation() -> None:
hidden = torch.randn(3, 4, 5, 2, dtype=torch.float32)
sigma = 0.4
power_exp = 3.0
spectrum = torch.fft.fftn(hidden, dim=(0, 1, 2))
gain = None
for axis in (0, 1, 2):
frequencies = torch.fft.fftfreq(hidden.shape[axis], dtype=torch.float32)
clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16)
axis_gain = (1.0 - sigma) * clean_power / ((1.0 - sigma) ** 2 * clean_power + sigma**2 + 1e-16)
shape = [1] * hidden.ndim
shape[axis] = hidden.shape[axis]
gain = axis_gain.reshape(shape) if gain is None else gain * axis_gain.reshape(shape)
assert gain is not None
gain = gain / gain.mean()
expected = torch.fft.ifftn(spectrum * gain, dim=(0, 1, 2)).real
torch.testing.assert_close(
apply_sea_filter(hidden, sigma=sigma, power_exp=power_exp),
expected,
)
def test_indicator_distance_and_linear_extrapolation() -> None:
previous = [torch.ones(2, 2)]
current = [torch.full((2, 2), 1.5)]
assert indicator_distance(current, previous) == pytest.approx(0.5)
assert indicator_distance([torch.ones(3)], previous) == float("inf")
history = [
(0, torch.full((2, 2), 2.0)),
(2, torch.full((2, 2), 6.0)),
]
torch.testing.assert_close(
extrapolate_residual(history, step=3, order=1),
torch.full((2, 2), 8.0),
)
torch.testing.assert_close(
extrapolate_residual(history, step=3, order=0),
torch.full((2, 2), 6.0),
)
quadratic_history = [
(0, torch.zeros(2, 2)),
(1, torch.ones(2, 2)),
(2, torch.full((2, 2), 4.0)),
]
torch.testing.assert_close(
extrapolate_residual(quadratic_history, step=3, order=2),
torch.full((2, 2), 9.0),
)
def test_hook_skips_middle_steps_and_forces_endpoints() -> None:
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=4)
hook = _apply_test_hook(
transformer,
metadata,
SeaCacheConfig(threshold=100.0, max_consecutive_cached=2),
)
hook.refresh(transformer)
for step, timestep in enumerate((1000, 750, 500, 250)):
metadata.step = step
metadata.sigma = timestep / 1000
_run_step(transformer, timestep, 1.0 - step * 0.01, hook=hook)
assert hook.full_count == 2
assert hook.skip_count == 2
assert [step for step, _ in hook.state_manager._states["cond"].history] == [0, 3]
def test_parameter_sharded_hook_skips_when_all_shard_ranks_agree(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.distributed import parallel_state
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
hook = _apply_test_hook(transformer, metadata)
hook._parameter_sharded = True
fs_group = object()
reduced_groups: list[object] = []
monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(
parallel_state,
"get_fs_group",
lambda: SimpleNamespace(world_size=2, device_group=fs_group),
)
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
def all_reduce(decision, *, op, group):
assert op == torch.distributed.ReduceOp.MAX
reduced_groups.append(group)
monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
_run_step(transformer, 1000, 1.0, hook=hook)
metadata.step = 1
metadata.sigma = 0.5
_run_step(transformer, 500, 0.99, hook=hook)
assert hook.full_count == 1
assert hook.skip_count == 1
assert reduced_groups == [fs_group, fs_group]
def test_parameter_sharded_peer_forces_full_compute(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.distributed import parallel_state
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
hook = _apply_test_hook(transformer, metadata)
hook._parameter_sharded = True
monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(
parallel_state,
"get_fs_group",
lambda: SimpleNamespace(world_size=2, device_group=object()),
)
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
reduce_count = 0
def all_reduce(decision, *, op, group):
nonlocal reduce_count
del op, group
reduce_count += 1
if reduce_count == 2:
decision.fill_(1)
monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
_run_step(transformer, 1000, 1.0, hook=hook)
metadata.step = 1
metadata.sigma = 0.5
_run_step(transformer, 500, 0.99, hook=hook)
assert hook.full_count == 2
assert hook.skip_count == 0
assert hook.state_manager._states["cond"].accumulated_distance == 0.0
def _uneven_cfg_sharded_worker(
rank: int,
init_method: str,
result_queue: torch.multiprocessing.Queue,
) -> None:
from vllm_omni.diffusion.distributed import parallel_state
torch.distributed.init_process_group(
"gloo",
init_method=init_method,
rank=rank,
world_size=4,
timeout=datetime.timedelta(seconds=10),
)
try:
fs_groups = [
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
for ranks in ([0, 1], [2, 3])
]
sp_groups = [
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
for ranks in ([0, 1], [2, 3])
]
cfg_groups = [
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
for ranks in ([0, 2], [1, 3])
]
fs_group = next(group for ranks, group in fs_groups if rank in ranks)
sp_group = next(group for ranks, group in sp_groups if rank in ranks)
cfg_group = next(group for ranks, group in cfg_groups if rank in ranks)
parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
parallel_state._SP = SimpleNamespace(world_size=2, device_group=sp_group)
hook = SeaCacheRootHook(SeaCacheConfig())
hook._parameter_sharded = True
hook._synchronize_compute(False, torch.device("cpu"))
if rank in (0, 1):
hook._synchronize_compute(False, torch.device("cpu"))
gathered = [torch.zeros(1) for _ in range(2)]
torch.distributed.all_gather(
gathered,
torch.tensor([rank], dtype=torch.float32),
group=cfg_group,
)
result_queue.put((rank, [int(value.item()) for value in gathered]))
finally:
parallel_state._FS = None
parallel_state._SP = None
torch.distributed.destroy_process_group()
def test_parameter_sharded_hook_handles_uneven_cfg_branch_dispatch() -> None:
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
result_queue = manager.Queue()
torch.multiprocessing.spawn(
_uneven_cfg_sharded_worker,
args=(get_distributed_init_method("seacache_uneven_cfg_"), result_queue),
nprocs=4,
)
results = sorted(result_queue.get(timeout=2) for _ in range(4))
assert results == [
(0, [0, 2]),
(1, [1, 3]),
(2, [0, 2]),
(3, [1, 3]),
]
def _hybrid_sp_worker(
rank: int,
init_method: str,
result_queue: torch.multiprocessing.Queue,
) -> None:
from vllm_omni.diffusion.distributed import parallel_state
torch.distributed.init_process_group(
"gloo",
init_method=init_method,
rank=rank,
world_size=4,
timeout=datetime.timedelta(seconds=10),
)
try:
fs_groups = [
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
for ranks in ([0, 1], [2, 3])
]
ulysses_groups = [
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
for ranks in ([0, 1], [2, 3])
]
ring_groups = [
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
for ranks in ([0, 2], [1, 3])
]
sp_group = torch.distributed.new_group(
[0, 1, 2, 3],
timeout=datetime.timedelta(seconds=10),
)
fs_group = next(group for ranks, group in fs_groups if rank in ranks)
ulysses_group = next(group for ranks, group in ulysses_groups if rank in ranks)
ring_group = next(group for ranks, group in ring_groups if rank in ranks)
parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
parallel_state._SP = SimpleNamespace(
world_size=4,
device_group=sp_group,
ulysses_world_size=2,
ulysses_group=ulysses_group,
ring_world_size=2,
ring_group=ring_group,
)
compute = rank == 3
non_sharded_hook = SeaCacheRootHook(SeaCacheConfig())
non_sharded_result = non_sharded_hook._synchronize_compute(compute, torch.device("cpu"))
sharded_hook = SeaCacheRootHook(SeaCacheConfig())
sharded_hook._parameter_sharded = True
sharded_result = sharded_hook._synchronize_compute(compute, torch.device("cpu"))
result_queue.put((rank, non_sharded_result, sharded_result))
finally:
parallel_state._FS = None
parallel_state._SP = None
torch.distributed.destroy_process_group()
def test_hook_synchronizes_full_hybrid_sequence_parallel_group() -> None:
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
result_queue = manager.Queue()
torch.multiprocessing.spawn(
_hybrid_sp_worker,
args=(get_distributed_init_method("seacache_hybrid_sp_"), result_queue),
nprocs=4,
)
results = sorted(result_queue.get(timeout=2) for _ in range(4))
assert results == [
(0, True, True),
(1, True, True),
(2, True, True),
(3, True, True),
]
def test_parameter_sharded_hook_fails_open_without_distributed_state(
monkeypatch: pytest.MonkeyPatch,
) -> None:
hook = SeaCacheRootHook(SeaCacheConfig())
hook._parameter_sharded = True
monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
assert hook._synchronize_compute(False, torch.device("cpu")) is True
def test_hook_keeps_three_transfer_branches_separate() -> None:
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
hook = _apply_test_hook(transformer, metadata)
hook.refresh(transformer)
contexts = ("cond", "cond_no_control", "uncond")
for step, (timestep, value) in enumerate(((1000, 1.0), (500, 0.9))):
metadata.step = step
metadata.sigma = timestep / 1000
_run_step(transformer, timestep, value, hook=hook, context=contexts[0], control=0.5)
_run_step(transformer, timestep, value, hook=hook, context=contexts[1])
_run_step(transformer, timestep, value, hook=hook, context=contexts[2], control=0.5)
assert set(hook.state_manager._states) == set(contexts)
assert hook.full_count == 3
assert hook.skip_count == 3
assert len(hook.state_manager._states["cond"].previous_indicator) == 2
assert len(hook.state_manager._states["cond_no_control"].previous_indicator) == 1
def test_hook_fails_open_without_noisy_vision() -> None:
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
hook = _apply_test_hook(transformer, metadata)
hook.refresh(transformer)
all_clean = torch.zeros(1, 1, 2, 1, 1)
_run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
metadata.step = 1
metadata.sigma = 0.5
_run_step(transformer, 500, 0.9, hook=hook, noisy_frame_mask=all_clean)
assert hook.full_count == 0
assert hook.skip_count == 0
def test_conditioning_only_model_failure_runs_forward_once(monkeypatch: pytest.MonkeyPatch) -> None:
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
hook = _apply_test_hook(transformer, metadata)
hook.refresh(transformer)
forward_calls = 0
def fail_forward(hidden_gen: torch.Tensor) -> torch.Tensor:
nonlocal forward_calls
forward_calls += 1
raise RuntimeError("model forward failed")
monkeypatch.setattr(transformer, "_run_gen_layers", fail_forward)
all_clean = torch.zeros(1, 1, 2, 1, 1)
with pytest.raises(RuntimeError, match="model forward failed"):
_run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
assert forward_calls == 1
def test_hook_fails_open_without_explicit_context() -> None:
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=2)
hook = _apply_test_hook(transformer, metadata)
hook.refresh(transformer)
_run_step(transformer, 1000, 1.0)
assert hook.full_count == 0
assert hook.skip_count == 0
assert hook.state_manager._states == {}
def test_hook_uses_exact_sigma_callback(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.cache.seacache import hook as hook_module
transformer = TinyCosmos3Transformer()
metadata = SimpleNamespace(step=0, sigma=0.37, num_steps=2)
observed_sigmas: list[float] = []
original_filter = hook_module.apply_sea_filter
def recording_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float) -> torch.Tensor:
observed_sigmas.append(sigma)
return original_filter(hidden_states, sigma, power_exp)
monkeypatch.setattr(hook_module, "apply_sea_filter", recording_filter)
hook = _apply_test_hook(transformer, metadata)
hook.refresh(transformer)
_run_step(transformer, timestep=999, value=1.0, hook=hook)
assert observed_sigmas
assert all(sigma == pytest.approx(0.37) for sigma in observed_sigmas)
def test_backend_selector_and_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
class FSDPTinyCosmos3Transformer(TinyCosmos3Transformer):
pass
assert get_extractor(FSDPTinyCosmos3Transformer) is _extract_tiny_cosmos3_context
backend = get_cache_backend(
"sea_cache",
{
"sea_threshold": 0.4,
"sea_residual_order": 0,
},
)
assert isinstance(backend, SeaCacheBackend)
assert backend.config.sea_threshold == 0.4
pipeline = Cosmos3OmniDiffusersPipeline()
backend.enable(pipeline)
hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
assert isinstance(hook, SeaCacheRootHook)
assert callable(getattr(pipeline, "_cache_context_factory", None))
for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
assert not hasattr(pipeline.transformer, name)
pipeline._current_step_index = 0
pipeline._current_sigma = 1.0
pipeline._num_timesteps = 7
_run_step(pipeline.transformer, 1000, 1.0, hook=hook)
assert hook.full_count == 1
backend.refresh(pipeline, num_inference_steps=7)
assert hook.full_count == 0
assert hook.skip_count == 0
assert hook.state_manager._states == {}
def test_backend_uses_resolved_pipeline_metadata_not_refresh_argument(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
pipeline = Cosmos3OmniDiffusersPipeline()
backend = SeaCacheBackend(DiffusionCacheConfig())
backend.enable(pipeline)
hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
assert isinstance(hook, SeaCacheRootHook)
backend.refresh(pipeline, num_inference_steps=35)
pipeline._num_timesteps = 50
pipeline._current_step_index = 0
pipeline._current_sigma = 1.0
_run_step(pipeline.transformer, 1000, 1.0, hook=hook)
assert hook.full_count == 1
assert hook.num_inference_steps_callback is not None
assert hook.num_inference_steps_callback() == 50
pipeline._current_step_index = 1
pipeline._current_sigma = 0.98
_run_step(pipeline.transformer, 980, 0.99, hook=hook)
assert hook.full_count == 1
assert hook.skip_count == 1
def test_shared_config_defaults() -> None:
config = DiffusionCacheConfig()
assert config.sea_threshold == 0.25
assert config.sea_residual_order == 1
assert config.sea_max_consecutive_cached == 2
assert config.sea_power_exp == 3.0
+5 -5
View File
@@ -68,12 +68,12 @@ class BaseExtractorTest(ABC):
pass
@abstractmethod
def get_module(self):
def get_module(self, module, /):
"""Return model module instance."""
pass
@abstractmethod
def get_sample_inputs(self):
def get_sample_inputs(self, sample_inputs, /):
"""Return sample inputs for model."""
pass
@@ -737,10 +737,10 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
monkeypatch,
):
"""Strict SP must not mix local TeaCache state with gathered block rows."""
from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
from vllm_omni.diffusion.distributed import parallel_state
from vllm_omni.diffusion.layers import indexed_modulation
seq_len = sample_inputs["x"].shape[1]
local_len = seq_len // 2
@@ -844,13 +844,13 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
)
monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
original_indexed_scale_shift = indexed_modulation.indexed_scale_shift_
def capture_indexed_scale_shift(hidden, shift, scale, indices):
captured_modulation_indices.append(indices.clone())
return original_indexed_scale_shift(hidden, shift, scale, indices)
monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
monkeypatch.setattr(indexed_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
local_img_pos = torch.tensor([0, 1])
local_audio_pos = torch.tensor([2, 3])
@@ -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)
@@ -440,29 +440,86 @@ def test_transformer_rejects_prompt_tuning():
BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
def _native_to_checkpoint_name(name: str) -> str:
"""Inverse of ``load_weights`` remapping: native param name -> diffusers name.
_QKV_FANOUT = {
"to_qkv": ("to_q", "to_k", "to_v"),
"img_to_qkv": ("img_to_q", "img_to_k", "img_to_v"),
"instruct_to_qkv": ("instruct_to_q", "instruct_to_k", "instruct_to_v"),
}
def _native_to_checkpoint_weights(name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]:
"""Inverse of ``load_weights`` remapping: native param -> diffusers weights.
The diffusers checkpoint stores fused projections as separate matrices, so a
single merged native param (QKV or FFN gate/up) fans out into several
checkpoint weights. Returns ``(diffusers_name, value)`` pairs whose values,
when fed through ``load_weights``, reassemble into the original ``param``.
- ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
- promoted joint-attention projections move back under ``.processor.``.
"""
q_size = NUM_HEADS * HEAD_DIM
kv_size = NUM_KV_HEADS * HEAD_DIM
# Fused QKV projections split back into per-matrix q/k/v weights.
for token, (q_name, k_name, v_name) in _QKV_FANOUT.items():
marker = f".{token}."
if marker in name:
q, k, v = param[:q_size], param[q_size : q_size + kv_size], param[q_size + kv_size :]
results = []
for sub_name, value in ((q_name, q), (k_name, k), (v_name, v)):
ckpt_name = name.replace(marker, f".{sub_name}.")
if token != "to_qkv":
# Joint-attention projections live under `.processor.` upstream.
ckpt_name = ckpt_name.replace(".img_instruct_attn.", ".img_instruct_attn.processor.")
results.append((ckpt_name, value))
return results
# Fused FFN gate/up splits into linear_1 (gate) / linear_3 (input).
if ".gate_up_proj." in name:
inner = param.shape[0] // 2
return [
(name.replace(".gate_up_proj.", ".linear_1."), param[:inner]),
(name.replace(".gate_up_proj.", ".linear_3."), param[inner:]),
]
# Default: 1:1 with the existing diffusers name promotions.
if ".to_out." in name:
name = name.replace(".to_out.", ".to_out.0.")
for proj in (
"img_to_q",
"img_to_k",
"img_to_v",
"instruct_to_q",
"instruct_to_k",
"instruct_to_v",
"instruct_out",
"img_out",
):
for proj in ("instruct_out", "img_out"):
token = f".img_instruct_attn.{proj}."
if token in name:
name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
break
return name
return [(name, param)]
def test_transformer_exposes_stacked_params_mapping():
"""The packed -> sub-layer mapping must be discoverable from the module tree.
``diffusion/lora/loader.py`` and the quantized weight loaders read
``stacked_params_mapping`` off the model, so it has to exist before (and
independently of) ``load_weights``.
"""
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
_BOOGU_STACKED_PARAMS_MAPPING,
BooguImageTransformer2DModel,
)
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
assert tuple(model.stacked_params_mapping) == _BOOGU_STACKED_PARAMS_MAPPING
# A per-instance copy, so consumers cannot mutate the module constant.
assert model.stacked_params_mapping is not _BOOGU_STACKED_PARAMS_MAPPING
# Every entry is a (param, shard, shard_id) triple.
for param_name, shard_name, shard_id in model.stacked_params_mapping:
assert param_name.startswith(".") and shard_name.startswith(".")
assert shard_id in {"q", "k", "v", 0, 1}
# The mapping targets exactly the fused projections this port creates.
mapped_leaves = {param.strip(".").split(".")[-1] for param, _, _ in model.stacked_params_mapping}
assert mapped_leaves == {"to_qkv", "img_to_qkv", "instruct_to_qkv", "gate_up_proj"}
def test_transformer_load_weights_round_trip():
@@ -473,24 +530,59 @@ def test_transformer_load_weights_round_trip():
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
native_params = dict(model.named_parameters())
# Build synthetic diffusers-named weights (one per native parameter).
checkpoint_weights = {}
# Build synthetic diffusers-named weights. Fused native projections fan out
# into the separate matrices the diffusers checkpoint stores.
expected: dict[str, torch.Tensor] = {}
checkpoint_weights: dict[str, torch.Tensor] = {}
for native_name, param in native_params.items():
checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
full = torch.randn_like(param)
expected[native_name] = full
for ckpt_name, value in _native_to_checkpoint_weights(native_name, full):
checkpoint_weights[ckpt_name] = value
# The remapping must be a bijection over the parameter set.
assert len(checkpoint_weights) == len(native_params)
# Every checkpoint name is unique; merged params fan out to >=1 weight.
assert len(checkpoint_weights) >= len(native_params)
loaded = model.load_weights(list(checkpoint_weights.items()))
# No missing / unexpected parameters.
assert loaded == set(native_params.keys())
# Values landed on the right parameters (TP=1: weight_loader copies verbatim).
# Values landed on the right parameters (TP=1: weight_loader copies verbatim;
# fused shards are placed at their q/k/v / gate/up offsets).
reloaded = dict(model.named_parameters())
for native_name in native_params:
expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
assert torch.allclose(reloaded[native_name], expected)
assert torch.allclose(reloaded[native_name], expected[native_name])
def test_transformer_load_weights_rejects_partial_fused_params():
"""A fused parameter must not count as loaded until all its shards arrive.
The loader compares parameter *names*, so reporting ``to_qkv`` complete
after a single ``to_q`` would let a checkpoint carrying only ``to_q`` start
up with the ``k``/``v`` slices left uninitialized.
"""
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
BooguImageTransformer2DModel,
)
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
native_params = dict(model.named_parameters())
complete_name = "noise_refiner.0.attn.to_qkv.weight"
partial_name = "noise_refiner.0.feed_forward.gate_up_proj.weight"
assert complete_name in native_params
assert partial_name in native_params
weights = list(_native_to_checkpoint_weights(complete_name, native_params[complete_name]))
# Only the gate half of the FFN, so its fused parameter stays incomplete.
gate_ckpt, gate_value = _native_to_checkpoint_weights(partial_name, native_params[partial_name])[0]
weights.append((gate_ckpt, gate_value))
loaded = model.load_weights(weights)
assert complete_name in loaded
assert partial_name not in loaded
def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
@@ -513,12 +605,15 @@ def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
try:
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
native_params = dict(model.named_parameters())
loaded_name = next(iter(native_params))
checkpoint_name = _native_to_checkpoint_name(loaded_name)
# A non-fused parameter, so one checkpoint entry loads it completely.
loaded_name = "x_embedder.weight"
assert loaded_name in native_params
(checkpoint_name, checkpoint_value), *_ = _native_to_checkpoint_weights(loaded_name, native_params[loaded_name])
checkpoint_value = torch.randn_like(checkpoint_value)
loaded = model.load_weights(
[
(checkpoint_name, torch.randn_like(native_params[loaded_name])),
(checkpoint_name, checkpoint_value),
("unexpected.weight", torch.ones(1)),
]
)
@@ -1,11 +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 json
import sys
import types
from contextlib import contextmanager
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any
@@ -85,6 +86,7 @@ class StubScheduler:
flow_shift: float = 1.0,
) -> None:
self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
self.sigmas = self.timesteps.float() / 10
self.config = SimpleNamespace(
num_train_timesteps=1000,
flow_shift=flow_shift,
@@ -113,6 +115,7 @@ class StubScheduler:
else:
assert num_inference_steps is not None
self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device)
self.sigmas = self.timesteps.float()
def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
del kwargs
@@ -250,7 +253,7 @@ def passthrough_progress_bar(iterable):
@pytest.fixture(autouse=True)
def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
module: Any = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
module.is_guardrails_enabled = lambda od_config, sampling_params=None: False
module.ensure_initialized = lambda od_config: None
module.check_text_safety = lambda text: None
@@ -283,6 +286,8 @@ def make_cosmos3_pipeline():
pipeline.is_edge_model = False
pipeline._guidance_scale = None
pipeline._num_timesteps = None
pipeline._current_step_index = None
pipeline._current_sigma = None
pipeline._cosmos3_branch_caches = None
pipeline._cache_dit_requires_paired_cfg = False
pipeline._sound_tokenizer = None
@@ -300,7 +305,7 @@ def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
def make_sampling_params(**overrides: Any) -> SimpleNamespace:
values = {
values: dict[str, Any] = {
"height": None,
"width": None,
"num_frames": None,
@@ -814,7 +819,7 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
t_list = [1.0, 0.75, 0.5, 0.25]
scheduler_dir = tmp_path / "scheduler"
scheduler_dir.mkdir()
scheduler_config = {"_class_name": scheduler_class_name}
scheduler_config: dict[str, Any] = {"_class_name": scheduler_class_name}
if expected_distilled:
scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list}
(scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config))
@@ -1531,7 +1536,7 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke
fused_add_rms_norm: list[str]
custom_op: list[str]
fake_kernel = types.ModuleType("vllm.config.kernel")
fake_kernel: Any = types.ModuleType("vllm.config.kernel")
fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig
monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel)
@@ -2154,6 +2159,47 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
def test_diffuse_publishes_exact_seacache_metadata_and_cfg_contexts(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
pipeline.scheduler.sigmas = torch.tensor([0.91, 0.17])
observations: list[tuple[str, int | None, float | None, int | None]] = []
class RecordingHook:
@contextmanager
def cache_context(self, name: str):
sigma = pipeline.current_sigma
observations.append(
(
name,
pipeline.current_step_index,
None if sigma is None else float(sigma),
pipeline.num_timesteps,
)
)
yield
pipeline._cache_context_factory = RecordingHook().cache_context
pipeline.diffuse(
latents=torch.zeros(1, 2, 1, 1, 1),
timesteps=torch.tensor([900, 100]),
cond_ids=_ids(2),
cond_mask=_mask(),
uncond_ids=_ids(1),
uncond_mask=_mask(),
guidance_scale=3.0,
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
)
assert observations == [
("cond", 0, pytest.approx(0.91), 2),
("uncond", 0, pytest.approx(0.91), 2),
("cond", 1, pytest.approx(0.17), 2),
("uncond", 1, pytest.approx(0.17), 2),
]
assert pipeline.current_step_index is None
assert pipeline.current_sigma is None
def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
pipeline._use_session_state = True
@@ -2213,6 +2259,53 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
torch.testing.assert_close(result, torch.full_like(latents, 254.0))
def test_diffuse_transfer_uses_named_seacache_contexts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
pipeline = make_cosmos3_pipeline()
pipeline.scheduler.sigmas = torch.tensor([0.42])
contexts: list[tuple[str, int | None, float | None, int | None]] = []
class RecordingHook:
@contextmanager
def cache_context(self, name: str):
sigma = pipeline.current_sigma
contexts.append(
(
name,
pipeline.current_step_index,
None if sigma is None else float(sigma),
pipeline.num_timesteps,
)
)
yield
pipeline._cache_context_factory = RecordingHook().cache_context
latents = torch.zeros(1, 2, 1, 1, 1)
velocity_mask = torch.ones(1, 1, 1, 1, 1)
pipeline.diffuse_transfer(
latents=latents,
timesteps=torch.tensor([7]),
cond_ids=_ids(2),
cond_mask=_mask(),
uncond_ids=_ids(1),
uncond_mask=_mask(),
guidance_scale=3.0,
control_guidance=1.5,
control_guidance_interval=None,
control_latents=[torch.zeros_like(latents)],
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
velocity_mask=velocity_mask,
condition_latents=torch.zeros_like(latents),
)
assert contexts == [
("cond", 0, pytest.approx(0.42), 1),
("cond_no_control", 0, pytest.approx(0.42), 1),
("uncond", 0, pytest.approx(0.42), 1),
]
assert pipeline.current_step_index is None
assert pipeline.current_sigma is None
def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
pipeline._use_session_state = True
@@ -2583,7 +2676,7 @@ def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -
class TestForwardRouting:
def _install_forward_stubs(self, pipeline):
captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []}
captured: dict[str, Any] = {"diffuse_calls": [], "prepare_calls": []}
def fake_format(
prompt,
@@ -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
@@ -513,6 +513,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
assert tuple(output.shape) == (1, 2, 1, 2, 2)
def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.cache.teacache.extractors import extract_cosmos3_context
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
class TrackingNorm(nn.Module):
def __init__(self) -> None:
super().__init__()
self.calls = 0
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
self.calls += 1
return hidden_states + 5.0
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
model = transformer_cosmos3.Cosmos3VFMTransformer(
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
)
norm = TrackingNorm()
model.norm_moe_gen = norm
captured: dict[str, torch.Tensor] = {}
def run_gen_layers(hidden_gen: torch.Tensor, **kwargs) -> torch.Tensor:
del kwargs
captured["input"] = hidden_gen.detach().clone()
return hidden_gen + 2.0
monkeypatch.setattr(model, "_run_gen_layers", run_gen_layers)
forward_kwargs = {
"hidden_states": torch.zeros(1, 2, 1, 2, 2),
"timestep": torch.tensor([1.0]),
"text_ids": torch.tensor([[1, 2]], dtype=torch.long),
"text_mask": torch.ones(1, 2, dtype=torch.long),
"video_shape": (1, 2, 2),
"fps": 24.0,
}
full_output = model(**forward_kwargs)
assert norm.calls == 1
norm.calls = 0
ctx = extract_cosmos3_context(model, **forward_kwargs)
execution_input = ctx.hidden_states.detach().clone()
execution_output = ctx.run_transformer_blocks()[0]
residual = execution_output - execution_input
extracted_output = ctx.postprocess(execution_output)
assert norm.calls == 1
expected_residual = ((captured["input"] + 2.0) + 5.0) - captured["input"]
assert torch.equal(residual, expected_residual)
assert torch.equal(extracted_output, full_output)
norm.calls = 0
def fail_if_gen_layers_run(*args, **kwargs):
del args, kwargs
raise AssertionError("SeaCache hit unexpectedly executed GEN layers")
monkeypatch.setattr(model, "_run_gen_layers", fail_if_gen_layers_run)
cached_ctx = extract_cosmos3_context(model, **forward_kwargs)
cached_output = cached_ctx.postprocess(cached_ctx.hidden_states + residual)
assert norm.calls == 0
assert torch.equal(cached_output, full_output)
for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
assert not hasattr(model, name)
def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
class CountingNorm(nn.Module):
def __init__(self, original: nn.Module) -> None:
super().__init__()
self.original = original
self.calls = 0
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
self.calls += 1
return self.original(hidden_states)
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
model = transformer_cosmos3.Cosmos3VFMTransformer(
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
)
norm = CountingNorm(model.norm_moe_gen)
model.norm_moe_gen = norm
output = model(
hidden_states=torch.zeros(1, 2, 1, 2, 2),
timestep=torch.tensor([1.0]),
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
text_mask=torch.ones(1, 2, dtype=torch.long),
video_shape=(1, 2, 2),
fps=24.0,
)
assert norm.calls == 1
assert tuple(output.shape) == (1, 2, 1, 2, 2)
def test_qwen_gen_mlp_remains_gated() -> None:
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
@@ -2711,3 +2711,125 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
assert "from .pipeline import" in lingbot_init
assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init
assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init
# ---------------------------------------------------------------------------
# a prompt is encoded once, whichever path asks for it
# ---------------------------------------------------------------------------
def _counting_text_encoder(pipeline) -> list[str]:
"""Give ``pipeline`` its real encode_prompt over stubs that record each encode."""
del pipeline.encode_prompt # _pipeline() replaces it; these tests need the real one
encoded: list[str] = []
def tokenizer(texts, **kwargs):
del kwargs
encoded.extend(texts)
return SimpleNamespace(
input_ids=torch.arange(512).view(1, 512),
attention_mask=torch.ones(1, 512, dtype=torch.long),
)
def text_encoder(input_ids, attention_mask):
del input_ids, attention_mask
return SimpleNamespace(last_hidden_state=torch.full((1, 512, 8), float(len(encoded))))
pipeline.tokenizer = tokenizer
pipeline.text_encoder = text_encoder
return encoded
def test_the_text_encoder_runs_once_per_distinct_prompt() -> None:
module = _load_pipeline_module()
pipeline = _pipeline(module)
encoded = _counting_text_encoder(pipeline)
def encode(prompt, *, length=512, dtype=torch.float32):
return pipeline.encode_prompt(prompt, max_sequence_length=length, dtype=dtype)
first = encode("move through the room")
assert encode("move through the room") is first
# The key is the text the tokenizer actually sees, so spacing alone is not new.
assert encode(" move through the room ") is first
assert encoded == ["move through the room"]
encode("enter the snowy valley")
# dtype and sequence length change the tensor, so each is its own entry.
encode("move through the room", dtype=torch.float16)
encode("move through the room", length=256)
assert encoded == [
"move through the room",
"enter the snowy valley",
"move through the room",
"move through the room",
]
def test_the_prompt_encode_cache_is_bounded() -> None:
module = _load_pipeline_module()
pipeline = _pipeline(module)
encoded = _counting_text_encoder(pipeline)
size = module._PROMPT_EMBEDS_CACHE_SIZE
for index in range(size + 1):
pipeline.encode_prompt(f"prompt {index}", max_sequence_length=512, dtype=torch.float32)
assert len(pipeline._prompt_embeds_cache) == size
# The oldest entry was the one evicted, so it is the one encoded again.
pipeline.encode_prompt("prompt 0", max_sequence_length=512, dtype=torch.float32)
pipeline.encode_prompt(f"prompt {size}", max_sequence_length=512, dtype=torch.float32)
assert encoded == [f"prompt {index}" for index in range(size + 1)] + ["prompt 0"]
def test_realtime_ticks_reuse_the_encode_and_still_invalidate_on_a_prompt_change(
monkeypatch: pytest.MonkeyPatch,
) -> None:
module = _load_pipeline_module()
pipeline = _pipeline(module)
pipeline._ar_height = 16
pipeline._ar_width = 16
pipeline._ar_diffusion_kv_state = object()
encoded = _counting_text_encoder(pipeline)
invalidations: list[bool] = []
def ar_text_caches(prompt_embeds, *, invalidate):
del prompt_embeds
invalidations.append(invalidate)
return [SimpleNamespace()]
monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches)
monkeypatch.setattr(
pipeline,
"_generate_block",
lambda **kwargs: torch.randn((1, 16, 3, 2, 2), generator=kwargs["generator"]),
)
for chunk_index in range(3):
pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=chunk_index))))
switched = _prompt()
switched["prompt"] = "enter the snowy valley"
pipeline(
_request(
sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=3, prompt=switched["prompt"])),
prompt=switched,
)
)
assert encoded == ["move through the room", "enter the snowy valley"]
# A new prompt must still drop the cross-attention K/V built from the old one.
assert invalidations == [False, False, False, True]
def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None:
module = _load_pipeline_module()
pipeline = _pipeline(module, transformer=_RecordingTransformer())
pipeline._ar_height = 16
pipeline._ar_width = 16
encoded = _counting_text_encoder(pipeline)
states = [_stepwise_state(request_id=request_id, num_frames=21) for request_id in ("req-a", "req-b")]
for state in states:
with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
pipeline.prepare_encode(state)
assert encoded == ["move through the room"]
assert states[0].prompt_embeds is states[1].prompt_embeds
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
import torch
from tests.diffusion.models.magi2.test_native_packing import (
_longer_text_tensors,
_sampler_tensors,
_tiny_model,
_tiny_sampler,
)
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def test_layer_regions_bracket_the_eager_kernels() -> None:
moe_layer, dense_layer = _tiny_model().block.layers
assert moe_layer.region_methods == ("_attention_input", "_moe_input", "_moe_output")
assert dense_layer.region_methods == ("_attention_input", "_dense_output")
def test_compile_regions_fullgraph_matches_eager_bitwise() -> None:
model = _tiny_model()
sampler = _tiny_sampler(model)
steps = [
sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
]
expected = [sampler.forward(step) for step in steps]
torch._dynamo.reset()
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
first = sampler.forward(steps[0])
graphs = torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before
with torch._dynamo.config.patch(error_on_recompile=True):
second = sampler.forward(steps[1])
for (video, audio), (video_ref, audio_ref) in zip((first, second), expected, strict=True):
assert torch.equal(video, video_ref)
assert torch.equal(audio, audio_ref)
# pre adapter, post adapter, three MoE-layer regions, two dense-layer regions
assert graphs == 7
def test_pipeline_setup_compile_forwards_config_dynamic() -> None:
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
calls: list[dict[str, object]] = []
class _Transformer(torch.nn.Module):
def compile_regions(self, **compile_kwargs: object) -> None:
calls.append(compile_kwargs)
pipeline = object.__new__(Magi2Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = OmniDiffusionConfig(
model="sand-ai/MAGI-2-preview",
model_class_name="Magi2Pipeline",
diffusion_compile_dynamic=False,
)
pipeline.transformer = _Transformer()
pipeline.setup_compile()
assert calls == [{"fullgraph": True, "dynamic": False, "options": {"emulate_precision_casts": True}}]
def test_pipeline_setup_compile_warns_when_granularity_is_not_regional(caplog: pytest.LogCaptureFixture) -> None:
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
calls: list[dict[str, object]] = []
class _Transformer(torch.nn.Module):
def compile_regions(self, **compile_kwargs: object) -> None:
calls.append(compile_kwargs)
pipeline = object.__new__(Magi2Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = OmniDiffusionConfig(
model="sand-ai/MAGI-2-preview",
model_class_name="Magi2Pipeline",
diffusion_compile_granularity="full",
)
pipeline.transformer = _Transformer()
pipeline.setup_compile()
assert calls == [{"fullgraph": True, "dynamic": True, "options": {"emulate_precision_casts": True}}]
assert "diffusion_compile_granularity='full'" in caplog.text
def test_compiled_graphs_are_shared_across_layers_of_one_kind() -> None:
model = _tiny_model(moe_layers=2)
sampler = _tiny_sampler(model)
step = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([900.0]), cfg_config=CFGConfig())
expected = sampler.forward(step)
torch._dynamo.reset()
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
actual = sampler.forward(step)
assert torch.equal(actual[0], expected[0])
# Two MoE layers and two dense layers still compile one graph per region
# kind: pre adapter, post adapter, three MoE regions, two dense regions.
assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
def test_compiled_regions_survive_new_requests_without_recompiling() -> None:
model = _tiny_model()
sampler = _tiny_sampler(model)
t = torch.tensor([900.0])
first_request = sampler.prepare_model_input(**_sampler_tensors(0), t=t, cfg_config=CFGConfig())
second_request = sampler.prepare_model_input(**_sampler_tensors(1), t=t, cfg_config=CFGConfig())
longer_request = sampler.prepare_model_input(**_longer_text_tensors(2), t=t, cfg_config=CFGConfig())
expected = [sampler.forward(request) for request in (first_request, second_request, longer_request)]
torch._dynamo.reset()
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
actual = [sampler.forward(first_request)]
with torch._dynamo.config.patch(error_on_recompile=True):
actual.append(sampler.forward(second_request))
actual.append(sampler.forward(longer_request))
for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
assert torch.equal(video, video_ref)
assert torch.equal(audio, audio_ref)
assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
@@ -0,0 +1,77 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
import torch
from tests.diffusion.models.magi2.test_native_packing import (
_longer_text_tensors,
_sampler_tensors,
_tiny_model,
_tiny_sampler,
)
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
pytestmark = [
pytest.mark.diffusion,
pytest.mark.core_model,
*hardware_marks(res={"cuda": "L4"}, num_cards=1),
pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
]
def _cuda_steps(sampler, tensors=_sampler_tensors):
timesteps = (torch.tensor([900.0]), torch.tensor([600.0]), torch.tensor([300.0]))
return [
sampler.prepare_model_input(
**{name: value.cuda() for name, value in tensors(seed).items()},
t=t,
cfg_config=CFGConfig(),
)
for seed, t in enumerate(timesteps)
]
def test_inductor_regions_match_eager_in_deterministic_mode(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("MAGI2_DETERMINISTIC", "1")
model = _tiny_model(params_dtype=torch.bfloat16).cuda()
sampler = _tiny_sampler(model)
steps = _cuda_steps(sampler)
expected = [sampler.forward(step) for step in steps]
new_request_step = _cuda_steps(sampler)[0]
longer_text_step = _cuda_steps(sampler, _longer_text_tensors)[0]
expected_new_request = sampler.forward(new_request_step)
expected_longer_text = sampler.forward(longer_text_step)
torch._dynamo.reset()
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
# The pipeline compiles with inductor emulating eager's intermediate
# low-precision casts, so fused bf16 chains keep eager's rounding.
model.compile_regions(fullgraph=True, dynamic=True, options={"emulate_precision_casts": True})
first = sampler.forward(steps[0])
with torch._dynamo.config.patch(error_on_recompile=True):
compiled = [first, sampler.forward(steps[1]), sampler.forward(steps[2])]
repeated = [sampler.forward(step) for step in steps]
# A second request repacks the same shapes; a third changes the
# text length and therefore every token count.
compiled_new_request = sampler.forward(new_request_step)
compiled_longer_text = sampler.forward(longer_text_step)
# Emulated precision casts restore eager's rounding boundaries; what is
# left is one-ulp transcendental noise that some Triton backends carry.
for (video, audio), (video_ref, audio_ref) in zip(compiled, expected, strict=True):
torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
for (video, audio), (video_again, audio_again) in zip(compiled, repeated, strict=True):
assert torch.equal(video, video_again)
assert torch.equal(audio, audio_again)
for (video, audio), (video_ref, audio_ref) in (
(compiled_new_request, expected_new_request),
(compiled_longer_text, expected_longer_text),
):
torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
@@ -0,0 +1,179 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import os
import tempfile
from contextlib import ExitStack, contextmanager
from unittest.mock import patch
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.distributed import init_device_mesh
from torch.distributed.fsdp import MixedPrecisionPolicy
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dlo_module
from tests.diffusion.models.magi2.test_native_distributed_parity import (
_current_group,
_new_groups,
_patched_groups,
)
from tests.diffusion.models.magi2.test_native_packing import _tiny_config, _tiny_model, _tiny_sampler
from tests.diffusion.offloader.helpers import patch_offload_runtime
from vllm_omni.diffusion.distributed.hsdp import shard_model
from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
from vllm_omni.diffusion.offloader.base import OffloadConfig
from vllm_omni.diffusion.offloader.config import OffloadStrategy
_WORLD_SIZE = 4
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@contextmanager
def _cpu_offload_runtime():
with pytest.MonkeyPatch.context() as monkeypatch:
patch_offload_runtime(monkeypatch, dlo_module.current_omni_platform, synchronize=True)
yield
def _enable_dlo(
model: Magi2PreviewTransformer,
*,
dp_group: dist.ProcessGroup | None,
dp_size: int,
dp_rank: int,
) -> None:
pipeline = Magi2Pipeline.__new__(Magi2Pipeline)
nn.Module.__init__(pipeline)
pipeline.transformer = model
backend = dlo_module.DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dp_size=dp_size,
dlo_use_allgather=dp_size > 1,
),
torch.device("cpu"),
)
backend.dp_group = dp_group
backend.rank = dp_rank
backend.enable(pipeline)
assert backend.enabled
def _enable_hsdp(model: Magi2PreviewTransformer) -> None:
mesh = init_device_mesh("cpu", mesh_shape=(1, _WORLD_SIZE), mesh_dim_names=("replicate", "shard"))
ignored = {
parameter for name in model._hsdp_ignored_modules for parameter in model.get_submodule(name).parameters()
}
shard_model(
model,
mesh=mesh,
mp_policy=MixedPrecisionPolicy(cast_forward_inputs=False),
hsdp_shard_conditions=model._hsdp_shard_conditions,
ignored_params=ignored,
)
model.requires_grad_(False)
def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
# 14 tokens per CFG branch, so the packed sequence splits evenly across
# two Ulysses ranks; gloo cannot gather uneven shards.
generator = torch.Generator(device="cpu").manual_seed(seed)
return {
"latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
"audio_latent": torch.randn(1, 5, 4, generator=generator),
"txt_feat": torch.randn(1, 3, 4, generator=generator),
"null_txt_feat": torch.randn(1, 3, 4, generator=generator),
}
def _steps(sampler):
return [
sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
]
def _run(model: Magi2PreviewTransformer, *, compile_regions: bool) -> list[tuple[torch.Tensor, torch.Tensor]]:
sampler = _tiny_sampler(model)
steps = _steps(sampler)
if not compile_regions:
return [sampler.forward(step) for step in steps]
torch._dynamo.reset()
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
outputs = [sampler.forward(steps[0])]
with torch._dynamo.config.patch(error_on_recompile=True):
outputs.append(sampler.forward(steps[1]))
return outputs
def _assert_same(actual, expected, *, exact: bool, label: str) -> None:
for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
if exact:
assert torch.equal(video, video_ref), label
assert torch.equal(audio, audio_ref), label
else:
torch.testing.assert_close(video, video_ref, atol=2e-5, rtol=2e-4, msg=label)
torch.testing.assert_close(audio, audio_ref, atol=2e-5, rtol=2e-4, msg=label)
def _distributed_worker(rank: int, rendezvous: str) -> None:
torch.set_num_threads(1)
dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=_WORLD_SIZE)
try:
singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
with _patched_groups(singleton, singleton):
oracle = _tiny_model()
checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()]
expected = _run(oracle, compile_regions=False)
sp2_groups = _new_groups(((0, 1), (2, 3)))
dp2_groups = _new_groups(((0, 2), (1, 3)))
sp_group = _current_group(rank, sp2_groups)
dp_ranks, dp_process_group = next((ranks, group) for ranks, group in dp2_groups if rank in ranks)
def dlo_rank_local(model: Magi2PreviewTransformer) -> None:
_enable_dlo(model, dp_group=None, dp_size=1, dp_rank=0)
def dlo_allgather(model: Magi2PreviewTransformer) -> None:
_enable_dlo(model, dp_group=dp_process_group, dp_size=len(dp_ranks), dp_rank=dp_ranks.index(rank))
variants = (
("dlo_rank_local_sp2", sp_group, dlo_rank_local),
("dlo_allgather_dp2sp2", sp_group, dlo_allgather),
("hsdp4", singleton, _enable_hsdp),
)
for name, ulysses_group, enable in variants:
with ExitStack() as stack:
stack.enter_context(_patched_groups(singleton, ulysses_group))
stack.enter_context(_cpu_offload_runtime())
outputs = {}
for compile_regions in (False, True):
model = Magi2PreviewTransformer(_tiny_config())
assert model.load_weights(checkpoint) == set(model.state_dict())
enable(model)
outputs[compile_regions] = _run(model, compile_regions=compile_regions)
_assert_same(outputs[False], expected, exact=False, label=f"{name} eager vs single-rank oracle")
_assert_same(outputs[True], outputs[False], exact=True, label=f"{name} compiled vs eager")
dist.barrier()
finally:
dist.destroy_process_group()
@pytest.mark.skipif(
not dist.is_available() or not dist.is_gloo_available(),
reason="requires torch.distributed gloo",
)
def test_compiled_regions_match_eager_under_dlo_and_hsdp() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
mp.spawn(_distributed_worker, args=(rendezvous,), nprocs=_WORLD_SIZE, join=True)
@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
import torch
from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
Magi2MHCConfig,
Magi2MoEConfig,
Magi2PreviewConfig,
)
from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm
from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE
from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
from vllm_omni.diffusion.models.magi2.preview_data_proxy import (
Magi2DataProxy,
Magi2PreviewDataProxyConfig,
)
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def _tiny_config(params_dtype: torch.dtype = torch.float32, *, moe_layers: int = 1) -> Magi2PreviewConfig:
# The leading layers are multimodal with MoE, the rest single-modality
# dense, so every layer belongs to one of two kinds.
moe_indices = tuple(range(moe_layers))
return Magi2PreviewConfig(
num_layers=2 * moe_layers,
hidden_size=16,
head_dim=8,
num_query_groups=2,
video_in_channels=4,
audio_in_channels=4,
text_in_channels=4,
intermediate_factor=2,
multimodal_layers=moe_indices,
params_dtype=params_dtype,
mhc=Magi2MHCConfig(num_streams=2),
moe=Magi2MoEConfig(
num_heads=2,
num_experts=4,
top_k=2,
expert_intermediate_size=8,
shared_expert_intermediate_size=8,
modality_shared_expert_intermediate_size=8,
layers=moe_indices,
),
)
def _tiny_model(
seed: int = 11,
params_dtype: torch.dtype = torch.float32,
*,
moe_layers: int = 1,
) -> Magi2PreviewTransformer:
model = Magi2PreviewTransformer(_tiny_config(params_dtype, moe_layers=moe_layers))
generator = torch.Generator(device="cpu").manual_seed(seed)
with torch.no_grad():
for parameter in model.parameters():
parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
for module in model.modules():
if isinstance(module, MultiModalityRMSNorm):
module.weight.zero_()
elif isinstance(module, Magi2MultiHeadMoE):
module.router.expert_bias.zero_()
module.router.expert_bias_ema.zero_()
return model
def _tiny_sampler(model: torch.nn.Module) -> Magi2PreviewSampler:
return Magi2PreviewSampler(model, Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=8)))
def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
generator = torch.Generator(device="cpu").manual_seed(seed)
return {
"latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
"audio_latent": torch.randn(1, 5, 4, generator=generator),
"txt_feat": torch.randn(1, 3, 4, generator=generator),
"null_txt_feat": torch.randn(1, 2, 4, generator=generator),
}
def _longer_text_tensors(seed: int) -> dict[str, torch.Tensor]:
tensors = _sampler_tensors(seed)
generator = torch.Generator(device="cpu").manual_seed(seed + 100)
tensors["txt_feat"] = torch.randn(1, 4, 4, generator=generator)
return tensors
def test_prepare_model_input_keeps_lengths_on_host() -> None:
sampler = Magi2PreviewSampler(torch.nn.Identity())
model_input = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([500.0]), cfg_config=CFGConfig())
assert model_input.audio_feat_len == [5, 5]
assert model_input.txt_feat_len == [3, 2]
assert model_input.ref_audio_feat_len == [0, 0]
assert model_input.ref_video_feat_len == [0, 0]
positive, negative = Magi2PreviewSampler._split_cfg_model_input(model_input)
assert positive.txt_feat_len == [3]
assert negative.txt_feat_len == [2]
def test_pre_adapter_embeds_directly_in_checkpoint_dtype() -> None:
model = _tiny_model(params_dtype=torch.bfloat16)
packed = torch.randn(6, 4)
indices = torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 5])
with torch.no_grad():
hidden = model.pre_adapter(packed, *indices)
assert hidden.dtype == torch.bfloat16
assert hidden.shape == (6, model.config.virtual_width)
@@ -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")
File diff suppressed because it is too large. Load diff
@@ -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]
)
@@ -6,7 +6,7 @@ import torch
from vllm.triton_utils import HAS_TRITON
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
from vllm_omni.diffusion.layers.indexed_modulation import (
_MAX_1D_GRID_SIZE,
_iter_row_chunks,
_launch_row_chunks,
@@ -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()
@@ -186,9 +186,8 @@ 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,
}
@@ -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)
@@ -0,0 +1,221 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import math
import sys
import types
from typing import Any
import pytest
import torch
from vllm_omni.diffusion.attention.backends.abstract import (
AttentionMetadata,
PackedPaddingMetadata,
VideoTokenLayout,
VideoTokenSpan,
)
from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import (
MiniMaxH3VSAImpl,
_get_h3_tile_metadata,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
extra: dict[str, Any] = {
"vsa_h3_prefix_segments": prefix_segments,
}
if gate is not None:
extra["gate_compress"] = gate
return AttentionMetadata(
packed_padding=packed_padding,
video_layout=VideoTokenLayout(
used_len=sum(prefix_segments) + math.prod(video_shape),
video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
),
extra=extra,
)
def _fake_block_sparse_kernel(monkeypatch, calls):
"""Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
calls["block_map"] = block_map.clone()
calls["variable_block_sizes"] = variable_block_sizes.tolist()
calls["q_shape"] = tuple(q.shape)
blocks = block_map.shape[-1]
assert blocks * 64 == q.shape[2]
# Expand the per-block map to a per-token mask so the stand-in honours
# the sparsity contract the real kernel implements, and drop the rows
# each partial block pads with, which the real kernel skips through
# variable_block_sizes.
token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
token_mask = token_mask & within.reshape(1, 1, 1, -1)
out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
return out.to(q.dtype), None
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
setattr(module, "block_sparse_attn", block_sparse_attn)
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
def _h3_impl(**backend_kwargs):
kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
kwargs.update(backend_kwargs)
return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
calls: dict[str, Any] = {}
_fake_block_sparse_kernel(monkeypatch, calls)
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5, 70), (2, 4, 4)
seq_len = sum(prefix_segments) + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
assert output.shape == query.shape
# topk covers every video block, so the block map is dense and the tiled
# kernel must reproduce plain attention over the same rows.
reference = torch.nn.functional.scaled_dot_product_attention(
*(x.float().transpose(1, 2) for x in (query, query, query))
).transpose(1, 2)
torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
# 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
assert calls["variable_block_sizes"] == [5, 64, 6, 32]
def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
calls: dict[str, Any] = {}
_fake_block_sparse_kernel(monkeypatch, calls)
impl = _h3_impl(topk=1)
prefix_segments, video_shape = (5,), (8, 4, 4)
seq_len = 5 + 8 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
# 1 prefix block + 2 video blocks. The odd count is padded to an even one
# for the paired-CTA kernel contract, and the Triton route slices that
# transport-only partner off again before launching.
block_map = calls["block_map"]
assert block_map.shape[-1] == 3
assert calls["variable_block_sizes"] == [5, 64, 64]
assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
assert block_map[..., :1].all(), "prefix keys are exempt from selection"
assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
# The same odd geometry must still untile back onto the original rows.
dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
reference = torch.nn.functional.scaled_dot_product_attention(
*(x.float().transpose(1, 2) for x in (query, query, query))
).transpose(1, 2)
torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
_fake_block_sparse_kernel(monkeypatch, {})
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5,), (2, 4, 4)
seq_len = 5 + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
zero_gate = torch.zeros_like(query)
with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
live_gate = torch.full_like(query, 0.5)
with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
assert not torch.allclose(with_live_gate.float(), without_gate.float())
def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
_fake_block_sparse_kernel(monkeypatch, {})
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5,), (2, 4, 4)
valid = 5 + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
padding = PackedPaddingMetadata(
q_length=valid,
kv_length=valid,
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
)
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
assert output.shape == query.shape
assert torch.count_nonzero(output[:, valid:]) == 0
def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
(5, 70, 9), (5, 6, 6), torch.device("cpu")
)
assert sizes[:4].tolist() == [5, 64, 6, 9]
assert prefix_blocks == 4
assert video_blocks == 8
assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
@pytest.mark.parametrize("fallback_on_error", [False, True])
def test_h3_kernel_failure_preserves_fallback_policy(monkeypatch, fallback_on_error):
impl = _h3_impl(fallback_on_error=fallback_on_error)
query = torch.randn(1, 37, 2, 8)
metadata = _h3_metadata((5,), (2, 4, 4))
def fail(*args, **kwargs):
raise RuntimeError("kernel failure")
monkeypatch.setattr(impl, "_forward_h3", fail)
monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
if fallback_on_error:
output = impl.forward_cuda(query, query, query, metadata)
expected = impl.sdpa_fallback.forward_cuda(query, query, query, metadata)
torch.testing.assert_close(output, expected, rtol=0, atol=0)
else:
with pytest.raises(RuntimeError, match="kernel failure"):
impl.forward_cuda(query, query, query, metadata)
def test_h3_accelerator_failure_is_not_swallowed(monkeypatch):
impl = _h3_impl(fallback_on_error=True)
query = torch.randn(1, 37, 2, 8)
def fail(*args, **kwargs):
raise torch.AcceleratorError("device failure")
monkeypatch.setattr(impl, "_forward_h3", fail)
with pytest.raises(torch.AcceleratorError, match="device failure"):
impl.forward_cuda(query, query, query, _h3_metadata((5,), (2, 4, 4)))
def test_h3_without_layout_delegates_to_shared_backend(monkeypatch):
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSAImpl
impl = _h3_impl()
query = torch.randn(1, 37, 2, 8)
metadata = AttentionMetadata()
expected = torch.empty_like(query)
calls = []
def forward(self, q, k, v, attn_metadata):
calls.append((self, q, k, v, attn_metadata))
return expected
monkeypatch.setattr(FastVideoVSAImpl, "forward_cuda", forward)
assert impl.forward_cuda(query, query, query, metadata) is expected
assert len(calls) == 1
assert all(actual is wanted for actual, wanted in zip(calls[0], (impl, query, query, query, metadata)))
@@ -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
+4 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from contextlib import contextmanager
@@ -33,7 +36,7 @@ class StubScheduler:
self.config = SimpleNamespace(num_train_timesteps=1000)
self.set_timesteps_calls: list[tuple[int, torch.device]] = []
def set_timesteps(self, num_steps: int, device: torch.device) -> None:
def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
self.set_timesteps_calls.append((num_steps, device))
@@ -0,0 +1,97 @@
{
"source_url": "https://github.com/Wan-Video/Wan2.2/blob/42bf4cfaa384bc21833865abc2f9e6c0e67233dc/wan/utils/fm_solvers_unipc.py",
"source_sha256": "0dec8c7ed17f6f2049275c6848113314da6ccec1c8db5bdc89df43c05c6038d9",
"git_head": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc",
"num_train_timesteps": 1000,
"use_dynamic_shifting": false,
"reproduction": "Load FlowUniPCMultistepScheduler from source_url; construct with num_train_timesteps=1000, shift=1.0, use_dynamic_shifting=False; call set_timesteps(40, device='cpu', shift=12.0); record timesteps.tolist() and sigmas.tolist().",
"constructor_shift": 1.0,
"runtime_shift": 12.0,
"steps": 40,
"timesteps": [
999,
997,
995,
993,
990,
988,
985,
982,
979,
976,
972,
969,
965,
961,
956,
952,
947,
941,
936,
929,
922,
915,
907,
898,
888,
877,
865,
852,
837,
819,
799,
776,
749,
717,
678,
631,
571,
492,
386,
235
],
"sigmas": [
0.9999166131019592,
0.9977826476097107,
0.9955461621284485,
0.9931996464729309,
0.9907346963882446,
0.988142192363739,
0.9854119420051575,
0.9825326800346375,
0.9794917702674866,
0.9762753844261169,
0.9728676676750183,
0.9692512154579163,
0.9654061794281006,
0.961310088634491,
0.9569376707077026,
0.95225989818573,
0.9472436308860779,
0.9418508410453796,
0.9360373020172119,
0.9297519326210022,
0.9229347705841064,
0.9155154228210449,
0.9074103832244873,
0.8985199332237244,
0.888724148273468,
0.8778772950172424,
0.865800678730011,
0.8522725701332092,
0.8370144367218018,
0.8196719288825989,
0.7997865080833435,
0.7767546772956848,
0.7497654557228088,
0.717703104019165,
0.6789888143539429,
0.6313128471374512,
0.5711563229560852,
0.49288034439086914,
0.38684695959091187,
0.23510953783988953,
0.0
],
"implementation": "Wan-Video/Wan2.2 wan/utils/fm_solvers_unipc.py"
}
@@ -1,15 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from PIL import Image
from torch import nn
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
Wan22I2VPipeline,
get_wan22_i2v_post_process_func,
@@ -50,7 +51,7 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
def _make_i2v_sampling(**overrides):
values = {
values: dict[str, object] = {
"height": 16,
"width": 16,
"num_frames": 5,
@@ -275,16 +276,18 @@ def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
assert first_frame_mask.shape == (2, 1, 2, 2, 2)
def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
@pytest.mark.parametrize("solver", ["unipc", "euler"])
@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch, solver: str, shift: float) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_i2v_pipeline(expand_timesteps=True)
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=5.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 5.0
pipeline.scheduler = build_wan_scheduler(solver, shift)
pipeline.od_config = SimpleNamespace(flow_shift=shift)
pipeline._sample_solver = solver
pipeline._flow_shift = shift
pipeline.boundary_ratio = 0.875
pipeline.has_image_encoder = False
pipeline._guidance_scale = None
@@ -327,6 +330,8 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
generator=gen_a,
latents=latents_a,
num_outputs_per_prompt=2,
num_inference_steps=50,
extra_args={"sample_solver": solver, "flow_shift": shift},
),
),
SimpleNamespace(
@@ -336,12 +341,25 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
generator=gen_b,
latents=latents_b,
num_outputs_per_prompt=2,
num_inference_steps=50,
extra_args={"sample_solver": solver, "flow_shift": shift},
),
),
]
)
outputs = pipeline.forward(batch)
if solver == "unipc":
sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
torch.testing.assert_close(
pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
)
assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
else:
reference = build_wan_scheduler("euler", shift)
reference.set_timesteps(50, device="cpu")
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
assert prepare_call["batch_size"] == 4
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
@@ -10,7 +10,7 @@ import torch
from torch import nn
from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, build_wan_scheduler
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -19,6 +19,12 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
def _cpu_platform(monkeypatch):
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
class _StubTransformer(nn.Module):
@property
def dtype(self) -> torch.dtype:
@@ -51,10 +57,10 @@ class _StubScheduler:
def __init__(self, timesteps: list[int]) -> None:
self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
self.config = SimpleNamespace(num_train_timesteps=1000)
self.set_timesteps_calls: list[tuple[int, torch.device]] = []
self.set_timesteps_calls: list[tuple[int, torch.device, float | None]] = []
def set_timesteps(self, num_steps: int, device: torch.device) -> None:
self.set_timesteps_calls.append((num_steps, device))
def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
self.set_timesteps_calls.append((num_steps, device, shift))
@contextmanager
@@ -184,7 +190,24 @@ def test_forward_delegates_denoising_to_diffuse(
assert captured["boundary_timestep"] == pytest.approx(875.0)
assert captured["latent_condition"] is None
assert captured["first_frame_mask"] is None
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"), 5.0)]
@pytest.mark.parametrize("solver", ["unipc", "euler"])
def test_forward_passes_request_shift_without_mutating_scheduler_config(solver: str) -> None:
pipeline = _make_pipeline()
pipeline.diffuse = lambda **kwargs: kwargs["latents"]
for shift in (3.0, 12.0, 5.0):
sampling = _make_sampling(num_inference_steps=5, extra_args={"sample_solver": solver, "flow_shift": shift})
request = OmniDiffusionRequest(prompt="prompt", request_id="schedule", sampling_params=sampling)
pipeline.forward(DiffusionRequestBatch(requests=[request]))
reference = build_wan_scheduler(solver, shift)
if solver == "unipc":
reference.set_timesteps(5, device="cpu", shift=shift)
assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
else:
reference.set_timesteps(5, device="cpu")
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
@@ -7,6 +10,7 @@ import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
Wan22S2VPipeline,
_make_clip_generators,
@@ -14,11 +18,32 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.mark.parametrize("configured_shift", [None, 12.0])
def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | None) -> None:
module = "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v"
config = SimpleNamespace(model="unused", flow_shift=configured_shift, enable_diffusion_pipeline_profiler=False)
def init_components(pipeline, *args):
pipeline.vae = SimpleNamespace(config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8))
with (
patch(f"{module}.get_local_device", return_value=torch.device("cpu")),
patch(f"{module}._resolve_model_path", return_value="unused"),
patch(f"{module}._is_diffusers_format", return_value=True),
patch.object(Wan22S2VPipeline, "_init_diffusers_format", init_components),
patch.object(Wan22S2VPipeline, "setup_diffusion_pipeline_profiler"),
):
pipeline = Wan22S2VPipeline(od_config=config)
assert pipeline._flow_shift == (3.0 if configured_shift is None else configured_shift)
assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
assert pipeline.scheduler.sigma_max == float(np.float32(0.999))
def _make_s2v_sampling(**overrides):
values = {
values: dict[str, object] = {
"height": 16,
"width": 16,
"num_frames": 8,
@@ -211,7 +236,8 @@ def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
pipeline.forward(batch)
def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
@pytest.mark.parametrize("shift", [3.0, 12.0])
def test_s2v_forward_batches_request_local_inputs_and_splits_outputs(shift: float) -> None:
pipeline = object.__new__(Wan22S2VPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
@@ -227,7 +253,8 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
enable_cpu_offload=False,
parallel_config=SimpleNamespace(use_hsdp=False),
)
pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
pipeline._flow_shift = shift
pipeline.scheduler = FlowUniPCMultistepScheduler(shift=1.0)
pipeline.vae_scale_factor_spatial = 8
pipeline.resolution_divisor = 16
pipeline.motion_frames = 7
@@ -262,6 +289,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
"multi_modal_data": {"image": image, "audio": audio_a},
},
sampling_params=_make_s2v_sampling(
num_inference_steps=5,
num_outputs_per_prompt=2,
generator=generators_a,
latents=latents_a,
@@ -275,6 +303,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
"multi_modal_data": {"image": image, "audio": audio_b},
},
sampling_params=_make_s2v_sampling(
num_inference_steps=5,
num_outputs_per_prompt=2,
generator=generators_b,
latents=latents_b,
@@ -287,6 +316,12 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
platform.is_available.return_value = False
outputs = pipeline.forward(batch)
expected = np.linspace(float(np.float32(0.999)), 0.0, 6)[:-1]
expected = shift * expected / (1.0 + (shift - 1.0) * expected)
torch.testing.assert_close(
pipeline.scheduler.sigmas, torch.tensor(np.append(expected, 0.0), dtype=torch.float32), rtol=0, atol=0
)
assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
assert len(outputs) == 2
assert outputs[0].output[0].shape[0] == 2
assert outputs[1].output[0].shape[0] == 2
@@ -519,6 +554,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
mock_scheduler.timesteps = torch.linspace(999, 0, 5)
mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),))
pipeline.scheduler = mock_scheduler
pipeline._flow_shift = 3.0
# -- Bind methods from the real class --
pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline)
@@ -0,0 +1,54 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Compare actual Wan schedules to native Wan outputs and independent sigma arithmetic."""
import json
from pathlib import Path
import numpy as np
import numpy.typing as npt
import pytest
import torch
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
FASTWAN_DMD_SCHEDULER_SHIFT,
build_wan_scheduler,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_wan_unipc_matches_captured_native_full_schedule() -> None:
# The fixture pins the official Wan source and records how to regenerate it.
golden = json.loads((Path(__file__).parent / "fixtures/native_unipc_shift12.json").read_text())
scheduler = build_wan_scheduler("unipc", golden["runtime_shift"])
for steps in (golden["steps"], 20, golden["steps"]):
scheduler.set_timesteps(steps, device="cpu", shift=golden["runtime_shift"])
if steps == golden["steps"]:
assert scheduler.timesteps.tolist() == golden["timesteps"]
torch.testing.assert_close(scheduler.sigmas, torch.tensor(golden["sigmas"]), rtol=0, atol=0)
@pytest.mark.parametrize("shift", [1.0, 3.0, 5.0, 12.0])
@pytest.mark.parametrize("steps", [20, 40])
def test_wan_unipc_shifts_unmodified_training_endpoints(shift: float, steps: int) -> None:
# Native Wan: float32 training sigma_max=(1-1/1000), sigma_min=0;
# interpolate in float64, apply the runtime shift, then cast sigma to FP32.
unshifted = np.linspace(float(np.float32(0.999)), 0.0, steps + 1)[:-1]
shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted)
expected_timesteps = (shifted * 1000).astype(np.int64)
expected_sigmas: npt.NDArray[np.float32] = np.append(shifted, 0.0).astype(np.float32)
scheduler = build_wan_scheduler("unipc", shift)
scheduler.set_timesteps(steps, device="cpu", shift=shift)
assert scheduler.config.shift == scheduler.config["shift"] == 1.0
torch.testing.assert_close(scheduler.timesteps, torch.from_numpy(expected_timesteps), rtol=0, atol=0)
torch.testing.assert_close(scheduler.sigmas, torch.from_numpy(expected_sigmas), rtol=0, atol=0)
def test_wan_euler_keeps_dmd_shift_semantics() -> None:
scheduler = build_wan_scheduler("euler", FASTWAN_DMD_SCHEDULER_SHIFT)
scheduler.set_timesteps(3, device="cpu")
# The unchanged Euler branch starts at sigma=1, with shift=8 for DMD.
expected = torch.tensor([1.0, 16.0 / 17.0, 0.8, 0.0])
torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=1e-7)
torch.testing.assert_close(scheduler.timesteps, expected[:-1] * 1000, rtol=0, atol=1e-4)
@@ -1,15 +1,17 @@
# 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 types import SimpleNamespace
import numpy as np
import pytest
import torch
from PIL import Image
from torch import nn
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import build_wan_scheduler
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
Wan22VACEPipeline,
create_vace_transformer_from_config,
@@ -33,6 +35,8 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
)
enable_layerwise_offload: bool = True
post_init_calls: int = 0
_resolved_diffusion_offload: object = None
_diffusion_offload_flags_materialized: bool = False
def __post_init__(self) -> None:
self.post_init_calls += 1
@@ -41,7 +45,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
resolved_offload = object()
config._resolved_diffusion_offload = resolved_offload
config._diffusion_offload_flags_materialized = True
captured = {}
captured: dict[str, object] = {}
def capture_base_init(self, *, od_config, prefix="") -> None:
captured.update(od_config=od_config, prefix=prefix)
@@ -51,6 +55,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
Wan22VACEPipeline(od_config=config, prefix="stage")
forwarded = captured["od_config"]
assert isinstance(forwarded, RuntimeConfig)
assert forwarded is not config
assert config.flow_shift is None
assert forwarded.flow_shift == 3.0
@@ -75,7 +80,7 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
def _make_vace_sampling(**overrides):
values = {
values: dict[str, object] = {
"height": 16,
"width": 16,
"num_frames": 5,
@@ -123,7 +128,7 @@ def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None:
captured = {}
captured: dict[str, object] = {}
class FakeVACETransformer:
def __init__(self, **kwargs) -> None:
@@ -203,7 +208,9 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
torch.testing.assert_close(result, torch.ones_like(latents))
def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
@pytest.mark.parametrize("solver", ["unipc", "euler"])
@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch, solver: str, shift: float) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
@@ -211,10 +218,10 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
pipeline = _make_vace_pipeline()
pipeline.transformer.vace_patch_embedding = None
pipeline.transformer_2 = None
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 3.0
pipeline.scheduler = build_wan_scheduler(solver, shift)
pipeline.od_config = SimpleNamespace(flow_shift=shift)
pipeline._sample_solver = solver
pipeline._flow_shift = shift
pipeline.boundary_ratio = None
pipeline._guidance_scale = None
pipeline._num_timesteps = None
@@ -247,6 +254,8 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
generator=gen_a,
latents=latents_a,
num_outputs_per_prompt=2,
num_inference_steps=50,
extra_args={"sample_solver": solver, "flow_shift": shift},
),
),
SimpleNamespace(
@@ -256,12 +265,25 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
generator=gen_b,
latents=latents_b,
num_outputs_per_prompt=2,
num_inference_steps=50,
extra_args={"sample_solver": solver, "flow_shift": shift},
),
),
]
)
outputs = pipeline.forward(batch)
if solver == "unipc":
sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
torch.testing.assert_close(
pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
)
assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
else:
reference = build_wan_scheduler("euler", shift)
reference.set_timesteps(50, device="cpu")
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
assert prepare_call["batch_size"] == 4
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
@@ -39,7 +39,8 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage
from vllm_omni.engine.stage_init_utils import StageMetadata
from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -42,7 +42,8 @@ from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
from vllm_omni.engine.stage_client import StagePoolClient
from vllm_omni.engine.stage_init_utils import StageMetadata
from vllm_omni.engine.stage_pool import StagePool
from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
from vllm_omni.outputs import OmniRequestOutput
Loaded 100 of 371 files, more files were not shown because too many files have changed in this diff. Show more