Compare commits

..
Author SHA1 Message Date
Gao Han 62aeffd15c Merge branch 'main' into fix/tp-local-kv-recv-consensus 2026-08-01 23:23:52 +08:00
tzhouamandClaude Fable 5 bac3eafc18 review: scope KV receive consensus to the TP group, compare full KV state
Address both P1 review comments:

- The exchange now runs on the stage's tensor-model-parallel group
  (vLLM parallel state) instead of pt.world. World also holds DP
  replicas and PP stages working on different requests, so the old
  group could mix verdicts across requests and — worse — deadlock if
  one replica skipped the exchange (dummy request) while another
  entered it. Within the TP group every skip guard is symmetric.

- Ranks now exchange (primary_received, sorted applied-KV roles)
  instead of a bare flag, using the same field enumeration as
  apply/broadcast, so partially-received CFG companion KV
  (cfg_text/cfg_img) diverging between ranks is caught too. Verdict is
  identical-state-everywhere: any divergence clears KV on every rank;
  symmetric absence is kept because it is consistent (Bagel's empty
  unconditional branch).

Local e2e rerun of test_bagel[tp_vae_patch_parallel_2] passed and the
new path fired on a real divergence (one rank primary+metadata, the
other only cfg_kv_request_ids) — both ranks fell back together.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-08-01 16:29:29 +08:00
Gao Han 0eda607077 Merge branch 'main' into fix/tp-local-kv-recv-consensus 2026-08-01 10:50:23 +08:00
tzhouamandClaude Fable 5 6d767a9a5b [BugFix] Make pure-TP KV receive fallback all-or-nothing across ranks
Fixes the recurring Bagel tp_vae_patch_parallel_2 NCCL ALLREDUCE
timeout (#5627, previously #5487): in pure TP the transfer topology is
LOCAL on every rank — each rank fetches its own per-rank KV key with an
independent recv_timeout and distribute_kv_cache is a no-op. A receive
timeout returns silently, so a single rank that missed its shard fell
back alone into the pipeline's no-KV prefill (Bagel
forward_cache_update_text through the TP-sharded LM) and entered TP
allreduces its peers never joined. The hung [~20, 3584] allreduce
(NumelIn=71680) blocked 600s until the NCCL watchdog killed the worker,
after which every request failed with "DiffusionExecutor is closed".
PR #5482 fixed the orchestrator-level CFG bundle race but never touched
this lower KV-transfer layer, which is why the signature recurred at
its own merge commit.

After the per-rank receive, pure-TP ranks now exchange their receive
flags (rank 0 collects and broadcasts the verdict over the stage world
group's CPU object plane) and act on it together: if any rank missed,
ranks that did receive drop their applied KV — cleared symmetrically
with what apply/broadcast set — so every rank takes the identical no-KV
prefill path and the TP collective schedule can never diverge. The skip
guards (topology, need_recv_cache, dummy warmup id) depend only on
values identical across ranks, so ranks always agree on whether the
exchange happens; a rank whose connector failed to initialize still
participates and reports a miss instead of deadlocking the exchange.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-07-31 11:08:13 +08:00
263 changed files with 3409 additions and 27990 deletions

No files matched your search

+4 -22
View File
@@ -427,12 +427,6 @@ steps:
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · MiniMax H3 T2VA Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_t2va_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2V · Other Accuracy Test"
timeout_in_minutes: 180
commands:
@@ -463,38 +457,26 @@ steps:
- group: ":card_index_dividers: Diffusion Test"
key: nightly-diffusion-test-group
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
steps:
- label: ":full_moon: Diffusion · Distributed Test with L4"
- label: ":full_moon: Diffusion · Distributed Test"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda' --run-level "full_model"
mirror_hardwares: l4_4
- label: ":full_moon: Diffusion · Distributed Test with H100"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model"
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion · Offloader Test"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/offloader -m "full_model and cuda"
mirror_hardwares: l4_1
- label: ":full_moon: Diffusion · Quantization Test with H100 · Single-GPU"
- label: ":full_moon: Diffusion · Quantization Test with H100"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion · MiniMax H3 FP8 Quality Test with H100 · 2-GPU"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py -m 'full_model and cuda and H100' --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion · Quantization Test with L4"
timeout_in_minutes: 60
commands:
-82
View File
@@ -95,80 +95,6 @@ steps:
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
- group: ":card_index_dividers: Diffusion I2I Model Test"
key: nightly-HunyuanImage3-x2i-test-group
depends_on: upload-nightly-pipeline
if: >-
build.env("NIGHTLY") == "1" ||
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
- label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
timeout_in_minutes: 180
mirror_hardwares: a3_npu_4
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HUNYUAN_IMAGE3_MODEL: "tencent/HunyuanImage-3.0-Instruct"
HUNYUAN_IMAGE3_DEVICES: "0,1,2,3"
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and A3" --run-level full_model
- group: ":card_index_dividers: Diffusion Benchmark Perf Test"
key: nightly-diffusion-benchmark-group
depends_on: upload-nightly-pipeline
if: >-
build.env("NIGHTLY") == "1" ||
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
- label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 CFGP=2"
timeout_in_minutes: 180
mirror_hardwares: a3_npu_4
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_TOKEN: "${HF_TOKEN}"
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
- label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 SP=2"
timeout_in_minutes: 180
mirror_hardwares: a3_npu_4
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_TOKEN: "${HF_TOKEN}"
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
- label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=4"
timeout_in_minutes: 180
mirror_hardwares: a3_npu_4
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_TOKEN: "${HF_TOKEN}"
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
- group: ":card_index_dividers: Diffusion X2V Model Test"
key: nightly-diffusion-x2v-group
depends_on: upload-nightly-pipeline
@@ -177,14 +103,6 @@ steps:
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2v-test"
steps:
- label: ":full_moon: Diffusion X2V · Function Test"
mirror_hardwares: a2b3_npu_4
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
+1 -1
View File
@@ -309,7 +309,7 @@ For Omni or custom models, create:
Required updates:
1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
2. `docs/user_guide/diffusion/cpu_offload.md` — if CPU offload supported (add to supported models table)
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
+4 -4
View File
@@ -16,15 +16,15 @@
/.claude/ @hsliuustc0106
# Tests
/tests/ @yenuo26 @NickCao
/tests/ @yenuo26
# CI
/.buildkite/ @yenuo26 @congw729 @NickCao
/.buildkite/ @yenuo26 @congw729
# Engine and model_executor
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/worker/ @tzhouam @gcanlin
@@ -37,7 +37,7 @@
# Omni and TTS support
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin
# Benchmarks
/benchmarks/accuracy/ @david6666666
+8 -5
View File
@@ -16,10 +16,14 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*Latest News* 🔥
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
---
@@ -51,13 +55,12 @@ vLLM-Omni is flexible and easy to use with:
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- OpenAI-compatible API server
- Full-duplex realtime serving with streaming audio input and output (experimental)
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
- **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
## Getting Started
+3 -7
View File
@@ -47,7 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
- **Generate Image** for text-to-image and image-to-image tasks
- **Generate Video** for text-to-video, first-frame/image-to-video, and reference-conditioned video
- **Generate Video** for text-to-video and image-to-video tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
@@ -105,7 +105,7 @@ You can configure per-stage sampling parameters for multi-stage models.
> [!TIP]
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
### Text-to-video and image-to-video generation (e.g., Wan, MiniMax-H3)
### Text-to-video and image-to-video generation (e.g., Wan)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
@@ -117,11 +117,7 @@ You can configure per-stage sampling parameters for multi-stage models.
</p>
> [!TIP]
> Connect a **frame** image for first-frame / image-to-video (e.g., Wan I2V, MiniMax-H3 FL2VA).
>
> For reference-conditioned generation (MiniMax-H3 Ref2VA), connect a **Video References** node instead.
>
> Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect.
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
### TTS (e.g., Qwen TTS series)
-6
View File
@@ -15,13 +15,11 @@ from .comfyui_vllm_omni.nodes import (
VLLMOmniDiffusionSampling,
VLLMOmniGenerateImage,
VLLMOmniGenerateVideo,
VLLMOmniMiniMaxH3Params,
VLLMOmniQwenTTSParams,
VLLMOmniRemoteLoRA,
VLLMOmniSamplingParamsList,
VLLMOmniTTS,
VLLMOmniUnderstanding,
VLLMOmniVideoReferences,
VLLMOmniVoiceClone,
VLLMOmniWanParams,
)
@@ -34,7 +32,6 @@ NODE_CLASS_MAPPINGS = {
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
"VLLMOmniTTS": VLLMOmniTTS,
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
"VLLMOmniVideoReferences": VLLMOmniVideoReferences,
# === Params ===
"VLLMOmniARSampling": VLLMOmniARSampling,
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
@@ -42,7 +39,6 @@ NODE_CLASS_MAPPINGS = {
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
"VLLMOmniWanParams": VLLMOmniWanParams,
"VLLMOmniMiniMaxH3Params": VLLMOmniMiniMaxH3Params,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -53,7 +49,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"VLLMOmniUnderstanding": "Multimodality Understanding",
"VLLMOmniTTS": "TTS (Text to Speech)",
"VLLMOmniVoiceClone": "TTS Voice Cloning",
"VLLMOmniVideoReferences": "Video References",
# === Params ===
"VLLMOmniARSampling": "AR Sampling Params",
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
@@ -61,7 +56,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"VLLMOmniRemoteLoRA": "LoRA",
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
"VLLMOmniWanParams": "Wan Video Params",
"VLLMOmniMiniMaxH3Params": "MiniMax-H3 Video Params",
}
WEB_DIRECTORY = "./web"
@@ -10,9 +10,7 @@ from .utils.types import (
AudioFormat,
AutoregressionSamplingParams,
DiffusionSamplingParams,
MiniMaxH3ModelSpecificParams,
QwenTTSModelSpecificParams,
VideoReferences,
WanModelSpecificParams,
)
from .utils.validators import (
@@ -165,8 +163,7 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
"num_frames": ("INT", {"default": 41, "min": 1}),
},
"optional": {
"frame": ("IMAGE",),
"references": ("VIDEO_REFERENCES",),
"image": ("IMAGE",),
"sampling_params": ("SAMPLING_PARAMS",),
"lora": ("REMOTE_LORA",),
"model_params": ("VIDEO_PARAMS",),
@@ -177,15 +174,6 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
RETURN_NAMES = ("video",)
FUNCTION = "generate"
@classmethod
def VALIDATE_INPUTS(cls, url, model, frame=None, references=None, **_kwargs) -> str | Literal[True]:
base = super().VALIDATE_INPUTS(url, model)
if base is not True:
return base
if frame is not None and references is not None:
return "Provide only one of frame or references, not both."
return True
async def generate(
self,
url: str,
@@ -196,8 +184,7 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
fps: int,
num_frames: int,
negative_prompt: str | None = None,
frame: torch.Tensor | None = None,
references: dict | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | list[dict] | None = None,
model_params: dict | None = None,
lora: dict | None = None,
@@ -220,16 +207,14 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
if sampling_params is not None:
sampling_params.pop("type", None) # internal fields
# model_params["type"] is also an internal field, but is used in api_client
# if model_params is not None:
# model_params.pop("type", None) # internal fields
if model_params is not None:
model_params.pop("type", None) # internal fields
client = VLLMOmniClient(url)
output = await client.generate_video(
model=model,
prompt=prompt,
frame=frame, # frame present => fl2va / Wan I2V
references=references,
image=image, # image present => i2v, absent => t2v
width=width,
height=height,
num_frames=num_frames,
@@ -749,78 +734,3 @@ class VLLMOmniWanParams:
def get_params(self, **kwargs):
return (WanModelSpecificParams(kwargs),)
class VLLMOmniMiniMaxH3Params:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"audio_flow_shift": (
"FLOAT",
{"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.1},
),
"flow_shift": (
"FLOAT",
{"default": 12.0, "min": 0.0, "max": 100.0, "step": 0.1},
),
}
}
RETURN_TYPES = ("VIDEO_PARAMS",)
RETURN_NAMES = ("MiniMax-H3 video params",)
FUNCTION = "get_params"
CATEGORY = "vLLM-Omni/Video Params"
def get_params(self, **kwargs):
params = MiniMaxH3ModelSpecificParams(kwargs)
params["type"] = "minimax_h3"
return (params,)
class VLLMOmniVideoReferences:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {},
"optional": {
"image_1": ("IMAGE",),
"image_2": ("IMAGE",),
"audio_1": ("AUDIO",),
"audio_2": ("AUDIO",),
"video_1": ("VIDEO",),
"video_2": ("VIDEO",),
},
}
RETURN_TYPES = ("VIDEO_REFERENCES",)
RETURN_NAMES = ("references",)
FUNCTION = "get_references"
CATEGORY = "vLLM-Omni"
def get_references(
self,
image_1: torch.Tensor | None = None,
image_2: torch.Tensor | None = None,
audio_1: AudioInput | None = None,
audio_2: AudioInput | None = None,
video_1: VideoInput | None = None,
video_2: 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
return (refs,)
@@ -24,7 +24,6 @@ from .format import (
image_tensor_to_base64,
image_tensor_to_png_bytes,
video_to_base64,
video_to_bytes,
)
from .logger import get_logger, pretty_printer
from .models import lookup_model_spec
@@ -64,11 +63,7 @@ async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get",
class VLLMOmniClient:
def __init__(
self,
base_url: str,
timeout: float | None = None,
poll_interval: float = 5.0,
max_poll_duration: float = 60 * 30,
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
):
self.base_url = base_url
self.timeout = aiohttp.ClientTimeout(total=timeout)
@@ -266,17 +261,12 @@ class VLLMOmniClient:
num_frames: int,
fps: int,
negative_prompt: str | None = None,
frame: torch.Tensor | None = None,
references: dict | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | None = None,
model_params: dict | None = None,
lora: dict | None = None,
**extra_params,
**extra_body,
) -> VideoInput:
if frame is not None and references is not None:
raise ValueError("Provide only one of frame or references, not both.")
# === regular payload fields ===
form = aiohttp.FormData()
form.add_field("model", model)
form.add_field("prompt", prompt)
@@ -289,91 +279,23 @@ class VLLMOmniClient:
if sampling_params is not None:
for k, v in sampling_params.items():
form.add_field(k, str(v))
if model_params is not None:
for k, v in model_params.items():
form.add_field(k, str(v))
if lora is not None:
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
if extra_body:
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
# === 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:
raise ValueError(
"Invalid references combination. Supported modes: "
"(1) one or more videos only, or (2) exactly one image and one audio."
)
else:
video_task = "t2va"
if input_reference_image is not None:
if image is not None:
image_filename = "image.png" # Required for multipart form
form.add_field(
"input_reference",
image_tensor_to_png_bytes(input_reference_image, image_filename),
image_tensor_to_png_bytes(image, image_filename),
filename=image_filename,
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)
model_params.pop("type", None)
spec, _ = lookup_model_spec(model)
params_builder = spec.get("params_builder") if spec else None
if params_builder is not None:
form_fields = params_builder(
model_params or {},
extra_params={**extra_params, "task": video_task},
)
for k, v in form_fields.items():
form.add_field(k, v if isinstance(v, str) else str(v))
else:
if model_params is not None:
for k, v in model_params.items():
form.add_field(k, str(v))
if extra_params:
form.add_field("extra_params", json.dumps(extra_params, ensure_ascii=False))
async with aiohttp.ClientSession(timeout=self.timeout) as session:
# Start the video generation job
url = f"{self.base_url}/videos"
@@ -1,12 +1,7 @@
import json
import re
from typing import Any
from .logger import get_logger
from .types import Modality, ModelMode, Spec
logger = get_logger(__name__)
def _bagel_payload_preprocessor(payload: dict) -> dict:
try:
@@ -34,30 +29,6 @@ def _qwen25_payload_preprocessor(payload: dict) -> dict:
return payload
def _minimaxh3_params_builder(
model_params: dict[str, Any],
*,
extra_params: dict[str, Any],
) -> dict[str, Any]:
"""Build multipart form fields for MiniMax-H3 from model_params + routed task."""
params = dict(model_params)
params.pop("type", None)
form_fields: dict[str, Any] = {}
merged_extra_params: dict[str, Any] = dict(extra_params or {})
if "flow_shift" in params:
form_fields["flow_shift"] = params.pop("flow_shift")
for key in ("audio_flow_shift",):
if key in params:
merged_extra_params[key] = params.pop(key)
if params:
logger.warning("Unused MiniMax-H3 model params ignored: %s", sorted(params))
if merged_extra_params:
form_fields["extra_params"] = json.dumps(merged_extra_params, ensure_ascii=False)
return form_fields
_MODEL_PIPELINE_SPECS: dict[str, Spec] = {
r"BAGEL-7B-MoT": {
"stages": [
@@ -100,21 +71,6 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
}
],
},
r"MiniMax-H3(/FL2VA|/Ref2VA)?": {
"stages": ["diffusion"],
"modes": [
{
"mode": ModelMode.VIDEO_GENERATION,
"input_modalities": [
Modality.TEXT,
Modality.IMAGE,
Modality.VIDEO,
Modality.AUDIO,
],
}
],
"params_builder": _minimaxh3_params_builder,
},
}
# Convert dict keys to regex patterns
MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {}
@@ -124,10 +80,12 @@ del _MODEL_PIPELINE_SPECS
def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]:
normalized = model.rstrip("/").lstrip("/")
lookup_key = normalized.split("/", 1)[-1]
try:
last_component = model.rstrip("/").rsplit("/", 1)[-1]
except IndexError:
last_component = model
for pattern, spec in MODEL_PIPELINE_SPECS.items():
if pattern.search(lookup_key):
if pattern.search(last_component):
return spec, pattern.pattern
return None, None
@@ -143,9 +101,6 @@ if __name__ == "__main__":
"Qwen/Qwen3-Omni-30B-A3B-Instruct",
"Custom/Path/UnknownModel-Instruct",
"Not/Matching/Anything",
"/path/to/MiniMax-H3/FL2VA",
"/path/to/MiniMax-H3/Ref2VA",
"MiniMaxAI/MiniMax-H3",
]
test_payload = {"messages": [{"role": "user", "content": "prompt"}]}
@@ -157,13 +112,7 @@ if __name__ == "__main__":
if preprocessor := spec.get("payload_preprocessor"):
result = preprocessor(test_payload)
print(f"✓ {path:<40} → {result}")
elif params_builder := spec.get("params_builder"):
result = params_builder(
{"audio_flow_shift": 3.0, "flow_shift": 12.0, "type": "minimax_h3"},
extra_params={"task": "t2va"},
)
print(f"✓ {path:<40} → {result}")
else:
print(f"✓ {path:<40} → No preprocessor/params_builder")
print(f"✓ {path:<40} → No preprocessor")
else:
print(f"✗ {path:<40} → No match")
@@ -23,14 +23,6 @@ class WanModelSpecificParams(dict):
pass
class MiniMaxH3ModelSpecificParams(dict):
pass
class VideoReferences(dict):
pass
class ModelMode(Enum):
IMAGE_GENERATION = auto()
VIDEO_GENERATION = auto()
@@ -51,11 +43,9 @@ class ModelModeSpec(TypedDict):
PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]]
ParamsBuilder: TypeAlias = Callable[..., dict[str, Any]]
class Spec(TypedDict):
stages: list[Literal["diffusion", "autoregression"]]
modes: list[ModelModeSpec]
payload_preprocessor: NotRequired[PayloadPreprocessor]
params_builder: NotRequired[ParamsBuilder]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 161 KiB

After

Width:  |  Height:  |  Size: 129 KiB

@@ -1 +1,513 @@
{"id":"3aea1d08-7de5-4997-ba71-b2f66ad72a21","revision":0,"last_node_id":24,"last_link_id":35,"nodes":[{"id":13,"type":"SaveVideo","pos":[1315.729153459605,-137.92038948394355],"size":[374,313.7692307692308],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":32},{"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/ComfyUI","auto","auto"]},{"id":10,"type":"MarkdownNote","pos":[183.74636797766794,-571.2529158046766],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both *text-to-video*, *frame-to-video*, and *reference-to-video* generation. The node will create corresponding payload depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[452.07189605801125,-569.7242536425791],"size":[270,314],"flags":{},"order":1,"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":[29]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["rap.png","image"]},{"id":23,"type":"LoadImage","pos":[9.206110246057657,-7.6955291509170785],"size":[270,314],"flags":{},"order":2,"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":[34]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["white-cat.jpg","image"]},{"id":22,"type":"MarkdownNote","pos":[-9.783113411416666,-148.8790300038605],"size":[297.3012007349086,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Reference to video (r2v)","properties":{},"widgets_values":["For r2v tasks, connect this Video Reference block to the main Generate Video block. Then disconnect the \"frame\" input above."],"color":"#432","bgcolor":"#000"},{"id":24,"type":"LoadAudio","pos":[317.9929742385805,-141.44821500065095],"size":[282.798828125,136],"flags":{},"order":4,"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":[35]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["nothing-beats-a-jet2-holiday.mp3",null,""]},{"id":21,"type":"VLLMOmniVideoReferences","pos":[376.6129172897595,105.41557277556292],"size":[156.9380859375,126],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"image_1","name":"image_1","shape":7,"type":"IMAGE","link":34},{"localized_name":"image_2","name":"image_2","shape":7,"type":"IMAGE","link":null},{"localized_name":"audio_1","name":"audio_1","shape":7,"type":"AUDIO","link":35},{"localized_name":"audio_2","name":"audio_2","shape":7,"type":"AUDIO","link":null},{"localized_name":"video_1","name":"video_1","shape":7,"type":"VIDEO","link":null},{"localized_name":"video_2","name":"video_2","shape":7,"type":"VIDEO","link":null}],"outputs":[{"localized_name":"references","name":"references","type":"VIDEO_REFERENCES","links":null}],"properties":{"Node name for S&R":"VLLMOmniVideoReferences"},"widgets_values":[]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[83.24620208113915,422.08543477542605],"size":[284.205078125,226],"flags":{},"order":5,"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_scLine truncated
{
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
"revision": 0,
"last_node_id": 14,
"last_link_id": 28,
"nodes": [
{
"id": 10,
"type": "MarkdownNote",
"pos": [
227.99306819575278,
-231.24143306069843
],
"size": [
240.3326478922474,
136.3505820791881
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Task and Input",
"properties": {},
"widgets_values": [
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 4,
"type": "LoadImage",
"pos": [
496.31859627609606,
-229.71277089860084
],
"size": [
270,
314
],
"flags": {},
"order": 1,
"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": [
25
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"cute-cat.jpg",
"image"
]
},
{
"id": 8,
"type": "VLLMOmniDiffusionSampling",
"pos": [
465.5140218220927,
148.80072646828987
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 2,
"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": [
27
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1,
1,
false,
false,
1992,
"randomize"
]
},
{
"id": 14,
"type": "VLLMOmniWanParams",
"pos": [
518.6005521845747,
428.4455663719664
],
"size": [
270,
106
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "guidance_scale_2",
"name": "guidance_scale_2",
"type": "FLOAT",
"widget": {
"name": "guidance_scale_2"
},
"link": null
},
{
"localized_name": "boundary_ratio",
"name": "boundary_ratio",
"type": "FLOAT",
"widget": {
"name": "boundary_ratio"
},
"link": null
},
{
"localized_name": "flow_shift",
"name": "flow_shift",
"type": "FLOAT",
"widget": {
"name": "flow_shift"
},
"link": null
}
],
"outputs": [
{
"localized_name": "Wan video params",
"name": "Wan video params",
"type": "VIDEO_PARAMS",
"links": [
28
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniWanParams"
},
"widgets_values": [
4,
0.875,
5
]
},
{
"id": 12,
"type": "VLLMOmniGenerateVideo",
"pos": [
827.2566336087867,
-73.77449831827545
],
"size": [
400,
346
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 25
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 27
},
{
"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": 28
},
{
"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": "num_frames",
"name": "num_frames",
"type": "INT",
"widget": {
"name": "num_frames"
},
"link": null
}
],
"outputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"links": [
26
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateVideo"
},
"widgets_values": [
"http://localhost:8000/v1",
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
"Cartoonish",
900,
1200,
12,
41
]
},
{
"id": 13,
"type": "SaveVideo",
"pos": [
1257.5441369501657,
-44.12956394419863
],
"size": [
374,
372.34896826553796
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"link": 26
},
{
"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": "COMBO",
"widget": {
"name": "codec"
},
"link": null
}
],
"outputs": [],
"properties": {},
"widgets_values": [
"video/ComfyUI",
"auto",
"auto"
]
}
],
"links": [
[
25,
4,
0,
12,
0,
"IMAGE"
],
[
26,
12,
0,
13,
0,
"VIDEO"
],
[
27,
8,
0,
12,
1,
"SAMPLING_PARAMS"
],
[
28,
14,
0,
12,
3,
"VIDEO_PARAMS"
]
],
"groups": [
{
"id": 1,
"title": "Input",
"bounding": [
213.1706010087147,
-313.20095750667554,
560.8246471437027,
407.4246532472182
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.1469075819486894,
"offset": [
101.7632950847089,
492.29121889347806
]
}
},
"version": 0.4
}
-1
View File
@@ -45,7 +45,6 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
"hyperpyyaml>=1.2.0" \
"onnx>=1.16.0" \
"librosa>=0.10.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 && \
+27 -44
View File
@@ -68,12 +68,13 @@ nav:
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
- VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
- CPU Offloading: user_guide/diffusion/cpu_offload.md
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
- LoRA: user_guide/diffusion/lora.md
- Custom Pipeline: features/custom_pipeline.md
- Execution Modes: user_guide/diffusion/execution_modes.md
- Request-Level Batching: user_guide/diffusion/request_batching.md
- Step Execution: user_guide/diffusion/step_execution.md
- Quantization:
- Quantization: user_guide/quantization/overview.md
- Overview: user_guide/quantization/overview.md
- Online Quantization: user_guide/quantization/online.md
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
- FP8 W8A8: user_guide/quantization/fp8.md
@@ -101,49 +102,31 @@ nav:
- CI: contributing/ci
- Design Documents:
- design/index.md
- Architecture Documents:
- Architecture Overview: design/architecture_overview.md
- design/architecture_overview.md
- Feature Design:
- Runtime and Stage Execution:
- Disaggregated Inference: design/feature/disaggregated_inference.md
- Async Chunk: design/feature/async_chunk.md
- Async Diffusion Output: design/feature/async_diffusion_output.md
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
- Automatic Prefix Caching: design/feature/prefix_caching.md
- Communication:
- OmniConnector Implementations:
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
- Mooncake Transfer Engine Connector: design/feature/omni_connectors/mooncake_transfer_engine_connector.md
- Mori Transfer Engine Connector: design/feature/omni_connectors/mori_transfer_engine_connector.md
- Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md
- Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md
- Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
- Diffusion Acceleration:
- Parallelism:
- CFG-Parallel: design/feature/cfg_parallel.md
- Expert Parallel: design/feature/expert_parallel.md
- Hybrid Sharded Data Parallel (HSDP): design/feature/hsdp.md
- Pipeline Parallel: design/feature/pipeline_parallel.md
- Sequence Parallel: design/feature/sequence_parallel.md
- Tensor Parallel: design/feature/tensor_parallel.md
- VAE Patch Parallelism: design/feature/vae_parallel.md
- Attention Backends:
- Skip-Softmax: design/feature/skip_softmax.md
- Quantization: design/feature/quantization.md
- Cache-DiT: design/feature/cache_dit.md
- TeaCache: design/feature/teacache.md
- Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
- Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
- Infrastructure and Performance:
- Prometheus Metrics: design/metrics.md
- Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
- design/feature/disaggregated_inference.md
- design/feature/ray_based_execution.md
- design/feature/omni_connectors/
- design/feature/prefix_caching.md
- design/feature/cfg_parallel.md
- design/feature/expert_parallel.md
- design/feature/sequence_parallel.md
- design/feature/tensor_parallel.md
- design/feature/vae_parallel.md
- design/feature/hsdp.md
- design/feature/cache_dit.md
- design/feature/teacache.md
- design/feature/skip_softmax.md
- design/feature/async_chunk.md
- design/feature/vae_parallel.md
- design/feature/diffusion_step_execution.md
- design/feature/diffusion_request_level_batching.md
- design/feature/diffusion_continuous_batching.md
- Module Design:
- Runtime Modules:
- AR Module: design/module/ar_module.md
- DIT Module: design/module/dit_module.md
- Orchestration:
- Entrypoint Module: design/module/entrypoint_module.md
- AsyncOmni Architecture: design/module/async_omni_architecture.md
- design/module/ar_module.md
- design/module/dit_module.md
- design/module/entrypoint_module.md
- design/module/async_omni_architecture.md
- Docs Guide: contributing/DOCS_GUIDE.md
- API Reference:
- api/README.md
+2 -3
View File
@@ -52,13 +52,12 @@ vLLM-Omni is flexible and easy to use with:
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- OpenAI-compatible API server
- Full-duplex realtime serving with streaming audio input and output (experimental)
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- Omni-modality models (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
- TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
- Diffusion models — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
- Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
For more information, checkout the following:
Binary file not shown.

Before

Width:  |  Height:  |  Size: 274 KiB

After

Width:  |  Height:  |  Size: 281 KiB

+7 -1
View File
@@ -18,12 +18,18 @@ Specify the port:
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
```
Load a custom deploy configuration with `--deploy-config`:
For a migrated model, load a custom deploy configuration with `--deploy-config`:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --deploy-config /path/to/deploy_config.yaml
```
The deprecated `--stage-configs-path` flag is retained for models that still use the legacy `stage_args` schema:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --stage-configs-path /path/to/legacy_stage_config.yaml
```
## bench
Run benchmark tests for online serving throughput.
+2 -1
View File
@@ -8,6 +8,7 @@ The stage-based CLI is designed for deployments that require launching each pipe
- For **migrated models** that utilize the bundled deployment YAML configurations located in
`vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
automatically loads the bundled deployment configuration.
- For custom legacy `stage_args` YAMLs, pass the file with `--stage-configs-path`.
Example: Initializing Stage 0 (Orchestrator and API Server):
The commands below show a common device mapping where Stage 0 uses GPU 0 and
@@ -31,7 +32,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--omni-master-port 26000
```
When utilizing a custom deployment YAML, append `--deploy-config /path/to/override.yaml` to each command execution.
When utilizing a custom deployment YAML based on the new schema, append `--deploy-config /path/to/override.yaml` to each command execution. Conversely, for legacy models, substitute this parameter with `--stage-configs-path /path/to/stage_configs.yaml`.
In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.
However, under the **stage-based CLI** paradigm, where each process strictly encapsulates a single stage, it is recommended to specify tuning parameters directly via discrete command-line flags for the respective stage, rather than constructing a composite `--stage-overrides` JSON string.
-1
View File
@@ -40,7 +40,6 @@ Committers have write access and merge rights. They typically have deep expertis
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
- [@NickCao](https://github.com/NickCao): CI, Tests, Entrypoint, and TTS Model Support
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
+1 -1
View File
@@ -2,7 +2,7 @@
This section lists the most common options for running vLLM-Omni.
For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration guide](https://docs.vllm.ai/en/v0.26.0/configuration/index.html).
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
Currently, the main options are maintained by stage configs for each model.
+15
View File
@@ -34,8 +34,12 @@ Two additional `StrategySpec` slots — `layer_hook_specs` and `kernel_specs`
| Flag | Description |
|------|-------------|
| `--strategy-config PATH` | Path to a `strategy.yaml`. Only takes effect on registry-based deploy paths (`--deploy-config` or the bundled default deploy YAML). Its derived sizing is overlaid onto the merged stages *before* CLI overrides. |
| `--stage-configs-path PATH` | **Legacy.** Loads the old `stage_args` YAML schema for un-migrated models. Passing `--strategy-config` together with this path is silently inapplicable — vLLM-Omni emits a warning and the strategy is ignored (see `vllm_omni/entrypoints/utils.py`, lines 372–380). |
| `--omni-lb-policy POLICY` | Orchestrator-wide load-balancer policy. Strategy-derived from any `stage_replica` axis; an explicit CLI value wins and must match the derived one or `AsyncOmniEngine` raises `Conflicting load-balancer policy` at construction. |
!!! warning
`--strategy-config` cannot be used with `--stage-configs-path` (the legacy deploy path). If your model has not been migrated to the registry-based pipeline yet, the strategy is silently dropped with a warning. Migrate to a registry-based model or use `--deploy-config` to apply a strategy.
## Sub-schemas
### `axis`
@@ -212,6 +216,17 @@ Strategy keys MUST be `model_stage` *names* (strings), the same labels stage con
!!! warning "Mixing CLI overrides with a strategy on the same axis"
A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--tensor-parallel-size 2`, `--stage-overrides '{"0": {"tensor_parallel_size": 2}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:
```text
--strategy-config (<path>) was provided but model '<model>' resolves via the
legacy stage_configs YAML path, which does not support composable-parallel
strategies; the strategy is ignored. Use a registry-based model to apply it.
```
Migrate the model to the new deploy schema (see `configuration/stage_configs.md`) before applying a strategy.
!!! warning "Conflicting `omni_lb_policy` across stages"
`omni_lb_policy` is a single pipeline-wide knob. If two `stage_replica` axes in different stages derive different policies, `apply_strategy_specs` raises:
+1 -1
View File
@@ -158,7 +158,7 @@ runtime:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/qwen3_omni_pd.yaml
--stage-configs-path /path/to/qwen3_omni_pd.yaml
```
## Operational Notes
+16 -8
View File
@@ -1,9 +1,9 @@
# Pipeline and deploy configurations
# Deploy and legacy stage configurations
In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run.
In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema.
!!! note
Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when `--deploy-config` is omitted. The model registry resolves the right pipeline and deploy YAML by `model_type`.
Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. Custom legacy `stage_args` YAMLs are still accepted via `--stage-configs-path`.
## Deploy configuration schema
@@ -86,9 +86,10 @@ stages:
| Flag | Description |
|------|-------------|
| `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
| `--deploy-config PATH` | Load a new-schema deploy YAML. It is mutually exclusive with `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
| `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
| `--stage-configs-path` | **Deprecated.** Load a legacy `stage_args` YAML. New-schema YAMLs passed to this flag are auto-detected for compatibility, but should use `--deploy-config`. It is mutually exclusive with `--deploy-config` and will be removed in a future release. |
### Stage-Based CLI Paradigm
@@ -118,7 +119,8 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--omni-master-port 26000
```
When instantiating a custom deployment YAML, append the `--deploy-config /path/to/override.yaml` directive to all node invocations.
When instantiating a custom deployment YAML conforming to the updated schema, append the `--deploy-config /path/to/override.yaml` directive
to all node invocations. For legacy architectures (e.g., BAGEL) configured via deprecated `stage_args:` schemas, continue to specify the relevant configuration via `--stage-configs-path /path/to/config.yaml`.
In the context of standard initialization architectures, utilizing the `--stage-overrides` parameter operates as the optimal methodology
for delineating stage-specific tuning from the CLI interface:
@@ -218,7 +220,8 @@ Therefore, as a core part of vLLM-Omni, a model's pipeline and deployment config
To override specific parameters, explicitly inject the customized configuration schema
in both online and offline instantiation flows. Use the `--deploy-config` flag
when loading a deploy configuration.
when loading a deploy configuration, reserving the `--stage-configs-path` parameter
exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
Examples:
@@ -233,6 +236,11 @@ For online serving:
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
```
Legacy online serving:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
!!! important
We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers.
@@ -387,8 +395,8 @@ that support request-level batching with `step_execution` disabled.
Use this together with `max_num_seqs > 1` for bursty serving traffic. `0`
disables admission waiting and preserves the lowest first-request latency.
For diffusion execution and batching tuning, see
[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md).
For diffusion request-level batching tuning, see
[Request-Level Batching](../user_guide/diffusion/request_batching.md).
Default: `0.0`
@@ -1,6 +1,6 @@
**Scope**
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
**Test Case Design**
@@ -44,11 +44,9 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_ttft_ms": [500, 800],
"mean_audio_ttfp_ms": [2000, 3500],
"mean_audio_rtf": [0.25, 0.35]
}
"mean_ttft_ms": [500, 800],
"mean_audio_ttfp_ms": [2000, 3500],
"mean_audio_rtf": [0.25, 0.35]
}
}
]
@@ -144,6 +142,10 @@ pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
# Optional baseline assertion (default off)
pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --assert-baseline \
--test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
```
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
@@ -200,4 +202,4 @@ You can add any benchmark running parameters you need here. For all optional par
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
| request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
| max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). |
| baseline | object | No | see below | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS.
| baseline | object | No | see above | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).
@@ -653,7 +653,7 @@ python end2end.py --model your-org/your-model-name --modality text2img --prompts
Mirror BAGEL’s online serving setup:
- Server launcher: `examples/online_serving/your_model_name/run_server.sh`
- Wrap `vllm serve ... --omni --port ...` (and `--deploy-config ...` if needed)
- Wrap `vllm serve ... --omni --port ...` (and `--stage-configs-path ...` if needed)
- Client: `examples/online_serving/your_model_name/openai_chat_client.py`
- Send requests to `POST /v1/chat/completions`
- Support multimodal inputs (e.g., base64 image) if your model needs it
@@ -755,8 +755,7 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
### Step Execution
See the detailed
[Diffusion Continuous Batching design guide](../../design/feature/diffusion_continuous_batching.md).
See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
Use this only when your pipeline can be split into stable request-scoped and
step-scoped phases. The reference implementation is
@@ -770,8 +769,9 @@ step-scoped phases. The reference implementation is
Do not enable `step_execution=True` until those four methods are implemented
and validated against the request-level path.
The same design guide covers the experimental batched step-wise path used when
`max_num_seqs > 1`.
If you want the pipeline to work with the experimental batched step-wise path
(`max_num_seqs > 1`), also see:
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
If you expose this in example scripts or recipes, keep it opt-in. Surface
runtime features like `step_execution` as optional flags instead of silently
@@ -824,7 +824,7 @@ omni = Omni(model="your-model",
### CPU Offload
See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload.md)
See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md)
vLLM-Omni provides two offloading strategies to reduce GPU memory usage:
+2 -2
View File
@@ -48,8 +48,8 @@ vllm_omni/deploy/
```
New in-tree models should register their `PipelineConfig` in
`vllm_omni/config/pipeline_registry.py`; use `--deploy-config` for custom
deployment YAMLs.
`vllm_omni/config/pipeline_registry.py`; use `--stage-configs-path` only for
custom legacy `stage_args` YAMLs.
## Step-by-Step Implementation
+3 -2
View File
@@ -449,8 +449,9 @@ Key configuration fields:
!!! note
New in-tree models should define frozen topology in `models/<model>/pipeline.py`,
register it in `vllm_omni/config/pipeline_registry.py`, and put deployment
knobs in `vllm_omni/deploy/<model>.yaml`. Use `--deploy-config` to load a
custom deployment YAML.
knobs in `vllm_omni/deploy/<model>.yaml`. Legacy `stage_args` YAMLs are still
accepted for custom configs via `--stage-configs-path`, but should not be used
for new bundled defaults.
### Batch mode
@@ -1,232 +1,165 @@
# Diffusion Continuous Batching
# Continuous Batching for Step-Wise Diffusion
This document describes the unified diffusion execution architecture, including
request execution, request-level batching, step execution, continuous batching,
and output delivery. For user-facing configuration and CLI examples, see
[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md).
!!! warning "Experimental Feature"
This feature is experimental. It currently applies only to native
diffusion pipelines running with `step_execution=True`.
This document describes the batching extension built on top of
[Diffusion Step Execution](diffusion_step_execution.md). The base
step-execution contract is unchanged. The batching work is mainly in the
scheduler and runner layers.
## Why It Helps
Step-wise execution breaks a long denoise loop into scheduler-visible units.
That gives the runtime a place to admit other compatible requests between
steps instead of waiting for an entire request to finish.
This matters most in low-MFU or bursty serving scenarios:
- one request's denoise step may not fully saturate the GPU
- several compatible requests can share the same denoise forward pass
- throughput and device utilization can improve without changing request-local
scheduler state
This is **not** a guaranteed single-request latency win. The main benefit is
usually higher utilization and better throughput when the workload contains
multiple in-flight compatible requests.
## Overview
`DiffusionEngine` selects an execution mode from `step_execution` and uses
`max_num_seqs` as that mode's scheduler capacity:
With continuous batching enabled:
| Configuration | Engine mode | Scheduler | Execution |
|---|---|---|---|
| `step_execution=False`, `max_num_seqs=1` | `REQUEST_BATCH` | `RequestScheduler` | One complete request-level `forward()` |
| `step_execution=False`, `max_num_seqs>1` | `REQUEST_BATCH` | `RequestScheduler` | One fused `forward()` over compatible requests |
| `step_execution=True`, `max_num_seqs=1` | `STEP_BATCH` | `StepScheduler` | One request advanced one denoise step per scheduler tick |
| `step_execution=True`, `max_num_seqs>1` | `STEP_BATCH` | `StepScheduler` | Compatible requests advanced together in step waves |
- the scheduler may keep multiple compatible requests active at the same time
- the runner packs request-local step state into one `InputBatch`
- `denoise_step()` runs on that batch
- `step_scheduler()` and `post_decode()` still run per request
Serial request execution is handled inside `REQUEST_BATCH`; it is not a
separate engine mode. Similarly, continuous batching is the multi-request
configuration of `STEP_BATCH`, not an independent execution mode.
The current implementation is conservative:
The engine performs the following initialization:
- only compatible requests are batched together
- per-request progress and completion remain independent
1. Resolve model-specific pre- and post-processing hooks.
2. Select `REQUEST_BATCH` or `STEP_BATCH`.
3. Construct the configured `DiffusionExecutor`.
4. Construct `RequestScheduler` or `StepScheduler`, unless a scheduler was
explicitly injected.
5. Bind `execute_batch` or `execute_step`.
Here, "continuous batching" means the step-wise path enabled by
`step_execution=True`. Request-mode `DiffusionRequestBatch` is static
request-level batching for one full pipeline `forward()` call; it does not
admit or remove requests between denoise steps.
`DiffusionEngine.make_engine()` owns startup warmup after engine construction.
This keeps construction, backend selection, and warmup failure cleanup in one
lifecycle.
## Enablement
## Unified Output Stream
Use `--step-execution` as the feature gate, then increase `--max-num-seqs`
above `1` if you want batching:
All asynchronous diffusion requests use the same per-request output stream,
regardless of whether the caller exposes streaming results to the user.
```python
async for output in engine.step_streaming(request):
...
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--step-execution \
--max-num-seqs 8
```
When a request is admitted, the engine creates an
`asyncio.Queue[DiffusionOutput]` keyed by its request ID. The busy loop sends
both intermediate chunks and final results through that queue:
`--max-num-seqs 1` keeps the step-wise path without enabling batching.
```text
request
-> scheduler.add_request()
-> scheduler.schedule()
-> execute_batch() or execute_step()
-> scheduler.update_from_output()
-> per-request output queue
-> step_streaming()
```
For a reproducible replay flow using the bundled serving benchmark, see the
Qwen-Image replay commands in
[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md)
and
[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md).
Streaming callers forward every yielded output. Non-streaming callers consume
the same generator to completion and return only its final output. This avoids
maintaining separate future-based and queue-based result paths.
## Scheduler
`DiffusionEngine.step()` and
`async_add_req_and_wait_for_response()` remain deprecated compatibility
wrappers. New integrations should consume `step_streaming()` or
`async_add_req_and_stream_response()`.
The scheduler derives its batch capacity from `max_num_seqs` through
`max_num_running_reqs`.
### Completion and Cancellation
Batch admission is gated by
[`StepBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
the core correctness rule for batching: requests are only co-batched when they
share the same denoise tensor contract.
Output delivery and scheduler lifecycle are related but separate:
There are three important details:
- A terminal scheduler result always finalizes request state.
- A disconnected or cancelled consumer removes only its delivery queue.
- The scheduler may still finish and clean up the request after its queue has
been removed.
- Engine shutdown sends an error output to all remaining streams.
- `num_inference_steps` is not part of the key, so requests with different
total step counts can still share a batch
- requests also do not need to be at the same current denoise progress; active
requests can continue batching even when their current step indices diverge
- admission is still FIFO, so an incompatible request at the head of the
waiting queue blocks later compatible requests
This separation prevents cancelled consumers from leaving active request IDs in
the scheduler.
Today that compatibility rule is still shape-sensitive. `height`, `width`,
`num_frames`, and CFG-related fields remain part of the key, so different
resolutions or incompatible guidance settings do **not** co-batch yet. The
key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
targeting different adapters or scales run in separate batches and the
worker can activate exactly one adapter per step.
## Request-Batch Execution
The scheduler batching unit is one logical `OmniDiffusionRequest`. In the
step-wise path, runtime tensor batching is represented as `StepInputBatch`. For
request-mode prompt semantics, see
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
Request-batch execution runs the complete diffusion pipeline for each scheduler
wave. Each `OmniDiffusionRequest` remains an independent logical request with
its own prompt, sampling parameters, seed, request ID, output, error, and abort
state.
## Runner
With `max_num_seqs=1`, `execute_batch()` handles the serial request path. With a
larger value, a pipeline can opt into fused execution by declaring:
The runner keeps persistent per-request execution state in
[`StepRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
while the scheduler owns a separate lightweight
[`SchedulerRequestState`](gh-file:vllm_omni/diffusion/sched/interface.py)
for queueing and lifecycle tracking.
```python
supports_request_batch = True
```
For each step, the runner builds an
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
active request states:
Its `forward()` method must accept `DiffusionRequestBatch` and return one
`DiffusionOutput` per request. The runner validates the result count and maps
outputs back to their original request IDs with `BatchRunnerOutput`.
- prompt embeddings and masks are normalized and padded
- dynamic tensors such as `latents` and `timesteps` are gathered each step
- buffers are reused when batch composition stays the same
Current in-tree request-batch implementations include the Flux, LTX-2, SD3, and
Qwen-Image pipelines. Treat the pipeline capability flag in the source as
authoritative because support continues to expand.
The step-wise batched path is:
### Request-Batch Data Flow
- `DiffusionSchedulerOutput` carries scheduled request IDs and payloads.
- `DiffusionRequestBatch` provides the pipeline-facing static batch.
- `BatchRunnerOutput` carries per-request results back to the engine.
Request-local setup, such as seed handling and output/error mapping, remains
per request. Cache refresh, homogeneous LoRA activation, and
`pipeline.forward(req_batch)` can be shared by the batch.
### Compatibility and Admission
`RequestBatchSamplingParamsKey` contains fields that affect the tensor contract,
including shape, guidance, output count, and LoRA identity. Only compatible
requests can share a fused forward call.
Admission is FIFO and conservative. An incompatible request at the front of the
waiting queue can prevent later compatible requests from joining the current
wave.
`request_batch_max_wait_ms` optionally delays the first schedule operation of a
new wave so bursty arrivals can accumulate. It applies only when:
- the pipeline supports request batching;
- step execution is disabled;
- no requests are currently running; and
- the configured wait is greater than zero.
The wait ends when the queue reaches capacity, stabilizes briefly, reaches its
deadline, or the engine stops.
## Step-Batch Execution
Step execution exposes denoising progress to the scheduler. A supporting
pipeline implements four stateful operations:
| Operation | Responsibility |
|---|---|
| `prepare_encode(state)` | Validate input, encode prompts, initialize latents and timesteps, and create request-local scheduler state |
| `denoise_step(input_batch, *, states=...)` | Run one denoise forward for the scheduler-provided request states |
| `step_scheduler(state, noise_pred)` | Update latents and advance request progress |
| `post_decode(state)` | Decode and postprocess a completed request or output boundary |
Persistent request state lives in `StepRequestState`. Pipeline-specific fields
that do not belong in the shared contract should be stored in `state.extra`.
Queueing and lifecycle metadata remain in the scheduler's request state.
Current native pipelines that explicitly enable step execution include
Qwen-Image, HunyuanImage3, and Helios. Step execution alone does not imply
continuous-batching support: Qwen-Image accepts batched step states.
HunyuanImage3 accepts batched step states only when its resolved self-attention
backend is `TORCH_SDPA`; otherwise it rejects groups larger than one request.
Configure `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or
`diffusion_attention_config.default.backend=TORCH_SDPA` when
`max_num_seqs>1`. See the
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
for its validated configuration. Helios supports only a single active step
request and must use `max_num_seqs=1`.
### Continuous Batching
When `max_num_seqs>1`, `StepScheduler` can keep multiple compatible requests
active. The runner gathers their state into `InputBatch`, performs one batched
denoise forward, then applies scheduler updates per request:
1. Receive transferred KV payloads, then run `prepare_encode()` for newly
admitted requests.
1. Run `prepare_encode()` for newly admitted requests.
2. Build or refresh `InputBatch`.
3. Run one batched `denoise_step(input_batch, states=states)`.
4. Slice noise predictions back to each request.
5. Run each request's `step_scheduler()`.
6. Run `post_decode()` at a chunk boundary or request completion.
7. Scatter updated latents back into persistent request state.
3. Run one batched `denoise_step(input_batch)`.
4. Slice the batched `noise_pred` back per request.
5. Run per-request `step_scheduler()`.
6. Run `post_decode()` only for requests that finished denoising.
7. Scatter updated latents back into persistent request state with
[`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py).
`StepBatchSamplingParamsKey` protects the batched tensor contract. Requests can
have different total step counts and current step indices, but shape-sensitive,
CFG-sensitive, and LoRA fields must be compatible. FIFO head-of-line blocking
also applies here.
This keeps the shared work limited to the denoise forward pass while preserving
request-local scheduler state and outputs.
Chunk-capable pipelines may emit intermediate results. Final-only step
pipelines emit when the request finishes even if `streaming_output=True`.
## Engine
## Model Author Guidelines
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides
the background loop and async add-request path needed for multiple requests to
accumulate in the scheduler.
When converting a request-level pipeline:
When `step_execution=True`, the engine routes work through the step-wise
executor path. The continuous batching behavior is defined by scheduler-side
compatibility gating and runner-side `StepInputBatch` packing.
- Reuse the same helpers as `forward()` to avoid behavior drift.
- Copy request-scoped scheduler state into `state.scheduler`.
- Advance `state.step_index` only in `step_scheduler()`.
- Keep model forward work in `denoise_step()`.
- Keep latent mutation in `step_scheduler()`.
- Keep final decoding equivalent to the tail of `forward()`.
- Store masks, condition latents, and other request-local tensors in the
request state rather than on the shared pipeline.
## Current Limitations
Before setting `supports_step_execution = True`, validate output parity for the
same seed and sampling parameters, scheduler isolation across concurrent
requests, abort cleanup, progress reporting, and all supported CFG paths.
- Experimental feature; use `max_num_seqs=1` for the conservative single-request
step path.
- Only native pipelines that already support `step_execution=True`.
- Only homogeneous batches keyed by `StepBatchSamplingParamsKey` are supported.
- `cache_backend`, KV transfer, and other request-mode extras are not wired
into the batched step-wise path yet.
- Future work can relax the current same-shape restriction with richer
heterogeneous batching policies such as bucketing or padded execution for
different resolutions.
Before setting `supports_request_batch = True`, validate the batch input
contract, one-output-per-request result shape, request identity and error
mapping, seeded concurrency, LoRA compatibility, and tensor IPC.
## Related Files
## Limitations
- Request batching requires an explicit pipeline capability when
`max_num_seqs>1`.
- Step execution requires the four-method stateful pipeline contract.
- Both batching paths currently require homogeneous compatibility keys.
- FIFO scheduling can reduce batching opportunities.
- `request_batch_max_wait_ms` trades first-request latency for burst
coalescing.
- All diffusion cache backends are currently unsupported in step mode.
- KV transfer is supported for newly admitted step requests; some other
request-mode extras remain unsupported.
- Step continuous batching remains experimental; use `max_num_seqs=1` for the
conservative step path.
## Related Implementation
- Engine: [`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
- Model contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
- Scheduler interface: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
- Request scheduler: [`vllm_omni/diffusion/sched/request_scheduler.py`](gh-file:vllm_omni/diffusion/sched/request_scheduler.py)
- Step scheduler: [`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
- Runner: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
- Request batch: [`vllm_omni/diffusion/worker/request_batch.py`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
- Step input batch: [`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
- Scheduler base:
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
- Scheduler interface:
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
- Step scheduler:
[`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
- Runner:
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
- Input batch:
[`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
- Tests:
[`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py)
@@ -0,0 +1,175 @@
# Request-Level Batching for Diffusion
This document describes the request-mode batching path for diffusion pipelines.
For end-user enablement and tuning, see
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
This is separate from
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
Request-level batching runs one full pipeline `forward()` over a static batch of
compatible requests. Step-wise continuous batching admits work between denoise
steps when `step_execution=True`.
## Why It Helps
The request-level design avoids coupling several logical requests to one request
object. This keeps request identity, abort/error handling, and per-request
metadata unambiguous while still allowing one fused pipeline forward pass for
bursty or concurrent traffic.
## Overview
With request-level batching enabled:
- each `OmniDiffusionRequest` contains one `prompt` and one `request_id`
- the scheduler groups compatible waiting requests into one scheduler wave
- `DiffusionRequestBatch` wraps the scheduled requests for pipeline `forward()`
- batch-capable pipelines return `list[DiffusionOutput]`, one output per
request
- `BatchRunnerOutput` maps each result back to its original `request_id`
Pipelines opt in with `supports_request_batch = True` and a `forward()` method
that accepts `DiffusionRequestBatch` and returns `list[DiffusionOutput]`.
Pipelines that do not opt in keep the existing per-request execution path.
## Enablement
Request-level batching is the request-mode path, so `step_execution` must remain
disabled. Increase `max_num_seqs` above `1` to let the scheduler keep multiple
compatible requests active:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--max-num-seqs 4
```
For bursty online ingress, `request_batch_max_wait_ms` can add a bounded
admission wait before the first `schedule()` of a scheduler wave:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--max-num-seqs 4 \
--request-batch-max-wait-ms 20
```
`request_batch_max_wait_ms=0` disables this wait and is the default.
## Request Contract
`OmniDiffusionRequest` represents one logical request. It owns one prompt,
sampling parameters, request id, and request-local metadata. Runtime batches are
formed by the scheduler and represented separately from the request payload.
Runtime batching is represented by:
- [`DiffusionSchedulerOutput`](gh-file:vllm_omni/diffusion/sched/interface.py)
for scheduled request ids and request payloads
- [`DiffusionRequestBatch`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
for the pipeline-facing request batch
- [`BatchRunnerOutput`](gh-file:vllm_omni/diffusion/worker/utils.py) for
per-request results
`DiffusionRequestBatch` intentionally exposes compatibility properties such as
`prompts`, `sampling_params`, `request_id`, and `kv_sender_info` so migrated
pipelines can stay close to upstream code while using a batch-aware contract.
## Scheduler
The scheduler derives its capacity from `max_num_seqs` through
`max_num_running_reqs`. It exposes waiting/running queue counters so the engine
can decide whether admission wait is useful before scheduling a new wave.
Batch compatibility is controlled by
[`RequestBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key
contains shape-sensitive and guidance-sensitive fields, including output count
and LoRA identity. Requests with incompatible shapes, CFG settings, output
counts, LoRA adapters, or LoRA scales are kept in separate batches.
Admission is conservative:
- the scheduler only batches compatible requests
- FIFO ordering is preserved
- an incompatible request at the head of the waiting queue blocks later
compatible requests
## Engine
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) resolves
request-batch capability during initialization from the configured pipeline
class, including custom pipeline classes.
The capability check uses the pipeline class attribute
`supports_request_batch = True`. Pipelines that set this attribute must implement
a request-batch-compatible `forward()` contract and return one
`DiffusionOutput` per request; the runner validates that return shape at runtime.
When the selected pipeline is batch-capable and `step_execution=False`, request
mode routes scheduler waves through the batch executor path. Otherwise it keeps
the per-request executor path.
The optional admission wait runs only when:
- request batching is supported
- `step_execution=False`
- `request_batch_max_wait_ms > 0`
- no requests are currently running
The wait exits early when the waiting queue reaches capacity, when the queue is
stable for a short window, when the deadline expires, or when the engine stops.
## Executor And Runner
The executor exposes two request-mode entries:
- `execute_request`: one worker call per scheduled request
- `execute_batch`: one worker call for the whole `DiffusionSchedulerOutput`
On the batch path, the worker builds a `DiffusionRequestBatch` and runs the
pipeline once. Request-local setup remains per request:
- KV transfer metadata
- random generator and seed handling
- request output/error/abort mapping
Shared batch setup happens once per batch when possible:
- cache refresh
- LoRA activation for the homogeneous adapter key
- pipeline `forward(req_batch)`
Large tensor IPC still uses the shared-memory packing path. The packer traverses
both normal `RunnerOutput.result` wrappers and nested batch results so batched
outputs do not fall back to pickle IPC for tensor payloads.
## Current Limitations
- Only pipelines that declare the request-batch contract use fused batch
execution.
- Batches are homogeneous under `RequestBatchSamplingParamsKey`; heterogeneous resolution or
incompatible guidance settings do not co-batch yet.
- FIFO scheduling can reduce batching opportunities when an incompatible
request is at the front of the queue.
- `request_batch_max_wait_ms` improves burst coalescing but can add latency to
the first request in a scheduler wave. Keep it small for latency-sensitive
serving.
- Step-wise continuous batching is documented separately and only applies when
`step_execution=True`.
## Related Files
- Request object and request batch:
[`vllm_omni/diffusion/request.py`](gh-file:vllm_omni/diffusion/request.py)
- Scheduler interface:
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
- Scheduler base:
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
- Engine:
[`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
- Worker runner:
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
- Executor interface:
[`vllm_omni/diffusion/executor/abstract.py`](gh-file:vllm_omni/diffusion/executor/abstract.py)
- Tests:
[`tests/diffusion/test_diffusion_engine.py`](gh-file:tests/diffusion/test_diffusion_engine.py)
@@ -0,0 +1,131 @@
# Adding Step Execution Support for Diffusion Pipelines
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
and contributors implementing `step_execution=True` support for a diffusion
pipeline.
For end-user enablement, supported models, and current limitations, see
[Step Execution](../../user_guide/diffusion/step_execution.md).
This document describes the base step-execution contract only. For the
experimental batching policy layered on top of the step-wise path, see
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
## Current Support Scope
`step_execution` is **not** a generic diffusion toggle. It only works for
pipelines that implement the segmented stateful contract in
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
This page is intentionally author-facing. Treat runtime enablement
(`step_execution=True` in Python or `--step-execution` in serving) as an
opt-in user knob layered on top of the implementation contract below.
Current in-tree support:
| Pipeline | Example models | Step execution |
|----------|----------------|----------------|
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
Current engine/runtime limitations:
- Continuous batching with `max_num_seqs > 1` is experimental and documented in
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
Keep `max_num_seqs=1` for the conservative single-request step path.
- `cache_backend` is not supported in step mode.
- Request-mode extras such as KV transfer are not wired into step mode yet.
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
## Execution Contract
Step mode is driven by four pipeline methods plus the shared mutable request
state object:
- `prepare_encode(state)`: one-time request preparation.
- `denoise_step(state)`: compute the noise prediction for the current step.
- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
- `post_decode(state)`: decode the final output after denoising is complete.
The state lives in
[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
as `StepRequestState`. Store request-scoped tensors there, or use
`state.extra` for model-specific fields that do not justify extending the core
dataclass.
The worker-side step loop lives in
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
1. `prepare_encode()` runs once for a new request.
2. `denoise_step()` runs every scheduler tick.
3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
## Recommended Split
When converting an existing request-level `forward()` pipeline, keep the split
strict and mechanical:
| Request-level phase | Stepwise method | What belongs there |
|---------------------|-----------------|--------------------|
| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
| VAE decode / postprocess | `post_decode()` | Final decode only |
Keep the stepwise path reusing the same helpers as the request-level path
whenever possible. Reimplementing the denoise loop from scratch is the easiest
way to introduce behavioral drift.
## Qwen-Image Reference
[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
is the reference implementation and is split correctly for the current
contract:
- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
with `forward()`.
- `prepare_encode()` deep-copies `self.scheduler` **after**
`prepare_timesteps()` so request-specific scheduler state is isolated.
- `denoise_step()` reuses `_build_denoise_kwargs()` plus
`predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
non-CFG behavior stay identical to the request-level path.
- `step_scheduler()` only calls
`scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
and increments `state.step_index`.
- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
the normal `forward()` path.
That decomposition is the target pattern for future models.
## Rules For New Pipelines
- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
`state.scheduler` during `prepare_encode()`.
- Do not mutate `state.step_index` inside `denoise_step()`. Only
`step_scheduler()` should advance the step.
- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
- If the request-level pipeline has condition latents, masks, or edit-specific
tensors, store them in `state` or `state.extra`, not in global pipeline
attributes.
- Preserve CFG behavior by sharing the same helper path used by `forward()`.
- Keep `post_decode()` equivalent to the tail of `forward()`.
## Validation Checklist
Before marking a pipeline as `supports_step_execution = True`, verify:
- Stepwise output matches request-level output for the same seed and sampling params.
- Per-request scheduler state is isolated across concurrent requests.
- Abort during denoise does not leak cached state.
- `step_index` reported by `RunnerOutput` matches the scheduler progress.
- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
## Related Files
- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
@@ -1,156 +0,0 @@
# Distributed Layerwise Offload
This document describes distributed layerwise offload (DLO) for diffusion
models. DLO keeps only a small number of DiT blocks on the accelerator and
streams the remaining blocks from host memory. The distributed backend can
either shard those host-side weights across an existing parallel group or keep
the standard loader's rank-local weights and avoid an additional collective.
For user-facing commands, see the
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
## Status
DLO is implemented for multi-device diffusion execution. The default
AllGather path is the primary path for DP and SP deployments. The
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
for standard-loader sharding, workstation bring-up, and systems where an
additional DLO collective is undesirable, but it does not reduce host weight
storage across ranks.
The compatibility matrix below describes the current implementation. The
unit-level guards are covered, but not every parallelism combination has a
full model-and-hardware end-to-end test.
## Design
### DLO consumes the existing parallel topology
DLO does not create a new DP, TP, or SP topology. It reads the configured
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
standard distributed groups have been initialized.
The DLO weight-sharding group is selected as follows:
1. Use the existing DP group when `data_parallel_size > 1`.
2. When DP is one and SP is greater than one, use the SP group.
3. Otherwise, run rank-locally without a DLO process group.
TP is deliberately not used as DLO's AllGather group. HSDP has its own
parameter-sharding lifecycle and is not allowed to be sharded a second time by
DLO's AllGather path.
### AllGather path
With the default `dlo_use_allgather=True`, each rank stores approximately
`1 / group_size` of each streamable block in pinned host memory. The next
block's shard is copied to a device buffer and reconstructed with
`all_gather_into_tensor` on a communication stream while the current block is
executing.
```text
Compute: [Block N] [Block N+1] [Block N+2]
H2D: [shard N+1] [shard N+2]
AllGather: [full N+1] [full N+2]
Buffers: [current slot] [prefetch slot] [current slot]
```
![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif)
The backend uses two shared device buffers, so accelerator weight residency is
bounded by the largest streamed blocks rather than the complete model.
When DP is greater than one, the engine can process one request per DP rank in
the same denoising wave. Because AllGather is a collective, all participating
requests must take the same execution path at every denoising step.
### Rank-local path without DLO AllGather
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
one. The regular model loader remains responsible for preparing each rank's
weights, including TP-local tensors or HSDP-managed parameters. DLO then
streams those rank-local tensors block by block using H2D copies only.
This mode means:
- DP still provides independent replicas, but DLO does not shard weights
across DP ranks.
- SP still performs its normal activation/attention collectives, but DLO does
not shard weights across SP ranks.
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
DLO's additional weight AllGather is disabled.
- Pure DP deployments keep a full host-side model copy per rank, subject to
shared OS page-cache behavior.
- The scheduler does not require a synchronized DP request wave for DLO.
This path is intentionally not implemented by reusing the AllGather mmap
loader. It relies on the standard loader so model-specific TP/HSDP transforms
remain intact.
## Parallelism compatibility
| Parallelism | DLO + AllGather | DLO without AllGather |
|---|---|---|
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
### Combined dimensions
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
than one; SP continues to use its own sequence-parallel group. If DP is one,
the SP group becomes DLO's sharding group in AllGather mode.
- **DP + TP/SP without AllGather:** standard model loading defines the
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
collective.
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
materialization and synchronization.
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
configuration.
## Request and loading constraints
AllGather DP multi-concurrency requires:
- explicit `num_inference_steps`;
- the same `num_inference_steps` for all requests in a wave; and
- identical request arguments that affect the collective execution path.
The no-AllGather path does not impose these DLO-specific synchronized-wave
requirements.
The mmap loader is used only by the supported DLO+AllGather path when the
model has a compatible checkpoint layout. Online quantization is incompatible
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
quantization. Models that do not meet the mmap requirements use the regular
loader path.
## Validation coverage
Current source-level validation includes:
- HSDP + DLO + AllGather rejection;
- HSDP + DLO without AllGather acceptance at configuration level;
- TP rejection in the DLO+AllGather mmap path;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
tests.
The highest-value missing coverage is end-to-end numerical comparison against
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
on the target CUDA/NCCL or CANN/HCCL hardware.
## Recommendations
- Use **DP + DLO AllGather** for the supported throughput and host-memory
scaling path.
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
not the goal.
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
the expectation of higher host-memory use and lower validation confidence.
- Prefer **HSDP alone** for production HSDP deployments until the combined
HSDP + DLO no-AllGather path has broader end-to-end coverage.
@@ -1,386 +0,0 @@
# Async Omni Output Materialization
## Table of Contents
1. [Overview](#overview)
2. [Performance](#performance)
3. [Architecture](#architecture)
4. [Enabling the Feature](#enabling-the-feature)
5. [Compatibility and Fallbacks](#compatibility-and-fallbacks)
6. [Related Files](#related-files)
## Overview
Async Omni output materialization moves CPU-side construction of
`OmniModelRunnerOutput` out of the autoregressive (AR) decode critical path.
It complements [Async Chunk](async_chunk.md): async chunk pipelines partial
outputs between stages, while async output materialization prevents the
construction of those partial outputs from blocking the next decode step.
An AR stage must return sampled token IDs quickly because the scheduler needs
them for the next decode step. However, an Omni stage can also produce hidden
states, multimodal tensors, streaming inter-stage payloads, and connector
metadata. Building these payloads can require device-to-host (D2H) copies,
tensor slicing, flattening, and Python object construction.
Without this feature, all of that work runs inline in `sample_tokens()`:
```text
sample tokens
-> update decode state
-> copy hidden and multimodal outputs to CPU
-> build per-request payloads
-> build the streaming wire payload
-> collect connector signals
-> return sampled tokens
-> launch the next decode step
```
With async output materialization, only work required by the next decode step
stays on the critical path:
```text
sample tokens
-> update decode state
-> snapshot output state and start asynchronous D2H copies
-> register sampled tokens
-> launch the next decode step
background output path
-> wait for the payload snapshot
-> build per-request payloads
-> build the streaming wire payload
-> collect connector signals
-> construct OmniModelRunnerOutput
```
This allows payload construction for step `N` to overlap with GPU decode work
for step `N + 1`. It does not change model computation or generated output.
!!! note "Platform and validation scope"
The performance and correctness validation in this document covers CUDA
and ROCm. This is a validation scope, not a runtime platform guard: XPU
inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so those
backends may enter the asynchronous path when its other guards pass. XPU
and MUSA have not been validated for this optimization. Ascend NPU uses a
separate `NPUARModelRunner` path that still constructs
`OmniModelRunnerOutput` and drains connector output synchronously before
wrapping the result as an asynchronous runner output.
<p align="center">
<img
alt="Decode step gap before and after async output materialization"
src="../../../source/architecture/qwen3-omni-async-output-step-gap.svg"
width="100%"
>
</p>
<p align="center">
<em>
Figure 1: Moving payload construction off the decode path reduces the
observed gap between consecutive Talker steps from about 2.8 ms to 41 µs.
Adapted from the
<a href="https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization">
Qwen3-Omni optimization blog
</a>.
</em>
</p>
For Qwen3-Omni, the feature is enabled for both AR stages:
- **Thinker**: Defers hidden-state and multimodal payload materialization.
- **Talker**: Runs its decode-state postprocess eagerly, then defers codec
payload materialization. Because Code2Wav consumes codec codes rather than
Talker hidden states, the Talker also skips unnecessary hidden-state D2H.
- **Code2Wav**: Does not use this path because it is a generation stage rather
than an AR stage.
Qwen3-TTS uses the same mechanism in its AR Talker stage. The Talker updates
the decode state it needs for the next step before returning, then constructs
the codec payload for Code2Wav in the background.
## Performance
In the controlled Qwen3-Omni optimization sweep, enabling async output
materialization on top of CUDA Graph and async chunk produced the following
results at concurrency 64:
| Configuration | Request throughput | Mean audio TTFP | Mean audio RTF |
|---|---:|---:|---:|
| CUDA Graph + async chunk | 9.3 req/s | 655 ms | 0.63 |
| + async output materialization | 11.3 req/s | 631 ms | 0.47 |
| Change | **+22%** | **-4%** | **-25%** |
The sweep used `Qwen/Qwen3-Omni-30B-A3B-Instruct`, 640 Seed-TTS English
prompts, and one replica each for Thinker, Talker, and Code2Wav. Async output
materialization recovers throughput by removing CPU payload construction from
the decode critical path without regressing time to first audio packet (TTFP).
See the
[Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization)
for the complete benchmark methodology and results.
The implementation was also validated against the synchronous safe-copy path.
The text and audio output hashes matched, confirming that overlap does not
change the generated result. See
[PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476) for the
correctness and profiling details.
## Architecture
### Output Lifecycle
`OmniAsyncGPUModelRunnerOutput` extends vLLM's
`AsyncGPUModelRunnerOutput`. It preserves the existing asynchronous
sampled-token feedback while adding a background builder for the complete Omni
output.
The output lifecycle is:
1. `GPUARModelRunner.sample_tokens()` samples tokens and performs the
bookkeeping required by the next decode step.
2. The runner snapshots step-local metadata such as request IDs, token spans,
scheduler output, and query start locations.
3. CUDA payload tensors are cloned before reusable model or CUDA Graph buffers
can be overwritten.
4. The cloned payload is copied to pinned CPU memory on a dedicated CUDA
stream. A CUDA event records when the copy is ready.
5. `OmniAsyncGPUModelRunnerOutput` starts a background thread that waits for
the payload event and builds `OmniModelRunnerOutput`.
6. The AR runner immediately registers the sampled-token CPU copy with the
input batch and returns, allowing async scheduling to advance.
7. When the engine calls `get_output()`, it joins the background builder,
propagates any builder exception, and finalizes sampled tokens and logprobs
through the upstream output implementation.
```text
GPU / decode path Background output path
forward + sample
|
+-- clone output tensors
+-- enqueue D2H copy --------------> wait for D2H event
+-- register sampled tokens |
+-- return async output +-- slice per request
| +-- build multimodal payloads
next decode step +-- partition streaming payloads
+-- build the wire payload
+-- drain connector signals
+-- build OmniModelRunnerOutput
```
### Safe State Snapshots
The next scheduler step can mutate runner state while the background builder is
still active. Most step-local builder inputs are therefore detached snapshots.
The snapshotted state includes:
- scheduler token counts and speculative-token metadata
- request IDs and request-to-batch-index mappings
- sampled token IDs, logprobs, and prompt logprobs
- query start locations and scheduled-token spans
- hidden states and multimodal output tensors
- KV-connector and encoder-cache outputs captured for the step
CUDA tensor cloning is required because CUDA Graph and model output buffers can
be reused by the next forward pass. The dedicated copy stream and pinned host
buffers make the D2H transfer asynchronous; the snapshot retains its cloned
CUDA sources until the transfer event completes.
### Live Runner State and Ownership
The complete builder input is not a detached snapshot. Connector output state
remains runner-owned: `get_omni_connector_output()` drains live, per-cycle
connector signals after the streaming payload is built. Correctness therefore
depends on that builder being the only consumer that drains connector signals
for its output cycle. Receive-side connector state written by the background
receiver is coordinated by the connector mixin's `_lock`.
`OmniAsyncGPUModelRunnerOutput.get_output()` joins the builder and is the
completion and exception boundary before the materialized output is consumed.
The shared output builder also contains a full-payload accumulation branch
that reads `self.requests`. That branch is not reachable when async Omni output
materialization is enabled: this feature requires `async_chunk`, while
`should_accumulate_full_payload_output()` returns `False` whenever
`async_chunk` is enabled. The feature path partitions each step's payload into
streaming inter-stage and client handoffs instead.
### Output Builder
The background builder performs the work that previously ran inline:
- resolves which requests need downstream payloads
- converts hidden and multimodal tensors into per-request CPU payloads
- applies model-specific payload processing
- partitions per-step payloads into inter-stage and client streams
- creates the tensor-only `multimodal_outputs` wire payload
- drains live connector readiness signals after building the streaming payload
- constructs the final `OmniModelRunnerOutput`
Input-side connector operations remain synchronous with model execution. In
particular, receiving inputs and flushing previously completed connector
outputs are not deferred because they affect scheduler-visible request state.
### Qwen3-Omni Stage Behavior
| Stage | Async output behavior | Reason |
|---|---|---|
| Thinker | Snapshots hidden states and multimodal outputs; builds the downstream payload in the background | The Talker needs the payload, but the next Thinker decode step only needs sampled-token feedback |
| Talker | Runs lightweight postprocess eagerly; snapshots codec outputs; omits hidden states from the downstream payload | `hidden_states.last` is needed by the next Talker step, while Code2Wav only needs codec codes |
| Code2Wav | Uses the normal generation-stage output path | Code2Wav is not executed by `GPUARModelRunner` |
The Qwen3-TTS Talker follows the same pattern as the Qwen3-Omni Talker:
postprocess runs eagerly because the next decode step needs its updated state,
while codec payload construction runs asynchronously. Its Code2Wav stage also
uses the normal generation-stage path.
## Enabling the Feature
There is no separate `async_omni_output` command-line or YAML option. The
feature is selected automatically for model stages that opt in when the runtime
conditions are safe. On the validated CUDA and ROCm configurations, using the
bundled deployment profile is enough to turn it on for the supported
Qwen3-Omni and Qwen3-TTS pipelines. The selection logic has no CUDA/ROCm-only
guard; see [Compatibility and Fallbacks](#compatibility-and-fallbacks) for the
other platform paths. The same profile does not enable background Omni
materialization on Ascend NPU.
### Example 1: Qwen3-Omni
Start the full Qwen3-Omni pipeline:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
--omni \
--port 8091
```
The model registry automatically loads
`vllm_omni/deploy/qwen3_omni_moe.yaml`. Its relevant settings are:
```yaml
async_chunk: true
stages:
- stage_id: 0 # Thinker
enable_prefix_caching: false
# async_scheduling defaults to true for AR stages
- stage_id: 1 # Talker
enable_prefix_caching: false
# async_scheduling defaults to true for AR stages
- stage_id: 2 # Code2Wav
enable_prefix_caching: false
async_scheduling: false
```
With this profile, async output materialization activates automatically for:
- **Stage 0, Thinker**: Defers hidden-state and multimodal payload
construction.
- **Stage 1, Talker**: Defers codec payload construction after eagerly
updating its decode state.
Stage 2 does not use this feature because Code2Wav is not an AR stage.
### Example 2: Qwen3-TTS
Start Qwen3-TTS with any checkpoint that uses the Qwen3-TTS Talker pipeline.
For example:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--omni \
--port 8091
```
The model registry automatically loads `vllm_omni/deploy/qwen3_tts.yaml`. Its
relevant settings are:
```yaml
async_chunk: true
stages:
- stage_id: 0 # Talker
async_scheduling: true
enable_prefix_caching: false
- stage_id: 1 # Code2Wav
enable_prefix_caching: false
```
Async output materialization activates only for Stage 0, the AR Talker. The
Talker keeps the state required by its next decode step on the GPU, skips an
unnecessary hidden-state D2H copy, and builds the codec payload for Code2Wav in
the background. Stage 1 uses the generation-stage output path.
The same behavior applies to the Base and VoiceDesign Qwen3-TTS checkpoints
because they use the same Talker implementation.
!!! warning
Do not pass `--no-async-chunk` or enable prefix caching when you want this
optimization. Either change causes the runner to fall back to synchronous
output construction. Enabling async chunk alone cannot activate the feature
for a model that has not opted into async Omni output.
!!! note "Prefix cache compatibility"
Async Omni output materialization and Omni prefix caching cannot currently
run together. `_should_use_async_omni_output()` returns `False` whenever
`self.omni_prefix_cache` is present, so prefix caching selects synchronous
output materialization. Supporting both features requires snapshotting or
otherwise synchronizing prefix-cache merge and update state before the
background output path can consume it safely.
## Compatibility and Fallbacks
`GPUARModelRunner` uses the async output path only when all of the following
runtime conditions hold:
| Requirement | Reason |
|---|---|
| AR async scheduling is enabled | The optimization relies on the scheduler advancing while the prior output is materialized |
| `async_chunk` is enabled | The feature targets incremental downstream Omni payloads |
| The model stage opts in with `use_async_omni_output` | Models must declare that their output lifecycle is safe to defer |
| Omni prefix cache is disabled | Prefix-cache merge and update ordering currently requires synchronous materialization |
| Speculative decoding is disabled | Speculative output state is not included in this deferred path |
| Routed-expert output is disabled | Routed-expert extraction currently requires the synchronous path |
| Postprocess is absent or explicitly runs eagerly | State needed by the next decode step must be updated before the runner returns |
These checks are evaluated per stage on runners based on `GPUARModelRunner`.
An unsupported combination does not prevent serving; it only falls back to
synchronous output construction for that stage. CUDA and ROCm are the validated
platforms. XPU inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so
they may select the async path when these checks pass, but this optimization is
not yet validated on either backend.
Ascend NPU does not use this selection path. `NPUARModelRunner.sample_tokens()`
fully constructs `OmniModelRunnerOutput`, calls
`get_omni_connector_output()`, and only then creates the asynchronous wrapper,
so Omni payload materialization remains synchronous.
!!! note
`use_async_omni_output`,
`eager_omni_postprocess_before_async_output`, and
`omni_pooler_payload_include_hidden` are model implementation contracts,
not user-facing configuration fields.
## Related Files
- `vllm_omni/worker/gpu_ar_model_runner.py`: Async output object, safe tensor
snapshots, compatibility guards, and deferred Omni output builder.
- `vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`: Separate Ascend NPU
runner, where Omni output materialization remains synchronous.
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Thinker and
Talker opt-in behavior.
- `vllm_omni/deploy/qwen3_omni_moe.yaml`: Default Qwen3-Omni deployment
settings.
- `vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py`: Qwen3-TTS
Talker opt-in and eager postprocess behavior.
- `vllm_omni/deploy/qwen3_tts.yaml`: Default Qwen3-TTS deployment settings.
- `tests/worker/test_gpu_ar_model_runner.py`: Snapshot, guard, connector
ordering, and background error-propagation tests.
- [Async Chunk](async_chunk.md): Inter-stage chunking and scheduling design.
- [Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization):
Optimization context and controlled performance results.
- [PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476): Feature
implementation, validation, and review history.
-249
View File
@@ -1,249 +0,0 @@
# Quantization
This document describes the quantization architecture in vLLM-Omni and the
extension points for contributors. For user-facing configuration, supported
hardware, and method-specific instructions, see the
[quantization user guide](../../user_guide/quantization/overview.md).
## Goals and scope
The design has four goals:
1. Provide one `quantization_config` entry point for diffusion-only,
multi-stage omni/TTS, and multi-stage diffusion models.
2. Reuse vLLM's quantization registry and `QuantizationConfig` contract where
a vLLM backend already provides the required behavior.
3. Allow different components of one pipeline to use different quantization
configurations.
4. Preserve the distinction between online quantization of a BF16/FP16
checkpoint and loading a pre-quantized checkpoint with serialized scales.
This document covers configuration resolution, component routing, model
integration, backend extension, and validation. Kernel implementations and
user-facing support matrices are outside its scope.
## Architecture
The quantization path is split across the following layers:
| Layer | Responsibility | Source |
| --- | --- | --- |
| Public API | Exposes `build_quant_config`, component routing, and backend registration | `vllm_omni/quantization/__init__.py` |
| Factory | Normalizes method names, selects Omni overrides or vLLM registry backends, detects ModelOpt configs, and reconciles checkpoint metadata | `vllm_omni/quantization/factory.py` |
| Component router | Resolves a quantization config from a layer prefix using longest-prefix matching | `vllm_omni/quantization/component_config.py` |
| Method configs | Implement `QuantizationConfig` and select the platform-specific linear method | `vllm_omni/quantization/*_config.py` |
| Diffusion configuration | Stores the canonical runtime config and propagates checkpoint configuration | `vllm_omni/diffusion/data.py` |
| Worker and model integration | Passes the resolved config into vLLM's model configuration and model layers | `vllm_omni/diffusion/worker/diffusion_worker.py` and `vllm_omni/diffusion/models/` |
The runtime contract remains vLLM's `QuantizationConfig`:
```text
user/CLI or model config
|
v
OmniDiffusionConfig.quantization_config
|
+--> build_quant_config()
| |
| +--> Omni override
| +--> vLLM registry backend
| +--> ComponentQuantizationConfig
|
+--> checkpoint metadata reconciliation
|
+--> diffusion worker / vLLM config
|
+--> model layers call get_quant_method(layer, prefix)
```
The model layer receives either a method implementation or `None`. Returning
`None` leaves that layer unquantized.
## Configuration resolution
`build_quant_config()` accepts the following forms:
| Input | Resolution |
| --- | --- |
| `None` or `"none"` | Disable quantization. |
| Method string, such as `"fp8"` | Resolve an Omni override first, then the vLLM quantization registry. |
| Flat dictionary with `method` or `quant_method` | Normalize the method and construct its `QuantizationConfig`. |
| Per-component dictionary | Construct a `ComponentQuantizationConfig`. |
| Existing `QuantizationConfig` | Pass through without rebuilding it. |
Method aliases are normalized case-insensitively, with `-` and `_` treated as
equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
The factory uses Omni-specific overrides for methods whose diffusion behavior
or platform dispatch is not provided by the vLLM registry. The current override
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through
vLLM's registry when their configuration is compatible.
## Per-component routing
Multi-stage models do not necessarily want one quantization policy for every
module. A component map expresses that policy without requiring each pipeline
to construct separate vLLM configurations:
```python
quantization_config = {
"transformer": {"method": "fp8"},
"vae": None,
"default": None,
}
```
`ComponentQuantizationConfig` applies the following rules:
1. Component keys are interpreted as runtime layer-prefixes.
2. The longest matching prefix wins.
3. `None` disables quantization for the matching component.
4. `default`, when present, handles prefixes with no explicit match.
5. A matching child config delegates to its normal
`get_quant_method(layer, prefix)` implementation.
Prefix matching happens after the model's weight-name mapping. New model
integrations must verify the actual runtime prefixes, especially when a
`WeightsMapper` changes checkpoint names. A prefix mismatch can silently route
a layer to the default configuration, so component-routing tests should cover
both positive matches and fall-through behavior.
## Online and pre-quantized checkpoints
The same method name can represent two different loading paths:
- **Online**: the source checkpoint contains BF16/FP16 weights; the method
creates quantized weights or scales while the model is loaded.
- **Pre-quantized**: the checkpoint contains quantized weights and the metadata
required by the method; the loader consumes those serialized tensors.
Method configuration owns the mode flag, for example
`is_checkpoint_int8_serialized` or `is_checkpoint_mxfp8_serialized`. The
configuration's `get_quant_method()` selects the corresponding linear method.
For pipelines with transformer-specific `config.json` files, such as cascade
models, `resolve_quant_config_from_disk()` reconciles the active configuration
with each transformer's checkpoint metadata. Its invariants are:
- If no active configuration exists, valid checkpoint metadata can be used for
auto-detection.
- A method mismatch is an error rather than an implicit conversion.
- Serialized-checkpoint flags rebuild an online configuration into the
matching offline configuration.
- Checkpoint-specific `ignored_layers` are preserved by rebuilding the active
configuration when they differ.
- AutoRound MXFP8 metadata (`data_type="mx_fp"`) selects its offline path.
This prevents a model from interpreting serialized weights with an online
loader or silently applying the wrong method to one transformer in a
multi-transformer pipeline.
## Model and layer integration
Diffusion model constructors receive the resolved configuration and pass it to
quantizable vLLM layers. Each method configuration is responsible for:
- identifying supported layers in `get_quant_method()`;
- creating the correct parameter shapes and loaders;
- processing weights after loading when conversion is required; and
- applying the platform-specific quantized operation at runtime.
Model code must keep precision-sensitive or unsupported modules unquantized.
The shared `safe_quant_config()` helper is used by model integrations for
normalization and modulation layers where a generic FP8 configuration would be
unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
are not implicitly quantized just because the main transformer is quantized;
they require explicit component routing and method support.
Weight quantization and quantized KV cache are separate paths. The former is
resolved through `quantization_config`; runtime attention/KV-cache quantization
is configured by the diffusion attention configuration and must be validated
independently.
## Platform and parallelism boundaries
Platform checks belong in the method implementation rather than in the common
factory. This keeps configuration parsing portable while allowing CUDA, NPU,
and XPU methods to select different kernels or reject unsupported execution at
model construction time. The user guide is the source of truth for the current
hardware support matrix.
Quantization methods must also respect the weight layout produced by the
configured parallelism strategy:
- Tensor-parallel layers may receive partitioned weight shapes and loaders.
- Sequence-parallel execution changes activation movement but does not by
itself change the quantization configuration.
- Data-parallel replicas build their own method state; quantization metadata
must be available consistently on every rank.
- `WeightsMapper` transformations must be applied before matching ignored-layer
lists or component prefixes.
Distributed layerwise offload has an additional boundary: its sharded mmap
AllGather path is not compatible with online quantization because the path
expects already materialized checkpoint weights. Use the rank-local
`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
combining these features. See the [DLO design document](distributed_layerwise_offload.md)
for the full compatibility matrix.
## Adding a quantization backend
Contributors should follow this sequence:
1. Check whether an existing vLLM backend can provide the required weight
format and kernel behavior. Prefer reuse over a duplicate Omni backend.
2. If Omni-specific behavior is required, implement a `QuantizationConfig` and
its `QuantizeMethodBase`/linear method in `vllm_omni/quantization/`.
3. Implement `get_name()`, supported activation dtypes, minimum capability,
checkpoint parsing, ignored-layer handling, and `get_quant_method()`.
4. Add a lazy factory override in `factory.py`. Keep optional platform
dependencies out of package import paths and normalize aliases to one
canonical method name.
5. If the method is pre-quantized, define the checkpoint metadata and the
serialized loading path before adding an online path.
6. Add model-specific component routing only when the model contains stages
that must use different policies.
7. Document user-visible configuration in the quantization user guide and keep
implementation rationale here.
The public factory registration hook,
`register_quantization_override(method, builder)`, is intended for integrations
that need to register an Omni-compatible builder without changing the vLLM
registry.
## Validation
Validation should cover the complete path, not only whether a config object can
be instantiated:
| Area | Required checks |
| --- | --- |
| Factory | Strings, aliases, flat dictionaries, invalid inputs, passthrough, and component maps |
| Method config | Online/offline selection, ignored layers, platform guards, and weight loading |
| Checkpoint metadata | Auto-detection, method mismatch errors, serialized flags, and per-transformer metadata |
| Model integration | Correct quantization scope, unquantized components, and parallel weight layouts |
| Numerical quality | Same-seed BF16 versus quantized output comparison at representative resolutions and step counts |
| Performance | Peak device memory, generation latency, and any backend-specific throughput metric |
The unit tests live under `tests/diffusion/quantization/`. The trajectory
similarity tool in
`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py`
provides a repeatable quality and memory comparison for diffusion models.
Method support and feature combinations remain model- and hardware-dependent;
passing factory tests alone does not establish production support.
## Design invariants
- `quantization_config` is the single runtime source of truth after config
normalization.
- A checkpoint method mismatch fails loudly instead of risking weight
corruption.
- Components without an explicit supported quantization path remain in their
original precision.
- Optional backend dependencies are imported lazily.
- Method implementations own platform-specific capability checks and kernels.
- User-facing support claims require model, hardware, numerical, and memory
validation in addition to configuration-level tests.
@@ -0,0 +1,59 @@
# Distributed utils
This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
## 1. Installation
```bash
pip install "ray[default]"
```
## 2. Ray Utils
The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
* **Multi-node deployment**: Running pipeline stages across different physical machines.
* **Resource management**: Efficient GPU/CPU allocation.
### 2.1 Basic Usage
To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
**Command Line Example:**
```bash
vllm serve Qwen/Qwen2.5-Omni-7B \
--omni \
--port 8091 \
--worker-backend ray \
--ray-address auto
```
### 2.2 Cluster Setup
**Step 1: Start Head Node**
Run this on your primary machine:
```bash
ray start --head --port=6399
```
**Step 2: Connect Worker Nodes**
Run this on each worker machine:
```bash
ray start --address=<HEAD_NODE_IP>:6399
```
> **Tip**: For a complete cluster setup script, refer to the vLLM example:
> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
### 2.3 Distributed Connector Support
When running on Ray, the system automatically adapts its communication strategy:
* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
* **SHM payload routing**: `SharedMemoryConnector` always stores payloads in `/dev/shm`.
### 2.4 Internal Helpers
* **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
## 3. Troubleshooting
* **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
* **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 195 KiB

+6 -53
View File
@@ -1,8 +1,6 @@
# Design Documents
This section contains design documents and architecture specifications for
vLLM-Omni. The sidebar groups documents by the system concern they describe,
while the source paths remain stable so existing links continue to work.
This section contains design documents and architecture specifications for vLLM-Omni.
## Architecture Documents
@@ -10,63 +8,18 @@ while the source paths remain stable so existing links continue to work.
## Feature Design Documents
### Runtime and stage execution
- [Disaggregated Inference](feature/disaggregated_inference.md)
- [Async Chunk](feature/async_chunk.md)
- [Async Diffusion Output](feature/async_diffusion_output.md)
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
- [Ray-based Execution](feature/ray_based_execution.md)
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
- [Request-Level Batching for Diffusion](feature/diffusion_request_level_batching.md)
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
### Communication
#### OmniConnector implementations
- [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md)
- [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md)
- [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md)
- [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md)
- [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md)
- [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md)
### Diffusion acceleration
#### Parallelism
- [CFG-Parallel](feature/cfg_parallel.md)
- [Expert Parallel](feature/expert_parallel.md)
- [Hybrid Sharded Data Parallel (HSDP)](feature/hsdp.md)
- [Pipeline Parallel](feature/pipeline_parallel.md)
- [Sequence Parallel](feature/sequence_parallel.md)
- [Tensor Parallel](feature/tensor_parallel.md)
- [VAE Patch Parallelism](feature/vae_parallel.md)
#### Attention Backends
- [Skip-Softmax](feature/skip_softmax.md)
#### Quantization
- [Quantization](feature/quantization.md)
- [Cache-DiT](feature/cache_dit.md)
- [TeaCache](feature/teacache.md)
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
## Infrastructure and Performance
## Infrastructure Design Documents
- [Prometheus Metrics](metrics.md)
- [Speech Generation Performance Optimizations](qwen3_omni_tts_performance_optimization.md)
## Module Design Documents
### Runtime Modules
- [AR Module](module/ar_module.md)
- [DIT Module](module/dit_module.md)
### Orchestration
- [Entrypoint Module](module/entrypoint_module.md)
- [AsyncOmni Architecture (Qwen3-Omni Example)](module/async_omni_architecture.md)
+2 -1
View File
@@ -252,7 +252,8 @@ while True:
batch of compatible independent requests for one full pipeline forward when
the pipeline supports request-level batching; the step-wise path can also
admit or remove compatible requests between denoise steps. See
[Diffusion Continuous Batching](../feature/diffusion_continuous_batching.md).
[Request-Level Batching](../feature/diffusion_request_level_batching.md) and
[Continuous Batching for Step-Wise Diffusion](../feature/diffusion_continuous_batching.md).
---
@@ -419,7 +419,8 @@ Notes:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
--omni \
--port 8091
--port 8091 \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
```
#### 3) Key config knobs
@@ -493,7 +494,7 @@ The default config (`qwen3_tts.yaml`) enables the full optimization stack:
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--omni \
--port 8000 \
--no-async-chunk
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml
```
#### 3) Key config knobs
@@ -1,8 +1,5 @@
# Installation
!!! important
vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
vLLM-Omni supports the following hardware platforms:
- [GPU](gpu.md)
@@ -5,7 +5,7 @@
# --8<-- [end:requirements]
# --8<-- [start:set-up-using-python]
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.26.x release line uses vLLM 0.26.x.
vLLM-Omni depends vLLM. So please follow instructions below mainly for vLLM.
!!! note
PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
```bash
DOCKER_BUILDKIT=1 docker build \
-f docker/Dockerfile.cuda \
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.22.1 \
-t vllm-omni-cuda .
```
@@ -53,20 +53,16 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
```bash
# Pin vLLM and vLLM-Ascend to the v0.26 release line
git clone -b v0.26.0 https://github.com/vllm-project/vllm.git
cd vllm
# Pin vLLM version to 0.18.0
git clone -b v0.18.0 https://github.com/vllm-project/vllm.git
VLLM_TARGET_DEVICE=empty pip install -v -e .
cd ..
git clone -b releases/v0.26.0rc https://github.com/vllm-project/vllm-ascend.git
cd vllm-ascend
git clone -b v0.18.0rc1 https://github.com/vllm-project/vllm-ascend.git
pip install -v -e .
cd ..
# Install vLLM-Omni from the latest main branch
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
cd /vllm-workspace/vllm-omni
pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
+3 -2
View File
@@ -87,8 +87,9 @@ if __name__ == "__main__":
!!! info
For diffusion request batching, step execution, and streaming controls, see
[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md).
For diffusion request-level batching controls such as `max_num_seqs` and
`request_batch_max_wait_ms`, see
[Request-Level Batching](../user_guide/diffusion/request_batching.md).
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
-1
View File
@@ -44,7 +44,6 @@ th {
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
@@ -1,150 +0,0 @@
<svg xmlns="http://www.w3.org/2000/svg" width="1024" height="600" viewBox="0 0 1024 600" role="img" aria-labelledby="title desc">
<title id="title">Qwen3-Omni Talker step gap before and after asynchronous output materialization</title>
<desc id="desc">Before the optimization, synchronous D2H copy, per-request payload construction, streaming payload handoff, and connector-signal drain separate consecutive Talker forward and sample steps by about 2.8 milliseconds. After the optimization, that Omni output work runs on a background builder while the next Talker step executes, leaving a boundary of about 41 microseconds for snapshot, D2H enqueue, state update, and sampled-token feedback.</desc>
<defs>
<marker id="time-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#64748b"/>
</marker>
<marker id="red-end" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="5" markerHeight="5" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#b91c1c"/>
</marker>
<marker id="red-start" viewBox="0 0 10 10" refX="2" refY="5" markerWidth="5" markerHeight="5" orient="auto">
<path d="M 10 0 L 0 5 L 10 10 z" fill="#b91c1c"/>
</marker>
<marker id="blue-end" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="6" markerHeight="6" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#2563eb"/>
</marker>
<style>
.canvas { fill: #ffffff; stroke: #e5e7eb; }
.title { font: 700 22px system-ui, -apple-system, Segoe UI, sans-serif; fill: #111827; }
.subtitle { font: 500 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.phase { font: 700 16px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1f2937; }
.phase-sub { font: 500 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.lane { font: 600 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #374151; }
.block-title { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; letter-spacing: .25px; fill: #ffffff; }
.block-op { font: 500 9.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.axis { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #94a3b8; }
.measure { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b91c1c; }
.inset-title { font: 700 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
.inset-step { font: 600 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.note-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e3a8a; }
.note-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
.stat-cap { font: 600 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.stat-big { font: 800 16px system-ui, -apple-system, Segoe UI, sans-serif; fill: #111827; }
.stat-sub { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.before-card { fill: #eef4ff; stroke: #c7dbf6; }
.after-card { fill: #eafaf0; stroke: #bfe7cf; }
.talker { fill: #16a34a; }
.sync-output { fill: #ea580c; }
.async-output { fill: #2563eb; }
.zoom-card { fill: #fff7f7; stroke: #f3caca; }
.note-card { fill: #eff6ff; stroke: #bfdbfe; }
.stat-blue { fill: #eef4ff; stroke: #c7dbf6; }
.stat-green { fill: #eafaf0; stroke: #bfe7cf; }
.stat-gray { fill: #f4f5f7; stroke: #e2e5ea; }
.guide { stroke: #e2e8f0; stroke-width: 1; }
.time-line { stroke: #64748b; stroke-width: 1.3; fill: none; marker-end: url(#time-arrow); }
.measure-line { stroke: #b91c1c; stroke-width: 1.4; marker-start: url(#red-start); marker-end: url(#red-end); }
.zoom-lead { stroke: #dc2626; stroke-width: 1.15; stroke-dasharray: 4 3; fill: none; }
.handoff { stroke: #2563eb; stroke-width: 1.3; stroke-dasharray: 3 2; fill: none; marker-end: url(#blue-end); }
</style>
</defs>
<rect x="14" y="14" width="996" height="572" rx="14" class="canvas"/>
<text x="44" y="52" class="title">Async output: what moves off the Talker decode path</text>
<text x="44" y="74" class="subtitle">Qwen3-Omni Talker stage &#183; consecutive decode steps in a fixed profiler window</text>
<!-- ===================== BEFORE ===================== -->
<rect x="44" y="94" width="144" height="110" rx="10" class="before-card"/>
<text x="60" y="120" class="phase">Before</text>
<text x="60" y="139" class="phase-sub">Omni output built</text>
<text x="60" y="156" class="phase-sub">synchronously on the</text>
<text x="60" y="173" class="phase-sub">AR critical path</text>
<text x="206" y="122" class="lane">AR path</text>
<line x1="300" y1="88" x2="300" y2="210" class="guide"/>
<rect x="300" y="104" width="190" height="54" rx="5" class="talker"/>
<text x="395" y="124" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N</text>
<text x="395" y="145" text-anchor="middle" class="block-op">forward &#8594; sample</text>
<rect x="490" y="104" width="250" height="54" rx="5" class="sync-output"/>
<text x="615" y="120" text-anchor="middle" class="block-title">SYNC OMNI OUTPUT &#183; STEP N</text>
<text x="615" y="138" text-anchor="middle" class="block-op">D2H &#8594; per-request slice / payload build</text>
<text x="615" y="152" text-anchor="middle" class="block-op">streaming handoff &#8594; connector-signal drain</text>
<rect x="740" y="104" width="220" height="54" rx="5" class="talker"/>
<text x="850" y="124" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N+1</text>
<text x="850" y="145" text-anchor="middle" class="block-op">forward &#8594; sample</text>
<line x1="490" y1="174" x2="740" y2="174" class="measure-line"/>
<text x="615" y="194" text-anchor="middle" class="measure">~2.8 ms GPU gap: output work is inline</text>
<!-- ===================== AFTER ===================== -->
<rect x="44" y="218" width="144" height="140" rx="10" class="after-card"/>
<text x="60" y="244" class="phase">After</text>
<text x="60" y="263" class="phase-sub">Only step feedback</text>
<text x="60" y="280" class="phase-sub">stays inline; payload</text>
<text x="60" y="297" class="phase-sub">work runs in the</text>
<text x="60" y="314" class="phase-sub">background</text>
<text x="206" y="254" class="lane">Talker GPU</text>
<text x="206" y="326" class="lane">Output builder</text>
<line x1="300" y1="212" x2="300" y2="370" class="guide"/>
<rect x="300" y="226" width="235" height="54" rx="5" class="talker"/>
<text x="418" y="246" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N</text>
<text x="418" y="267" text-anchor="middle" class="block-op">forward &#8594; sample</text>
<rect x="540" y="226" width="235" height="54" rx="5" class="talker"/>
<text x="658" y="246" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N+1</text>
<text x="658" y="267" text-anchor="middle" class="block-op">forward &#8594; sample</text>
<path d="M 535 280 L 548 298" class="handoff"/>
<rect x="540" y="298" width="360" height="58" rx="5" class="async-output"/>
<text x="720" y="315" text-anchor="middle" class="block-title">BACKGROUND OMNI OUTPUT &#183; STEP N</text>
<text x="720" y="334" text-anchor="middle" class="block-op">wait for D2H &#8594; per-request slice / payload build</text>
<text x="720" y="349" text-anchor="middle" class="block-op">streaming handoff &#8594; connector-signal drain</text>
<line x1="300" y1="366" x2="930" y2="366" class="time-line"/>
<text x="922" y="382" text-anchor="end" class="axis">time</text>
<!-- leads from the real 41 us boundary to the magnified inset -->
<line x1="535" y1="280" x2="485" y2="424" class="zoom-lead"/>
<line x1="540" y1="280" x2="525" y2="424" class="zoom-lead"/>
<!-- ===================== 41 µs BOUNDARY ===================== -->
<rect x="300" y="386" width="400" height="92" rx="10" class="zoom-card"/>
<text x="316" y="404" class="inset-title">Zoom &#183; Talker step boundary (magnified)</text>
<rect x="320" y="416" width="165" height="26" rx="4" class="talker"/>
<rect x="525" y="416" width="155" height="26" rx="4" class="talker"/>
<text x="403" y="433" text-anchor="middle" class="inset-step">Talker step N</text>
<text x="603" y="433" text-anchor="middle" class="inset-step">Talker step N+1</text>
<line x1="485" y1="452" x2="525" y2="452" class="measure-line"/>
<text x="505" y="469" text-anchor="middle" class="measure">~41 &#181;s</text>
<rect x="720" y="386" width="240" height="92" rx="10" class="note-card"/>
<text x="736" y="405" class="note-title">Inline at the ~41 &#181;s boundary</text>
<text x="736" y="424" class="note-text">&#8226; eager Talker state update</text>
<text x="736" y="441" class="note-text">&#8226; output snapshot + D2H enqueue</text>
<text x="736" y="458" class="note-text">&#8226; sampled-token feedback</text>
<!-- ===================== SUMMARY ===================== -->
<rect x="44" y="494" width="296" height="70" rx="9" class="stat-blue"/>
<text x="60" y="514" class="stat-cap">Step-to-step gap</text>
<text x="60" y="538" class="stat-big">2.8 ms &#8594; 41 &#181;s</text>
<text x="60" y="555" class="stat-sub">~68&#215; tighter</text>
<rect x="356" y="494" width="296" height="70" rx="9" class="stat-green"/>
<text x="372" y="514" class="stat-cap">Talker GPU</text>
<text x="372" y="538" class="stat-big">next step starts immediately</text>
<text x="372" y="555" class="stat-sub">decode workers stay busy</text>
<rect x="668" y="494" width="312" height="70" rx="9" class="stat-gray"/>
<text x="684" y="514" class="stat-cap">Background builder</text>
<text x="684" y="538" class="stat-big">materializes Omni output</text>
<text x="684" y="555" class="stat-sub">payload work overlaps the next Talker step</text>
</svg>

Before

Width:  |  Height:  |  Size: 9.8 KiB

@@ -14,8 +14,8 @@ The full set of backends and their platform defaults is in the **Backend Options
| Value | Notes |
|---|---|
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention and **SAGE** quantized attention — see below. Requires `flashinfer`. |
| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
@@ -222,30 +222,6 @@ Same prompt and seed across runs. `Total generation time` from `text_to_video.py
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
## Known Limitations
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
```bash
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model Lightricks/LTX-2 ...
```
### FlashAttention-4 on Blackwell
Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
```bash
pip install 'vllm-omni[fa4]'
```
`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
when it is unavailable. Version `4.0.0b18` is required; earlier beta
wheels had known JIT failures on Blackwell.
## Choosing a Backend Manually
@@ -184,15 +184,10 @@ class Flux2Transformer2DModel(nn.Module):
### How It Works
Distributed layerwise offloading extends single-GPU layerwise offloading to
multi-device deployments. In the default mode, each rank stores only
**1/dp_size** of the model weights in host memory; full layer weights are
reconstructed at runtime via **AllGather** on a dedicated communication stream,
overlapped with computation via a fixed double-buffer scheme. The
`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
weights independently and does not shard weights across ranks.
For the implementation design and compatibility matrix, see
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
multi-device deployments. Each DP rank stores only **1/dp_size** of the model
weights in host memory; full layer weights are reconstructed at runtime via
**AllGather** on a dedicated communication stream, overlapped with computation
via a fixed double-buffer scheme.
**Key features:**
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
@@ -224,7 +219,7 @@ vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4
# Without DLO AllGather (each rank streams standard-loader rank-local weights)
# Without AllGather (each rank loads full weights, no sharding)
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4 \
@@ -252,9 +247,9 @@ m = Omni(
| Flag | Description | Default |
|------|-------------|---------|
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
### How model weights are loaded (mmap path)
@@ -308,15 +303,9 @@ request while AllGather synchronizes only weight shards (request-independent).
### Limitations
- Online quantization (FP8) is incompatible with mmap loading — falls back to
regular loading only when the model is outside the DLO+AllGather mmap path;
online quantization with DLO+AllGather is rejected. Use
`--dlo-no-use-allgather` or disable online quantization.
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
weight-loader callbacks are bypassed. The no-AllGather path retains the
standard TP loader and is experimental.
- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
with no-AllGather is accepted at configuration level, but has limited
end-to-end validation.
regular `load_weights()` automatically
- Tensor Parallel is not supported (DLO uses DP-based sharding)
- HSDP + AllGather is rejected (would double-shard weights)
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
**Module Discovery**
@@ -349,7 +338,7 @@ OffloadBackend (base class)
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
│ (single-GPU, full weights on host)
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
(multi-GPU, sharded AllGather or rank-local mode)
(multi-GPU, 1/dp_size sharded weights + AllGather)
```
Factory function `get_offload_backend()` selects the appropriate backend based on
@@ -380,10 +369,4 @@ reasoner/generator components inside the model forward pass.
**Notes:**
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
- Distributed Layerwise Offloading reuses the layerwise block topology, but the
model and parallelism combination still needs validation. AllGather mode
additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
retains the standard loader's rank-local tensors. See the
[DLO design](../../design/feature/distributed_layerwise_offload.md) and
[Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
for usage examples.
- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required. See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.
@@ -1,252 +0,0 @@
# Diffusion Execution Modes
vLLM-Omni supports complete-request execution and step-wise diffusion
execution. Both modes use the same asynchronous engine output stream; the
configuration controls how work is scheduled and whether intermediate outputs
are exposed.
## Choose a Mode
| Goal | CLI configuration |
|---|---|
| Serial request execution | `--max-num-seqs 1` |
| Fused request-level batching | `--max-num-seqs N` |
| Single-request step execution | `--step-execution --max-num-seqs 1` |
| Step-wise continuous batching | `--step-execution --max-num-seqs N` |
| Chunked diffusion output | `--diffusion-streaming-output` |
`N` must be greater than `1` to allow batching. The selected pipeline must
support request-level batching or batched-step execution as appropriate;
single-request step support alone is not sufficient for step-wise continuous
batching.
## Request Execution
Request execution is the default when `--step-execution` is omitted. The
pipeline performs a complete `forward()` for each scheduler wave.
### Serial Requests
Use the conservative serial path with any request-mode pipeline:
```bash
vllm serve MODEL --omni \
--port 8091 \
--max-num-seqs 1
```
### Request-Level Batching
Set `max_num_seqs` above one to combine compatible independent requests into a
single pipeline forward:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--max-num-seqs 4
```
For bursty traffic, add a small admission window:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--max-num-seqs 4 \
--request-batch-max-wait-ms 20
```
`--request-batch-max-wait-ms 0` is the default. A nonzero value can improve
batch formation but adds up to that much latency before a new scheduler wave.
It has no effect in step mode.
Each prompt remains a separate logical request. Do not submit a top-level list
as one packed prompt. Concurrent serving requests are batched internally when
their shapes, guidance settings, output counts, and LoRA settings are
compatible.
Only pipelines that declare request-batch support accept
`max_num_seqs>1` in request mode. Unsupported pipelines fail during engine
initialization; use `max_num_seqs=1` instead.
## Step Execution
Step execution lets the scheduler advance and abort requests between denoise
steps:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--step-execution \
--max-num-seqs 1
```
Set a larger capacity to allow compatible requests to share denoise waves:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--step-execution \
--max-num-seqs 8
```
Step continuous batching is experimental. Start with `max_num_seqs=1` when
validating a model or debugging correctness, then increase it for
multi-request throughput.
Step execution is capability-based, not a generic switch for every diffusion
model. Qwen-Image supports step-wise continuous batching. HunyuanImage3 also
supports it, but only when its resolved self-attention backend is `TORCH_SDPA`;
set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure
`diffusion_attention_config.default.backend=TORCH_SDPA` before using
`--max-num-seqs >1`. See the
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
for its validated configuration. Helios supports single-request step execution only: use
`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's
documentation and source for the latest support status.
## Streaming Output
Use `--diffusion-streaming-output` for a pipeline that can produce intermediate
diffusion outputs:
```bash
vllm serve BestWishYsh/Helios-Distilled --omni \
--port 8000 \
--diffusion-streaming-output
```
Streaming output requires step execution. If
`--diffusion-streaming-output` is set without `--step-execution`, the engine
enables step execution automatically. Model initialization fails if the
pipeline does not implement step execution.
Chunk-capable pipelines emit intermediate and final outputs through the same
request stream. Final-only step pipelines emit only the final result even when
streaming output is enabled.
The normal non-streaming serving path uses the same internal stream but drains
it and returns only the final output. Users do not need to enable
`--diffusion-streaming-output` to benefit from the unified engine lifecycle.
## Send Requests
Execution-mode flags configure the server. They do not change the
OpenAI-compatible client request format. For example, after starting any of
the non-streaming configurations above:
```bash
curl -X POST http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"prompt": "a cat sitting on a windowsill",
"size": "1024x1024",
"num_inference_steps": 50,
"seed": 42
}' | jq -r '.data[0].b64_json' | base64 -d > cat.png
```
To create a request-level batch, send multiple independent requests
concurrently. Do not put several prompts into one `prompt` field:
```bash
for prompt in "a red fox" "a blue bird" "a green frog"; do
curl -s -X POST http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d "{\"prompt\":\"${prompt}\",\"size\":\"1024x1024\",\"seed\":42}" \
> /dev/null &
done
wait
```
The scheduler may batch these requests when their sampling parameters are
compatible. See the
[Image Generation API](../../serving/image_generation_api.md) for response
formats and additional client examples. Streaming models can use
model-specific streaming endpoints documented by their serving guide.
## Python API
The Python arguments mirror the CLI flags:
```python
from vllm_omni import Omni
omni = Omni(
model="Qwen/Qwen-Image",
max_num_seqs=4,
request_batch_max_wait_ms=20.0,
)
outputs = omni.generate(
[
"a cup of coffee on a table",
"a toy dinosaur on a sandy beach",
"a fox waking up in bed and yawning",
]
)
```
Each list item passed to `Omni.generate()` becomes an independent logical
request that the scheduler may batch with compatible items.
For step execution:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
step_execution=True,
max_num_seqs=1,
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
## Stage Configuration
Use the equivalent engine arguments in a deployment YAML:
```yaml
stage_args:
- stage_id: 0
stage_type: diffusion
engine_args:
step_execution: false
max_num_seqs: 4
request_batch_max_wait_ms: 20
```
For step execution, set `step_execution: true` and remove
`request_batch_max_wait_ms`.
## CLI Reference
| Flag | Default | Effect |
|---|---:|---|
| `--step-execution` | disabled | Select step-wise scheduling |
| `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity |
| `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode |
| `--diffusion-streaming-output` | disabled | Expose supported intermediate diffusion outputs and require step execution |
## Limitations and Troubleshooting
- Requests batch only when their compatibility-sensitive parameters match.
- Different LoRA adapters or scales run in separate batches.
- FIFO scheduling can cause an incompatible request to block later compatible
requests.
- All diffusion cache backends are unsupported in step mode. KV transfer is
supported for newly admitted step requests; some other request-mode extras
remain unsupported.
- If request-mode startup reports that the pipeline does not support batching,
use `--max-num-seqs 1`.
- If step-mode startup mentions `prepare_encode()`, `denoise_step()`,
`step_scheduler()`, or `post_decode()`, the pipeline does not implement the
required step contract.
For implementation details and model-author guidance, see
[Diffusion Continuous Batching](../../design/feature/diffusion_continuous_batching.md).
@@ -0,0 +1,119 @@
# Request-Level Batching
Request-level batching lets diffusion serving combine multiple compatible
logical requests into one pipeline `forward()` call. Each prompt remains a
separate request with its own `request_id`, sampling parameters, seed, output,
error, and abort state. The scheduler decides which requests can run together.
!!! warning "Prompt List Semantics"
Diffusion request-level batching does not support a top-level packed
list-prompt request. Submit multiple prompts as independent requests and let
the scheduler batch compatible in-flight requests. Multimodal payloads stay
inside a single prompt dict, for example
`{"prompt": "...", "multi_modal_data": {"image": image}}`.
## Enablement
Request-level execution is the default diffusion request mode when
`step_execution=False`. `max_num_seqs` controls the maximum number of compatible
requests in one scheduler wave:
- `max_num_seqs=1` is the serial request path.
- `max_num_seqs>1` allows fused request-level batching for pipelines that
declare request-batch support.
Increase `max_num_seqs` above `1` to allow the request scheduler to keep more
than one compatible request active:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--max-num-seqs 4
```
For bursty online traffic, you can also set a small admission wait window. This
lets the engine wait briefly before the first `schedule()` of a new wave so
nearby compatible requests can arrive and share the same fused forward pass:
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--max-num-seqs 4 \
--request-batch-max-wait-ms 20
```
`--request-batch-max-wait-ms 0` is the default and disables admission waiting,
so there is no added wait latency.
For deploy YAMLs, configure the diffusion stage engine args:
```yaml
stage_args:
- stage_id: 0
stage_type: diffusion
engine_args:
max_num_seqs: 4
request_batch_max_wait_ms: 20
```
## Compatibility
Only pipelines that declare request-batch support can use the fused
multi-request `pipeline.forward(batch)` path. Pipelines that do not support this
contract can still use request-level execution with `max_num_seqs=1`, which is
the serial request path. Setting `max_num_seqs>1` for a pipeline without
request-batch support fails early during engine initialization.
The scheduler batches only compatible requests. Compatibility is based on
shape-sensitive and guidance-sensitive sampling fields, including resolution,
frame count, CFG settings, output count, and LoRA identity. Requests with
different LoRA adapters or scales are kept in separate batches so the worker
activates one adapter per batch.
Request-level batching applies when `step_execution=False`. For the separate
step-wise runtime, see [Step Execution](step_execution.md).
## Tuning
- `max_num_seqs` caps the number of active compatible requests in one scheduler
wave. `1` means serial request execution; values above `1` enable fused
request-level batching when the pipeline supports it.
- `request_batch_max_wait_ms` is an upper bound on extra admission wait before a
new wave starts. Keep it small for latency-sensitive serving; values such as
`10` to `50` ms are a practical starting range for bursty HTTP ingress.
- `0` disables admission waiting and preserves the lowest first-request latency.
- FIFO ordering is conservative: an incompatible request at the front of the
waiting queue can block later compatible requests from joining the current
batch.
## Python API
When constructing `Omni`, pass the same engine arguments:
```python
from vllm_omni.entrypoints.omni import Omni
omni = Omni(
model="Qwen/Qwen-Image",
max_num_seqs=4,
request_batch_max_wait_ms=20.0,
)
outputs = omni.generate(
[
"a cup of coffee on a table",
"a toy dinosaur on a sandy beach",
"a fox waking up in bed and yawning",
]
)
```
`Omni.generate([...])` submits each list item as its own logical diffusion
request. The runtime may batch those requests internally when their sampling
parameters are compatible.
## For Contributors
For implementation details and model-author guidance, see
[Request-Level Batching for Diffusion](../../design/feature/diffusion_request_level_batching.md).
+100
View File
@@ -0,0 +1,100 @@
# Step Execution
Step execution is an opt-in diffusion execution mode enabled with
`step_execution=True` when constructing `Omni`.
It is not a generic diffusion toggle for every pipeline. Only pipelines that
implement the stepwise contract support it today.
## Quick Start
### Python API
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Qwen/Qwen-Image",
step_execution=True,
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(
num_inference_steps=50,
),
)
```
### Serving
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091 \
--step-execution \
--max-num-seqs 8
```
For serving, `--step-execution` enables the step-wise runtime. Continuous
batching only becomes relevant when `--max-num-seqs > 1`.
`max_num_seqs` controls the step scheduler capacity:
- `max_num_seqs=1` runs one request through the step-wise state machine.
- `max_num_seqs>1` allows compatible active requests to share step waves.
## Supported Pipelines
| Pipeline | Example models | Step execution |
|----------|----------------|----------------|
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
!!! warning "Experimental continuous batching"
When `--step-execution` is enabled and `max_num_seqs > 1` is configured,
the step-wise path can batch
compatible requests together. This is experimental. Requests with
incompatible sampling parameters are intentionally kept in separate batches,
and `max_num_seqs=1` remains the conservative single-request step path.
## Current Limitations
- Continuous batching under `step_execution` is experimental and only batches
compatible requests.
- `cache_backend` is not supported together with step execution.
- Unsupported pipelines fail early during model loading.
- Request-mode extras such as KV transfer are not wired into step mode yet.
- LoRA is supported in step mode, but each batch must use a single adapter:
requests with different `lora_request` or `lora_scale` are scheduled into
separate batches.
## When To Use It
Use step execution only when you specifically need the pipeline to run through
its stepwise request state machine. For normal diffusion inference, leave it
disabled unless your workflow depends on this mode.
For Qwen-Image online serving, the usual progression is:
- start with `--step-execution --max-num-seqs 1` if you only need the step-wise path
- increase `--max-num-seqs` after that if you want the experimental compatible-request batching behavior
If you are looking for general diffusion speedups, see
[Diffusion Features Overview](../diffusion_features.md).
## Troubleshooting
If model loading fails with a message mentioning `prepare_encode()`,
`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected
pipeline does not support step execution.
## For Model Authors
If you want to add step execution support to a new diffusion pipeline, see the
implementation guide:
[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).
If you also want that pipeline to participate in the experimental batched
step-wise path, see:
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
+13 -12
View File
@@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
| Method | Description | Best For |
|--------|-------------|----------|
| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
@@ -80,11 +80,12 @@ denoise steps.
| Method | Description | Best For |
|--------|-------------|----------|
| **[Diffusion Execution Modes](diffusion/execution_modes.md)** | Configures serial requests, request batching, step execution, continuous batching, and streaming output | Matching latency, throughput, cancellation, and output-delivery requirements |
| **[Request-Level Batching](diffusion/request_batching.md)** | Scheduler batches compatible independent diffusion requests into one pipeline forward pass | Bursty online serving and multi-request throughput |
| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control |
**Note:** Request-level batching and step execution are capability-based.
Consult the execution guide and the selected pipeline's documentation for
current support.
**Note:** Request-level batching is available for pipelines that declare the
request-batch forward contract. Step execution is currently supported by
QwenImagePipeline only. See [Supported Models](#supported-models) for details.
### Quantization Methods
@@ -104,8 +105,8 @@ The following tables show which models support each feature:
> Notes:
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md).
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
### ImageGen
@@ -159,7 +160,7 @@ The following tables show which models support each feature:
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **MiniMax-H3** | ❌ | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
**Frame Interpolation Support**
@@ -205,7 +206,7 @@ The following tables show which models support each feature:
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. CPU Offloading (Layerwise) supports single-card for now.
5. Using FP8-Quant as an example of qunatization methods.
6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
## Multi-Thread Weight Loading
@@ -274,9 +275,9 @@ Measured on NVIDIA H800:
**Memory Optimization:**
- **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage
- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
- **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
- **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
**Extensions:**
@@ -285,7 +286,7 @@ Measured on NVIDIA H800:
**Execution Modes:**
- **[Diffusion Execution Modes](diffusion/execution_modes.md)** - Configure request batching, step execution, continuous batching, and streaming output
- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
**Startup Optimization:**
@@ -184,7 +184,7 @@ if __name__ == "__main__":
!!! info
For diffusion request-level batching controls such as `max_num_seqs`, see
[Diffusion Execution Modes](../../diffusion/execution_modes.md).
[Request-Level Batching](../../diffusion/request_batching.md).
### Negative Prompts
@@ -40,7 +40,7 @@ bash run_server_stage_cli.sh --stage 0
bash run_server_stage_cli.sh --stage 1
```
To use a custom deploy YAML:
To use a custom deploy YAML (note: `--stage-configs-path` is deprecated in favor of `--deploy-config`):
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
@@ -11,8 +11,8 @@ memory, sampling params) live in `vllm_omni/deploy/glm_image.yaml`.
vllm serve zai-org/GLM-Image --omni --port 8091
```
The config system auto-detects the pipeline from the model's `model_index.json`, so
no manual `--deploy-config` is needed.
The config system auto-detects the pipeline from the model's `model_index.json` — no
manual `--stage-configs-path` or `--deploy-config` needed.
By default, stage 0 (AR) runs on GPU 0 and stage 1 (Diffusion) on GPU 1. To colocate
both stages on a single GPU, override per stage:
@@ -15,9 +15,9 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
```
To use a custom deploy config, launch the server with the command below:
If you have custom stage configs file, launch the server with command below
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
### Send Multi-modal Request
@@ -182,7 +182,7 @@ The script supports the following arguments:
- `--model`: Model name/path (default: Qwen/Qwen2.5-Omni-7B)
- `--server-port`: Port for vLLM server (default: 8091)
- `--gradio-port`: Port for Gradio demo (default: 7861)
- `--deploy-config`: Path to a custom deploy config YAML file (optional)
- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
- `--server-host`: Host for vLLM server (default: 0.0.0.0)
- `--gradio-ip`: IP for Gradio demo (default: 127.0.0.1)
- `--share`: Share Gradio demo publicly (creates a public link)
@@ -195,9 +195,9 @@ The script supports the following arguments:
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
```
To use a custom deploy config:
If you have custom stage configs file:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
**Step 2: Run the Gradio demo**
@@ -69,19 +69,19 @@ The behaviors and conclusions above are summarized from current fault injection
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
The following log is based on a real run where a `SIGTERM` signal was sent to the serve root process. Its command line has been normalized to the current `--deploy-config` option.
The following full log is from a real run where a `SIGTERM` signal was sent to the serve root process.
```text
Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
Server ready on 127.0.0.1:60675 (OmniServer startup took 206.023s)
OmniServer started successfully
[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
[reliability][process-kill] current_server_proc pid=557118 name=python3 cmdline=/workspace/.venv/bin/python3 -c from multiprocessing.resource_tracker import main;main(57)
[reliability][process-kill] current_server_proc pid=557119 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
[reliability][process-kill] current_server_proc pid=557122 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
[reliability][process-kill] current_server_proc pid=558201 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [async_omni_engine.py:2397] [AsyncOmniEngine] Shutting down Orchestrator
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:351] [Orchestrator] Received shutdown signal
-1
View File
@@ -93,7 +93,6 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
-3
View File
@@ -5,9 +5,6 @@ path. The same configuration entrypoint is used across diffusion-only models,
multi-stage omni/TTS models, and multi-stage diffusion models, but each model
type has a different quantization scope.
For the internal architecture and backend extension points, see the
[quantization design document](../../design/feature/quantization.md).
## Quantization Modes
| Mode | Guide | Description | Methods |
@@ -1,77 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline PersonaPlex inference through the vllm-omni fullduplex backend.
This is the runnable end-to-end proof for the integration: it drives a user WAV
through :class:`PersonaPlexEngine` + :class:`PersonaPlexSession` (our wrappers
around the Moshi/Mimi lockstep loop) and writes the agent's audio + inner-monologue
text. It mirrors ``python -m moshi.offline`` but goes through our ``FrameStepper``
seam, so a green run validates that the session/adapter plumbing is faithful.
Prereqs: HF access to ``nvidia/personaplex-7b-v1`` (``export HF_TOKEN=...``). The
stack is moshi-free (native temporal + depformer + streaming Mimi via
``transformers``); decoding is greedy.
Example:
python examples/offline_inference/personaplex/personaplex_offline.py \\
--input-wav input_assistant.wav --output-wav out.wav \\
--voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
"""
from __future__ import annotations
import argparse
import json
import numpy as np
from vllm_omni.experimental.fullduplex.personaplex import (
PersonaPlexConfig,
PersonaPlexEngine,
PersonaPlexSession,
)
from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--input-wav", required=True, help="User-side input WAV")
parser.add_argument("--output-wav", required=True, help="Agent audio output WAV")
parser.add_argument("--output-text", default=None, help="Optional JSON for agent text")
parser.add_argument("--voice-prompt", default="NATF2.pt", help="Voice clone (bundled .pt or .wav)")
parser.add_argument("--persona", default=DEFAULT_PERSONA, help="System role text")
parser.add_argument("--hf-repo", default="nvidia/personaplex-7b-v1")
parser.add_argument("--device", default="cuda")
parser.add_argument("--cpu-offload", action="store_true")
args = parser.parse_args()
import sphn # local import: only the runnable driver needs the audio I/O dep
config = PersonaPlexConfig(
hf_repo=args.hf_repo,
voice_prompt=args.voice_prompt,
persona=args.persona,
device=args.device,
cpu_offload=args.cpu_offload,
)
engine = PersonaPlexEngine(config).load()
session = PersonaPlexSession(engine, config)
session.open()
# Load exactly like the reference (moshi.lm.load_audio): native read then
# explicit resample, so greedy runs stay bit-faithful to `python -m moshi.offline`.
user_pcm, src_sr = sphn.read(args.input_wav)
user_pcm = sphn.resample(user_pcm, src_sample_rate=src_sr, dst_sample_rate=config.sample_rate)
audio, text = session.run(np.asarray(user_pcm[0], dtype=np.float32))
sphn.write_wav(args.output_wav, audio, config.sample_rate)
print(f"wrote {audio.shape[0] / config.sample_rate:.2f}s of agent audio to {args.output_wav}")
print(f"agent text: {text!r}")
if args.output_text:
with open(args.output_text, "w") as f:
json.dump({"text": text}, f, ensure_ascii=False)
if __name__ == "__main__":
main()
@@ -260,7 +260,7 @@ if __name__ == "__main__":
For diffusion pipelines, the input list is sliced into single-item requests
before feeding into the diffusion pipeline. For request-level batching
controls such as `max_num_seqs`, see
[Diffusion Execution Modes](../../../docs/user_guide/diffusion/execution_modes.md).
[Request-Level Batching](../../../docs/user_guide/diffusion/request_batching.md).
### Negative Prompts
+1 -1
View File
@@ -53,5 +53,5 @@ vllm-omni serve "$MODEL" \
--host 0.0.0.0 \
--port "$PORT" \
--trust-remote-code \
--deploy-config "$DEPLOY_YAML" \
--stage-configs-path "$DEPLOY_YAML" \
--omni
@@ -28,11 +28,8 @@ If image.command is set, uses that as a full override.
- {{ .Values.omniArgs.numGpus | quote }}
{{- end }}
{{- if .Values.omniArgs.stageConfigsPath }}
{{- fail "omniArgs.stageConfigsPath has been removed; use omniArgs.deployConfigPath" }}
{{- end }}
{{- if .Values.omniArgs.deployConfigPath }}
- "--deploy-config"
- {{ .Values.omniArgs.deployConfigPath | quote }}
- "--stage-configs-path"
- {{ .Values.omniArgs.stageConfigsPath | quote }}
{{- end }}
{{- if and .Values.omniArgs.cacheBackend (ne .Values.omniArgs.cacheBackend "none") }}
- "--cache-backend"
@@ -78,26 +78,6 @@ tests:
path: spec.template.spec.containers[0].command
content: "--enable-layerwise-offload"
- it: should include deploy config in command
set:
omniArgs:
deployConfigPath: "/configs/deploy.yaml"
asserts:
- equal:
path: spec.template.spec.containers[0].command[8]
value: "--deploy-config"
- equal:
path: spec.template.spec.containers[0].command[9]
value: "/configs/deploy.yaml"
- it: should reject removed stage config value
set:
omniArgs:
stageConfigsPath: "/configs/legacy.yaml"
asserts:
- failedTemplate:
errorMessage: "omniArgs.stageConfigsPath has been removed; use omniArgs.deployConfigPath"
- it: should use full command override when image.command is set
set:
image:
@@ -67,8 +67,8 @@ omniArgs:
enableCpuOffload: false
# -- Number of GPUs for diffusion inference (empty = auto)
numGpus:
# -- Path to deploy config file (empty = auto-detected from model)
deployConfigPath:
# -- Path to stage configs file (empty = auto-detected from model)
stageConfigsPath:
# -- Cache backend: none, tea_cache, or cache_dit
cacheBackend: "none"
# -- Default sampling params as a JSON string
@@ -3,7 +3,7 @@
#
# Prereq:
# Start a vllm-omni OpenAI server for MiniCPM-o 4.5 on :8099 (see the
# 8x4090 deploy config at vllm_omni/deploy/minicpmo_4_5_8x4090.yaml).
# 8x4090 stage config under vllm_omni/model_executor/stage_configs).
set -e
HERE="$(cd "$(dirname "$0")" && pwd)"
@@ -1,84 +0,0 @@
# PersonaPlex full-duplex online serving
Serve [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
(a Moshi-based full-duplex speech-to-speech model) with the native vLLM-Omni engine.
The server hosts the **official PersonaPlex web client** at `/` (auto-downloaded from
the model repo) and implements its WebSocket protocol at `/api/chat`, so you talk to the
model live in the browser — mic in, agent speech + inner-monologue text out.
> Experimental. Requires a GPU and Hugging Face access to the gated repo
> (`HF_TOKEN` with access to `nvidia/personaplex-7b-v1`).
## Start the server
```bash
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 bash run_server.sh
# or directly:
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
vllm_omni.experimental.fullduplex.personaplex.serving.server --port 8124
```
## Talk to it
- **Browser (recommended):** open `http://localhost:8124/` and allow the microphone.
Use headphones so the agent does not hear itself.
- **Headless:** stream a 24 kHz mono WAV and save the reply:
```bash
python duplex_client.py --url ws://localhost:8124/api/chat --input user.wav --out reply.wav
```
## Unified Realtime API
The default `vllm_omni/deploy/personaplex.yaml` also enables main's engine-owned
full-duplex control plane:
```bash
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
/path/to/personaplex-7b-v1 \
--omni \
--stage-configs-path vllm_omni/deploy/personaplex.yaml
```
Validate the actual `/v1/realtime?duplex=1` scheduler path with paced 24 kHz
PCM, two concurrent sessions, overflow admission, per-session slot recycling,
and non-silent output:
```bash
python tests/e2e/online_serving/personaplex_realtime_duplex.py \
--model /path/to/personaplex-7b-v1 \
--input-wav /path/to/speech.wav \
--output-dir /tmp/personaplex-realtime-duplex
```
This is distinct from the standalone `/api/chat` compatibility server above.
The unified endpoint advertises `supports_barge_in=false`: overlapping speech
is native model behavior, but destructive output interruption and model-state
rewind have not been validated for PersonaPlex.
## Protocol (`/api/chat`)
Binary WebSocket messages, first byte is a tag (identical to `moshi.server`):
| Direction | Tag | Payload |
| --- | --- | --- |
| server → client | `\x00` | ready handshake (after the system-prompt prefill) |
| client → server | `\x01` | Opus-encoded mic audio (24 kHz) |
| server → client | `\x01` | Opus-encoded agent audio |
| server → client | `\x02` | UTF-8 inner-monologue text |
Connect with query params `text_prompt` (persona) and `voice_prompt` (voice file).
A raw-PCM endpoint (`/v1/audio/duplex`, JSON `open` + float32 frames) is also available
for clients that do not want an Opus dependency.
## Notes
- **Run the client near the server.** Real-time 80 ms audio is sensitive to network
latency/jitter; over a high-latency remote link playback can stutter regardless of
engine speed. On localhost it is smooth.
- Greedy decoding by default. One conversation per server by default (a second
concurrent connection is rejected with WS close 1013 until the first ends);
pass `--batch-size N` to host N concurrent conversations on one engine
(elastic batching; connections beyond N are rejected until a slot frees).
- Full runbook: [`recipes/NVIDIA/PersonaPlex.md`](../../../recipes/NVIDIA/PersonaPlex.md).
@@ -1,114 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Headless client for the PersonaPlex duplex server (official Moshi protocol).
Streams a 24 kHz mono WAV to ``/api/chat`` exactly as the browser client does
(Opus over WebSocket, tags ``\\x00`` ready / ``\\x01`` audio / ``\\x02`` text),
collects the agent's reply audio + inner-monologue text, and writes a WAV.
The browser UI at ``http://<host>:<port>/`` is the primary demo; this script is
for headless testing / scripting.
python duplex_client.py --url ws://localhost:8124/api/chat \
--input user.wav --out reply.wav
Deps: ``websockets``, ``sphn``, ``soundfile``, ``numpy``.
"""
from __future__ import annotations
import argparse
import asyncio
import numpy as np
import soundfile as sf
import sphn
SR = 24000 # PersonaPlex / Mimi codec rate
async def converse(url: str, pcm: np.ndarray) -> tuple[np.ndarray, str]:
import websockets
reader = sphn.OpusStreamReader(SR) # decode agent audio
writer = sphn.OpusStreamWriter(SR) # encode our mic audio
text: list[str] = []
ready = asyncio.Event()
last = {"t": 0.0}
async with websockets.connect(url, max_size=None) as ws:
async def receive() -> None:
loop = asyncio.get_event_loop()
while True:
try:
msg = await ws.recv()
except Exception:
return
if not isinstance(msg, (bytes, bytearray)) or not msg:
continue
tag, payload = msg[0], msg[1:]
if tag == 0:
ready.set()
elif tag == 1:
reader.append_bytes(bytes(payload))
last["t"] = loop.time()
elif tag == 2:
text.append(payload.decode("utf8", errors="replace"))
recv_task = asyncio.create_task(receive())
await asyncio.wait_for(ready.wait(), timeout=120) # waits for system-prompt prefill
frame = 1920 # 80 ms
for i in range(0, len(pcm), frame):
chunk = pcm[i : i + frame]
if len(chunk) < frame:
chunk = np.concatenate([chunk, np.zeros(frame - len(chunk), dtype=np.float32)])
writer.append_pcm(chunk)
data = writer.read_bytes()
if data:
await ws.send(b"\x01" + data)
await asyncio.sleep(frame / SR) # pace at real time
loop = asyncio.get_event_loop()
last["t"] = loop.time()
while loop.time() - last["t"] < 2.5: # drain until the agent stops talking
await asyncio.sleep(0.1)
recv_task.cancel()
chunks: list[np.ndarray] = []
while True:
out = reader.read_pcm()
if out is None or out.shape[-1] == 0:
break
chunks.append(np.asarray(out, dtype=np.float32))
audio = np.concatenate(chunks) if chunks else np.zeros(0, dtype=np.float32)
return audio, "".join(text)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--url", default="ws://localhost:8124/api/chat")
ap.add_argument("--input", required=True, help="24 kHz mono WAV of the user turn")
ap.add_argument("--out", default="personaplex_reply.wav")
ap.add_argument("--tail", type=float, default=10.0, help="trailing silence (s) for the agent reply")
ap.add_argument("--persona", default="")
ap.add_argument("--voice", default="NATF2.pt")
args = ap.parse_args()
wav, sr = sf.read(args.input, dtype="float32")
if wav.ndim > 1:
wav = wav.mean(axis=1)
if sr != SR:
raise SystemExit(f"input must be {SR} Hz mono; got {sr} Hz. Resample it first.")
wav = np.concatenate([wav, np.zeros(int(args.tail * SR), dtype=np.float32)])
url = f"{args.url}?text_prompt={args.persona}&voice_prompt={args.voice}"
audio, text = asyncio.run(converse(url, wav))
sf.write(args.out, audio if audio.size else np.zeros(SR, dtype=np.float32), SR)
print(f"agent reply: {audio.size / SR:.2f}s -> {args.out}")
print(f"inner-monologue: {text.strip()!r}")
if __name__ == "__main__":
main()
@@ -1,15 +0,0 @@
#!/bin/bash
# PersonaPlex full-duplex online serving.
#
# Serves the OFFICIAL PersonaPlex web client at / (auto-downloaded from
# nvidia/personaplex-7b-v1) and the Moshi WebSocket protocol at /api/chat,
# backed by the native vLLM-Omni engine. Requires HF access to the gated repo.
HOST="${HOST:-0.0.0.0}"
PORT="${PORT:-8124}"
echo "Starting PersonaPlex duplex server on ${HOST}:${PORT}..."
echo "Open http://localhost:${PORT}/ for the official web client (run it near the server)."
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
--host "$HOST" --port "$PORT"
@@ -12,9 +12,9 @@ Please refer to [README.md](../../../README.md)
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
```
To use a custom deploy config, launch the server with the command below:
If you have custom stage configs file, launch the server with command below
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
### Send Multi-modal Request
@@ -173,7 +173,7 @@ The script supports the following arguments:
- `--model`: Model name/path (default: Qwen/Qwen2.5-Omni-7B)
- `--server-port`: Port for vLLM server (default: 8091)
- `--gradio-port`: Port for Gradio demo (default: 7861)
- `--deploy-config`: Path to a custom deploy config YAML file (optional)
- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
- `--server-host`: Host for vLLM server (default: 0.0.0.0)
- `--gradio-ip`: IP for Gradio demo (default: 127.0.0.1)
- `--share`: Share Gradio demo publicly (creates a public link)
@@ -186,9 +186,9 @@ The script supports the following arguments:
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
```
To use a custom deploy config:
If you have custom stage configs file:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
**Step 2: Run the Gradio demo**
@@ -13,7 +13,7 @@ set -e
MODEL="Qwen/Qwen2.5-Omni-7B"
SERVER_PORT=8091
GRADIO_PORT=7861
DEPLOY_CONFIG_PATH=""
STAGE_CONFIGS_PATH=""
SERVER_HOST="0.0.0.0"
GRADIO_IP="127.0.0.1"
GRADIO_SHARE=false
@@ -33,8 +33,8 @@ while [[ $# -gt 0 ]]; do
GRADIO_PORT="$2"
shift 2
;;
--deploy-config)
DEPLOY_CONFIG_PATH="$2"
--stage-configs-path)
STAGE_CONFIGS_PATH="$2"
shift 2
;;
--server-host)
@@ -56,7 +56,7 @@ while [[ $# -gt 0 ]]; do
echo " --model MODEL Model name/path (default: Qwen/Qwen2.5-Omni-7B)"
echo " --server-port PORT Port for vLLM server (default: 8091)"
echo " --gradio-port PORT Port for Gradio demo (default: 7861)"
echo " --deploy-config PATH Path to custom deploy config YAML file (optional)"
echo " --stage-configs-path PATH Path to custom stage configs YAML file (optional)"
echo " --server-host HOST Host for vLLM server (default: 0.0.0.0)"
echo " --gradio-ip IP IP for Gradio demo (default: 127.0.0.1)"
echo " --share Share Gradio demo publicly"
@@ -87,8 +87,8 @@ echo "=========================================="
# Build vLLM server command
SERVER_CMD=("vllm" "serve" "$MODEL" "--omni" "--port" "$SERVER_PORT" "--host" "$SERVER_HOST")
if [ -n "$DEPLOY_CONFIG_PATH" ]; then
SERVER_CMD+=("--deploy-config" "$DEPLOY_CONFIG_PATH")
if [ -n "$STAGE_CONFIGS_PATH" ]; then
SERVER_CMD+=("--stage-configs-path" "$STAGE_CONFIGS_PATH")
fi
# Function to cleanup on exit
@@ -336,7 +336,7 @@ vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
# or:
./moss_tts_nano/run_server.sh
```
The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--deploy-config`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--stage-configs-path`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
### Sending requests
```bash
-8
View File
@@ -71,14 +71,6 @@ plugins:
- search
- autorefs
- awesome-nav
- redirects:
redirect_maps:
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
user_guide/quantization/index.md: user_guide/quantization/overview.md
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
- glightbox
- git-revision-date-localized:
# exclude files
+2 -8
View File
@@ -70,7 +70,8 @@ dev = [
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests
"kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
"s3tokenizer==0.3.0", # for tts test
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
"step-audio2==1.0.0", # Step-Audio2 Token2Wav (flashcosyvoice / hyperpyyaml)
"stepaudio2-minicpmo", # for MiniCPM-o 4.5 talker vocoder (stepaudio2-minicpmo Token2wav)
]
demo = [
@@ -95,12 +96,6 @@ quack = [
"quack-kernels>=0.3.11",
]
# Optional Blackwell-native FlashAttention-4 backend. Keep this opt-in because
# the CuTe/CUTLASS dependency is CUDA 13 and Blackwell specific.
fa4 = [
"flash-attn-4[cu13]==4.0.0b18",
]
docs = [
"mkdocs>=1.5.0",
"mkdocs-api-autonav",
@@ -108,7 +103,6 @@ docs = [
"mkdocstrings-python",
"mkdocs-gen-files",
"mkdocs-awesome-nav",
"mkdocs-redirects",
"mkdocs-glightbox",
"mkdocs-git-revision-date-localized-plugin",
"mkdocs-minify-plugin",
+1 -1
View File
@@ -332,6 +332,6 @@ bundled offline CLI do not currently expose `sigmas`.
- <https://huggingface.co/Lightricks/LTX-2.3>
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
- [Request batching](../../docs/user_guide/diffusion/request_batching.md)
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
- [I2V offline example](../../examples/offline_inference/image_to_video/README.md)
-61
View File
@@ -1,61 +0,0 @@
# MiniMax-H3 on RTX 5090
This recipe uses BF16 weights, tiled VAE decode, tensor parallelism where a
second GPU is available, and distributed layerwise offload (DLO). It is a
memory-first serving configuration; lower resident counts reduce HBM use and
increase CPU-to-GPU transfer time.
## Capacity requirements
| Resource | One RTX 5090 | Two RTX 5090s |
| --- | ---: | ---: |
| GPU HBM | 32 GiB | 32 GiB per GPU |
| Checkpoint storage | 135 GiB per partition | 135 GiB per partition |
| Available system RAM | 200 GiB minimum | 200 GiB minimum |
| Recommended system RAM | 384 GiB | 384 GiB |
`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one
server at a time. DLO keeps rank-local weights in pinned host memory; increasing
`--dlo-resident-layers` improves latency but does **not** reduce host RAM in the
current implementation because resident layers retain pinned CPU master copies.
> **Modular H3:** after #5720 lands, preserve this recipe's one-partition
> behavior with `--task-type fl2va` or `--task-type ref2va`.
## One RTX 5090: 1344x768, 5 seconds
Use 12 resident DiT layers. A 50-step B300 allocation test with this exact
single-rank topology peaked at 26.50 GiB; re-measure peak HBM on the target
card before increasing the resident count.
```bash
CUDA_VISIBLE_DEVICES=0 vllm serve /path/to/MiniMax-H3/FL2VA \
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
--num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
--usp 1 --ring 1 --vae-patch-parallel-size 1 \
--vae-parallel-mode tile --vae-use-tiling \
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
--dlo-resident-layers 12 --enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
## Two RTX 5090s: 1344x768, 5 seconds
Use TP2 and 20 resident DiT layers. The two-rank B300 capacity run peaked at
27,726 MiB per rank for this shape and 50 steps. This is a memory/correctness
proxy, not a consumer-GPU latency claim.
```bash
CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/MiniMax-H3/FL2VA \
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
--num-gpus 2 --tensor-parallel-size 2 --text-encoder-tp-size 2 \
--usp 1 --ring 1 --vae-patch-parallel-size 2 \
--vae-parallel-mode tile --vae-use-tiling \
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
--dlo-resident-layers 20 --enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
For Ref2VA, stop the FL2VA server and restart the same command with
`/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video count and prompt length
can increase activation memory; begin with one request at a time.
-147
View File
@@ -1,147 +0,0 @@
# MiniMax H3 on Moore Threads MUSA
> Joint video and audio generation on Moore Threads GPUs
## Summary
- Vendor: MiniMaxAI
- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
- Tasks: T2VA, FL2VA, and Ref2VA
- Mode: OpenAI-compatible `/v1/videos` HTTP serving
- Hardware: 4x Moore Threads MTT S5000 for the validated Ref2VA profile
- Maintainer: Community
This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for MUSA environments.
MiniMax H3 has two independently served checkpoint partitions:
- `FL2VA` serves text-to-video+audio (`t2va`) and
first-frame-to-video+audio (`fl2va`).
- `Ref2VA` serves image/audio or video-reference generation (`ref2va`).
## Prerequisites
### Checkpoint
Only download the partition needed by the server. Downloading the whole model
repository also retrieves duplicate shared components and the original
top-level pipeline layout.
For T2VA and FL2VA:
```bash
python -m pip install modelscope
export MODEL_ROOT=/path/to/MiniMax-H3
modelscope download MiniMax/MiniMax-H3 \
--local_dir "${MODEL_ROOT}" \
--max-workers 16 \
--include 'FL2VA/**'
```
For Ref2VA, download `Ref2VA/**` instead.
### Environment
Install a compatible PyTorch/torch-musa, torchada, MATE, and vLLM-MUSA
stack before installing vLLM-Omni. Importing vLLM-Omni should report that the
`musa` vLLM platform plugin is active.
Install vLLM-Omni from a checkout containing MiniMax H3 support:
```bash
python -m pip install -e .
```
MiniMax H3 reference inputs and MP4 output require `soundfile`, `ffmpeg`, and
`ffprobe`. The Python dependency is installed by vLLM-Omni; install the two
executables with the operating system package manager and verify them:
```bash
python -c 'import soundfile; print(soundfile.__version__)'
ffmpeg -version
ffprobe -version
```
When torchaudio cannot load TorchCodec, vLLM-Omni automatically falls back to
soundfile. Formats that libsndfile cannot read are demuxed through ffmpeg.
## Start a server
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
and FL2VA, or to `Ref2VA` for Ref2VA.
The validated Ref2VA configuration uses tensor parallelism across four MTT
S5000 GPUs and offloads inactive model components to CPU:
```bash
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8091
MUSA_VISIBLE_DEVICES=0,1,2,3 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 4 \
--tensor-parallel-size 4 \
--vae-patch-parallel-size 4 \
--enable-cpu-offload \
--vae-use-tiling \
--diffusion-attention-backend FLASH_ATTN
```
On MUSA, `FLASH_ATTN` selects the MATE/FlashAttention-3 implementation. Do
not install the CUDA-only `fa4` optional dependency.
For request examples and parameter definitions, see
[MiniMax-H3.md](MiniMax-H3.md#http-api-examples).
## MUSA validation
The following weight-independent gates have been validated on one MTT S5000:
- packed variable-length attention selects MATE FlashAttention-3 and returns
finite BF16 output;
- MiniMax H3 RoPE casts CPU FP64 position metadata to FP32 before device
arithmetic and returns finite output;
- the video-VAE seeded RNG context is deterministic on MUSA and restores both
CPU and MUSA RNG states;
- the soundfile fallback preserves stereo channel layout, FP32 samples, and
the native sample rate when torchaudio/TorchCodec is unavailable.
Ref2VA image+audio generation was also validated end to end on one MTT S5000
with CPU offload and VAE tiling. The smoke used a 448x256 reference image, a
four-second 32 kHz stereo audio reference, two denoising steps, and seed 42.
The synchronous HTTP request completed in 50.6 seconds and returned an MP4
with the following decoded properties:
- 107 video frames at 448x256 and 24 FPS;
- non-static video (`pixel_std=13.07`, `temporal_delta=0.49`);
- finite stereo audio at 32 kHz (`audio_rms=0.0458`).
Official Ref2VA video-reference serving was validated with the expanded
MiniMax prompt at 1344x768, 5 seconds, seed 0, and 50 steps. A true TP4 run on
four MTT S5000 GPUs returned HTTP 200 in 638.6 seconds and produced 124 H.264
frames at 24 FPS with AAC stereo audio at 32 kHz.
Keep `--vae-use-tiling` enabled for this serving profile.
## Known limitations
- Keep Ring Attention at degree 1. The current Ring path does not preserve
MiniMax H3 packed padding boundaries.
- H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
- Each server process loads only one checkpoint partition.
- H3 currently executes one generation request per diffusion batch.
- FP8 quantization has not been enabled for MiniMax H3.
- MP3, M4A, MP4, and reference-video audio fallback requires `ffmpeg` on
`PATH`; WAV inputs can be read directly through soundfile.
## Additional resources
- [MiniMax-H3.md](MiniMax-H3.md) — full model and API guide
- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — Ascend NPU deployment guide
- [Supported models](../../docs/models/supported_models.md)
- [Video API](../../docs/serving/videos_api.md)
-162
View File
@@ -1,162 +0,0 @@
# MiniMax H3 on Ascend NPU
> Joint video and audio generation with text, first/last keyframes, and mixed
> image/video/audio references — Ascend NPU deployment guide
## Summary
- Vendor: MiniMaxAI
- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
- Tasks: T2VA, FL2VA, and Ref2VA
- Mode: OpenAI-compatible `/v1/videos` HTTP serving
- Hardware: Atlas 800I A3
- Maintainer: Community
This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for Ascend NPU
environments. Differences from the GPU path:
- Runs on a CPU-only (aarch64) PyTorch build with `torch_npu`; no CUDA
runtime is present.
- Audio loading does **not** require TorchCodec (whose aarch64 wheels are
built against CUDA torch and fail to load on CPU-only builds). vLLM-Omni
automatically falls back to soundfile / ffmpeg for wav/mp3/m4a/mp4 audio
inputs.
## Prerequisites
### Checkpoint
Same as the GPU recipe — Hugging Face access approval is required:
```bash
hf auth login
export MODEL_ROOT=/path/to/MiniMax-H3
hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
```
### Environment
- Ascend driver & firmware: 25.5.1
- CANN toolkit: 9.0.1
- Python: 3.12
- PyTorch: 2.10.0+cpu
- torch_npu: 2.10.0.post2
- Install vLLM-Omni from a checkout with MiniMax H3 support:
```bash
uv venv
source .venv/bin/activate
uv pip install -e .
```
Install the **mindie-sd** operator library to enable Ascend-optimized fused
operators (`adalayernorm`, etc.):
```bash
git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
# Comment out the tik_ops build step (not needed for this use case)
sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh
python setup.py bdist_wheel
cd dist
pip install mindiesd-*.whl
```
- `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
reference-video preparation and MP4 output.
- Reference-video decoding uses `decord` when available and falls back to
PyAV otherwise.
- Audio inputs do not require TorchCodec: wav/mp3/m4a/mp4 files are loaded
at native sample rate through the soundfile / ffmpeg fallback.
## Start a server
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
and FL2VA requests, or to `Ref2VA` for Ref2VA requests.
### Multi-NPU: 768P validated configuration
Validated on eight NPUs of an Atlas 800I A3 server with Ulysses sequence
parallelism degree 8, text-encoder tensor parallelism degree 8, native tiled
VAE patch parallelism degree 8, and layerwise offload:
```bash
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export PORT=9098
export MODEL="${MODEL_ROOT}/FL2VA"
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800
export PYTHONDONTWRITEBYTECODE=1
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 8 \
--usp 8 \
--ring 1 \
--text-encoder-tp-size 8 \
--enable-layerwise-offload \
--vae-parallel-mode tile \
--vae-use-tiling \
--vae-patch-parallel-size 8 \
--diffusion-attention-backend FLASH_ATTN
```
Do not add `--enforce-eager`. The first request includes regional
compilation; warm the server once before measuring steady-state latency.
H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
To serve Ref2VA, stop the FL2VA server and restart with:
```bash
export MODEL="${MODEL_ROOT}/Ref2VA"
```
## HTTP API examples
The request format is identical to the GPU recipe; see
[MiniMax-H3.md § HTTP API examples](MiniMax-H3.md#http-api-examples).
Use the validated 768P shapes (e.g. `width=1344 height=768`) on NPU.
## Key parameters
Same as the GPU recipe; see
[MiniMax-H3.md § Key parameters](MiniMax-H3.md#key-parameters).
The validated resolution on NPU is 768P (e.g. 1344x768).
## Validated NPU evidence
Measured on an Atlas 800I A3 server (8x NPU) with CANN 9.0.1,
PyTorch 2.10.0+cpu, and torch_npu 2.10.0.post2, using the multi-NPU
configuration above:
| Workload | Configuration |
|----------|---------------|
| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
These measurements describe the validated shapes rather than a general
throughput guarantee.
## Known limitations
- Each server process loads only one checkpoint partition.
- H3 currently executes one generation request per diffusion batch.
- The first regional-compile request is a warmup and should not be included
in steady-state performance measurements.
- The official H3 input matrix and media limits are documented in the [GPU
recipe](MiniMax-H3.md#official-input-matrix-and-limits); this NPU path uses
the same HTTP request contract.
- VAE patch parallelism requires size 1 or the full DiT group size and
supports the H3 native `tile` mode only.
## Additional resources
- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
- [Supported models](../../docs/models/supported_models.md)
- [Video API](../../docs/serving/videos_api.md)
-568
View File
@@ -1,568 +0,0 @@
# MiniMax H3
> Joint video and audio generation with text, first/last keyframes, and
> mixed image/video/audio references
## Summary
- Vendor: MiniMaxAI
- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
- Tasks: T2VA, FL2VA, and Ref2VA
- Mode: OpenAI-compatible `/v1/videos` HTTP serving
- Maintainer: Community
MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its
checkpoint has two independently served partitions:
- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
(`fl2va`)
- `Ref2VA`: up to 9 images, 3 videos, and 3 audio references in supported
image/video/audio combinations (`ref2va`; audio-only is rejected)
The generated MP4 contains H.264 video and synchronized stereo audio.
## Prerequisites
The checkpoint requires Hugging Face access approval. Download it and point
`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA`
subdirectories:
```bash
hf auth login
export MODEL_ROOT=/path/to/MiniMax-H3
hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
```
Install vLLM-Omni from the checkout containing MiniMax H3 support. The
two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need
FlashAttention-4. Install the optional dependency only for the four-GPU
B300/GB200 `FLASH_ATTN` profile:
```bash
uv venv
source .venv/bin/activate
uv pip install -e '.[fa4]'
```
`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
reference-video preparation and MP4 output.
## Start a server
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
### Memory and storage requirements
Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each
H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of
BF16 safetensors (about **135 GiB** on disk). Keeping both partitions locally
therefore needs roughly **270 GiB** of model storage, although the examples
serve only one partition at a time.
CPU offload and distributed layerwise offload reduce GPU residency; they do
not make the model weights disappear. With `--dlo-no-use-allgather`, each
worker retains its standard-loader rank-local weights in host memory, including
pinned CPU buffers used for H2D streaming. Use at least **200 GiB available
system RAM** before starting the two-GPU recipe; a **384 GiB host is
recommended** to leave room for the OS, CUDA/PyTorch allocations, request
inputs, and filesystem cache. Do not run the FL2VA and Ref2VA servers at the
same time on a host sized for this minimum.
The consumer-GPU profiles below are HBM budgets only. They still require the
host-RAM budget above.
### Single GPU: accuracy and memory first
The single-GPU configuration uses model-level CPU offload.
This matches the accuracy-qualified reference path and prevents the Qwen3-VL
encoder and DiT from being resident on the GPU at the same time.
```bash
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8091
CUDA_VISIBLE_DEVICES=0 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 1 \
--enable-cpu-offload \
--diffusion-attention-backend FLASH_ATTN
```
Use a GPU with enough memory for the active H3 component and enough system RAM
for the offloaded components. CPU offload reduces GPU memory pressure but adds
PCIe/NVLink transfer latency.
### Two 24/32 GB GPUs: TP2 distributed layerwise offload
For two PCIe consumer GPUs, combine TP2 with distributed layerwise offload
(DLO). The standard loader first creates the rank-local TP shard. DLO keeps
that shard in pinned host memory and streams the 30 tail DiT blocks through a
shared two-buffer window without a DP AllGather. The first 20 DiT blocks are
copied to the GPUs once per denoise stage, reused by every sampling step, and
released before VAE decode so the decoder can reuse their HBM.
```bash
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8091
CUDA_VISIBLE_DEVICES=0,1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 2 \
--tensor-parallel-size 2 \
--usp 1 \
--ring 1 \
--text-encoder-tp-size 2 \
--vae-patch-parallel-size 2 \
--vae-parallel-mode tile \
--vae-use-tiling \
--enable-distributed-layerwise-offload \
--dlo-no-use-allgather \
--dlo-resident-layers 20 \
--enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
Use the profile that matches the per-GPU memory capacity:
| Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
|---|---:|---:|---:|---|---|---|
| `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
| `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
This topology uses all available parallel capacity: TP2 shards both the DiT
and text encoder, `--dlo-no-use-allgather` streams each rank's local TP shard
without reconstructing full blocks, and VAE patch parallelism splits tiled
decode across both GPUs. cuDNN attention is selected explicitly for the RTX
consumer path; the server stays eager to avoid an unqualified compile path.
The resident count changes placement and transfer frequency only; it does not
quantize or change the BF16/FP32 denoise math. Re-measure peak memory before
increasing it on a different request shape.
### RTX 5090 target-hardware validation
At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
2 x RTX 5090 without OOM:
| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
|---:|---:|---:|---:|---|
| 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
This is a single end-to-end validation run, not a warmed multi-run latency
benchmark. The sampled `nvidia-smi` peak is also not a CUDA allocator
high-water mark. The environment used vLLM 0.26.0, vLLM-Omni
`0.26.1.dev14+gae6577ea`, and PyTorch 2.11.0+cu130. The
[run record](https://github.com/lishunyang12/vllm-omni-rankings/blob/dcd06d7e83cb069842535918c0169ee9f3f29ba0/scripts/%E5%BE%AE%E4%BF%A1%E5%9B%BE%E7%89%87_20260805000034_86_237.png)
captures the environment, output contract, elapsed time, and sampled peak.
Before the target run, both profiles were exercised on two B300 ranks as an
allocation and correctness proxy. At 1344x768, 124 frames, and 50 steps, the
20-layer profile peaked at 27,726 MiB per rank. At 1024x576, the 12-layer
profile peaked at 18,888 MiB per rank in a 5-step capacity run. The resident
and fully streamed placements produced identical decoded video-frame and audio
hashes for the same shape, step count, prompt, and seed. The B300 result does
not establish RTX 4090 PCIe latency; treat the 4090 profile as a conservative
starting point until it is measured on that GPU.
To run T2VA, FL2VA, image+audio Ref2VA, and two-video Ref2VA in order, validate
every MP4's H.264/AAC streams, and retain live server and GPU-memory logs:
```bash
RUN_ROOT=/path/to/run-root \
MODEL_ROOT=/path/to/MiniMax-H3 \
GPU_IDS=0,1 \
PROFILE=rtx5090 \
bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh
```
The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for
`PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default.
### Four GPUs: measured best-practice throughput
The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
- no CPU or layerwise offload;
- Ulysses sequence parallelism degree 4;
- native tiled VAE patch parallelism degree 4;
- regional `torch.compile` for the repeated DiT blocks;
- CuTe FlashAttention-4 through the `FLASH_ATTN` backend, with Ring and TP
left at 1.
```bash
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8091
CUDA_VISIBLE_DEVICES=0,1,2,3 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 4 \
--usp 4 \
--ring 1 \
--vae-patch-parallel-size 4 \
--vae-parallel-mode tile \
--vae-use-tiling \
--diffusion-attention-backend FLASH_ATTN
```
On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
`Using CuTe FlashAttention-4 on Blackwell` before recording measurements.
Do not add `--enforce-eager` to this performance configuration. The first
request includes regional compilation; warm the server once before measuring
steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
1. The H3 VAE supports its native `tile` mode, not
`spatial_shard_height` or `spatial_shard_width`.
### Text encoder tensor parallelism
The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
default fully resident on the DiT main rank. On multi-GPU no-offload runs that
rank becomes the peak-memory hotspot. Add `--text-encoder-tp-size N` to shard
the encoder across the first `N` DiT ranks (the encoder is implemented with
vLLM-style tensor-parallel layers and runs with distributed collectives over
its own encoder process group):
```bash
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8091
CUDA_VISIBLE_DEVICES=0,1,2,3 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 4 \
--usp 4 \
--ring 1 \
--text-encoder-tp-size 4 \
--vae-patch-parallel-size 4 \
--vae-parallel-mode tile \
--vae-use-tiling \
--diffusion-attention-backend FLASH_ATTN
```
`N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so
valid values on a 4-GPU server are 1, 2, and 4 (1, 2, 4, 8 on 8 GPUs). The
encoder TP rank set is the first `N` DiT ranks; on 4 GPUs
`--text-encoder-tp-size 4` shards the encoder 4-way, dropping the DiT main
rank's no-offload peak by roughly `(N-1)/N` of the ~51.5 GB encoder while the
other ranks each gain `~51.5/N` GB. The encoder output remains identical to
the reference path within bf16 rounding: every encoder rank all-reduces the
row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is
replicated on every rank.
To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU
command with:
```bash
export MODEL="${MODEL_ROOT}/Ref2VA"
```
### Online FP8 quantization
MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
uses dynamic activation scaling during inference. By default, attention and
MLP linears in the token refiner and main DiT blocks, the condition
projection, and all AdaLN projections use FP8. Patch, timestep, and final
projections remain FP32; the text encoder and VAEs are unchanged.
Add this option to an existing H3 server command:
```bash
--quantization fp8
```
Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
resolves the `transformer` component before constructing the DiT, so names do
not start with `transformer.`. Entries are exact runtime linear prefixes; a
parent name such as `blocks.0.attn` does not exclude its children.
Eligible names are the `attn.qkv_proj`, `attn.out_proj`, `mlp.fc1`, and
`mlp.fc2` children under `token_refiner.blocks.<0-1>` or `blocks.<0-49>`.
The other eligible names are `condition_proj`,
`blocks.<0-49>.adaln_proj.linear`, and `final_layer.adaln_proj.linear`.
For example, keep the first main block's attention projections in BF16 with:
```bash
--diffusion-quantization-config \
'{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
```
The structured option replaces `--quantization fp8`. Online FP8 is currently
incompatible with H3 layerwise offload because the offload path produces a
weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
parallelism and VAE tiling instead.
## HTTP API examples
The following requests use the synchronous endpoint so the returned body can
be saved directly as an MP4. The asynchronous `POST /v1/videos` endpoint can
also be used when job polling is preferred.
All four tasks use 24 FPS, 50 sigma points, seed values from the validated
workloads, and the checkpoint-reference video/audio flow shifts of 12 and 3.
Decimal durations are passed through `extra_params`.
Set the endpoint once:
```bash
export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
```
### 1. T2VA: text to video and audio
Run this request against the `FL2VA` partition:
```bash
curl -sS -X POST "${API_URL}" \
-F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
-F 'width=1344' \
-F 'height=768' \
-F 'aspect_ratio=16:9' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0}' \
-o t2va.mp4
```
### 2. FL2VA: first frame to video and audio
Run this request against the `FL2VA` partition. When width and height are
omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short
edge.
```bash
export FIRST_FRAME=/path/to/fl2va_first_frame.png
curl -sS -X POST "${API_URL}" \
-F 'prompt=A man stands beside a yellow car at night. The car drives away; he follows it with his eyes and begins singing sadly, with synchronized voice and city ambience.' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=2101' \
-F 'extra_params={"task":"fl2va","duration":8.7,"audio_flow_shift":3.0}' \
-F "input_reference=@${FIRST_FRAME};type=image/png" \
-o fl2va.mp4
```
Use the same `FL2VA` partition for the official tail-keyframe forms. A single
image with `frame_indices=[-1]` conditions the last frame; two ordered images
with `frame_indices=[0,-1]` condition the first and last frames:
```bash
export LAST_FRAME=/path/to/fl2va_last_frame.png
export FIRST_FRAME=/path/to/fl2va_first_frame.png
curl -sS -X POST "${API_URL}" \
-F 'prompt=The subject moves naturally from the first image to the last image.' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=2102' \
-F 'extra_params={"task":"fl2va","duration":8.7,"frame_indices":[0,-1],"audio_flow_shift":3.0}' \
-F "input_references=@${FIRST_FRAME};type=image/png" \
-F "input_references=@${LAST_FRAME};type=image/png" \
-o fl2va_first_last.mp4
```
### 3. Ref2VA: image-only, image/audio, or mixed references
Run these requests against the `Ref2VA` partition. Image-only Ref2VA omits
`audio_reference`; adding one or more audio references is optional. The typed
fields accept one object or an ordered JSON list. `audio_reference` accepts an
HTTP(S) URL or a `data:` URL.
```bash
python -m http.server 8092 \
--bind 127.0.0.1 \
--directory /path/to/reference_assets
```
Then submit an image-only request from another terminal:
```bash
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
curl -sS -X POST "${API_URL}" \
-F 'prompt=A white cat sits on a beige couch and slowly looks toward the camera.' \
-F 'aspect_ratio=adaptive' \
-F 'short_edge=768' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=3100' \
-F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
-F "input_reference=@${REF_IMAGE};type=image/png" \
-o ref2va_image_only.mp4
```
An image-plus-audio request is:
```bash
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
export AUDIO_URL=http://127.0.0.1:8092/ref2va_audio.mp3
curl -sS -X POST "${API_URL}" \
-F 'prompt=A white cat with black mustache and eyebrow markings sits on a beige couch, lip-syncing precisely to the complete reference audio before shifting from confusion to deadpan speechlessness.' \
-F 'width=1344' \
-F 'height=768' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=3101' \
-F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
-F "input_reference=@${REF_IMAGE};type=image/png" \
-F "audio_reference={\"audio_url\":\"${AUDIO_URL}\"}" \
-o ref2va_image_audio.mp4
```
The requested duration should cover the complete audio. If `duration` is
shorter, the reference soundtrack is truncated to the generated clip.
### 4. Ref2VA: video, separate audio, and mixed references
Run this request against the `Ref2VA` partition. Repeat the
`input_references` multipart field once per source video. H3 consumes the
videos in form order and preserves their original soundtracks during
conditioning.
```bash
export SUBJECT_VIDEO=/path/to/green_screen_subject.mp4
export BACKGROUND_VIDEO=/path/to/fairytale_background.mov
curl -sS -X POST "${API_URL}" \
-F 'prompt=Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2. Match the background motion to the character actions and relight the character to fit the scene.' \
-F 'width=1344' \
-F 'height=768' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=3101' \
-F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
-F "input_references=@${SUBJECT_VIDEO};type=video/mp4" \
-F "input_references=@${BACKGROUND_VIDEO};type=video/quicktime" \
-o ref2va_video_video.mp4
```
The server stores uploaded references only for the lifetime of the request and
deletes temporary files after generation. A video may use its embedded
soundtrack, a separate `audio_reference`, or both. To send a mixed multipart
request, repeat `input_references` for each image, video, or audio file; the
server classifies them by MIME type and preserves the per-type order.
Reference videos must be MP4/MOV with H.264/H.265 video, optional AAC/MP3
audio, 2–15 seconds each, and at most 15 seconds combined. They may still be
longer than the generated clip. Use `start_time_seconds` to select a
synchronized segment; for multiple typed video references, pass one value per
video in `extra_params.start_time_seconds`.
Reference images accept JPG/JPEG, PNG, WEBP, HEIC, or HEIF up to 30 MiB. Standalone
audio references accept WAV or MP3 up to 15 MiB, with 2–15 seconds per file and
at most 15 seconds combined.
## Official input matrix and limits
| Task | Supported references | Limits |
|------|----------------------|--------|
| T2VA | text only | prompt must be non-empty |
| FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` |
| Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference |
The H3 output contract is 4–15 seconds at 24 FPS, stereo 32 kHz audio, and a
32-pixel canvas multiple. T2VA requires one named output ratio from `21:9`,
`16:9`, `4:3`, `1:1`, `3:4`, or `9:16`. FL2VA always follows the first input
image's ratio and ignores a generic `aspect_ratio` override. Ref2VA defaults to
`16:9`; `adaptive` and SGLang's `auto` spelling are accepted aliases for that
default. `short_edge` controls the 768-pixel canvas and must be `768`.
`num_outputs_per_prompt` accepts 1–10 and derives each output seed as
`seed + output_index`. The asynchronous endpoint returns all
outputs; the synchronous raw-MP4 endpoint returns the first output when more
than one is requested.
## Key parameters
| Parameter | Recommended value | Notes |
|-----------|-------------------|-------|
| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition |
| `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
| `fps` | `24` | H3 output FPS is fixed |
| `num_inference_steps` | `50` | Matches the reference accuracy workloads |
| `flow_shift` | `12` | Video sigma shift |
| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
| `seed` | Task-specific | Use a fixed value for reproducibility |
| `aspect_ratio` | Task-specific | T2VA requires a named ratio; FL2VA follows the input image; Ref2VA defaults to `16:9` |
| `short_edge` | `768` | H3 shape policy requires exactly 768 when `width`/`height` are omitted |
| `num_outputs_per_prompt` | `1` | 1–10; async API returns every output |
| `start_time_seconds` | `0` | Reference-video segment start; use a list in `extra_params` for multiple videos |
| `width`, `height` | Multiples of 32 | Output aspect ratio must be between 1:4 and 4:1 |
## ComfyUI Frontend
Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 service. The ComfyUI frontend can run in a separate environment or machine. Refer to [vLLM-Omni ComfyUI Integration](../../docs/features/comfyui.md) for details.
## Validated four-GPU evidence
The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
excluded warmup followed by three requests.
| Workload | Configuration | Observed result |
|----------|---------------|-----------------|
| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
| Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
These measurements describe the validated shapes rather than a general
throughput guarantee. Multi-video Ref2VA is much slower because the two
reference videos expand both the Qwen3-VL vision sequence and the packed DiT
attention sequence.
## Validated FP8 evidence
With eager DiT/text-encoder TP2 and VAE tiling, the 384x672, 107-frame,
10-step quality case measured LPIPS 0.1156 (limit 0.20), PSNR 23.6316 dB,
audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
reduction.
## Known limitations
- Each server process loads only one checkpoint partition.
- H3 currently executes one generation request per diffusion batch.
- The first regional-compile request is a warmup and should not be included in
steady-state performance measurements.
- Online FP8 is not compatible with layerwise offload.
- 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.
- VAE patch parallelism requires size 1 or the full DiT group size and supports
the H3 native `tile` mode only.
## Additional resources
- [Supported models](../../docs/models/supported_models.md)
- [Video API](../../docs/serving/videos_api.md)
- [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md)
-174
View File
@@ -1,174 +0,0 @@
# PersonaPlex
> Live full-duplex speech-to-speech (talk to the model in your browser, it listens
> while it speaks) with persona and voice control, on the native vLLM-Omni engine.
## Summary
- Vendor: NVIDIA
- Model: [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
(gated; a Moshi finetune)
- Task: Full-duplex speech-to-speech. 24 kHz mic audio in, agent speech +
inner-monologue text out, in 80 ms lockstep. Persona (role text) and voice
(zero-shot clone) are set per session.
- Mode: Live duplex WebSocket server with the official browser client, single or
multi-session (`--batch-size`), plus an offline WAV-in / WAV-out example
- Maintainer: [`@linyueqian`](https://github.com/linyueqian)
## When to use this recipe
Use this recipe to run `nvidia/personaplex-7b-v1` as a live voice agent on one GPU.
The integration is **moshi-free**: the temporal transformer, depformer, and
streaming Mimi codec are all native vLLM-Omni modules (Mimi runs on
`transformers.MimiModel`), so no vendored fork is installed. Decoding is greedy,
matching the reference implementation frame for frame on the golden replays used
as the acceptance gate.
PersonaPlex is the first Moshi-class (pure-lockstep) model in the experimental
full-duplex framework (`vllm_omni/experimental/fullduplex/`), and the first with
elastic multi-session batching: one engine hosts `--batch-size` concurrent
conversations, callers join or leave on any 80 ms tick without disturbing live
slots (greedy isolation is bit-exact).
## References
- Serving package:
[`vllm_omni/experimental/fullduplex/personaplex/`](../../vllm_omni/experimental/fullduplex/personaplex/)
- Native model modules (temporal / depformer / streaming Mimi):
[`vllm_omni/model_executor/models/personaplex/`](../../vllm_omni/model_executor/models/personaplex/)
- Online example (server + headless client):
[`examples/online_serving/personaplex/`](../../examples/online_serving/personaplex/)
- Offline example:
[`examples/offline_inference/personaplex/personaplex_offline.py`](../../examples/offline_inference/personaplex/personaplex_offline.py)
- Staged offline pipeline (talker -> Mimi code2wav, async-chunk streaming):
[`vllm_omni/deploy/personaplex.yaml`](../../vllm_omni/deploy/personaplex.yaml)
- Integration PR: [vllm-project/vllm-omni#4771](https://github.com/vllm-project/vllm-omni/pull/4771)
- Framework context: [RFC #3745](https://github.com/vllm-project/vllm-omni/issues/3745)
(duplex adapter patterns), [#1335](https://github.com/vllm-project/vllm-omni/issues/1335)
(full-duplex target), [PR #3907](https://github.com/vllm-project/vllm-omni/pull/3907)
(fullduplex core contracts this adapter conforms to)
- Upstream model card:
[`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
## Why it is different from MiniCPM-o / JoyVL duplex
| | MiniCPM-o 4.5 (#3907) | JoyVL | **PersonaPlex (this)** |
|---|---|---|---|
| Cadence | 1 s chunk groups | ~1 fps frames | **80 ms lockstep (12.5 Hz)** |
| Turn control | learned `⟨listen⟩`/`⟨speak⟩` | `</silence>`/`</response>` | **none, pure lockstep** |
| Per step | variable-length token group | text decision | **1 user frame in -> 1 agent frame + 1 text token** |
| Barge-in | at chunk boundary | n/a | **native (model always hears the user)** |
| Session state | chunk-group KV | per-tick HTTP | **persistent ring KV + streaming Mimi state** |
This is the lockstep ("parallel-frame joint") shape of the duplex adapter
patterns: the adapter declares `DuplexCapability.continuous = True`, and
`core.DuplexRuntime` runs ONE eternal response that consumes input frames as they
arrive and drains on close. The flag is a small, default-off lifecycle mode, so
turn-based adapters (JoyVL, MiniCPM-o) are unaffected.
## Architecture (Moshi RQ-Transformer, all native)
- **Mimi codec** 24 kHz, 12.5 Hz, 1920 samples/frame, 8 active codebooks
(card 2048). Streaming encode + decode via `transformers.MimiModel` with the
PersonaPlex reference checkpoint weights.
- **Temporal transformer** (Helium backbone) 4096-d / 32 layers / 32 heads,
sliding window 3000 frames (a 4-minute rolling context) held in a ring KV cache
with per-slot recycle.
- **Depformer** 1024-d / 6 layers / 16 heads, autoregressive over 16 codebooks
per frame (8 vocoded), reset each frame.
- Per-frame token column `[B, 17, 1]`: row 0 = inner-monologue text,
rows 1-8 = agent audio, rows 9-16 = user audio.
- Persona + voice are injected once at session open through the same lockstep
step: voice clone forces the agent stream from reference-audio Mimi codes;
persona forces the text stream from `<system> ... <system>` tokens.
## Hardware Support
Verified on GPU (Hopper-class). The serving path is plain PyTorch eager plus the
native modules, so other CUDA GPUs with enough memory are expected to work.
## GPU
### 1x 141 GB Hopper-class GPU (H20-class, verified)
#### Environment
- OS: Linux
- Python: 3.10+ (CI covers 3.11 / 3.12)
- vLLM-Omni: PR #4771 branch or current `main` once merged
- Extra dependency for the live server: `pip install sphn` (Opus framing;
`aiohttp` already ships with vLLM)
- `export HF_TOKEN=...` with access to the gated
`nvidia/personaplex-7b-v1` repo (accept the license on the model page)
First run auto-downloads from the model repo: `model.safetensors`,
`tokenizer_spm_32k_3.model`, `voices.tgz`, and the Mimi reference checkpoint;
`kyutai/mimi` is fetched via `transformers` for the codec module graph.
#### Command
Live duplex server (official browser client served at `/`):
```bash
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
--port 8124 --voice NATF2.pt
```
Multi-session on the same GPU (elastic batching, callers join/leave any tick):
```bash
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
--port 8124 --batch-size 4
```
Then open `http://localhost:8124/` and allow the microphone (use headphones so
the agent does not hear itself). Headless alternative:
```bash
python examples/online_serving/personaplex/duplex_client.py \
--url ws://localhost:8124/api/chat --input user.wav --out reply.wav
```
Offline WAV-in / WAV-out (no server):
```bash
python examples/offline_inference/personaplex/personaplex_offline.py \
--input-wav question.wav --output-wav reply.wav --output-text reply.json \
--voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
```
#### Verification
```bash
# GPU-free contract tests (stubbed FrameStepper)
pytest tests/e2e/features/fullduplex/test_personaplex_adapter.py \
tests/e2e/features/fullduplex/test_personaplex_session.py \
tests/e2e/features/fullduplex/test_personaplex_batched.py -q
# GPU e2e: batched isolation + slot recycle (coherence check needs a real
# question WAV via PPLEX_QUESTION_WAV)
pytest tests/e2e/features/fullduplex/test_personaplex_elastic.py -q
```
A green offline run producing a coherent spoken reply (and inner-monologue text
in `reply.json`) is the quickest end-to-end check.
#### Notes
- Realtime budget is 80 ms/frame. Verified eager per-tick latency at
`--batch-size 4` is ~70-74 ms on this hardware class, i.e. all four
conversations stay realtime; single-session has comfortable headroom.
- Decoding is greedy only (temperature/top-k knobs are intentionally not
exposed): greedy is what the parity gates pin against the reference
implementation, and sampling amplifies sub-bit numeric drift into divergence.
- With `--batch-size N`, connections beyond N are rejected until a slot frees;
a recycled slot is bit-exact with a fresh engine for the same inputs.
- Voice prompts: bundled `voices.tgz` names (`NATF*`/`NATM*` natural,
`VARF*`/`VARM*` varied) or a path to your own `.pt`/`.wav`.
- Run the browser client near the server. The 80 ms cadence is sensitive to
network jitter; over a high-latency link playback can stutter regardless of
engine speed. On localhost it is smooth.
- Sliding window is 3000 frames (~4 min). The ring KV recycles beyond that;
very long sessions keep running with a rolling context.
-2
View File
@@ -45,8 +45,6 @@ recipes/
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
+4 -5
View File
@@ -164,11 +164,10 @@ curl -s http://localhost:8091/v1/chat/completions \
PR [#2495](https://github.com/vllm-project/vllm-omni/pull/2495) adds
performance CI configs for the same DiT-only settings. The CI step is
currently opt-in (gated by `RUN_HUNYUAN_IMAGE3_PERF=1`) with `soft_fail`
enabled, intended for initial data collection. Perf runners no longer
compare metrics against JSON baselines at assert time; the `baseline`
blocks in the configs are reference-only and are copied into result JSON
for reporting. Regression gates can be reintroduced once enough nightly
data has been collected.
enabled, intended for initial data collection. Performance assertions are
skipped (`skip-performance-assertion: true`); the baseline values in the
JSON configs are reference-only and will be promoted to regression gates
once enough nightly data has been collected.
The user-facing equivalent is to launch one of the CLI commands above and
generate 1024x1024 images with 50 denoising steps.
+10 -18
View File
@@ -56,14 +56,10 @@ vllm serve /path/to/Cosmos3-Super \
--data-parallel-size 4
```
### Without DLO AllGather (rank-local weights, no DLO sharding)
### Without AllGather (full weights per rank, no sharding)
Use `--dlo-no-use-allgather` when you want each rank to stream the standard
loader's rank-local weights independently — no DLO AllGather synchronization,
no concurrent request requirement, but no DLO weight sharding and higher host
memory use. In pure DP, this is equivalent to independent full-weight
layerwise-offload replicas; with TP or HSDP, each rank retains the standard
loader's local tensor layout.
Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
no AllGather synchronization, no concurrent request requirement, but N× host memory.
```bash
vllm serve /path/to/Cosmos3-Nano \
@@ -158,9 +154,9 @@ wait
| Flag | Description | Default |
|------|-------------|---------|
| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
| `--data-parallel-size N` | Number of DP ranks; DLO weight sharding + concurrent requests when AllGather is enabled | `1` |
| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
| `--dlo-no-use-allgather` | Stream standard-loader rank-local weights independently (no DLO sharding, no DLO AllGather) | `false` |
| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
## Mode comparison
@@ -170,7 +166,7 @@ wait
| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Standard-loader rank-local weights | 2 blocks | 0.4× HSDP | No DLO AllGather overhead, high CPU |
| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
## Memory expectations
@@ -189,11 +185,7 @@ wait
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
use `--dlo-no-use-allgather` or disable online quantization
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
`--dlo-no-use-allgather` retains the standard TP loader and should be treated
as an experimental compatibility path
- HSDP + AllGather is rejected (would double-shard weights). HSDP with
`--dlo-no-use-allgather` is accepted at configuration level, but is not yet a
broadly validated production combination
- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
`load_weights()` automatically
- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
- HSDP + AllGather is rejected (would double-shard weights)
-150
View File
@@ -423,156 +423,6 @@ ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_
once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
unknown keys for the model are dropped.
## ROCm
### 1x MI350X (Online serving)
#### Environment
- OS: Ubuntu 22.04+
- Python: 3.12+
- Driver / runtime: ROCm, HIP 7.2, **gfx942 or gfx950** (validated on gfx950 / MI350X,
which exposes 252 GiB of HBM to PyTorch)
- PyTorch: 2.11.0 (ROCm build)
- vLLM version: match the repository requirements from your current checkout
- vLLM-Omni version or commit: use the commit you are deploying from
No ROCm-specific patch or flag is needed: Cosmos3 runs on the stock code path.
#### Command
The NVIDIA command works unchanged. Guardrails need the gated
`nvidia/Cosmos-1.0-Guardrail` repo, so the quick path is to disable them (you are
then responsible for license compliance):
```bash
vllm serve nvidia/Cosmos3-Nano \
--omni \
--no-guardrails \
--host 0.0.0.0 --port 8000 \
--init-timeout 1800
```
For 720p video, add `--vae-use-tiling`: it cuts peak VRAM by ~68% for ~13%
latency, which is the difference between needing a 120 GiB card and a 40 GiB one.
`--quantization fp8` and `--enable-layerwise-offload` are also supported; see the
Notes for the measured cost of each. For extra GPUs use `--ulysses-degree N` or
`--tensor-parallel-size N` (not validated on ROCm yet).
#### Verification
Use the same `curl` requests as the CUDA section above. The ROCm images do not
ship `ffprobe`, so the CUDA section's `ffprobe` line does not work here. `imageio`
and `pyav` are already installed and report the same thing:
```bash
python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print(im.size, im.mode)"
python -c "import imageio.v3 as iio; print(iio.improps('cosmos3_t2v.mp4', plugin='pyav').shape)"
```
Expect `(1024, 1024) RGB` and `(189, 720, 1280, 3)` for the requests above.
Re-running the same request with the same seed reproduces a byte-identical mp4,
so an md5 comparison across two runs also works as a smoke check.
#### Notes
- **Warm up once per output shape before timing anything.** The first generation
of a shape pays a one-time cost that is cached on disk and reused by later
processes in the same container: a ~250 s aiter attention-kernel JIT build plus
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
**~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 |
| `--quantization fp8` | 150 s (−6.6%) | 107 GiB (−11%) | 82 GiB (−14%) | weights (online, no calibration) |
Reserved is the caching-allocator high water mark and runs ~2–25 GiB above
allocated; allocated is the figure to compare against another platform.
- **`--vae-use-tiling` is by far the biggest lever for video.** Of the 95 GiB
allocated by default only ~35 GiB is weights and fixed overhead; the remaining
**60 GiB is activations from decoding all 189 frames in one piece**, which is
why fp8 and offload cannot touch it and why tiling, which leaves under 5 GiB of
them, wins so much. `--vae-use-slicing` is a no-op for a single video request:
peak memory is bit-identical to the default.
- **`--enable-layerwise-offload` is close to free on this part:** per-block DiT
compute is long enough to hide the weight transfer, so on a card this size it
buys headroom for concurrency or higher resolution rather than the
speed-for-memory trade it is on small cards.
- **Attention backend:** the numbers above are aiter FlashAttention, not SDPA.
ROCm picks aiter when it is installed and the encoded device capability
satisfies `90 < major * 10 + minor < 100` (gfx942 and gfx950); otherwise it
falls back to `TORCH_SDPA` and says so only at debug level, which is easy to
miss. If you are reporting numbers, assert the choice instead of trusting the
default: `RocmOmniPlatform.get_diffusion_attn_backend_cls(None, 128)` returns
the resolved path.
- **Determinism is per-configuration:** same seed and same flags reproduce a
byte-identical mp4 across processes, but different configurations do not match
each other — including offload, which nominally only relocates weights, so do
not use it as a bit-exact regression baseline. Quality was not evaluated for
any configuration.
### 1x MI350X (Offline generation)
#### Environment
Same as the online serving section above.
#### Command
The standard task examples run unchanged:
```bash
# Text-to-image
python examples/offline_inference/text_to_image/text_to_image.py \
--model nvidia/Cosmos3-Nano \
--prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
--negative-prompt "blurry, distorted, low quality" \
--height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
--extra-body '{"flow_shift": 3.0, "guardrails": false}' \
--output cosmos3_t2i.png
# Text-to-video
python examples/offline_inference/text_to_video/text_to_video.py \
--model nvidia/Cosmos3-Nano \
--prompt "A robot arm is cleaning a plate in the kitchen." \
--negative-prompt "blurry, distorted, low quality" \
--height 720 --width 1280 --num-frames 189 --fps 24 \
--num-inference-steps 35 --guidance-scale 6.0 --flow-shift 10.0 \
--extra-body '{"max_sequence_length": 4096, "guardrails": false,
"use_resolution_template": false, "use_duration_template": false}' \
--output cosmos3_t2v.mp4
```
Both scripts print `Total generation time` and `Worker peak GPU memory
(reserved)`, which is where the latency and reserved figures in the Notes above
come from. Add `--vae-use-tiling`, `--quantization fp8` or
`--enable-layerwise-offload` to reproduce the other rows. The allocated column
needs `VLLM_LOGGING_LEVEL=DEBUG`, which enables a per-request line reporting
reserved, allocated and pool overhead together.
#### Verification
Same as the online serving section above.
#### Notes
- Run each command twice and read the second run; the first pays the one-time
compilation described above. If you run in a container, keep the container
alive between runs (`docker run --rm` throws the JIT cache away and you pay
the ~250 s kernel build every time).
- Latency and memory figures are identical to the online serving section — they
were taken from this path, because a single process has no client/server
timing ambiguity.
## NPU
### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.1 MiB

-18
View File
@@ -81,21 +81,6 @@ def _from_pipeline_key(
)
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
omni_config = _from_pipeline_key("personaplex")
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
deploy_path = tmp_path / "personaplex-turn.yaml"
deploy_path.write_text("session_mode: turn\n", encoding="utf-8")
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
pipeline = _resolve_pipeline_or_skip(model_type)
@@ -120,7 +105,6 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
engine_args = legacy_stage.yaml_engine_args
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
@@ -395,7 +379,6 @@ def test_runtime_config_fields_match_rfc_runtime_scope():
def test_sub_config_fields_match_rfc_scopes():
assert {f.name for f in fields(OmniStageModelConfig)} == {
"active_stream_window",
"duplex_max_sessions",
"enable_sleep_mode",
"default_sampling_params",
"subtalker_sampling_params",
@@ -457,7 +440,6 @@ def test_sub_config_fields_match_rfc_scopes():
"ulysses_mode",
"cfg_parallel_size",
"vae_patch_parallel_size",
"text_encoder_tp_size",
"vae_parallel_mode",
"use_hsdp",
"mask_sp_padding",
@@ -3,8 +3,7 @@
from __future__ import annotations
from collections import defaultdict
from collections.abc import Callable
from types import MethodType, SimpleNamespace
from types import SimpleNamespace
import numpy as np
import pytest
@@ -13,22 +12,9 @@ from vllm.v1.outputs import LogprobsLists
from vllm.v1.request import RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# Shared update_from_output helpers moved onto OmniSchedulerMixin; bind them on
# the lightweight stub so direct method calls still exercise the real paths.
_MIXIN_UPDATE_HELPERS = (
"_remove_stopped_requests_from_queues",
"_handle_failed_kv_load_outputs",
"_aggregate_kv_connector_stats",
"_publish_kv_cache_events",
"_attach_finished_request_sets",
"_attach_scheduler_stats",
"_capture_omni_connector_output",
)
def _logprobs(
token_ids: list[list[int]],
@@ -119,14 +105,18 @@ class _RequestQueue(list):
self.remove(request)
def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
"""Build a minimal stub for OmniARScheduler.update_from_output()."""
def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
"""Spec decode can sample tokens past a stop; the trailing tokens are
trimmed by _update_request_with_output, and the emitted logprob rows must
be trimmed with them (upstream slices logprobs after the trim)."""
request = _Request("req")
request.num_computed_tokens = 8
scheduler = SimpleNamespace(
perf_metrics=None,
connector=None,
chunk_transfer_adapter=None,
requests={request.request_id: request for request in requests},
running=list(requests),
requests={request.request_id: request},
running=[request],
waiting=_RequestQueue(),
skipped_waiting=_RequestQueue(),
structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
@@ -139,42 +129,7 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
finished_req_ids_dict=defaultdict(set),
kv_cache_manager=SimpleNamespace(take_events=lambda: None),
kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
recompute_kv_load_failures=False,
)
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.make_spec_decoding_stats = lambda *args, **kwargs: None
scheduler.make_stats = lambda *args, **kwargs: None
return scheduler
def _bind_request_lifecycle(
scheduler: SimpleNamespace,
*,
update_request: Callable,
handle_stopped: Callable | None = None,
) -> None:
def free_request(request):
scheduler.requests.pop(request.request_id, None)
scheduler.finished_req_ids.add(request.request_id)
scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
# vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
return None, None
scheduler._update_request_with_output = update_request
scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
scheduler._handle_stopped_request = handle_stopped or (lambda _request: True)
scheduler._free_request = free_request
def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
"""Spec decode can sample tokens past a stop; the trailing tokens are
trimmed by _update_request_with_output, and the emitted logprob rows must
be trimmed with them (upstream slices logprobs after the trim)."""
request = _Request("req")
request.num_computed_tokens = 8
scheduler = _make_scheduler_stub([request])
def update_request_trimming(req, token_ids):
# Mirror vllm Scheduler._update_request_with_output: the first sampled
@@ -184,7 +139,19 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
req.status = RequestStatus.FINISHED_STOPPED
return token_ids, True
_bind_request_lifecycle(scheduler, update_request=update_request_trimming)
def free_request(req):
scheduler.requests.pop(req.request_id, None)
scheduler.finished_req_ids.add(req.request_id)
scheduler.finished_req_ids_dict[req.client_index].add(req.request_id)
return None, None
scheduler._update_request_with_output = update_request_trimming
scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
scheduler._handle_stopped_request = lambda _request: True
scheduler._free_request = free_request
scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
scheduler.make_stats = lambda *args, **kwargs: None
scheduler._capture_omni_connector_output = lambda _output: None
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"req": 3},
@@ -223,14 +190,45 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
"""A bad runner response must not bring down unrelated batch requests."""
bad = _Request("bad")
good = _Request("good")
scheduler = _make_scheduler_stub([bad, good])
scheduler = SimpleNamespace(
perf_metrics=None,
connector=None,
chunk_transfer_adapter=None,
requests={bad.request_id: bad, good.request_id: good},
running=[bad, good],
waiting=_RequestQueue(),
skipped_waiting=_RequestQueue(),
structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
transfer_triggered_requests=set(),
active_kv_transfers=set(),
pending_stop_after_extraction=set(),
waiting_for_transfer_free=set(),
_new_prompt_len_snapshot={},
finished_req_ids=set(),
finished_req_ids_dict=defaultdict(set),
kv_cache_manager=SimpleNamespace(take_events=lambda: None),
kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
)
update_calls: list[str] = []
def update_request(request, token_ids):
update_calls.append(request.request_id)
return token_ids, False
_bind_request_lifecycle(scheduler, update_request=update_request)
def free_request(request):
scheduler.requests.pop(request.request_id, None)
scheduler.finished_req_ids.add(request.request_id)
scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
# vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
return None, None
scheduler._update_request_with_output = update_request
scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
scheduler._handle_stopped_request = lambda _request: True
scheduler._free_request = free_request
scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
scheduler.make_stats = lambda *args, **kwargs: None
scheduler._capture_omni_connector_output = lambda _output: None
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"bad": 1, "good": 1},
@@ -110,9 +110,6 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
session.status = RequestStatus.RUNNING
session.resumable = True
session.num_computed_tokens = len(session.prompt_token_ids)
# schedule() has counted this step's token as in flight; update_from_output()
# must settle it before handling the segment boundary.
session.num_in_flight_tokens = 1
sched = MagicMock()
sched.requests = {session.request_id: session}
@@ -121,20 +118,12 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
is_done_receiving_chunks=lambda _request_id: True,
segment_finished_requests={session.request_id},
)
def rearm_resumable_request(request: Request) -> bool:
request.status = RequestStatus.WAITING
return False
sched._handle_stopped_request.side_effect = rearm_resumable_request
sched._handle_stopped_request.return_value = False
sched.running = [session]
sched.waiting = MagicMock()
sched.skipped_waiting = MagicMock()
sched.structured_output_manager.should_advance.return_value = False
# Async scheduling can observe the segment boundary in the next schedule()
# before this model output is applied and defer the same request for finish.
# update_from_output() must not re-arm the resumable request twice.
sched._pending_finish_reqs = [session]
sched._pending_finish_reqs = []
sched.recompute_kv_load_failures = False
sched.connector = None
sched.kv_cache_manager.take_events.return_value = None
@@ -166,8 +155,6 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
output = outputs[session.client_index].outputs[0]
assert output.finish_reason is not None
assert output.is_segment_finished is True
sched._handle_stopped_request.assert_called_once_with(session)
assert sched._pending_finish_reqs == []
def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:
@@ -1,90 +0,0 @@
from collections import defaultdict
from types import SimpleNamespace
import pytest
from vllm.v1.engine import FinishReason
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
from vllm_omni.core.sched.output import OmniChunkRecvHandle
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _Scheduler(OmniSchedulerMixin):
pass
def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
calls = []
scheduler = _Scheduler()
scheduler.waiting = ["waiting"]
scheduler.running = ["running"]
scheduler.requests = {"request": object()}
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
scheduler.chunk_transfer_adapter = SimpleNamespace(
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
("process", waiting, running, scheduler_requests)
),
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
("restore-chunks", waiting, running, scheduler_requests)
),
)
scheduler.input_coordinator = SimpleNamespace(
restore_queues=lambda waiting: calls.append(("restore-full", waiting))
)
scheduler._process_pending_omni_inputs("ar")
scheduler._restore_omni_wait_queues()
assert calls == [
("consume", "ar"),
("timeouts",),
("process", scheduler.waiting, scheduler.running, scheduler.requests),
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
("restore-full", scheduler.waiting),
]
@pytest.mark.parametrize(
("synthesize_abort_outputs", "expected_finish_reason"),
[(False, None), (True, FinishReason.ABORT)],
)
def test_finished_request_attachment_keeps_ar_abort_policy_explicit(
synthesize_abort_outputs,
expected_finish_reason,
):
scheduler = _Scheduler()
scheduler.finished_req_ids_dict = defaultdict(set, {2: {"req-finished"}})
outputs = {}
scheduler._attach_finished_request_sets(
outputs,
synthesize_abort_outputs=synthesize_abort_outputs,
)
assert outputs[2].finished_requests == {"req-finished"}
if expected_finish_reason is None:
assert outputs[2].outputs == []
else:
assert outputs[2].outputs[0].finish_reason == expected_finish_reason
assert scheduler.finished_req_ids_dict == {}
def test_chunk_receive_handle_carries_minimal_registration_fields():
handle = OmniChunkRecvHandle(request_id="req", external_req_id="external")
assert handle.request_id == "req"
assert handle.external_req_id == "external"
def test_output_helper_preserves_required_nan_counter_default():
scheduler = _Scheduler()
request = SimpleNamespace(
request_id="req-output",
trace_headers=None,
take_events=lambda: [],
)
output = scheduler._make_omni_engine_output(request, new_token_ids=[])
assert output.num_nans_in_logits == 0
-316
View File
@@ -1,6 +1,3 @@
import math
from typing import NamedTuple
import pytest
import torch
@@ -320,103 +317,6 @@ def test_deferred_multimodal_cache_write_commits_on_completion():
assert torch.equal(rows[8:16], expected)
def test_discard_deferred_mm_outputs_drops_pending_chunks():
"""discard_deferred_mm_outputs is the abort path: a request that
leaves without a cache commit must not leak its staged chunks into
the cache if commit_deferred_mm_outputs is later called for it."""
cache = get_omni_pcache()
mm_key = "codes.audio"
block_table = torch.tensor(
[
[2, 3, 4, 5],
[6, 7, 8, 9],
],
dtype=torch.long,
)
input_batch = MockInputBatch(
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
block_table=block_table,
)
cache.update_omni_tensor_prefix_cache(
hidden_states=None,
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
num_tokens_unpadded=4,
slot_mapping=torch.arange(8, 12),
skip_mm_cache_keys={mm_key},
)
chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
cache.stage_deferred_mm_outputs(
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
input_batch=input_batch,
multimodal_outputs={mm_key: chunk},
num_scheduled_tokens={"req1": 4, "req2": 0},
deferred_mm_cache_keys={mm_key},
)
assert "req1" in cache._deferred_mm_outputs
cache.discard_deferred_mm_outputs("req1")
assert "req1" not in cache._deferred_mm_outputs
# A later commit call for the now-discarded req_id must be a no-op,
# not a resurrection of the dropped chunk.
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
assert torch.all(rows[8:12] == 0)
def test_deferred_mm_outputs_request_id_reused_after_discard():
"""A request ID reused after discard_deferred_mm_outputs must start
clean: the new occupant's committed data must not be mixed with the
previous (discarded) occupant's staged chunks."""
cache = get_omni_pcache()
mm_key = "codes.audio"
block_table = torch.tensor(
[
[2, 3, 4, 5],
[6, 7, 8, 9],
],
dtype=torch.long,
)
input_batch = MockInputBatch(
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
block_table=block_table,
)
cache.update_omni_tensor_prefix_cache(
hidden_states=None,
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
num_tokens_unpadded=4,
slot_mapping=torch.arange(8, 12),
skip_mm_cache_keys={mm_key},
)
stale_chunk = torch.full((4, 2), -1, dtype=torch.long)
cache.stage_deferred_mm_outputs(
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
input_batch=input_batch,
multimodal_outputs={mm_key: stale_chunk},
num_scheduled_tokens={"req1": 4, "req2": 0},
deferred_mm_cache_keys={mm_key},
)
cache.discard_deferred_mm_outputs("req1")
# "req1" is reused by a brand-new request with different data.
fresh_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
cache.stage_deferred_mm_outputs(
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
input_batch=input_batch,
multimodal_outputs={mm_key: fresh_chunk},
num_scheduled_tokens={"req1": 4, "req2": 0},
deferred_mm_cache_keys={mm_key},
)
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
assert torch.equal(rows[8:12], fresh_chunk)
def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
cache = get_omni_pcache()
mm_key = "codes.audio"
@@ -699,219 +599,3 @@ def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker):
assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim])
assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs)
@pytest.mark.parametrize(
"passthrough_shape",
[(5,), (5, 3)],
ids=["1d_token_metadata", "2d_feature_tensor"],
)
def test_get_merged_multimodal_outputs_slices_passthrough_tensor_per_request(passthrough_shape, mocker):
"""A passthrough key that is itself a real tensor (not the string/list
passthrough already covered above) must be sliced per request through
to_payload_element's tensor branch, under the same mixed cache
hit/miss + uneven scheduled-length batch as test_get_merged_multimodal_outputs.
Covers both a 1D token-aligned metadata tensor (e.g. ref_code_len,
codec_streaming -- shape (total_tokens,)) and a 2D per-token feature
tensor (shape (total_tokens, feat_dim)), since W2's batching could
easily special-case one and silently mishandle the other.
"""
feat_dims = {"foo": 100, "bar": 100}
cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN)
orig_num_tokens_unpadded = 8
slot_offset = 8
orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
orig_mm_outputs = {
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
}
cache.update_omni_tensor_prefix_cache(
hidden_states=None,
multimodal_outputs=orig_mm_outputs,
num_tokens_unpadded=orig_num_tokens_unpadded,
slot_mapping=orig_slot_mapping,
num_tokens_padded=None,
)
num_new_toks_req1 = 3
num_new_toks_req2 = 2
cache.add_prefix_cached_new_req_id("req1")
num_scheduled_tokens = {
"req1": num_new_toks_req1,
"req2": num_new_toks_req2,
}
new_mm_outputs = {}
for mm_key in cache.mm_cache_keys:
new_mm_outputs[mm_key] = torch.rand(
(num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]),
dtype=DTYPE,
)
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
assert passthrough_shape[0] == total_scheduled_tokens
passthrough_tensor = torch.arange(int(torch.prod(torch.tensor(passthrough_shape))), dtype=DTYPE).reshape(
passthrough_shape
)
new_mm_outputs["passthrough_tensor"] = passthrough_tensor
input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
mocker.patch(
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
new=fake_get_cached_block_ids,
)
merged_mm_outputs = cache.get_merged_multimodal_states(
query_start_loc=[0, num_new_toks_req1],
input_batch=input_batch,
multimodal_outputs=new_mm_outputs,
num_scheduled_tokens=num_scheduled_tokens,
)
assert "passthrough_tensor" not in cache.mm_cache_keys
passthrough_out = merged_mm_outputs["passthrough_tensor"]
# req1 (cache hit) gets tokens [0:3), req2 (cache miss) gets tokens [3:5) --
# a passthrough tensor is sliced by scheduled position, not merged with
# any cached history (unlike the mm_cache_keys tensors above).
assert torch.equal(passthrough_out["req1"], passthrough_tensor[0:num_new_toks_req1])
assert torch.equal(passthrough_out["req2"], passthrough_tensor[num_new_toks_req1:total_scheduled_tokens])
class MergeHarness(NamedTuple):
"""A primed cache plus the arguments get_merged_multimodal_states needs."""
cache: OmniTensorPrefixCache
input_batch: MockInputBatch
num_scheduled_tokens: dict[str, int]
query_start_loc: list[int]
total_scheduled_tokens: int
new_mm_outputs: dict[str, object]
def build_mixed_hit_miss_harness(mocker, num_new_toks_req1: int = 3, num_new_toks_req2: int = 2) -> MergeHarness:
"""Two requests over one cache key: req1 hits 8 cached tokens, req2 misses,
and their scheduled lengths differ -- the same batch shape as
test_get_merged_multimodal_outputs, which is what makes a per-request
boundary error observable. Callers add their passthrough keys to
``new_mm_outputs`` before calling get_merged_multimodal_states.
"""
cache = get_omni_pcache_with_mm_tensors({"foo": 100}, seq_len=DEFAULT_SEQ_LEN)
orig_num_tokens_unpadded = 8
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
cache.update_omni_tensor_prefix_cache(
hidden_states=None,
multimodal_outputs={
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
},
num_tokens_unpadded=orig_num_tokens_unpadded,
slot_mapping=torch.arange(8, 8 + orig_num_tokens_unpadded),
num_tokens_padded=None,
)
cache.add_prefix_cached_new_req_id("req1")
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
mocker.patch(
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
new=fake_get_cached_block_ids,
)
return MergeHarness(
cache=cache,
input_batch=MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])),
num_scheduled_tokens={"req1": num_new_toks_req1, "req2": num_new_toks_req2},
query_start_loc=[0, num_new_toks_req1],
total_scheduled_tokens=total_scheduled_tokens,
new_mm_outputs={
key: torch.rand((total_scheduled_tokens, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
},
)
def test_get_merged_multimodal_outputs_passes_per_request_tensor_list_through(mocker):
"""The passthrough shape qwen3-tts actually emits: a nested ``codes``/``meta``
dict whose ``codes.ref`` holds one tensor per request
(qwen3_tts_talker.make_omni_output), carried through to_payload_element's
pass_lists_through branch beside 1D int metadata that IS sliced per request.
Each request must receive the list *whole*. Per-request indexing happens
later, in the runner's _unwrap_lists; a W2 that split the list here would
make that second index select a different request's reference codes. The
existing list coverage passes a list of strings, which cannot catch a
tensor-specific regression, and no other test covers an int dtype.
"""
harness = build_mixed_hit_miss_harness(mocker)
ref_codes = [torch.tensor([11, 12], dtype=torch.long), torch.tensor([21], dtype=torch.long)]
ref_code_len = torch.arange(harness.total_scheduled_tokens, dtype=torch.int32)
harness.new_mm_outputs["codes"] = {"ref": ref_codes}
harness.new_mm_outputs["meta"] = {"ref_code_len": ref_code_len}
merged = harness.cache.get_merged_multimodal_states(
query_start_loc=harness.query_start_loc,
input_batch=harness.input_batch,
multimodal_outputs=harness.new_mm_outputs,
num_scheduled_tokens=harness.num_scheduled_tokens,
)
assert "codes" not in harness.cache.mm_cache_keys
for req_id in ("req1", "req2"):
req_refs = merged["codes"][req_id]["ref"]
assert len(req_refs) == len(ref_codes)
for got, want in zip(req_refs, ref_codes):
assert torch.equal(got, want)
assert got.dtype == torch.long
# Per-request clones: the branch exists so one request's downstream mutation
# cannot reach another request's payload or the source list.
merged["codes"]["req1"]["ref"][0] += 100
assert torch.equal(merged["codes"]["req2"]["ref"][0], torch.tensor([11, 12], dtype=torch.long))
assert torch.equal(ref_codes[0], torch.tensor([11, 12], dtype=torch.long))
# Token-aligned 1D metadata is still split per request, unlike the list.
req1_len = harness.num_scheduled_tokens["req1"]
assert torch.equal(merged["meta"]["req1"]["ref_code_len"], ref_code_len[:req1_len])
assert torch.equal(merged["meta"]["req2"]["ref_code_len"], ref_code_len[req1_len:])
assert merged["meta"]["req1"]["ref_code_len"].dtype == torch.int32
@pytest.mark.parametrize(
"passthrough_shape",
[(3, 4), (1, 4), (7,)],
ids=["shorter_2d", "single_row_2d", "longer_1d"],
)
def test_get_merged_multimodal_outputs_broadcasts_non_token_aligned_passthrough(passthrough_shape, mocker):
"""A passthrough tensor whose first dimension is not the scheduled-token
count is request-invariant: to_payload_element hands every request the whole
tensor, cloned, rather than slicing it.
W2 batching most plausibly breaks by assuming every passthrough tensor is
token-aligned; that assumption would slice these and give each request a
truncated view, or an empty one once start exceeds the first dimension.
"""
harness = build_mixed_hit_miss_harness(mocker)
assert passthrough_shape[0] != harness.total_scheduled_tokens
passthrough = torch.arange(math.prod(passthrough_shape), dtype=DTYPE).reshape(passthrough_shape)
expected = passthrough.clone()
harness.new_mm_outputs["speaker_embedding"] = passthrough
merged = harness.cache.get_merged_multimodal_states(
query_start_loc=harness.query_start_loc,
input_batch=harness.input_batch,
multimodal_outputs=harness.new_mm_outputs,
num_scheduled_tokens=harness.num_scheduled_tokens,
)
assert "speaker_embedding" not in harness.cache.mm_cache_keys
out = merged["speaker_embedding"]
for req_id in ("req1", "req2"):
assert out[req_id].shape == passthrough_shape
assert torch.equal(out[req_id], expected)
out["req1"] += 100
assert torch.equal(out["req2"], expected)
assert torch.equal(passthrough, expected)
-119
View File
@@ -1,119 +0,0 @@
import pytest
import torch
from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
# `cpu` is the lane selector, not a hardware claim: the only Buildkite steps that
# reach tests/ select `core_model and cpu`, and they run on GPU hardware (l4_1).
# Marking this `cuda` collects it nowhere. Real hardware need is the skipif below.
pytestmark = [
pytest.mark.core_model,
pytest.mark.cpu,
pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device"),
]
NUM_BLOCKS = 10
BLOCK_SIZE = 4
HIDDEN_SIZE = 2
DTYPE = torch.float32
def get_omni_pcache() -> OmniTensorPrefixCache:
"""Build an OmniTensorPrefixCache, but don't init mm tensors."""
return OmniTensorPrefixCache(
num_blocks=NUM_BLOCKS,
block_size=BLOCK_SIZE,
hidden_size=HIDDEN_SIZE,
hs_dtype=DTYPE,
)
def test_drain_ready_async_writes_is_a_noop_before_the_event_is_ready(monkeypatch):
"""drain_ready_async_writes() must not touch the cache while the
scheduled D2H's event has not completed, and must consume it exactly
once after it has.
Real GPU timing (e.g. via torch.cuda._sleep to stall the copy stream)
was tried and found unreliable on this box: cross-stream wait_stream/
wait_event did not reliably delay event.query() from reporting ready
immediately, likely a virtualization/driver quirk of this cloud GPU.
Forcing the event's query() result directly tests the branching logic
in drain_ready_async_writes/_consume_pending_write deterministically,
independent of real CUDA scheduling behavior.
"""
cache = get_omni_pcache()
mm_key = "codes.audio"
cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4)
num_tokens = 4
slot_mapping_gpu = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda")
mm_gpu = {
mm_key: torch.arange(num_tokens * HIDDEN_SIZE, dtype=torch.float32, device="cuda").reshape(
num_tokens, HIDDEN_SIZE
)
}
cache.schedule_async_write(
hidden_states_gpu=None,
multimodal_outputs_gpu=mm_gpu,
slot_mapping_gpu=slot_mapping_gpu,
num_tokens_unpadded=num_tokens,
num_tokens_padded=num_tokens,
)
pending_event = cache._pending_write.event
rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE)
monkeypatch.setattr(pending_event, "query", lambda: False)
assert cache.drain_ready_async_writes() == 0
assert torch.all(rows[8 : 8 + num_tokens] == 0)
assert cache._pending_write is not None
monkeypatch.setattr(pending_event, "query", lambda: True)
assert cache.drain_ready_async_writes() == 1
assert torch.equal(rows[8 : 8 + num_tokens], mm_gpu[mm_key].cpu())
assert cache._pending_write is None
# Nothing pending anymore: a second drain is a no-op, not a re-scatter.
assert cache.drain_ready_async_writes() == 0
def test_schedule_async_write_consumes_previous_pending_write_first():
"""Scheduling step N+1's write must first consume (scatter) step N's
pending write, even if the caller never explicitly called
drain_ready_async_writes() in between -- otherwise step N's data would
be silently dropped when overwritten by step N+1's pending write."""
cache = get_omni_pcache()
mm_key = "codes.audio"
cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4)
num_tokens = 4
first_slots = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda")
first_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 1.0, device="cuda")}
second_slots = torch.arange(12, 12 + num_tokens, dtype=torch.int64, device="cuda")
second_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 2.0, device="cuda")}
cache.schedule_async_write(
hidden_states_gpu=None,
multimodal_outputs_gpu=first_gpu,
slot_mapping_gpu=first_slots,
num_tokens_unpadded=num_tokens,
num_tokens_padded=num_tokens,
)
torch.accelerator.synchronize()
# This call's internal _consume_pending_write() must flush the FIRST
# write before the second one is scheduled; the caller never drained.
cache.schedule_async_write(
hidden_states_gpu=None,
multimodal_outputs_gpu=second_gpu,
slot_mapping_gpu=second_slots,
num_tokens_unpadded=num_tokens,
num_tokens_padded=num_tokens,
)
torch.accelerator.synchronize()
cache.drain_ready_async_writes()
rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE)
assert torch.all(rows[8:12] == 1.0)
assert torch.all(rows[12:16] == 2.0)
+146 -8
View File
@@ -1,15 +1,16 @@
import copy
import json
import os
import re
import subprocess
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Any
import pytest
from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import modify_stage_config
@@ -118,6 +119,32 @@ def create_unique_server_pytest_params(
]
def create_benchmark_pytest_params(
benchmark_configs: list[dict[str, Any]],
server_to_benchmark_mapping: dict[str, dict],
) -> list[Any]:
"""Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks."""
marks_by_name = _marks_by_test_name(benchmark_configs)
params: list[Any] = []
seen: set[str] = set()
for config in benchmark_configs:
test_name = config["test_name"]
if test_name in seen:
continue
seen.add(test_name)
params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
for idx, id_suffix in enumerate(id_suffixes):
params.append(
pytest.param(
(test_name, idx),
marks=marks_by_name.get(test_name, []),
id=f"{test_name}-{id_suffix}",
)
)
return params
def create_paired_benchmark_pytest_params(
server_entries: list[tuple[Any, str]],
params_by_test_name: dict[str, list[Any]],
@@ -330,6 +357,10 @@ def supports_video_generation(model_name: str) -> bool:
return any(key in lower for key in ("wan", "video", "i2v", "t2v"))
def supports_chat_generation(model_name: str) -> bool:
return not supports_video_generation(model_name)
def parse_stage_devices(stage_config_path: str) -> str:
text = Path(stage_config_path).read_text(encoding="utf-8")
raw_devices: list[str] = re.findall(r"^\s*devices:\s*\"?([0-9,\s]+)\"?\s*$", text, flags=re.MULTILINE)
@@ -358,6 +389,48 @@ def assert_fault_exception(exc: Exception, error_keywords: tuple[str, ...]) -> N
assert any(key in text for key in error_keywords), f"unexpected error under fault injection: {exc}"
def wait_chat_request_ready(host: str, port: int, model: str, timeout_sec: int = 180) -> None:
"""Poll a minimal chat request until success."""
deadline = time.time() + timeout_sec
payload = json.dumps(
{
"model": model,
"messages": [{"role": "user", "content": "Say hello in one short sentence."}],
"stream": False,
"modalities": ["text"],
}
)
last_error: str | None = None
while time.time() < deadline:
try:
status, body = post_chat_completions_raw(host, port, payload)
if status == 200:
return
last_error = f"http={status} body={body[:200]!r}"
except Exception as exc: # noqa: BLE001
last_error = str(exc)
time.sleep(2)
raise TimeoutError(f"runtime-teardown warmup request did not succeed within {timeout_sec}s: {last_error}")
def assert_no_extra_worker_processes(
baseline_pids: set[int],
worker_pattern: str,
timeout_sec: int = 60,
) -> None:
"""Ensure no extra worker PIDs remain after teardown."""
deadline = time.time() + timeout_sec
while time.time() < deadline:
current = set(list_remote_process_pids_by_pattern(worker_pattern))
extra = current - baseline_pids
if not extra:
return
time.sleep(2)
current = set(list_remote_process_pids_by_pattern(worker_pattern))
extra = sorted(current - baseline_pids)
assert not extra, f"orphan worker processes remain after container teardown: {extra}"
def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]:
mapping = {}
for config in configs:
@@ -564,6 +637,57 @@ def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
return get_runtime_resource_label()
def resolve_baseline_value(
baseline_raw: Any,
*,
sweep_index: int | None,
max_concurrency: Any = None,
request_rate: Any = None,
) -> Any:
"""Pick the baseline threshold for this sweep step."""
if baseline_raw is None:
return 100000
if isinstance(baseline_raw, dict):
if max_concurrency is not None:
for key in (max_concurrency, str(max_concurrency)):
if key in baseline_raw:
return baseline_raw[key]
if request_rate is not None:
for key in (request_rate, str(request_rate)):
if key in baseline_raw:
return baseline_raw[key]
raise KeyError(
f"baseline dict has no key for max_concurrency={max_concurrency!r} "
f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
)
if isinstance(baseline_raw, (list, tuple)):
if sweep_index is None:
raise ValueError("list baseline requires sweep_index")
if not (0 <= sweep_index < len(baseline_raw)):
raise IndexError(f"baseline list len={len(baseline_raw)} has no index {sweep_index}")
return baseline_raw[sweep_index]
return baseline_raw
def _baseline_thresholds_for_step(
baseline_data: dict[str, Any],
*,
sweep_index: int | None = None,
max_concurrency: Any = None,
request_rate: Any = None,
) -> dict[str, Any]:
"""Resolve baseline config to one threshold per metric for this iteration."""
return {
metric_name: resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
request_rate=request_rate,
)
for metric_name, baseline_raw in baseline_data.items()
}
def run_benchmark(
args: list[str],
test_name: str,
@@ -572,17 +696,18 @@ def run_benchmark(
num_prompt: int,
*,
baseline_config: dict[str, Any] | None = None,
sweep_index: int | None = None,
request_rate: Any | None = None,
max_concurrency: Any | None = None,
random_input_len: Any | None = None,
random_output_len: Any | None = None,
resource_label: str | None = None,
) -> dict[str, Any]:
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the raw
``baseline`` object from the test config (when ``baseline_config`` is provided),
including any hardware-nested map (e.g. ``{"H100": {...}}``). Downstream report
tooling can select a hardware bucket when needed. If the benchmark exits without
writing a result file, ``result_omni_template.json`` is used as a fallback.
After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the resolved per-metric thresholds
(when ``baseline_config`` is provided). If the benchmark exits without writing a result file,
``result_omni_template.json`` is used as a fallback.
"""
current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
ri = _safe_filename_token(random_input_len)
@@ -649,8 +774,12 @@ def run_benchmark(
result = json.load(f)
if baseline_config:
# Persist the JSON baseline as-is (no hardware / sweep resolution).
result["baseline"] = copy.deepcopy(baseline_config)
result["baseline"] = _baseline_thresholds_for_step(
baseline_config,
sweep_index=sweep_index,
request_rate=request_rate,
max_concurrency=max_concurrency,
)
else:
result["baseline"] = {}
if random_input_len is not None:
@@ -671,3 +800,12 @@ def pytest_addoption(parser: pytest.Parser) -> None:
default=None,
help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
)
parser.addoption(
"--assert-baseline",
action="store_true",
default=False,
help=(
"When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config "
"(default: off)."
),
)
+60 -6
View File
@@ -15,10 +15,15 @@ from tests.dfx.conftest import (
get_runtime_resource_label,
is_diffusion_perf_config,
load_benchmark_configs,
resolve_baseline_value,
run_benchmark,
)
from tests.helpers.runtime import OmniServer
# Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline``
# (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are
# defined in the same module.
#
# Optional JSON field ``mark`` is applied as pytest marks via
# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
@@ -157,8 +162,36 @@ def benchmark_params(request):
}
def assert_result(result, num_prompt) -> None:
def assert_result(
result,
params,
num_prompt,
*,
assert_baseline: bool,
sweep_index: int | None = None,
max_concurrency: Any | None = None,
request_rate: Any | None = None,
) -> None:
assert result["completed"] == num_prompt, "Request failures exist"
if not assert_baseline:
return
baseline_data = params.get("baseline", {})
for metric_name, baseline_raw in baseline_data.items():
current_value = result[metric_name]
baseline_value = resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
request_rate=request_rate,
)
if "throughput" in metric_name:
if current_value <= baseline_value:
print(
f"ERROR: Throughput test results were below baseline: {metric_name}: {current_value} > {baseline_value}"
)
else:
if current_value >= baseline_value:
print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}")
@pytest.mark.benchmark
@@ -167,7 +200,7 @@ def assert_result(result, num_prompt) -> None:
paired_benchmark_params,
indirect=["omni_server", "benchmark_params"],
)
def test_performance_benchmark(omni_server, benchmark_params):
def test_performance_benchmark(omni_server, benchmark_params, request):
test_name = benchmark_params["test_name"]
params = benchmark_params["params"]
dataset_name = params.get("dataset_name", "")
@@ -179,6 +212,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
print(f"Running benchmark for model: {model}")
print(f"Benchmark parameters: {benchmark_params}")
assert_baseline = request.config.getoption("--assert-baseline", default=False)
resource_label = get_runtime_resource_label()
def to_list(value, default=None):
@@ -237,7 +271,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
break
# QPS / request-rate sweep
for qps, num_prompt in zip(qps_list, num_prompt_list):
for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
result = run_benchmark(
args=args,
@@ -246,14 +280,24 @@ def test_performance_benchmark(omni_server, benchmark_params):
dataset_name=dataset_name,
num_prompt=num_prompt,
baseline_config=params.get("baseline"),
sweep_index=i,
request_rate=qps,
max_concurrency=None,
random_input_len=params.get("random_input_len"),
random_output_len=params.get("random_output_len"),
resource_label=resource_label,
)
assert_result(result, num_prompt)
assert_result(
result,
params,
num_prompt,
assert_baseline=assert_baseline,
sweep_index=i,
request_rate=qps,
)
# concurrency test
for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
result = run_benchmark(
args=args,
@@ -262,8 +306,18 @@ def test_performance_benchmark(omni_server, benchmark_params):
dataset_name=dataset_name,
num_prompt=num_prompt,
baseline_config=params.get("baseline"),
sweep_index=i,
request_rate=None,
max_concurrency=concurrency,
random_input_len=params.get("random_input_len"),
random_output_len=params.get("random_output_len"),
resource_label=resource_label,
)
assert_result(result, num_prompt)
assert_result(
result,
params,
num_prompt,
assert_baseline=assert_baseline,
sweep_index=i,
max_concurrency=concurrency,
)
+103 -12
View File
@@ -13,6 +13,8 @@ A config JSON file may be passed via --test-config-file. If omitted, every ``*.j
pytest run_diffusion_benchmark.py -m "diffusion"
pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off).
Optional JSON field ``mark`` is applied as pytest marks on that case via
``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).
@@ -23,7 +25,6 @@ of all runs from cases in that file). Bulk load without ``--test-config-file`` u
the same per-file aggregation; ``-m`` only selects which cases run.
"""
import copy
import json
import os
import socket
@@ -273,6 +274,11 @@ def _write_result_record(record: dict[str, Any]) -> Path:
return target
def _append_to_aggregated_file(record: dict[str, Any]) -> None:
"""Backward-compatible wrapper; prefer :func:`_write_result_record`."""
_write_result_record(record)
_server_lock = threading.Lock()
# ---------------------------------------------------------------------------
@@ -672,7 +678,7 @@ def run_benchmark(
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
tmp_result_file = Path(tmp.name)
exclude_keys = {"baseline", "dataset", "task", "name"}
exclude_keys = {"baseline", "dataset", "task", "name", "skip-performance-assertion"}
cmd = [
sys.executable,
@@ -766,17 +772,12 @@ def run_benchmark(
completed = metrics.get("completed_requests", metrics.get("completed", 0))
failed = metrics.get("failed_requests", metrics.get("failed", 0))
# Persist the JSON baseline as-is (no hardware / sweep resolution).
baseline = copy.deepcopy(params.get("baseline") or {})
metrics["baseline"] = baseline
record: dict[str, Any] = {
"test_name": test_name,
"endpoint": endpoint,
"timestamp": timestamp,
"server_params": server_cfg.get("server_params"),
"benchmark_params": params,
"baseline": baseline,
"result": metrics,
"log_file": str(log_file),
"Model": model,
@@ -826,14 +827,70 @@ def run_benchmark(
# ---------------------------------------------------------------------------
def _resolve_baseline_value(
baseline_raw: Any,
*,
sweep_index: int | None,
max_concurrency: Any = None,
request_rate: Any = None,
) -> Any:
"""Pick the baseline threshold for this sweep step."""
if baseline_raw is None:
return 100000
if isinstance(baseline_raw, dict):
if max_concurrency is not None:
for key in (max_concurrency, str(max_concurrency)):
if key in baseline_raw:
return baseline_raw[key]
if request_rate is not None:
for key in (request_rate, str(request_rate)):
if key in baseline_raw:
return baseline_raw[key]
raise KeyError(
f"baseline dict has no key for max_concurrency={max_concurrency!r} "
f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
)
if isinstance(baseline_raw, (list, tuple)):
if sweep_index is None:
raise ValueError("sweep_index is required when baseline is a list or tuple")
return baseline_raw[sweep_index]
return baseline_raw
def assert_result(
result: dict[str, Any],
params: dict[str, Any],
num_prompts: int,
*,
sweep_index: int | None = None,
max_concurrency: Any = None,
request_rate: Any = None,
assert_baseline: bool = True,
) -> None:
"""Assert that the expected number of requests completed."""
"""Assert that benchmark metrics satisfy the configured baselines."""
completed = result.get("completed_requests", result.get("completed", 0))
assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
if not assert_baseline:
return
if params.get("skip-performance-assertion", False):
print("Skipping performance assertions.")
return
for metric, baseline_raw in params.get("baseline", {}).items():
current = result.get(metric)
assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}"
threshold = _resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
request_rate=request_rate,
)
if "throughput" in metric:
assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}"
else:
assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}"
def _default_benchmark_endpoint_for_task(task: str) -> str:
"""Return the default client-side benchmark endpoint for a diffusion task."""
@@ -862,6 +919,7 @@ def _build_run_params(
params: dict[str, Any],
*,
num_prompts: int,
sweep_index: int | None = None,
request_rate: Any | None = None,
max_concurrency: Any | None = None,
) -> dict[str, Any]:
@@ -873,6 +931,16 @@ def _build_run_params(
run_params["request-rate"] = request_rate
if max_concurrency is not None:
run_params["max-concurrency"] = max_concurrency
if "baseline" in params:
run_params["baseline"] = {
metric: _resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
request_rate=request_rate,
)
for metric, baseline_raw in params["baseline"].items()
}
return run_params
@@ -895,28 +963,36 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
sweep_runs: list[dict[str, Any]] = []
for request_rate, num_prompts in zip(request_rate_list, num_prompt_list):
for i, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
sweep_runs.append(
{
"params": _build_run_params(
params,
request_rate=request_rate,
num_prompts=num_prompts,
sweep_index=i,
),
"num_prompts": num_prompts,
"sweep_index": i,
"request_rate": request_rate,
"max_concurrency": None,
}
)
for max_concurrency, num_prompts in zip(max_concurrency_list, num_prompt_list):
for i, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
sweep_runs.append(
{
"params": _build_run_params(
params,
max_concurrency=max_concurrency,
num_prompts=num_prompts,
sweep_index=i,
request_rate="inf",
),
"num_prompts": num_prompts,
"sweep_index": i,
"request_rate": None,
"max_concurrency": max_concurrency,
}
)
@@ -929,6 +1005,9 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
num_prompts=default_num_prompts,
),
"num_prompts": default_num_prompts,
"sweep_index": None,
"request_rate": None,
"max_concurrency": None,
}
)
@@ -944,12 +1023,14 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
paired_benchmark_params,
indirect=["diffusion_server", "benchmark_params"],
)
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
"""Run the diffusion performance benchmark and verify request completion.
One server is started per unique parallel configuration (module scope).
For each server, all benchmark parameter sets defined in the config JSON
are executed sequentially; metrics are recorded to the aggregated result file.
Pass ``--assert-baseline`` to compare ``throughput_qps``, ``latency_mean``, etc. to the JSON ``baseline`` block.
"""
test_name = benchmark_params["test_name"]
params = benchmark_params["params"]
@@ -992,4 +1073,14 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
print(f"\n Aggregated results: {_aggregated_result_file_for_source(str(source))}")
print("=" * 60)
assert_result(result, sweep_run["num_prompts"])
assert_baseline = request.config.getoption("--assert-baseline", default=False)
assert_result(
result,
params,
sweep_run["num_prompts"],
sweep_index=sweep_run["sweep_index"],
max_concurrency=sweep_run["max_concurrency"],
request_rate=sweep_run["request_rate"],
assert_baseline=assert_baseline,
)
+16 -24
View File
@@ -40,12 +40,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"H100": {
"throughput_qps": 0.495,
"latency_mean": 2.0238,
"peak_memory_mb_max": 29393.4286,
"peak_memory_mb_mean": 29392.7143
}
"throughput_qps": 0.495,
"latency_mean": 2.0238,
"peak_memory_mb_max": 29393.4286,
"peak_memory_mb_mean": 29392.7143
}
}
]
@@ -92,12 +90,10 @@
"enable-negative-prompt": true,
"num-input-images": 1,
"baseline": {
"H100": {
"throughput_qps": 0.3451,
"latency_mean": 2.9031,
"peak_memory_mb_max": 30480.2857,
"peak_memory_mb_mean": 30479.2857
}
"throughput_qps": 0.3451,
"latency_mean": 2.9031,
"peak_memory_mb_max": 30480.2857,
"peak_memory_mb_mean": 30479.2857
}
}
]
@@ -143,12 +139,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"H100": {
"throughput_qps": 0.5048,
"latency_mean": 1.987,
"peak_memory_mb_max": 29380.0,
"peak_memory_mb_mean": 29380.0
}
"throughput_qps": 0.5048,
"latency_mean": 1.987,
"peak_memory_mb_max": 29380.0,
"peak_memory_mb_mean": 29380.0
}
}
]
@@ -195,12 +189,10 @@
"enable-negative-prompt": true,
"num-input-images": 1,
"baseline": {
"H100": {
"throughput_qps": 0.149,
"latency_mean": 6.7157,
"peak_memory_mb_max": 30574.0,
"peak_memory_mb_mean": 30574.0
}
"throughput_qps": 0.149,
"latency_mean": 6.7157,
"peak_memory_mb_max": 30574.0,
"peak_memory_mb_mean": 30574.0
}
}
]
@@ -24,12 +24,10 @@
"warmup-num-inference-steps": 28,
"enable-negative-prompt": true,
"baseline": {
"H100": {
"throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
"latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
"peak_memory_mb_max": 40500,
"peak_memory_mb_mean": 40500
}
"throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
"latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
"peak_memory_mb_max": 40500,
"peak_memory_mb_mean": 40500
}
}
]
@@ -24,12 +24,10 @@
"warmup-num-inference-steps": 28,
"enable-negative-prompt": true,
"baseline": {
"H100": {
"throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
"latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
"peak_memory_mb_max": 40500,
"peak_memory_mb_mean": 40500
}
"throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
"latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
"peak_memory_mb_max": 40500,
"peak_memory_mb_mean": 40500
}
}
]
@@ -48,11 +48,9 @@
"guardrails": false
},
"baseline": {
"H100": {
"throughput_qps": 1.107,
"latency_mean": 0.8929,
"peak_memory_mb_mean": 80000.0
}
"throughput_qps": 1.107,
"latency_mean": 0.8929,
"peak_memory_mb_mean": 80000.0
}
}
]
@@ -117,11 +115,9 @@
"use_duration_template": false
},
"baseline": {
"H100": {
"throughput_qps": 0.027,
"latency_mean": 36.6824,
"peak_memory_mb_mean": 30155.4
}
"throughput_qps": 0.027,
"latency_mean": 36.6824,
"peak_memory_mb_mean": 30155.4
}
}
]
@@ -185,11 +181,9 @@
"guardrails": false
},
"baseline": {
"H100": {
"throughput_qps": 0.027,
"latency_mean": 36.968,
"peak_memory_mb_mean": 30443.6
}
"throughput_qps": 0.027,
"latency_mean": 36.968,
"peak_memory_mb_mean": 30443.6
}
}
]
@@ -257,11 +251,9 @@
"condition_video_keep": "first"
},
"baseline": {
"H100": {
"throughput_qps": 0.027,
"latency_mean": 37.2117,
"peak_memory_mb_mean": 30135.6
}
"throughput_qps": 0.027,
"latency_mean": 37.2117,
"peak_memory_mb_mean": 30135.6
}
}
]
+15 -19
View File
@@ -46,14 +46,12 @@
"seed_tts_locale": "en",
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
223.3584
],
"median_audio_rtf": [
0.4393
]
}
"median_audio_ttfp_ms": [
223.3584
],
"median_audio_rtf": [
0.4393
]
}
},
{
@@ -72,17 +70,15 @@
"seed_tts_locale": "en",
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
276.575
],
"median_audio_rtf": [
0.5073
],
"audio_throughput": [
27.3728
]
}
"median_audio_ttfp_ms": [
276.575
],
"median_audio_rtf": [
0.5073
],
"audio_throughput": [
27.3728
]
}
}
]
@@ -5,8 +5,7 @@
{
"hardware_marks": {
"res": {
"cuda": "H100",
"npu": "A3"
"cuda": "H100"
},
"num_cards": 4
}
@@ -28,9 +27,16 @@
"async_chunk": false,
"trust_remote_code": true,
"connectors": {
"shared_mem_connector": {
"name": "SharedMemoryConnector",
"extra": {}
"rdma_connector": {
"name": "MooncakeTransferEngineConnector",
"extra": {
"host": "auto",
"zmq_port": 50051,
"protocol": "rdma",
"device_name": "",
"memory_pool_size": 17179869184,
"memory_pool_device": "cpu"
}
}
},
"stages": [
@@ -60,7 +66,7 @@
"need_send_cache": true
},
"output_connectors": {
"to_stage_1": "shared_mem_connector"
"to_stage_1": "rdma_connector"
},
"default_sampling_params": {
"temperature": 0.0,
@@ -90,7 +96,7 @@
},
"vae_use_tiling": true,
"input_connectors": {
"from_stage_0": "shared_mem_connector"
"from_stage_0": "rdma_connector"
},
"default_sampling_params": {
"num_inference_steps": 50,
@@ -133,14 +139,13 @@
"bot-task": "think_recaption",
"max-concurrency": 16,
"num-prompts": 96,
"skip-performance-assertion": true,
"baseline": {
"H100": {
"throughput_qps": 0.1874,
"latency_p99": 101.852,
"latency_mean": 79.0385,
"peak_memory_mb_max": 70000,
"peak_memory_mb_mean": 70000
}
"throughput_qps": 0.1874,
"latency_p99": 101.852,
"latency_mean": 79.0385,
"peak_memory_mb_max": 70000,
"peak_memory_mb_mean": 70000
}
}
]
Loaded 100 of 263 files, more files were not shown because too many files have changed in this diff. Show more