Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 1c3d59796f [Diffusion] Document DLO mmap model opt-in
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-15 06:57:38 +08:00
Hongsheng Liu 76e257b0f3 [Diffusion] Address rank-local mmap review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-14 17:12:05 +08:00
Hongsheng Liu 2956af1e9d [Diffusion] Add rank-local mmap storage for DLO
Decouple node-local checkpoint storage from the DLO AllGather transfer protocol. Compatible rank-local workers retain shared file-backed sources and stream complete blocks through bounded staging without changing request scheduling.

Refs #6195

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-14 14:38:06 +08:00
zhengjiaandCursor 5d50bdc2ab [Perf] Cap diffusion worker thread count before spawn (#6165)
Signed-off-by: zjli2013 <leezhengjiang@126.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-14 12:14:55 +08:00
Samit 8a5bf2fb3f [doc]: Update latest news section for verl-omni release (#6187)
Signed-off-by: Samit <285365963@qq.com>
2026-08-14 12:04:15 +08:00
596c16a550 [Kernel] Fuse Q/K RMSNorm and RoPE (#5990)
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-14 00:46:59 +00:00
Lucas(Ruijie) 90a08c0eda [Model][TTS][WIP] Add dots.tts (rednote-hilab) — continuous-AR 48kHz TTS (#4765)
Signed-off-by: Moore-Z <zhumoore9661996@gmail.com>
2026-08-13 17:58:33 +00:00
Sy03 dbc0dd6d85 [Bugfix][Higgs-Audio-V3] Disable XQA decode (#6068)
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Sysy <1370724210@qq.com>
2026-08-13 10:00:50 -07:00
Mu GuanLin 83d1f11449 [Feat] LTX Standard Two-Stage Pipeline (#5500)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-14 00:13:14 +08:00
Yukim1andHongsheng Liu 4b66f62e10 [Refactor] Remove internal stage_configs_path plumbing (#5741)
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 21:29:03 +08:00
Weiming Liao f41451f64c [NPU] Upgrade to v0.27.0 (#6096)
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
2026-08-13 20:01:01 +08:00
Hongsheng Liu bbe6ccc512 [Doc] Refresh community documentation (#6141)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 17:04:06 +08:00
Canlin Guo bbb709cfb5 [BugFix] Use MOSS-TTS-Local official sample params (#6156)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-13 08:56:26 +00:00
146 changed files with 9571 additions and 1589 deletions

No files matched your search

+8 -2
View File
@@ -456,7 +456,7 @@ steps:
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
@@ -477,10 +477,16 @@ steps:
commands:
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
+1
View File
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*Latest News* 🔥
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
- [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.
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0
ARG VLLM_ASCEND_TAG=v0.27.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0-a3
ARG VLLM_ASCEND_TAG=v0.27.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0
ARG VLLM_ASCEND_TAG=v0.27.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0-a3
ARG VLLM_ASCEND_TAG=v0.27.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+2 -1
View File
@@ -163,7 +163,8 @@ nav:
- CLI Reference: cli
- Community:
- Governance: community/governance.md
- community/*
- Contact Us: community/contact_us.md
- Community Meetings: community/meetings.md
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
- Forum: https://discuss.vllm.ai
+7 -3
View File
@@ -1,5 +1,9 @@
# Contact Us
- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature
- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues).
- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/).
- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md).
## Security Disclosures
Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request.
+3 -23
View File
@@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
## Diffusion Workload Division
## Reviewer Routing
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).
### Runtime-related
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
### Optimization
- Parallel: @RuixiangMa, @wtomin
- Attention backend: @ZJY0516
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
- Offload: @yuanheng-zhao
- Intel GPU support: @xuechendi
- Quantization: @lishunyang12, @david6666666
- Kernel fusion / communication-computation: @ZJY0516
### Docs & Test
- Diffusion docs: @RuixiangMa, @wtomin
## Meetings
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights.
## Committer Nomination Process
+8
View File
@@ -0,0 +1,8 @@
# Community Meetings
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth:
- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0)
- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0)
Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/).
-1
View File
@@ -1 +0,0 @@
# Meetups
-12
View File
@@ -1,12 +0,0 @@
# Volunteers for Bugfix and CI
We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
|----------|----------|----------|----------|----------|----------|
| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
We kindly welcome more contributors to fix bugs and contribute new features!
+2 -2
View File
@@ -3,7 +3,7 @@
Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
!!! note
We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. See [Community Meetings](../community/meetings.md) for current schedules, access details, agendas, and past notes.
## Getting Started
@@ -122,7 +122,7 @@ Repository skills guide the agent, but they do not replace contributor judgment,
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
!!! important
If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
Do not report suspected security vulnerabilities through a public issue, pull request, discussion, or Slack channel. Follow the [security disclosure instructions](../community/contact_us.md#security-disclosures) to arrange a private report.
## Pull Requests & Code Reviews
@@ -5,6 +5,8 @@ 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.
The rank-local mmap storage design is tracked in
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
For user-facing commands, see the
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
@@ -16,8 +18,10 @@ 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.
additional DLO collective is undesirable. For explicitly mmap-compatible
models in a TP1, non-HSDP topology, this path shares file-backed DiT checkpoint
pages across processes and keeps only bounded staging buffers private. TP,
HSDP, and other models retain standard-loader rank-local host storage.
The compatibility matrix below describes the current implementation. The
unit-level guards are covered, but not every parallelism combination has a
@@ -49,6 +53,10 @@ 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.
Multi-rank AllGather behavior is unchanged. When the effective DLO group size
is one, the backend now uses rank-local mmap staging; this closes a previous
loader/backend mismatch that could leave parameters on the `meta` device.
```text
Compute: [Block N] [Block N+1] [Block N+2]
H2D: [shard N+1] [shard N+2]
@@ -68,9 +76,18 @@ 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.
one and streams complete runtime blocks using H2D copies only. Host storage is
selected independently from that transfer protocol:
- For an explicitly mmap-compatible model with TP1 and no HSDP, each worker
retains immutable safetensors views backed by the node's shared OS page
cache. Two pinned host buffers, sized to the largest streamed block, stage
mmap pages for asynchronous H2D. Checkpoint-to-runtime adapters run while a
block is packed, so transformed weights do not become a persistent private
model copy.
- With TP, HSDP, online quantization, or a model without mmap support, the
regular loader remains responsible for preparing rank-local tensors. This
preserves TP/HSDP/quantization loader callbacks.
This mode means:
@@ -80,19 +97,26 @@ This mode means:
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.
- Pure DP deployments with a compatible mmap checkpoint share the streamed
DiT's file-backed pages. Other components and fallback paths may still keep
rank-private host weights.
- 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.
The mmap source remains immutable in both transfer modes. AllGather copies a
persistent `1 / group_size` shard from it and then closes the mappings;
rank-local mode keeps the mappings open and copies a complete block through
bounded staging without adding a collective.
The current implementation packs mmap views into a staging slot synchronously
on the host; the subsequent pinned H2D transfer overlaps with block compute.
This keeps host memory bounded but makes CPU packing bandwidth part of the
rank-local path's performance profile.
## 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. |
| **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 collective. Compatible TP1 models share file-backed DiT pages across replicas. |
| **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. |
@@ -122,11 +146,35 @@ AllGather DP multi-concurrency requires:
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.
The mmap loader is used when the model explicitly declares a compatible
checkpoint layout. It is independent of AllGather for TP1/non-HSDP execution.
Models can declare a per-parameter checkpoint-to-runtime adapter; the adapter
must preserve dtype and element count so it can be applied while packing one
staging block. MiniMax-H3 uses this mechanism for grouped QKV. Online
quantization and TP/HSDP no-AllGather execution use the regular loader path.
### Adding mmap support to another model
A new pipeline should opt in only after its ordinary loader contract can be
represented without materializing a private full-model copy:
1. Populate `weights_sources` with safetensors component sources and prefixes
that place every DiT checkpoint key in the pipeline parameter namespace.
2. If prefixed checkpoint keys already equal runtime parameter names, declare
`_supports_mmap_loading = True`. If names differ, provide
`_remap_ckpt_key` to map checkpoint names to runtime names (or return
`None` for intentionally ignored keys).
3. If a parameter needs a TP1 layout conversion, attach a callable
`mmap_weight_transform` to that parameter during model construction. The
transform must preserve dtype and element count because it runs while one
bounded block is packed.
4. Verify that mmap and the regular loader produce the same runtime tensors,
that every expected DiT tensor maps successfully, and that TP, HSDP, and
online-quantized configurations retain their standard-loader behavior.
Setting `_supports_mmap_loading = True` alone is not sufficient: explicit
opt-ins fail during initialization if any expected DiT tensor remains on the
`meta` device after checkpoint mapping.
## Validation coverage
@@ -136,6 +184,8 @@ Current source-level validation includes:
- HSDP + DLO without AllGather acceptance at configuration level;
- TP rejection in the DLO+AllGather mmap path;
- resident-layer requests requiring no-AllGather;
- rank-local mmap source retention, layout transforms, and bounded host
staging;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
tests.
@@ -150,7 +200,8 @@ on the target CUDA/NCCL or CANN/HCCL hardware.
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.
- Use **no-AllGather** when replicas must avoid DLO collectives. TP1 models with
mmap support can share DiT checkpoint pages; TP/HSDP and unsupported models
retain the higher-memory standard-loader path.
- Prefer **HSDP alone** for production HSDP deployments until the combined
HSDP + DLO no-AllGather path has broader end-to-end coverage.
+3 -1
View File
@@ -42,7 +42,8 @@ th {
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
| `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` | ✅︎ | ✅︎ | | | — |
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `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` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -58,6 +59,7 @@ th {
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+3 -6
View File
@@ -10,15 +10,12 @@ A: Refer to [GPU memory calculation and configuration](../configuration/gpu_memo
> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
A: Search the current [issues](https://github.com/vllm-project/vllm-omni/issues) first. If no existing issue matches a bug or CI failure, open a [bug report](https://github.com/vllm-project/vllm-omni/issues/new?template=400-bug-report.yml) with a minimal reproducer, environment details, and relevant logs. For technical or usage questions, use the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or the [vLLM Forum](https://discuss.vllm.ai/).
> Q: Does vLLM-Omni support AWQ or any other quantization?
A: AWQ is not currently listed as a validated vLLM-Omni path. Quantization
support depends on the model, pipeline stage, checkpoint format, and hardware
backend. See the current [quantization support matrix](../user_guide/quantization/overview.md)
for online methods, pre-quantized checkpoints, and model-level validation.
A: AWQ is available through the inherited vLLM quantization registry, and vLLM-Omni supports several additional methods. Actual support and validation vary by model, component, hardware, and checkpoint. See the [quantization overview](../user_guide/quantization/overview.md) for validated combinations and method-specific guides.
> Q: Does vLLM-Omni support multimodal streaming input and output?
A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!
A: Yes, for specific models and serving paths. Current capabilities include [streaming video input](../serving/video_stream_api.md), [streaming diffusion output](../user_guide/diffusion/execution_modes.md#streaming-output), and experimental full-duplex real-time audio described in the [project overview](../README.md). Check the model and API documentation because input and output support is model-specific.
+25 -13
View File
@@ -189,7 +189,10 @@ multi-device deployments. In the default mode, each rank stores only
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.
weights independently and does not shard weights across ranks. For explicitly
mmap-compatible TP1 models, the same mode can retain node-shared file-backed
DiT weights and use bounded per-worker host staging instead of a private full
DiT copy.
For the implementation design and compatibility matrix, see
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
@@ -201,8 +204,9 @@ For the implementation design and compatibility matrix, see
time, regardless of model size
- **DP multi-concurrency**: N concurrent requests processed in parallel
(AllGather only gathers weight shards, which are request-independent)
- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
page cache, eliminating O(dp_size × model_size) RSS during model creation
- **mmap weight storage**: compatible checkpoints use views into the shared OS
page cache. AllGather copies persistent rank shards; rank-local mode keeps
the views and packs complete blocks through two bounded host staging slots
- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
(CANN/HCCL) via vLLM-Omni's platform abstraction
@@ -224,7 +228,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 DLO AllGather (compatible TP1 models use node-shared mmap storage)
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4 \
@@ -258,20 +262,26 @@ m = Omni(
### How model weights are loaded (mmap path)
When DLO + AllGather is active, the offloader:
When DLO and an mmap-compatible model are active, the offloader:
1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
from the normally-created transformer
2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
initialization weights
3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
which return views into the OS page cache (shared across all ranks, 0 RSS)
4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
(e.g. Cosmos3's `time_embedder` → FP32)
which return views into the OS page cache without eagerly copying the full
tensor into anonymous process memory
4. **Applies declared per-parameter layout adapters while packing**, when the
checkpoint and TP1 runtime layouts differ (for example MiniMax-H3 grouped
QKV)
5. **Restores non-persistent buffers** from saved copies
This approach requires **zero model-specific code changes** — no pipeline or
transformer modifications are needed.
With AllGather, each rank copies only its persistent shard and then releases
the source mappings. Without AllGather, mappings remain open and every worker
owns two pinned staging slots sized to the largest streamed block. The latter
shares file-backed pages, not process-local Python objects, and does not alter
request routing or scheduling. Packing into the host slot is synchronous; the
pinned H2D copy can overlap with computation.
### OffloadPlan (declarative topology metadata)
@@ -313,7 +323,7 @@ request while AllGather synchronizes only weight shards (request-independent).
`--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.
standard TP loader instead of rank-local mmap 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.
@@ -373,6 +383,7 @@ reasoner/generator components inside the model forward pass.
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
| MiniMaxH3Pipeline | `MiniMaxAI/MiniMax-H3` (`FL2VA`) | `MiniMaxH3DiTModel` | - | ✓ | ✓ | `"blocks"` |
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
@@ -382,8 +393,9 @@ reasoner/generator components inside the model forward pass.
- 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
additionally requires a compatible checkpoint/mmap layout. For explicitly
compatible TP1 models, no-AllGather uses node-shared mmap sources; other
layouts retain 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.
+1 -1
View File
@@ -92,7 +92,7 @@ current support.
|--------|--------------|-------------|----------|
| **[FP8](quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on validated transformer stages | Memory reduction, inference speedup |
| **[INT8](quantization/int8.md)** | `quantization="int8"` | INT8 W8A8 on validated transformer stages | Memory reduction, broad GPU compatibility |
| **[BitsAndBytes W4](quantization/bitsandbytes.md)** | `quantization="bitsandbytes"` | Online NF4/FP4 weight-only quantization on validated CUDA transformer stages | Reducing transformer weight memory without preparing a separate checkpoint |
| **[GGUF](quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
## Supported Models
@@ -143,8 +143,7 @@ python text_to_video.py \
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
- `--quantization`: online quantization method, for example `fp8`; supported
pre-quantized checkpoints are normally auto-detected from checkpoint metadata.
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
- `--flow-shift`: scheduler flow_shift parameter.
- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
- `--lora-scale`: scale factor for LoRA weights.
@@ -90,8 +90,8 @@ vllm serve "Qwen/Qwen-Image" \
--quantization-config '{"method": "fp8"}'
```
Other vLLM-Omni quantization methods, such as `modelopt`, `mxfp4`, `mxfp8`,
serialized checkpoints, static FP8 configs, and layer-name skip lists
Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`,
`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists
such as `ignored_layers`, are intentionally not translated by this compatibility
shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or
a native vLLM-Omni pipeline for those cases.
+21 -44
View File
@@ -10,10 +10,6 @@ checkpoint's `config.json` and auto-detects
AutoRound is static quantization: no `--quantization` flag is needed at
inference time when the checkpoint already contains the quantization config.
The validation terms on this page follow the
[support levels](overview.md#support-levels) defined in the overview. A listed
checkpoint is not enough by itself to establish end-to-end model support.
## Hardware Support
| Device | Support |
@@ -25,57 +21,46 @@ checkpoint is not enough by itself to establish end-to-end model support.
| Intel XPU | ✅ |
| Ascend NPU | ❌ |
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. This table is method-level backend availability. AutoRound is an Intel
project; vLLM selects a compatible CUDA or XPU compute backend at load time.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. AutoRound is Intel-supported.
## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2)
| Model | Checkpoint | Scope | Scheme | Validation |
|-------|------------|-------|--------|------------|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | Validated |
| Qwen-Image | `INC4AI/Qwen-Image-AutoRound-W4A16` | Diffusion transformer | W4A16 | Validated with end-to-end generation and quality comparison |
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Both diffusion transformers | W4A16 | CI-backed |
| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Both diffusion transformers | W4A16 | CI-backed |
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | Validated; not in the scheduled Wan AutoRound job |
CUDA execution uses GPTQ-Marlin or another compatible vLLM AutoRound backend;
Intel XPU uses the Intel-supported backend. The text encoder and VAE stay in
BF16 for all entries above.
| Model | Checkpoint | Scope | Scheme | Backend |
|-------|------------|-------|--------|---------|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
### World Model (Cosmos3)
| Model | Checkpoint | Scope | Scheme | Validation |
|-------|------------|-------|--------|------------|
| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | Integrated; no model-specific end-to-end test is maintained in-tree |
| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | Validated with a manual serving run |
The Cosmos VAE and guardrail components stay in BF16.
| Model | Checkpoint | Scope | Scheme | Backend |
|-------|------------|-------|--------|---------|
| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
| Model | Checkpoint | Scope | Scheme | Validation |
|-------|------------|-------|--------|------------|
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Thinker language-model stage | W4A16 | Validated outside scheduled CI |
| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | Validated outside scheduled CI |
| Model | Checkpoint | Scope | Scheme | Backend |
|-------|------------|-------|--------|---------|
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Language-model stage | W4A16 | AutoRound |
| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | AutoRound |
| Qwen3-TTS | Not listed | TTS language-model stage | W4A16 | Not validated |
AutoRound support is checkpoint-driven. A model is supported when its
checkpoint uses a compatible INC/AutoRound config and the target stage maps to
vLLM-Omni's runtime module names.
For Qwen-Omni, only the thinker language model uses the AutoRound checkpoint.
Audio and vision encoders, the talker, and waveform-decoder stages stay in
BF16.
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
| Model | Checkpoint | Scope | Validation |
|-------|------------|-------|------------|
| GLM-Image | `Intel/GLM-Image-int4-AutoRound` | Diffusion transformer | CI-backed |
| BAGEL | Not listed | Checkpoint-defined diffusion or transformer stage | Not validated |
| Model | Scope | Status | Notes |
|-------|-------|--------|-------|
| GLM-Image | Diffusion transformer | ✅ | `Intel/GLM-Image-int4-AutoRound` |
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
## Configuration
@@ -134,14 +119,6 @@ The checkpoint should contain a config like:
At load time, vLLM-Omni builds an `OmniINCConfig`, remaps checkpoint block names
to runtime module names, and selects the matching vLLM compute backend.
Checkpoint auto-detection and successful layer construction establish the
**Integrated** level. Promote a model to **Validated** only after the named
checkpoint completes an end-to-end generation or multimodal request and its
output passes a model-appropriate sanity or quality check. Compare against the
BF16 baseline for quality-sensitive image and video paths. Promote it to
**CI-backed** only when that named checkpoint is selected by scheduled hardware
CI.
Example checkpoint creation:
```bash
+6 -14
View File
@@ -9,10 +9,6 @@ stay in BF16/FP16. No pre-quantized checkpoint is required.
This is an online (dynamic) path: load the normal HuggingFace checkpoint and
quantize during model loading via the `bitsandbytes` CUDA kernels.
The validation terms on this page follow the
[support levels](overview.md#support-levels). BitsAndBytes kernel support does
not imply that every diffusion transformer has been quality-validated.
## Hardware Support
| Device | Support |
@@ -25,7 +21,7 @@ not imply that every diffusion transformer has been quality-validated.
| Intel XPU | ❌ |
| Ascend NPU | ❌ |
Legend: `✅` backend available, `❌` unsupported. Non-CUDA platforms raise in
Legend: `✅` supported, `❌` unsupported. Non-CUDA platforms raise in
`get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`).
Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
@@ -34,11 +30,11 @@ Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
### Diffusion Model (Qwen-Image, Wan2.2)
| Model | HF models | Scope | Level | Recommendation |
|-------|-----------|-------|-------|----------------|
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Diffusion transformer | Validated | All heavy linear layers; sensitive embedders stay BF16 |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Diffusion transformer | Not validated | Compare with BF16 before enabling |
| Wan2.2 | Wan2.2 diffusion pipelines | One or both diffusion transformers | Not validated | Validate quality and cascade behavior before enabling |
| Model | HF models | CUDA | Mode | Recommendation |
|-------|-----------|:----:|------|----------------|
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling |
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs |
Other diffusion models may work if their transformer uses supported linear
layers, but they are not validated in this guide.
@@ -108,10 +104,6 @@ Multi-GPU tensor parallelism (`tensor_parallel_size` > 1) is not validated for
BitsAndBytes in diffusion models; each rank quantizes its own weight shard
independently.
Z-Image-Turbo is validated outside scheduled full-model BitsAndBytes CI. The
real CUDA kernel smoke tests cover quantization and matrix multiplication, but
do not establish model-level quality for unlisted architectures.
If quality regresses, use `ignored_layers` to keep sensitive projections in BF16
(for example `to_out` or `w2`).
+3 -1
View File
@@ -80,12 +80,14 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2)
### Diffusion Models
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
+10 -21
View File
@@ -13,11 +13,6 @@ components, such as the tokenizer, scheduler, text encoder, vision/audio
encoder, and VAE, on the base checkpoint unless a model-specific guide says
otherwise.
The validation terms on this page follow the
[support levels](overview.md#support-levels) defined in the overview. ModelOpt
format detection and checkpoint adaptation are shared runtime integrations;
each model and stage still requires its own end-to-end validation.
!!! note
ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
`--quantization fp8` for these checkpoints. The checkpoint
@@ -63,11 +58,11 @@ stay unquantized unless the model-specific recipe validates otherwise.
| Intel XPU | ⭕ |
| Ascend NPU | ❌ |
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt
NVFP4 and mixed FP8/NVFP4 diffusion checkpoints are currently validated on
Blackwell CUDA systems in the recipes below; other CUDA generations require
separate backend and quality validation.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
CUDA systems in the recipes below; other CUDA generations require separate
backend and quality validation.
## Model Type Support
@@ -84,9 +79,6 @@ separate backend and quality validation.
| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
| Wan2.2 | Not available | Diffusion transformer | Not validated |
The validated diffusion entries above have named checkpoint and recipe
evidence. They are not all exercised by scheduled full-model CI.
For full serving commands and benchmark context, see
[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
and
@@ -96,8 +88,8 @@ and
| Model | Scope | Status |
|-------|-------|--------|
| Qwen3-Omni | Thinker language-model stage (FP8) | Validated outside scheduled CI |
| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | CI-backed for the H100 Marlin fallback; native Blackwell path is recipe-validated |
| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path |
| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | Validated; see [Qwen3-Omni NVFP4 W4A4](#qwen3-omni-nvfp4-w4a4-thinker) below |
| Qwen3-TTS | TTS language-model stage | Not validated |
Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless
@@ -110,9 +102,9 @@ Qwen3-Omni-30B-A3B-Instruct thinker language model. The thinker text body
(attention + MoE experts) is quantized to NVFP4 with FP8 per-tensor input
scales; the audio encoder, vision encoder, talker, and code2wav stay in BF16.
| Variant | HF checkpoint | Hardware and validation |
|---------|---------------|-------------------------|
| W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ native FlashInfer FP4 GEMM; H100 Marlin fallback is CI-backed |
| Variant | HF checkpoint | Hardware |
|---------|---------------|----------|
| W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ (Blackwell, FlashInfer FP4 GEMM) |
Calibration uses ModelOpt `mtq.NVFP4_DEFAULT_CFG` with the `awq_clip` algorithm
on 1024 ultrachat samples chat-templated through the Qwen3-Omni tokenizer.
@@ -240,6 +232,3 @@ omni = Omni(
5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
ModelOpt path only when the checkpoint already contains ModelOpt quantized
weights and scales.
6. The Qwen3-Omni full-model CI smoke exercises the H100 Marlin fallback. A
scheduled Blackwell test is still required before treating the native
FlashInfer FP4 kernel path as CI-backed.
+11 -31
View File
@@ -3,17 +3,12 @@
## Overview
[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend compression
toolkit for producing pre-quantized model checkpoints. vLLM-Omni consumes its
output either through a compatible generic Ascend checkpoint path or through
the native MXFP8 and MXFP4 merge workflows documented in this guide.
toolkit for producing pre-quantized model checkpoints. In vLLM-Omni, these
checkpoints run through the Ascend/NPU path with `--quantization ascend`.
msModelSlim is static quantization: quantized weights are generated offline
before vLLM-Omni inference starts.
The validation terms on this page follow the
[support levels](overview.md#support-levels). A producer recipe or hardware
backend alone is an integration point, not model-level validation.
## Hardware Support
| Device | Support |
@@ -25,28 +20,23 @@ backend alone is an integration point, not model-level validation.
| Intel XPU | ❌ |
| Ascend NPU | ✅ |
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide.
## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2)
| Model | Base model | Scope | Level | Notes |
|-------|------------|-------|-------|-------|
| Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Integrated | Upstream msModelSlim provides a producer recipe; use the [MXFP8](mxfp8.md) or [MXFP4](mxfp4.md) workflow for vLLM-Omni's native merged formats |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Not validated | No supported checkpoint workflow is documented |
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Integrated | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation and validate against BF16 |
| Model | Base model | Scope | Hardware | Notes |
|-------|------------|-------|----------|-------|
| Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Ascend NPU | Upstream msModelSlim provides a Wan2.2 quantization recipe; vLLM-Omni inference validation is not listed |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Ascend NPU | Not validated in this guide |
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Ascend A2/A3 NPU | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation |
Public Hugging Face quantized weights are not available yet. Use the
[HunyuanImage-3.0 msModelSlim adaptation](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8)
to generate the checkpoint manually.
msModelSlim is also the producer used by the native offline MXFP8 and MXFP4
DualScale workflows. Those workflows preprocess and annotate each transformer
checkpoint and are auto-detected at inference; follow their method pages rather
than passing `--quantization ascend`.
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
| Model | Scope | Status | Notes |
@@ -63,32 +53,25 @@ than passing `--quantization ascend`.
## Configuration
For a compatible generic Ascend checkpoint, select the Ascend quantization
method explicitly:
Offline inference:
```bash
python text_to_image.py --model <compatible-ascend-checkpoint> --quantization ascend
python text_to_image.py --model <quantized-model-path> --quantization ascend
```
Online serving:
```bash
vllm serve <compatible-ascend-checkpoint> --omni --quantization ascend
vllm serve <quantized-model-path> --omni --quantization ascend
```
## Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `--quantization` | str | - | Use `ascend` only for a compatible generic Ascend-format checkpoint; omit it for native merged MXFP8/MXFP4 checkpoints |
| `--quantization` | str | - | Use `ascend` for msModelSlim-produced checkpoints |
| `model` | str | - | Path to the quantized checkpoint generated by Ascend tooling |
Do not pass `--quantization ascend` for a Wan checkpoint processed by the native
MXFP8 or MXFP4 merge tools. Their transformer metadata selects `mxfp8` or
`mxfp4_dualscale` automatically.
Example msModelSlim command for a Wan2.2 W8A8 checkpoint:
```bash
@@ -110,6 +93,3 @@ For HunyuanImage-3.0, use the Hunyuan-specific adaptation linked above.
tooling; it is not a load-time CUDA quantizer.
3. Keep the quantized checkpoint aligned with the same model architecture and
stage config used for BF16 inference.
4. No named full-model msModelSlim checkpoint is currently scheduled in NPU
CI. Treat new producer outputs as **Integrated** until an end-to-end BF16
comparison establishes **Validated** support.
+23 -52
View File
@@ -13,10 +13,6 @@ vLLM-Omni provides two quantization methods with different scale structures:
| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
The validation terms on this page follow the
[support levels](overview.md#support-levels). Current model guidance covers
Wan diffusion transformers only; the text encoder and VAE stay in BF16.
!!! tip "Recommended: `mxfp4_dualscale` offline"
For production deployments, use the `mxfp4_dualscale` offline mode with a
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
@@ -38,45 +34,26 @@ Wan diffusion transformers only; the text encoder and VAE stay in BF16.
## Hardware Support
| Device | Online `mxfp4` | Online `mxfp4_dualscale` | Offline `mxfp4_dualscale` |
|--------|-----------------|---------------------------|----------------------------|
| NVIDIA Blackwell GPU (SM 100+) | ❌ | ❌ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ❌ | ❌ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ❌ | ❌ | ❌ |
| AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ |
| Intel XPU | ❌ | ❌ | ❌ |
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ |
| Device | Support |
|--------|---------|
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
| AMD ROCm (gfx950 / MI355X) | ✅ |
| Intel XPU | ⭕ |
| Ascend NPU (Atlas 950 A5) | ✅ |
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. Backend availability does not imply full-model checkpoint validation.
!!! note "Why NVIDIA is marked unsupported"
This table describes the methods currently wired into vLLM-Omni, not every
low-precision format NVIDIA hardware could execute. `DiffusionMXFP4Config`
has runtime implementations only for Ascend and ROCm gfx950; on CUDA it
raises `NotImplementedError` because no CUDA MXFP4 linear method or
checkpoint loader is registered.
On NVIDIA, use [ModelOpt](modelopt.md) NVFP4 for the documented 4-bit
checkpoint path. NVFP4 and OCP MXFP4 have different checkpoint metadata,
scale layouts, and kernel contracts, so they are not interchangeable. A
native CUDA MXFP4 path would require a CUDA linear implementation, loader
contract, and model-level output validation.
The ROCm path dispatches online single-scale `mxfp4` to AITER on gfx950. It has
kernel and dispatch coverage, but this guide does not yet claim end-to-end
quality validation for a named Wan checkpoint on ROCm. DualScale and
pre-quantized MXFP4 checkpoints are not supported on ROCm.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
## Model Type Support
### Diffusion Model (Wan2.2)
| Model | Online (Ascend) | Offline (Ascend) | Validation and notes |
|-------|-----------------|------------------|----------------------|
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | Validated outside scheduled full-model offline NPU CI; both cascade transformers are quantized |
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | Validated outside scheduled full-model offline NPU CI; same cascade scheme as T2V-A14B |
| Wan2.2-TI2V-5B | ❌ | ❌ | Unsupported because W4A4 causes unacceptable accuracy loss |
| Model | Online | Offline | Notes |
|-------|--------|---------|-------|
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
quantization quality, not model compatibility — both work on cascade (A14B) and
@@ -106,13 +83,12 @@ by the merge script); there is no offline `mxfp4` single-scale format.
### `mxfp4` — Single-Scale Online Mode
Online mode requires no pre-processing. On Ascend, vLLM-Omni quantizes BF16
weights at load time using `npu_dynamic_mx_quant`. On ROCm gfx950, the same
`mxfp4` config dispatches to the AITER online kernel. A single block scale
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
calibration `mul_scale` is available. The config applies equally to
single-transformer and cascade (A14B) models, but named Wan model validation is
currently limited to the documented Ascend path.
calibration `mul_scale` is available. Applies equally to single-transformer
and cascade (A14B) models — both transformers in a cascade receive the same
quantization config automatically.
```python
from vllm_omni import Omni
@@ -135,10 +111,10 @@ vllm serve <your-model> --omni --quantization mxfp4
### `mxfp4_dualscale` — DualScale Online Mode
Ascend-only online DualScale mode computes both fine and coarse scales on the
fly from BF16 weights using `npu_dynamic_dual_level_mx_quant`. It applies
equally to single-transformer and cascade (A14B) models. Compared to `mxfp4`
online, DualScale provides better quantization accuracy at higher compute cost.
Online DualScale mode computes both fine and coarse scales on the fly from BF16
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
DualScale provides better quantization accuracy at higher compute cost.
The default configuration keeps the leading 5 transformer blocks in BF16
(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
@@ -365,11 +341,6 @@ unless they appear in `ignored_layers`.
quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
off compression vs. accuracy for precision-sensitive layers.
8. The native methods, merge mappings, and per-transformer cascade metadata
have unit coverage. The current Buildkite configuration does not schedule a
named full-model offline MXFP4 checkpoint on NPU, so this path is
**Validated**, not **CI-backed**.
## Adapting MXFP4 for a New Model
This section is aimed at developers who want to add MXFP4 support to a model
+11 -30
View File
@@ -15,45 +15,30 @@ This method supports three modes:
| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
The validation terms on this page follow the
[support levels](overview.md#support-levels). Current model guidance covers
Wan diffusion transformers only; the text encoder and VAE stay in BF16.
## Hardware Support
| Device | Online | Offline (Native) | Offline (AutoRound) |
|--------|--------|------------------|---------------------|
| NVIDIA Blackwell GPU (SM 100+) | ❌ | ❌ | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ❌ | ❌ | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ❌ | ❌ | ⭕ |
| AMD ROCm | ❌ | ❌ | ⭕ |
| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ |
| AMD ROCm | ⭕ | ⭕ | ⭕ |
| Intel XPU | ✅ | ❌ | ✅ |
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. Backend availability does not imply full-model checkpoint validation.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
The native `DiffusionMXFP8Config` currently dispatches only to Intel XPU and
Ascend NPU implementations. CUDA reaches `NotImplementedError`, so the NVIDIA
rows are marked unsupported even though this is an implementation gap rather
than a general claim about NVIDIA hardware capabilities. On NVIDIA, use
[FP8](fp8.md) for load-time quantization or [ModelOpt](modelopt.md) FP8 for
pre-quantized checkpoints. Adding native CUDA MXFP8 requires a registered CUDA
linear method, a compatible checkpoint contract, and model-level validation.
Native `mxfp8` is available online on Intel XPU and online or offline on
Ascend. Intel XPU only supports AutoRound MXFP8 for offline mode. The AutoRound
column is a separate checkpoint format and compute path.
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
## Model Type Support
### Diffusion Model (Wan2.2)
| Model | Mode | Scope | Validation |
|-------|------|-------|------------|
| Wan2.2-T2V-A14B | Online + Offline | MoE cascade: `transformer` + `transformer_2` | Validated outside scheduled full-model offline NPU CI |
| Wan2.2-I2V-A14B | Online + Offline | MoE cascade: `transformer` + `transformer_2` | Validated outside scheduled full-model offline NPU CI |
| Wan2.2-TI2V-5B | Online + Offline | Single diffusion transformer | Validated outside scheduled full-model offline NPU CI |
| Model | Mode | Notes |
|-------|------|-------|
| Wan2.2-T2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
| Wan2.2-I2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
| Wan2.2-TI2V-5B | Online + Offline | Single transformer |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -233,7 +218,3 @@ omni = Omni(model="<autoround-mxfp8-model>")
`--output-dir`), regenerate it with the current script. The old script
wrote a separate `quantization_config.json` that is not read by vLLM-Omni;
the current script injects the config directly into `transformer/config.json`.
4. The native loader, merge mappings, and Wan cascade config propagation have
unit coverage. The current Buildkite configuration does not schedule a
named full-model offline MXFP8 checkpoint on NPU, so this path is
**Validated**, not **CI-backed**.
+15 -27
View File
@@ -6,7 +6,7 @@ Online quantization means vLLM-Omni computes quantized weights and scales while
loading the model. Use it when you want memory savings without preparing a
separate quantized checkpoint.
This mode is different from pre-quantized checkpoint formats such as ModelOpt,
This mode is different from pre-quantized checkpoint formats such as GGUF,
AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
prepared before serving and are documented in their method-specific guides.
For MXFP8 and MXFP4, use this page for load-time quantization from BF16
@@ -15,22 +15,18 @@ by msModelSlim and the merge tools.
## Hardware Support
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | MXFP8 W8A8 | MXFP4 W4A4 |
|--------|----------|-----------|------------------|------------|------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ❌ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ❌ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ❌ | ❌ |
| AMD ROCm (gfx950) | ⭕ | ⭕ | ❌ | ❌ | ✅ |
| Intel XPU | ⭕ | ⭕ | ❌ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ |
| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 |
|--------|----------|-----------|------------|------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
Legend: `✅` backend available, `❌` unsupported, `⭕` not verified in this
guide. These entries describe method-level backend availability, not validation
for every model. FP8 on Ampere may use a weight-only path where available.
On ROCm gfx950, MXFP4 means online single-scale `mxfp4`; online and offline
`mxfp4_dualscale` are Ascend-only. The native MXFP8/MXFP4 configs do not
currently register CUDA implementations; use FP8 or ModelOpt FP8/NVFP4 on
NVIDIA instead.
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and
MXFP4 are documented for the Ascend NPU path.
## Model Type Support
@@ -40,9 +36,8 @@ NVIDIA instead.
|--------|-------|----------------|--------|
| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Z-Image-Turbo; Qwen-Image and Wan2.2 are not validated | Validated for Z-Image-Turbo |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Validated on Ascend NPU; the ROCm gfx950 backend is integrated but not model-validated; TI2V-5B is unsupported |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -64,7 +59,6 @@ from vllm_omni import Omni
omni_fp8 = Omni(model="<your-model>", quantization="fp8")
omni_int8 = Omni(model="<your-model>", quantization="int8")
omni_bnb = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="bitsandbytes")
omni_mxfp8 = Omni(model="<your-model>", quantization="mxfp8")
omni_mxfp4 = Omni(model="<your-model>", quantization="mxfp4")
omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale")
@@ -75,7 +69,6 @@ CLI:
```bash
vllm serve <your-model> --omni --quantization fp8
vllm serve <your-model> --omni --quantization int8
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --quantization bitsandbytes
vllm serve <your-model> --omni --quantization mxfp8
vllm serve <your-model> --omni --quantization mxfp4
vllm serve <your-model> --omni --quantization mxfp4_dualscale
@@ -92,19 +85,14 @@ config = build_quant_config({
})
```
Component routing is an extension mechanism, not automatic support for every
pipeline component. See [Quantization Scope](overview.md#quantization-scope)
before applying it to an encoder, VAE, or decoder.
## Parameters
| Parameter | Methods | Description |
|-----------|---------|-------------|
| `method` | FP8, Int8, BitsAndBytes, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"bitsandbytes"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
| `ignored_layers` | FP8, Int8, BitsAndBytes, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
| `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
| `quant_type` | BitsAndBytes | Weight format: `"nf4"` (default) or `"fp4"` |
| `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` |
## Validation and Notes
+58 -130
View File
@@ -16,123 +16,72 @@ For the internal architecture and backend extension points, see the
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
Online quantization starts from a normal BF16/FP16 checkpoint and repeats the
conversion on each fresh model load. It is the simplest path for experiments.
Pre-quantized checkpoints store the packed weights and scales produced by a
separate tool; they are easier to reuse across deployments and can carry
calibrated or mixed-precision policies that load-time conversion cannot
reproduce.
## Hardware Support
## Support Levels
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim |
|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
Quantization support is specific to a model, component, checkpoint format, and
hardware backend. The tables in this guide use the following levels:
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available.
| Level | Meaning |
|-------|---------|
| **CI-backed** | A named model and quantized checkpoint are exercised by scheduled full-model hardware CI. |
| **Validated** | A named model and checkpoint have end-to-end or quality-validation evidence, but the check is not necessarily scheduled in CI. |
| **Integrated** | The config, loader, quantized layers, or model wiring exists, but this guide has no completed end-to-end validation for the named model and checkpoint. |
| **Not validated** | There is no current support recommendation for this model/component combination. |
| **Unsupported** | The combination is intentionally rejected or known not to work. |
## Model Type Support
A method being available on a device does not make every model on that device
supported. Likewise, successfully detecting a checkpoint's
`quantization_config` establishes loader compatibility, not output quality.
### Diffusion Model (Qwen-Image, Wan2.2)
## Hardware × Quantization Method
These models run a diffusion transformer as the primary inference module. The
default quantization target is the transformer; tokenizer, scheduler, text
encoder, and VAE stay on the base checkpoint unless a method guide says
otherwise.
This matrix uses concrete inference paths as columns. Producer toolkits are not
separate methods: NVIDIA ModelOpt produces the ModelOpt FP8/NVFP4 formats,
Intel AutoRound produces AutoRound checkpoints, and Ascend msModelSlim produces
checkpoints consumed by the native MXFP or compatible Ascend paths.
| Method | Guide | Mode | Example models | Status |
|--------|-------|------|----------------|--------|
| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Online W4 weight-only | Z-Image-Turbo; Qwen-Image/Wan2.2 not validated | Validated for Z-Image-Turbo |
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
| Hardware | [Online FP8](fp8.md) | [Int8](int8.md) | [BitsAndBytes W4](bitsandbytes.md) | [ModelOpt FP8](modelopt.md) | [ModelOpt NVFP4](modelopt.md#supported-modelopt-checkpoint-formats) | [AutoRound W4A16](autoround.md) | [OCP MXFP8](mxfp8.md) | [OCP MXFP4](mxfp4.md) |
|----------|-----------------------|------------------|------------------------------------------|--------------------------------|-------------------------------------------------------------------------|--------------------------------------|--------------------------|--------------------------|
| NVIDIA Blackwell (SM 100+) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| NVIDIA Hopper (SM 90) | ✅ | ✅ | ✅ | ✅ | ✅ (Marlin fallback) | ✅ | ❌ | ❌ |
| NVIDIA Ada (SM 89) | ✅ | ✅ | ✅ | ✅ | ⭕ | ✅ | ❌ | ❌ |
| NVIDIA Ampere (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | ❌ |
| AMD ROCm (gfx950 for MXFP4) | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ❌ | ✅ (online only) |
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
Legend: `✅` backend available, `❌` rejected or no implementation, `⭕` not
verified in this guide. This matrix does not establish model-level support; use
the model-and-stage tables below for that.
These models combine an AR language model with audio, vision, talker, or TTS
stages. Quantization is scoped to the AR language-model stage when the
checkpoint contains a supported `quantization_config`; the non-AR stages stay
in BF16 unless the model guide explicitly adds support.
The Int8 column covers online and serialized Int8 paths. OCP MXFP8 is online on
Intel XPU and online/offline on Ascend; Intel's separate AutoRound MXFP8
checkpoint path is described in the [MXFP8 guide](mxfp8.md). OCP MXFP4 on AMD
means online single-scale `mxfp4` on gfx950, while MXFP4 DualScale is
Ascend-only. NVIDIA's NVFP4 and OCP MXFP4 have different checkpoint metadata,
scale layouts, and kernel contracts. Mixed ModelOpt FP8/NVFP4 checkpoints also
require a model-specific mixed-precision policy and validation. FP8 on Ampere
may use a weight-only path where available.
| Method | Guide | Scope | Example models | Status |
|--------|-------|-------|----------------|--------|
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
## Current Model and Stage Support
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
The current offline checkpoint guidance is transformer-first. A listed
checkpoint quantizes the stage named in the **Scope** column; it does not imply
that the model's encoders, VAE, decoder, or other stages are also quantized.
These models split generation across multiple stages. Quantization must be
attached to the intended stage rather than applied globally.
### Vendor and Runtime Ecosystems
The producer and runtime backend are related, but they are not the same support
claim. A producer can emit several checkpoint formats, and each format still
needs a compatible vLLM-Omni loader, model-stage wiring, and validation. A
runtime may also provide online quantization without an offline producer path.
| Ecosystem | Current paths | Current vLLM-Omni scope | Model evidence |
|-----------|---------------|-------------------------|----------------|
| NVIDIA | ModelOpt FP8, NVFP4, and mixed FP8/NVFP4 | Diffusion transformer or Qwen-Omni thinker | Qwen-Image and other named diffusion checkpoints are validated; Qwen3-Omni NVFP4 has H100 CI for the Marlin fallback |
| Intel | AutoRound W4A16 and AutoRound MXFP8 | Diffusion/world-model transformer or Qwen-Omni thinker | Wan A14B and GLM-Image W4A16 are CI-backed; Qwen-Image, Qwen-Omni, and Cosmos3-Super have validation outside scheduled CI |
| AMD ROCm | Online OCP MXFP4 through AITER on gfx950; no pre-quantized MXFP4 or AMD Quark workflow is integrated | Diffusion transformer linear layers | Integrated backend with dispatch and weight-allocation unit coverage; no named-model quality validation or scheduled gfx950 full-model CI |
| Ascend | msModelSlim output through native MXFP8/MXFP4 merge workflows or a compatible generic Ascend format | Wan diffusion transformers | Native Wan MXFP8/MXFP4 paths are validated outside scheduled full-model NPU CI; other model families remain integrated or unvalidated |
### Pre-Quantized Model Matrix
| Model family | Method | Scope | Level | Notes |
|--------------|--------|-------|-------|-------|
| Qwen-Image | [ModelOpt](modelopt.md) FP8 or mixed FP8/NVFP4 | Diffusion transformer | Validated | Named Qwen-Image 2512 checkpoints and recipes are available. Text encoder and VAE stay BF16. |
| Qwen-Image | [AutoRound](autoround.md) W4A16 | Diffusion transformer | Validated | Uses `INC4AI/Qwen-Image-AutoRound-W4A16`; text encoder and VAE stay BF16. |
| Wan2.2 I2V/T2V A14B | [AutoRound](autoround.md) W4A16 | Both diffusion transformers | CI-backed | The high-noise and low-noise DiTs are quantized; text encoder and VAE stay BF16. |
| Wan2.2 TI2V 5B | [AutoRound](autoround.md) W4A16 | Diffusion transformer | Validated | Named checkpoint available; not covered by the scheduled Wan AutoRound job. |
| Wan2.2 | [MXFP8](mxfp8.md) | One or both diffusion transformers | Validated | Native Ascend and AutoRound checkpoint formats are integrated; no scheduled full-model offline NPU job. |
| Wan2.2 I2V/T2V A14B | [MXFP4 DualScale](mxfp4.md) | Both diffusion transformers | Validated | Offline calibrated checkpoints are recommended; no scheduled full-model offline NPU job. |
| Qwen2.5-Omni | [AutoRound](autoround.md) W4A16 | Thinker language model | Validated | Encoders and output/audio stages stay BF16. |
| Qwen3-Omni | [ModelOpt](modelopt.md) FP8 | Thinker language model | Validated | Tested outside scheduled CI. |
| Qwen3-Omni | [ModelOpt](modelopt.md) NVFP4 | Thinker language model | CI-backed | H100 CI exercises the Marlin fallback; the native Blackwell FP4 path has separate recipe validation. |
| Qwen3-Omni | [AutoRound](autoround.md) W4A16 | Thinker language model | Validated | Audio encoder, vision encoder, talker, and Code2Wav stay BF16. |
| GLM-Image | [AutoRound](autoround.md) W4A16 | Diffusion transformer stage | CI-backed | Text-to-image and image-to-image full-model smoke coverage. |
| Cosmos3-Super | [AutoRound](autoround.md) W4A16 | World-model transformer | Validated | Manually validated serving path; VAE and other components stay BF16. |
| Cosmos3-Nano | [AutoRound](autoround.md) W4A16 | World-model transformer | Integrated | Named checkpoint is listed, but no model-specific end-to-end test is maintained in-tree. |
| MiniMax H3 | None | None | Not validated | No offline vendor checkpoint path is documented; online FP8 targets the DiT only. |
The method pages contain the complete checkpoint lists, including FLUX,
Z-Image, and HunyuanImage-3.0 variants.
### Online Quantization
| Model family | Methods | Scope | Level |
|--------------|---------|-------|-------|
| Qwen-Image | FP8, Int8 | Diffusion transformer | Validated |
| Z-Image | Int8, BitsAndBytes W4 | Diffusion transformer | Validated |
| Wan2.2 | MXFP8 | One or both diffusion transformers | Validated on the documented XPU/NPU paths |
| Wan2.2 A14B | MXFP4 | Both diffusion transformers | Validated on Ascend; the ROCm gfx950 online backend is integrated but not model-validated |
| MiniMax H3 | FP8 | DiT and optional reference DiT | Validated |
| Qwen3-Omni | Generic online methods, such as FP8 | Thinker language model only | Integrated; not model-validated as an online recommendation |
See [Online Quantization](online.md) for configuration and hardware details.
### Unlisted and Long-Tail Models
For models not listed above, generic quantized linear layers or compatible
checkpoint metadata are only an integration starting point. The default support
target is the model's primary transformer, DiT, or AR language-model stage.
Text/audio/vision encoders, VAEs, vocoders, decoders, and output stages remain
BF16 until a model-specific integration supplies quant-aware layers, checkpoint
mapping, end-to-end output validation, and backend coverage.
| Method | Guide | Scope | Example models | Status |
|--------|-------|-------|----------------|--------|
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated |
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
!!! note
"Online quantization" means vLLM-Omni computes the quantization data while
@@ -143,8 +92,8 @@ mapping, end-to-end output validation, and backend coverage.
### Diffusion Model (Qwen-Image, Wan2.2)
The default target is the diffusion transformer. `build_quant_config()` can
construct a component router:
The default target is the diffusion transformer. Component routing is available
through `build_quant_config()`:
```python
from vllm_omni.quantization import build_quant_config
@@ -155,19 +104,11 @@ config = build_quant_config({
})
```
!!! warning "Routing syntax is not a support claim"
A component dictionary only selects a quantization config by module prefix.
It does not make an encoder or VAE quantizable. The model must construct
that component from quantization-aware layers, and an offline checkpoint
must contain matching packed weights and scales. Current user guidance
keeps diffusion text encoders and VAEs in BF16.
| Component | Default quantized? | Notes |
|-----------|--------------------|-------|
| Diffusion transformer | Yes, for listed model/method pairs | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
| Second diffusion transformer | Yes, for listed Wan2.2 A14B methods | `transformer` and `transformer_2` read their own offline checkpoint metadata; arbitrary mixed methods are not currently validated |
| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
| VAE | No | Keep BF16 |
| VAE | No | Keep BF16; storage-only paths are method-specific |
| Scheduler/tokenizer | No | Loaded from the base model repository |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -180,18 +121,6 @@ config = build_quant_config({
| Talker or TTS stage | No | BF16 unless model-specific support is documented |
| Code2Wav | No | BF16 |
An explicit `ComponentQuantizationConfig` can be used by model developers to
route different configs to `audio_tower`, `visual`, and `language_model`, but
the supported Qwen-Omni checkpoints currently quantize only the thinker
language model.
### World Model (Cosmos3)
| Component | Default quantized? | Notes |
|-----------|--------------------|-------|
| World-model transformer | AutoRound checkpoints only | Cosmos3-Super is validated; Cosmos3-Nano is integrated but not end-to-end tested in-tree |
| VAE and guardrail components | No | BF16 |
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
| Component | Default quantized? | Notes |
@@ -212,7 +141,6 @@ build_quant_config("fp8")
build_quant_config({"method": "fp8", "activation_scheme": "static"})
build_quant_config("bitsandbytes")
build_quant_config("auto-round", bits=4, group_size=128)
# Component routing syntax; model-specific support is still required.
build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
build_quant_config(None)
```
@@ -59,7 +59,7 @@ def main(args):
omni = Omni(
model=args.model_name,
stage_configs_path=args.stage_configs_path,
deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
batch_timeout=args.batch_timeout,
@@ -167,10 +167,10 @@ def parse_args():
help="Sampling rate for audio loading (default: 16000).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to stage configs YAML file.",
help="Path to a deploy config YAML file.",
)
parser.add_argument(
"--log-stats",
@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
parser.add_argument(
"--stage-config-path",
"--deploy-config",
type=str,
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
help="Path to deploy config yaml.",
@@ -1394,8 +1394,8 @@ def main() -> None:
from vllm_omni.entrypoints.omni import Omni
stage_config_path = str(Path(args.stage_config_path).expanduser())
omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
deploy_config_path = str(Path(args.deploy_config).expanduser())
omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
sampling_params_list = [
SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
for _ in range(omni.num_stages)
@@ -73,7 +73,6 @@ def parse_args():
)
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
parser.add_argument(
"--stream",
action="store_true",
@@ -146,13 +145,7 @@ def main():
else:
bot_task = args.bot_task
if args.deploy_config is not None and args.stage_configs_path is not None:
raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
deploy_config = args.deploy_config
stage_configs_path = args.stage_configs_path
if deploy_config is None and stage_configs_path is None:
deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
omni_kwargs = {
"model": args.model,
@@ -168,10 +161,7 @@ def main():
if additional_config is not None:
omni_kwargs["additional_config"] = additional_config
if deploy_config is not None:
omni_kwargs["deploy_config"] = deploy_config
else:
omni_kwargs["stage_configs_path"] = stage_configs_path
omni_kwargs["deploy_config"] = deploy_config
omni = Omni(**omni_kwargs)
@@ -272,10 +262,7 @@ def main():
print(f" Modality: {args.modality}")
print(f" Prompt task: {task}")
print(f" Bot task: {bot_task}")
if deploy_config is not None:
print(f" Deploy config: {deploy_config}")
else:
print(f" Stage config: {stage_configs_path}")
print(f" Deploy config: {deploy_config}")
print(f" Num stages: {omni.num_stages}")
if args.modality in ("text2img", "img2img"):
print(f" Inference steps: {args.steps}")
@@ -452,12 +452,6 @@ def parse_args():
default=None,
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to a stage configs file (deprecated; prefer --deploy-config).",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -11,6 +11,13 @@ Examples:
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
# Rank-local mmap storage, with no DLO collective or DP wave scheduling.
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2-no-allgather
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
@@ -45,9 +52,9 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
parser.add_argument(
"--mode",
choices=("dlo-dp2", "request"),
choices=("dlo-dp2", "dlo-dp2-no-allgather", "request"),
required=True,
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
help="DLO DP2 (with or without AllGather) or ordinary non-DLO TP2 request mode",
)
parser.add_argument("--steps", type=int, default=2)
parser.add_argument("--duration", type=float, default=5.0)
@@ -74,14 +81,14 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
"stage_init_timeout": args.init_timeout,
"init_timeout": args.init_timeout,
}
if args.mode == "dlo-dp2":
if args.mode in {"dlo-dp2", "dlo-dp2-no-allgather"}:
common.update(
tensor_parallel_size=1,
data_parallel_size=2,
text_encoder_tp_size=1,
vae_patch_parallel_size=1,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
dlo_use_allgather=args.mode == "dlo-dp2",
dlo_resident_layers=0,
)
else:
@@ -484,10 +484,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to a stage configs file.",
help="Path to a deploy config YAML.",
)
parser.add_argument(
"--video-path",
@@ -232,15 +232,8 @@ def main(args):
raise ValueError(f"Unknown query type: {args.query_type}")
# Initialize vLLM-Omni with Step-Audio2
# Resolve deploy config path. Keep explicit stage_configs_path for legacy
# custom configs, but use bundled deploy configs by default.
deploy_config_path = None
stage_config_path = None
if args.stage_configs_path and args.deploy_config:
raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
if args.stage_configs_path:
stage_config_path = args.stage_configs_path
elif args.deploy_config:
# Resolve an explicit deploy config or use the bundled task default.
if args.deploy_config:
deploy_config_path = args.deploy_config
else:
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
@@ -262,10 +255,7 @@ def main(args):
"ray_address": args.ray_address,
"trust_remote_code": True,
}
if stage_config_path is not None:
omni_kwargs["stage_configs_path"] = stage_config_path
else:
omni_kwargs["deploy_config"] = deploy_config_path
omni_kwargs["deploy_config"] = deploy_config_path
omni_llm = Omni(**omni_kwargs)
@@ -418,12 +408,6 @@ def parse_args():
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -92,12 +92,6 @@ def parse_args() -> argparse.Namespace:
"meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
"stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -475,8 +469,6 @@ def main():
omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
if args.enforce_eager is not None:
omni_kwargs["enforce_eager"] = args.enforce_eager
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if use_nextstep:
@@ -515,8 +507,8 @@ def main():
print(f" Image size: {args.width}x{args.height}")
if args.lora_path:
print(f" LoRA: scale={args.lora_scale}")
if args.stage_configs_path:
print(f" stage-configs-path: {args.stage_configs_path}")
if args.deploy_config:
print(f" deploy-config: {args.deploy_config}")
print(f"{'=' * 60}\n")
# Build LoRA request when --lora-path is set
@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
@@ -430,6 +431,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
---
## dots.tts
Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage.
### Quick start
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
### Voice cloning
Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet.
### Streaming
The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example.
### Notes
- Output: 48 kHz mono WAV.
- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`).
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet.
- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap.
---
## IndexTTS-2 and IndexTTS-2.5
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
@@ -0,0 +1,125 @@
"""Offline dots.tts inference example (single-stage native AR pipeline).
Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml).
Reference checkpoint: rednote-hilab/dots.tts-soar.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
import soundfile as sf
import torch
from vllm_omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
REPO_ROOT = Path(__file__).resolve().parents[4]
SAMPLE_RATE = 48_000
def parse_args():
parser = TrackingArgumentParser(description="Offline dots.tts native AR inference")
parser.add_argument(
"--model",
type=str,
default="rednote-hilab/dots.tts-soar",
help="dots.tts model path or HuggingFace repo ID.",
)
parser.add_argument(
"--text",
type=str,
default="Hello, this is a test of dots TTS running on vLLM Omni.",
help="Text to synthesize.",
)
parser.add_argument(
"--output-dir",
type=str,
default="output_audio",
help="Directory for output WAV files.",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override the deploy config path. If unset, auto-loads "
"vllm_omni/deploy/dots_tts.yaml based on the HF model_type.",
)
return parser.parse_args()
def extract_audio(multimodal_output: dict) -> torch.Tensor:
"""Extract the final complete audio tensor from multimodal output.
The output processor concatenates per-step delta tensors under
``model_outputs``. Falls back to ``audio`` for backwards compat.
"""
audio = multimodal_output.get("model_outputs")
if audio is None:
audio = multimodal_output.get("audio")
if audio is None:
raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")
if isinstance(audio, list):
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
if not valid:
raise ValueError("Audio list is empty or all elements are None.")
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
return torch.as_tensor(audio).float().cpu().reshape(-1)
def main():
args = parse_args()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
engine = Omni(
model=args.model,
deploy_config=args.deploy_config,
)
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
build_dots_tts_prompt,
)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
prompt = build_dots_tts_prompt(
tokenizer=tokenizer,
text=args.text,
)
print(f"Model : {args.model}")
print(f"Text : {args.text}")
print(f"Output dir : {output_dir}")
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
t_start = time.perf_counter()
outputs = engine.generate([prompt])
elapsed = time.perf_counter() - t_start
request_output = outputs[0]
mm = request_output.outputs[0].multimodal_output
audio = extract_audio(mm)
duration = audio.numel() / SAMPLE_RATE
rtf = elapsed / duration if duration > 0 else float("inf")
output_path = output_dir / "output.wav"
sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV")
print(f"Saved : {output_path}")
print(f"Duration : {duration:.2f}s")
print(f"Inference : {elapsed:.2f}s")
print(f"RTF : {rtf:.3f}")
if __name__ == "__main__":
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
main()
@@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "fishaudio/s2-pro"
DEFAULT_STAGE_CONFIG = os.path.join(
DEFAULT_DEPLOY_CONFIG = os.path.join(
os.path.dirname(__file__),
"..",
"..",
@@ -152,10 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
else:
omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
return omni_kwargs
@@ -278,12 +275,6 @@ def parse_args():
default=None,
help="Deploy config YAML path (under vllm_omni/deploy/).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to stage configs YAML (legacy, prefer --deploy-config).",
)
parser.add_argument(
"--output-dir",
default="output_audio",
@@ -145,7 +145,7 @@ def main(args):
omni = Omni(
model=args.model,
stage_configs_path=deploy_config_path,
deploy_config=deploy_config_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
@@ -490,10 +490,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to a stage configs file.",
help="Path to a deploy config YAML.",
)
parser.add_argument(
"--audio-path",
+2
View File
@@ -73,6 +73,8 @@ plugins:
- awesome-nav
- redirects:
redirect_maps:
community/meetups.md: community/meetings.md
community/volunteers.md: community/contact_us.md
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
+72 -46
View File
@@ -4,11 +4,20 @@
## Pipelines
| Model | Checkpoint | Task | `--model-class-name` | Batching |
|---|---|---|---|---|
| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
| `--model-class-name` | Task | Required checkpoint repositories |
|---|---|---|
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
Repositories in the table are download units. A full pipeline repository
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
and tokenizer; an additional repository supplies LoRA or upsampler sidecars.
The runtime first searches the model root for sidecars under their official
filenames, then downloads them from the matching Lightricks Hub repository.
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
@@ -16,10 +25,10 @@ image selects I2V. Both one-stage repositories declare this class, so
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
or provide one initial image for I2V. As in the official LTX pipelines,
distilled mode is selected explicitly rather than inferred from the checkpoint
name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
merged distilled Transformer in both stages. Both entries support T2V and I2V;
select their class explicitly.
## API Migration
@@ -68,24 +77,28 @@ The model recipe also supplies the default negative prompt. Top-level
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
default to `121`.
## Distilled Two-Stage Defaults
## Two-Stage Defaults
| Parameter | Value |
|---|---:|
| Final width × height | 1536 × 1024 |
| Stage 1 width × height | 768 × 512 |
| Frames / frame rate | 121 / 24 |
| Stage 1 / Stage 2 denoise steps | 8 / 3 |
| Guidance | Fixed positive-only |
| Parameter | Ordinary | Full-distilled |
|---|---:|---:|
| Final width × height | 1536 × 1024 | 1536 × 1024 |
| Stage 1 width × height | 768 × 512 | 768 × 512 |
| Frames / frame rate | 121 / 24 | 121 / 24 |
| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 |
| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only |
The dimensions passed to the API are the final output dimensions and must be
divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
`8` when supplied.
API dimensions are final dimensions and must be divisible by 64. All LTX
requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or
LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so
`num_inference_steps`, when supplied, must be `8`. Both entries reject custom
sigmas and input latents.
Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
and automatically switches to dynamic LoRA when quantization is enabled.
## Serving
Start either checkpoint:
Start either one-stage checkpoint:
```bash
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
@@ -95,11 +108,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
```
Start the distilled checkpoint with its two-stage entry:
Start a two-stage checkpoint with its explicit class:
```bash
vllm serve rootonchair/LTX-2-19b-distilled --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# Ordinary LTX-2.3 two-stage
vllm serve diffusers/LTX-2.3-Diffusers --omni \
--model-class-name LTX2TwoStagePipeline \
--enable-layerwise-offload \
--stage-init-timeout 600
```
The same server handles T2V and I2V. A T2V request using the selected recipe's
@@ -135,8 +156,9 @@ together with `input_reference`.
## Guidance
LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
cross-modality guidance, and rescaling:
One-stage and ordinary Stage 1 support independent video/audio CFG,
spatio-temporal guidance (STG), cross-modality guidance, and rescaling.
Distilled stages and ordinary Stage 2 are fixed positive-only.
| Parameter | Default | Effect | Alias |
|---|---:|---|---|
@@ -172,9 +194,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
text CFG, STG, and cross-modality guidance passes; guidance rescaling is
applied after all pass predictions are gathered.
The default one-stage recipe has four Transformer passes per denoise step:
`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
configurations are therefore:
The default one-stage recipe and ordinary Stage 1 have four Transformer passes
per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod`
(cross-modality). The useful balanced configurations are therefore:
| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
|---:|---:|---:|---|
@@ -195,9 +217,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \
```
The total device count is the product of `cfg_parallel_size` and the other
configured parallel dimensions. Distilled pipelines use positive-only
guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
them.
configured parallel dimensions. Positive-only distilled phases do not benefit
from `cfg_parallel_size > 1`.
### Python API
@@ -242,22 +263,22 @@ noted below.
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. |
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. |
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. |
| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. |
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
@@ -273,13 +294,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`.
### Recipe-Specific Request Capabilities
| Override | One-stage | Distilled two-stage |
|---|---|---|
| Guidance | Supported | Fixed positive-only |
| Negative prompt/embeddings | Supported | Rejected |
| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
| Video/audio latents | Supported | Rejected |
| Override | One-stage | Ordinary two-stage | Distilled two-stage |
|---|---|---|---|
| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only |
| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected |
| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 |
| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules |
| Video/audio latents | Supported | Rejected | Rejected |
These capability checks apply equally to direct `forward` keywords and
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
@@ -309,6 +330,10 @@ bundled offline CLI do not currently expose `sigmas`.
`--max-num-seqs 1`.
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
modality guidance, and rescale-compatible prediction gathering.
- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
divisible request shape.
- Cache-DiT is one-stage only; multi-stage configurations are rejected.
## Operational Notes
@@ -323,14 +348,15 @@ bundled offline CLI do not currently expose `sigmas`.
request parameter.
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
- Current two-stage support is limited to the distilled checkpoint; official
LTX two-stage/HQ execution remains out of scope.
- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution
remains out of scope.
## References
- <https://huggingface.co/Lightricks/LTX-2>
- <https://huggingface.co/Lightricks/LTX-2.3>
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- <https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers>
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
+152
View File
@@ -0,0 +1,152 @@
# dots.tts
> Offline continuous-AR TTS at 48 kHz (rednote-hilab)
## Summary
- Vendor: rednote-hilab
- Model: `rednote-hilab/dots.tts-soar`
- Task: Text-to-speech, zero-shot synthesis only
- Mode: Offline end-to-end example only (not yet wired for online serving —
see [Known limitations](#known-limitations))
- Maintainer: Community
## When to use this recipe
Use this recipe as a known-good starting point for running
`rednote-hilab/dots.tts-soar` offline on vLLM-Omni on consumer-class GPUs.
dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM
+ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio.
It follows the same "vLLM-native base LM + side-path computation" pattern as
VoxCPM2 — single-stage pipeline
`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR
loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone
instead of MiniCPM4, and no FSQ / residual-LM stage.
This is an early integration (see [Known limitations](#known-limitations)
below before depending on it for anything beyond offline experimentation).
## References
- Offline end-to-end script:
[`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py)
- Example guide:
[`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts)
- Default deploy config:
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
- Talker / pipeline source:
[`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/)
- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts)
## Hardware Support
This recipe documents one tested 16 GB consumer-GPU configuration. Other
vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as
community validation lands.
## GPU
### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended)
dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++
speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The
bundled default config at
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`,
`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model
weights and ~0.3 GiB peak activation; the remainder of the configured
budget is available for KV cache. Total resident footprint at idle is
roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph
roadmap would be step 8's graph capture (not implemented yet; this
release runs `enforce_eager: true`, so it doesn't apply today).
#### Environment
- OS: Linux (WSL2)
- Python: 3.12
- Driver / runtime: NVIDIA driver 595.95
- torch: 2.11.0+cu130
- vLLM: 0.26.0
- vLLM-Omni: 0.22.1.dev (current `main`)
#### Command
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
The deploy config at
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
is loaded automatically by the model registry (HF `model_type=dots_tts`).
Pass `--deploy-config <path>` to override.
#### Verification
**T1 — offline zero-shot synthesis**:
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single
one-shot process (init → one `generate()` → exit), so the reported
numbers include engine init and first-request warmup, not just steady-state
per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`,
`RTF: 1.569`.
Weight-loading breakdown from the same run (all tensors matched, no
missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270
patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder.
Whisper transcription of the output matched the input text with no
dropped leading word (confirms the streaming-vocoder patch-boundary fix
described in [Known limitations](#known-limitations)).
#### Notes
- Output: 48 kHz mono WAV.
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default
used throughout this recipe. `rednote-hilab/dots.tts-base` shares the
same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf`
(MeanFlow, 2-4 step) is not supported — see below.
- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy
config are load-bearing, not just conservative defaults: with prefix
caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path
does not preserve the `sparse_audio` marker this model relies on to
route audio output correctly, and generation silently truncates to a
single ~160 ms patch. Do not override `enable_prefix_caching` for this
model until that framework-level gap is fixed.
## Known limitations
- **No online serving yet.** dots.tts is not registered in
`vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS
models in this repo use for `/v1/audio/speech`), so `vllm serve
rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible
speech endpoint does not know how to build request params for it yet.
Offline `Omni()` (as used in `end2end.py`) is the only supported path
today.
- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder
weights load and are exercised by `load_weights()`, but nothing in the
prompt builder consumes a reference audio yet — generation is zero-shot
only. `end2end.py` has no `--ref-audio`/`--ref-text` flags.
- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only
the fixed 10-step Euler DiT sampler used by `dots.tts-soar` /
`dots.tts-base` is implemented.
- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three
captured graphs (base LM decode, CFM solver, VAE decode) have no
dots.tts equivalent yet.
- **Concurrent requests do not scale.** Each request's 10-step DiT Euler
integration runs serially in the side path (no cross-request batching,
unlike voxcpm2's `enable_batched_cfm`). A community review of this
integration measured no throughput gain at `c=4` concurrent requests
versus `c=1`.
- **`SamplingParams.seed` does not control audio-generation randomness.**
The DiT's flow-matching noise is deterministically derived per-request
from a fixed internal seed (reproducible run-to-run, matching
voxcpm2's `deterministic_cfm_seed` convention) rather than from the
caller-supplied `seed` field — the same limitation voxcpm2 has today.
+4 -48
View File
@@ -6,7 +6,7 @@ Unit tests for StageConfigFactory and related classes.
import importlib
import warnings
from dataclasses import dataclass, fields
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
@@ -34,9 +34,8 @@ from vllm_omni.config.stage_config import (
load_deploy_config,
merge_pipeline_deploy,
pipeline_cfg_resolver,
strip_parent_engine_args,
)
from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -475,7 +474,7 @@ class TestStageConfigFactory:
cli_overrides = {
"gpu_memory_utilization": 0.9,
"model": "some_model", # Internal
"stage_configs_path": "/path", # Internal
"deploy_config": "/path", # Internal
"batch_timeout": 10, # Internal
}
@@ -483,7 +482,7 @@ class TestStageConfigFactory:
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "stage_configs_path" not in overrides
assert "deploy_config" not in overrides
assert "batch_timeout" not in overrides
def test_per_stage_override_excludes_internal_keys(self):
@@ -525,49 +524,6 @@ class TestStageResolutionHelpers:
assert "model" not in overrides
assert "parallel_config" not in overrides
def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
parent_fields = {f.name: f for f in fields(EngineArgs)}
filtered, overridden = strip_parent_engine_args(
{
"model": "some/model",
"stage_configs_path": "/tmp/stages.yaml",
"tensor_parallel_size": 4,
"worker_extension_cls": "some.Extension",
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
},
parent_fields=parent_fields,
keep_keys={"worker_extension_cls"},
strip_keys={"stage_configs_path"},
no_warn_keys={"model"},
)
assert filtered == {
"worker_extension_cls": "some.Extension",
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
}
assert overridden == ["tensor_parallel_size"]
def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
parent_fields = {f.name: f for f in fields(EngineArgs)}
filtered, overridden = strip_parent_engine_args(
{
"model": "some/model",
"stage_configs_path": "/tmp/stages.yaml",
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
},
parent_fields=parent_fields,
keep_keys={"allowed_local_media_path", "allowed_media_domains"},
strip_keys={"stage_configs_path"},
no_warn_keys={"model"},
)
assert filtered == {
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
}
assert overridden == []
class TestPipelineDiscovery:
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
+1
View File
@@ -471,6 +471,7 @@ def test_sub_config_fields_match_structured_scopes():
"trust_remote_code",
"dtype",
"attention_backend",
"attention_config",
"moe_backend",
"hf_overrides",
"limit_mm_per_prompt",
@@ -0,0 +1,82 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm.triton_utils import HAS_TRITON
pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
_HEAD_DIM = 128
_ROTARY_DIM = 96
_EPS = 1e-5
def _reference(q, k, q_weight, k_weight, rope_table):
q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
half = _ROTARY_DIM // 2
cos = rope_table[..., :half].unsqueeze(1)
sin = rope_table[..., half:].unsqueeze(1)
def apply(x):
first = x[..., :half]
second = x[..., half:_ROTARY_DIM]
return torch.cat(
(
first * cos - second * sin,
second * cos + first * sin,
x[..., _ROTARY_DIM:],
),
dim=-1,
)
return apply(q), apply(k)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
@pytest.mark.parametrize("seq_len", [1, 257, 1024])
def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
fused_qk_norm_rope,
)
torch.manual_seed(17)
heads = 14
qkv = torch.randn(
seq_len,
heads * _HEAD_DIM * 3,
device="cuda",
dtype=torch.bfloat16,
)
q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
seq_len,
heads,
_HEAD_DIM,
)
q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
expected_q, expected_k = _reference(
q,
k,
q_weight,
k_weight,
rope_table,
)
actual_q, actual_k = fused_qk_norm_rope(
q,
k,
q_weight,
k_weight,
rope_table,
_EPS,
)
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
@@ -23,7 +23,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
def test_prepare_video_latents_samples_directly_in_packed_token_space():
def test_prepare_video_latents_matches_official_values_and_token_major_layout():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
@@ -44,6 +44,7 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space():
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
assert actual.stride()[1:] == (1, actual.shape[1])
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
@@ -81,6 +82,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:]))
def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding():
pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1)
pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
generator_sp1 = torch.Generator().manual_seed(42)
generator_sp4 = torch.Generator().manual_seed(42)
latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp1,
)
latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp4,
)
torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0)
torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:]))
torch.testing.assert_close(
torch.randn(8, generator=generator_sp4),
torch.randn(8, generator=generator_sp1),
rtol=0,
atol=0,
)
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
@@ -137,7 +174,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
packed = latent_ops.pack_audio_latents(latents)
torch.testing.assert_close(padded[:, :10], packed[:, :10])
torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:]))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
@@ -12,6 +12,19 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0
sigma = mocker.MagicMock()
sigma.to.return_value.item.return_value = 0.5
actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma)
sigma.to.assert_called_once_with(torch.float32)
sigma.to.return_value.item.assert_called_once_with()
torch.testing.assert_close(actual, torch.tensor([2.0]))
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
@@ -199,6 +212,7 @@ class TestCFGParallelHelpers:
),
attention_kwargs=None,
audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
original_audio_num_frames=state.audio.shape[1],
)
actual_video, actual_audio = executor.predict_parallel_guidance(
@@ -0,0 +1,483 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the fixed LTX refinement-phase adapter."""
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter
from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import (
AdapterTarget,
parse_ltx_adapter,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE
from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_TWO_STAGE_RECIPE,
LTX_DISTILLED_ADAPTER_SLOT,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _target(
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
):
return AdapterTarget(
source_module=source_module,
module=module,
slice=target_slice,
a_key="a",
b_key="b",
a_shape=tuple(lora_a.shape),
b_shape=tuple(lora_b.shape),
)
def _wrap(
layer,
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
layer_fused=False,
):
target = _target(
lora_a,
lora_b,
source_module=source_module,
module=module,
target_slice=target_slice,
)
wrapper = ltx2_phase_adapter._PhaseAdapterLinear(
layer,
[target],
adapter_dtype=lora_a.dtype,
layer_fused=layer_fused,
)
wrapper.load_target(target, lora_a, lora_b)
wrapper.set_enabled(True)
return wrapper
class _LoRAAttention(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.to_qkv = torch.nn.Identity()
class _LoRATransformer(torch.nn.Module):
stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),)
def __init__(self) -> None:
super().__init__()
self.proj_in = torch.nn.Linear(4, 6, bias=False)
block = torch.nn.Module()
block.attn1 = _LoRAAttention()
self.transformer_blocks = torch.nn.ModuleList([block])
def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2),
},
adapter_path,
)
manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == {
("proj_in", "proj_in", None),
("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"),
}
def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)},
adapter_path,
)
with pytest.raises(ValueError, match="Missing paired LoRA A tensors"):
parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
def _pipeline(*, dtype, quantization_config):
return SimpleNamespace(
pipeline_recipe=LTX2_TWO_STAGE_RECIPE,
component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE,
transformer=object(),
od_config=SimpleNamespace(
model="unused",
lora_path=None,
dtype=dtype,
quantization_config=quantization_config,
),
)
@pytest.mark.parametrize(
("dtype", "quantization_config", "layer_fused"),
[
(torch.float32, object(), False),
(torch.bfloat16, None, True),
],
)
def test_ltx_phase_adapter_selects_execution_from_quantization(
monkeypatch,
dtype,
quantization_config,
layer_fused,
):
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
manifest = object()
captured = {}
def resolve_artifact(*_args, **kwargs):
assert not kwargs
return "adapter"
def parse(transformer, path):
assert transformer is pipeline.transformer
assert path == "adapter"
return manifest
class Runtime:
def __init__(self, transformer, parsed_manifest, **kwargs):
captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact)
monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse)
monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime)
runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
assert isinstance(runtime, Runtime)
assert captured == {
"transformer": pipeline.transformer,
"manifest": manifest,
"dtype": dtype,
"layer_fused": layer_fused,
}
def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch):
pipeline = _pipeline(dtype=torch.float16, quantization_config=None)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter")
with pytest.raises(ValueError, match="bfloat16"):
ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path):
class TinyTransformer(torch.nn.Module):
stacked_params_mapping = ()
def __init__(self) -> None:
super().__init__()
self.proj = torch.nn.Linear(2, 2, bias=False)
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]),
},
adapter_path,
)
transformer = TinyTransformer()
runtime = LTXPhaseAdapterRuntime(
transformer,
parse_ltx_adapter(transformer, str(adapter_path)),
dtype=torch.float32,
)
assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight"
assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"}
with torch.no_grad():
transformer.proj.base_layer.weight.copy_(torch.eye(2))
runtime.finalize()
x = torch.tensor([[1.0, 1.0]])
runtime.activate(None)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]]))
runtime.activate(LTX_DISTILLED_ADAPTER_SLOT)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]]))
def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion():
layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16)
base_weight = torch.tensor(
[
[0.5, -0.25, 0.125, 0.75],
[-0.375, 0.625, 0.25, -0.5],
[0.875, -0.125, -0.75, 0.375],
],
dtype=torch.bfloat16,
)
lora_a = torch.tensor(
[
[0.1015625, -0.203125, 0.3046875, -0.40625],
[0.5078125, -0.609375, 0.7109375, -0.8125],
],
dtype=torch.bfloat16,
)
lora_b = torch.tensor(
[[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]],
dtype=torch.bfloat16,
)
with torch.no_grad():
layer.weight.copy_(base_weight)
layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16))
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeRow(torch.nn.Module):
input_size = 4
output_size = 2
input_size_per_partition = 2
input_is_parallel = True
tp_rank = 0
tp_size = 2
reduce_results = True
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16))
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow)
reductions = []
def all_reduce(value):
reductions.append(value.clone())
return value + 1
monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce)
layer = FakeRow()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a[:, :2])
expected_weight.add_(layer.weight)
x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16)
expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias)
actual = wrapper(x)
assert len(reductions) == 1
assert torch.equal(reductions[0], expected_local)
assert torch.equal(actual, expected_local + 1)
def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
input_size = 3
output_size = 4
output_size_per_partition = 2
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16)
)
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
layer = FakeColumn()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b[2:], lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
pass
class FakeQKV(FakeColumn):
input_size = 4
head_size = 2
v_head_size = 2
total_num_heads = 2
total_num_kv_heads = 2
num_heads = 1
num_kv_heads = 1
num_kv_head_replicas = 1
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = True
return_bias = True
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16
)
self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8)
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV)
layer = FakeQKV()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(
layer,
lora_a,
lora_b,
source_module="to_k",
module="to_qkv",
target_slice="k",
layer_fused=True,
)
expected_weight = base_weight.clone()
expected_k = torch.matmul(lora_b[2:], lora_a)
expected_k.add_(base_weight[2:4])
expected_weight[2:4].copy_(expected_k)
x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16)
expected_output = torch.nn.functional.linear(x, expected_weight)
actual_output, returned_bias = wrapper(x)
assert torch.equal(actual_output, expected_output)
assert returned_bias is layer.bias
assert torch.equal(layer.weight, base_weight)
@pytest.mark.parametrize(
("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"),
[
(
lambda: type(
"FakeColumn",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 4,
"output_size_per_partition": 2,
"tp_rank": 1,
"gather_output": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[12963.0, 17284.0]]),
),
(
lambda: type(
"FakeRow",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 2,
"input_size_per_partition": 2,
"tp_rank": 1,
"reduce_results": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[2.0], [3.0]]),
torch.tensor([[8600.0, 12900.0]]),
),
(
lambda: type(
"FakeQKV",
(torch.nn.Module,),
{
"input_size": 4,
"head_size": 2,
"v_head_size": 2,
"total_num_heads": 2,
"total_num_kv_heads": 2,
"num_heads": 1,
"num_kv_heads": 1,
"num_kv_head_replicas": 1,
"tp_rank": 1,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)),
},
)(),
"k",
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]),
),
],
)
def test_ltx_phase_adapter_uses_local_tp_lora_slices(
monkeypatch,
layer_factory,
target_slice,
expected_a,
expected_b,
expected_output,
):
layer = layer_factory()
layer.register_buffer("device_anchor", torch.empty(0))
layer_type = type(layer)
if target_slice == "k":
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type)
elif "Column" in layer_type.__name__:
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type)
else:
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type)
wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice)
torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output)
expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size)
assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width
+362 -10
View File
@@ -13,13 +13,21 @@ import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.ltx2 import ltx2_components
from vllm_omni.diffusion.models.ltx2 import (
ltx2_components,
ltx2_transformer,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX2_DISTILLED_COMPONENT_PROFILE,
LTX2_TWO_STAGE_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
LTX23_DISTILLED_COMPONENT_PROFILE,
LTX23_TWO_STAGE_COMPONENT_PROFILE,
_install_connector_attention,
detect_ltx_model_version,
resolve_ltx_artifact,
resolve_ltx_component_profile,
)
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
@@ -33,6 +41,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTXGuidancePlan,
LTXGuidanceSpec,
LTXModalityGuidance,
_repeat_batch,
build_perturbation_kwargs,
combine_guided_x0,
)
@@ -40,18 +49,29 @@ from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX2_ONE_STAGE_RECIPE,
LTX2_TWO_STAGE_RECIPE,
LTX23_DISTILLED_TWO_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
LTX23_TWO_STAGE_RECIPE,
LTX_DEFAULT_NEGATIVE_PROMPT,
LTX_POSITIVE_ONLY_RECIPE,
resolve_ltx_pipeline_recipe,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import (
LTXRequestInputs,
validate_ltx_checkpoint,
validate_pipeline_request,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2DistilledPipeline,
LTX2TwoStagePipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -104,10 +124,22 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
)
assert issubclass(LTX2Pipeline, LTXRuntime)
assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
for pipeline_cls, profile in (
(LTX2Pipeline, LTX2_COMPONENT_PROFILE),
(LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
(LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
):
assert pipeline_cls._dit_modules == list(profile.dit_modules)
assert pipeline_cls._encoder_modules == list(profile.encoder_modules)
assert pipeline_cls._vae_modules == list(profile.vae_modules)
assert pipeline_cls._resident_modules == list(profile.resident_modules)
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
@@ -123,6 +155,72 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2"
def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
calls = []
def fake_download(**kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
)
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
expected = tmp_path / filename
expected.write_bytes(b"sidecar")
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config"),
[
("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}),
("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}),
],
)
def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config", "error"),
[
("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"),
("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"),
],
)
def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error):
with pytest.raises(ValueError, match=error):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
def test_ltx_rejects_advanced_uaa_before_component_initialization():
od_config = SimpleNamespace(
model="unused",
parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
)
with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
LTX2Pipeline(od_config=od_config)
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
@@ -144,16 +242,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
assert pipe.reports_stage_durations
def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
initialized = False
def stub_components(*_args, **_kwargs):
nonlocal initialized
initialized = True
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
with pytest.raises(ValueError, match="only one-stage LTX recipes"):
LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
cache_backend="cache_dit",
)
)
assert not initialized
def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace())
pipe = LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
lora_path=None,
model_config={},
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE
def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = LTX2DistilledPipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert pipe._phase_adapter is None
assert pipe.preserve_sp_padded_audio_duration
assert pipe.reports_stage_durations
def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts():
assert LTX2Pipeline.supports_request_batch
assert LTX2Pipeline.support_image_input
assert LTX2Pipeline.unified_text_image_entry
assert LTX2Pipeline.distributed_video_decode
assert not LTX2TwoStagePipeline.supports_request_batch
assert LTX2TwoStagePipeline.support_image_input
assert LTX2TwoStagePipeline.unified_text_image_entry
assert not LTX2DistilledPipeline.supports_request_batch
assert LTX2DistilledPipeline.support_image_input
assert LTX2DistilledPipeline.unified_text_image_entry
def test_ltx_two_stage_resolves_request_embedded_images():
pipe = object.__new__(LTX2TwoStagePipeline)
request_inputs = SimpleNamespace(latents=None)
image = torch.zeros(3, 8, 8)
req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}])
assert pipe._resolve_request_image(req, None, request_inputs) is image
def test_ltx_dmd2_entries_retain_their_task_boundaries():
assert not LTX2T2VDMD2Pipeline.support_image_input
assert LTX2I2VDMD2Pipeline.support_image_input
@@ -330,6 +516,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics():
assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",)
def test_ltx_positive_only_guidance_preserves_token_major_layout():
tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2)
repeated = _repeat_batch(tensor, 1)
assert repeated is tensor
assert repeated.stride() == tensor.stride()
def test_ltx_guidance_rescale_is_invariant_to_request_batching():
cond = torch.tensor(
[
@@ -380,6 +575,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction():
torch.testing.assert_close(actual, torch.zeros_like(cond))
def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens():
logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]])
logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]])
guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0)
expected = combine_guided_x0(
cond=logical_cond,
uncond_text=logical_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
)
padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1)
padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1)
actual = combine_guided_x0(
cond=padded_cond,
uncond_text=padded_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
rescale_token_count=2,
)
torch.testing.assert_close(actual[:, :2], expected)
def test_ltx_guidance_builds_per_sample_stg_and_modality_masks():
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1))
@@ -547,11 +768,19 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin)
assert issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe
assert "_forward_request" not in LTX2DistilledPipeline.__dict__
def test_ltx_two_stage_entry_is_a_thin_pipeline_variant():
assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__
assert "_forward_request" not in LTX2TwoStagePipeline.__dict__
assert "load_weights" not in LTX2TwoStagePipeline.__dict__
assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin)
assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline)
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
assert "_denoise_loop" not in LTX2Pipeline.__dict__
assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin)
@@ -594,7 +823,8 @@ def test_denoise_executor_owns_progress_and_interrupt():
torch.testing.assert_close(state.audio, torch.tensor(11.0))
def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
@@ -671,12 +901,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
prompt_context_sentinel = prompt_context
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
activated_slots = []
if pipeline_cls is LTX2TwoStagePipeline:
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
@@ -686,14 +919,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
assert len(phase_calls) == 2
assert phase_calls[1][2] is prompt_context_sentinel
if pipeline_cls is LTX2TwoStagePipeline:
assert activated_slots == [None, "ltx_distilled"]
torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls):
from vllm_omni.diffusion.models.ltx2 import ltx2_latents
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
@@ -750,6 +987,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only():
assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536)
def test_ltx23_full_distilled_route_uses_merged_weights_without_lora():
assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE
assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases)
assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None
assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls
assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == (
"ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
)
@pytest.mark.parametrize(
("request_inputs", "error"),
[
@@ -958,6 +1207,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting():
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
@pytest.mark.parametrize(
("recipe", "steps", "stg_block"),
[
(LTX2_TWO_STAGE_RECIPE, 40, 29),
(LTX23_TWO_STAGE_RECIPE, 30, 28),
],
)
def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block):
stage1, stage2 = recipe.phases
assert (recipe.width, recipe.height) == (1536, 1024)
assert recipe.num_inference_steps == steps
assert stage1.spatial_downscale == 2
assert stage1.adapter_slot is None
assert stage1.guidance.video.stg_blocks == (stg_block,)
assert stage1.guidance.audio.stg_blocks == (stg_block,)
assert stage2.input_transform == "spatial_upsample"
assert stage2.adapter_slot == "ltx_distilled"
assert stage2.guidance == LTXGuidanceSpec.positive_only()
assert stage2.num_inference_steps == 3
assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
# Official Stage 2 refines video only; its audio result is discarded.
assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0)
def test_ltx_two_stage_entries_select_the_official_adapter_slots():
phase_switches = []
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
for phase in LTX2_TWO_STAGE_RECIPE.phases:
ordinary_pipeline._enter_phase(phase)
assert phase_switches == [None, "ltx_distilled"]
distilled_pipeline = object.__new__(LTX2DistilledPipeline)
for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases:
distilled_pipeline._enter_phase(phase)
def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch):
class FakeColumn(torch.nn.Module):
def __init__(
self,
input_size,
output_size,
*,
bias,
gather_output,
return_bias,
quant_config,
prefix,
):
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.bias_enabled = bias
self.gather_output = gather_output
self.return_bias = return_bias
self.quant_config = quant_config
self.prefix = prefix
class FakeRow(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
class FakeAttention(torch.nn.Module):
def __init__(self, **kwargs):
super().__init__()
monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow)
monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention)
monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
attention = ltx2_transformer.LTX2Attention(
query_dim=16,
heads=4,
dim_head=4,
qk_norm="rms_norm",
apply_gated_attention=True,
pack_qkv=False,
prefix="transformer_blocks.0.attn1",
)
gate = attention.to_gate_logits
assert isinstance(gate, FakeColumn)
assert gate.input_size == 16
assert gate.output_size == 4
assert gate.gather_output is False
assert gate.return_bias is False
assert gate.quant_config is None
assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits"
class TestLTXRequestParsing:
def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -1432,6 +1773,7 @@ class TestPipelineComponents:
def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self):
assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
class TestLTX23DecodeConditioning:
@@ -1477,6 +1819,11 @@ class TestRegistryIntegration:
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline")
assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
"LTX2TwoStagePipeline",
)
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -1488,6 +1835,9 @@ class TestRegistryIntegration:
"LTX23ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2ImageToVideoTwoStagePipeline",
"LTX23TwoStagePipeline",
"LTX23ImageToVideoTwoStagePipeline",
}
assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
@@ -1497,6 +1847,7 @@ class TestRegistryIntegration:
expected = [
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1509,6 +1860,7 @@ class TestRegistryIntegration:
"model_class_name",
[
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1,13 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from torch import nn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
LTX2AudioVideoAttnProcessor,
LTX2VideoTransformer3DModel,
_LTX2ParallelAttention,
_make_rms_norm,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -66,3 +73,83 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
processor = LTX2AudioVideoAttnProcessor()
hidden_states = torch.zeros(1, 3, 4)
encoder_hidden_states = torch.zeros(1, 4, 4)
key_padding_mask = torch.tensor([[True, True, True, False]])
actual = processor._prepare_attention_mask(
SimpleNamespace(),
hidden_states,
encoder_hidden_states,
key_padding_mask,
batch_size=1,
sequence_length=4,
)
torch.testing.assert_close(actual, key_padding_mask[:, None, None, :])
def _fake_ltx_parallel_attention():
calls = []
class Backend:
@staticmethod
def get_name():
return "FLASH_ATTN"
@staticmethod
def supports_attention_mask():
return True
def run(name):
def forward(query, _key, _value, _metadata):
calls.append(name)
return torch.zeros_like(query)
return forward
layer = object.__new__(_LTX2ParallelAttention)
nn.Module.__init__(layer)
layer.attn_backend = Backend
layer.attention = SimpleNamespace(forward=run("native"))
layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
layer.backend_pref = "FLASH_ATTN"
return layer, calls
@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
layer, calls = _fake_ltx_parallel_attention()
query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_masked_self_attention_uses_sdpa():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True, True, True, False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_unmasked_cross_attention_keeps_native_backend():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
layer._run_local_attention(query, key, value, None)
assert calls == ["native"]
@@ -191,6 +191,30 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
]
def test_unquantized_qkv_declares_equivalent_mmap_transform(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "Attention", _FakeAttention)
arch = h3.MiniMaxH3DiTArchConfig(
hidden_size=1,
num_attention_heads=2,
attention_head_dim=1,
rope_inv_freq_len=1,
)
attention = h3.MiniMaxH3Attention(
arch,
quant_config=None,
prefix="blocks.0.attn",
)
transform = attention.qkv_proj.weight.mmap_weight_transform
checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)
assert transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
def test_pipeline_resolves_transformer_component_quant_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_resolve_component_quant_config,
@@ -30,7 +30,13 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
PinnedResidentLayerGroup,
)
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
from vllm_omni.diffusion.offloader.offload_plan import (
OffloadPlan,
get_offload_plan,
has_online_quantization,
should_select_dlo_mmap,
supports_mmap_loading,
)
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -48,6 +54,9 @@ class DummyEvent:
def record(self, _stream) -> None:
return None
def synchronize(self) -> None:
return None
@contextmanager
def dummy_stream(_stream):
@@ -350,6 +359,79 @@ class TestDistributedLayerwiseOffloadHook:
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime):
current_block = nn.Linear(2, 2, bias=False)
next_block = nn.Linear(2, 2, bias=False)
next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2))
source_storage = next_block.weight.untyped_storage().data_ptr()
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
next_block.weight.mmap_weight_transform_pending = True
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
rank_local_mmap=True,
)
hook.initialize_hook(current_block)
source = hook.cpu_sources[torch.float32][0]["tensor"]
assert source.untyped_storage().data_ptr() == source_storage
assert hook.cpu_shards == {}
assert next_block.weight.numel() == 0
hook.cpu_staging_buffers = [
{torch.float32: torch.empty(4)},
{torch.float32: torch.empty(4)},
]
hook.prefetch_layer(slot=0, non_blocking=False)
assert torch.equal(
next_block.weight,
torch.tensor([[2.0, 3.0], [0.0, 1.0]]),
)
assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2))
def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime):
hooks = []
for size in (4, 9):
current_block = nn.Linear(1, 1, bias=False)
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32))
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
shared_buffers=[None, None],
rank_local_mmap=True,
)
hook.initialize_hook(current_block)
hooks.append(hook)
resident_block = nn.Module()
resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32))
resident_group = PinnedResidentLayerGroup(
[resident_block],
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers(
hooks,
resident_group,
)
assert len(staging) == 2
assert staging[0][torch.float32].numel() == 12
assert staging[1][torch.float32].numel() == 12
class TestPinnedResidentLayerGroup:
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -386,6 +468,42 @@ class TestPinnedResidentLayerGroup:
group.offload()
def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)]
expected = []
source_storage = []
for index, block in enumerate(blocks):
block.weight.data.fill_(index + 1)
expected.append(block.weight.detach().clone())
source_storage.append(block.weight.untyped_storage().data_ptr())
group = PinnedResidentLayerGroup(
blocks,
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
for index, state in enumerate(group._states):
source = state["cpu_sources"][torch.float32][0]["tensor"]
assert source.untyped_storage().data_ptr() == source_storage[index]
assert state["cpu_shards"] == {}
assert len(group._cpu_staging_buffers) == 2
group.load()
for block, weight in zip(blocks, expected):
assert torch.equal(block.weight, weight)
blocks[0].weight.data.zero_()
group.offload()
group.load()
for block, weight in zip(blocks, expected):
assert torch.equal(block.weight, weight)
group.offload()
assert all(block.weight.numel() == 0 for block in blocks)
def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
module = nn.Linear(2, 2, bias=False)
@@ -779,6 +897,46 @@ class TestOffloadPlan:
with pytest.raises(Exception):
plan.block_attrs = {"x": ("y",)} # type: ignore
def test_explicit_component_sources_can_opt_in_to_mmap(self):
class Pipeline(nn.Module):
_supports_mmap_loading = True
weights_sources = [object()]
assert supports_mmap_loading(Pipeline())
def test_component_sources_do_not_implicitly_enable_mmap(self):
class Pipeline(nn.Module):
weights_sources = [object()]
assert not supports_mmap_loading(Pipeline())
def test_online_quantization_detection(self):
model = nn.Sequential(nn.Linear(1, 1))
assert not has_online_quantization(model)
model[0].quant_method = SimpleNamespace(uses_meta_device=True)
assert has_online_quantization(model)
@pytest.mark.parametrize(
("kwargs", "expected"),
[
({"dlo_use_allgather": True, "tensor_parallel_size": 2}, True),
({"dlo_use_allgather": False, "tensor_parallel_size": 1}, True),
({"dlo_use_allgather": False, "tensor_parallel_size": 2}, False),
({"dlo_use_allgather": False, "tensor_parallel_size": 1, "use_hsdp": True}, False),
({"dlo_use_allgather": False, "tensor_parallel_size": 1, "has_online_quant": True}, False),
],
)
def test_dlo_mmap_policy(self, kwargs, expected):
defaults = {
"mmap_supported": True,
"dlo_use_allgather": False,
"has_online_quant": False,
"tensor_parallel_size": 1,
"use_hsdp": False,
}
assert should_select_dlo_mmap(**(defaults | kwargs)) is expected
def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
class TokenRefiner(nn.Module):
def __init__(self):
@@ -827,6 +985,93 @@ class TestOffloadPlan:
class TestMmapValidation:
"""Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
def test_component_source_prefix_builds_model_namespace(self, tmp_path):
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
checkpoint_file = checkpoint_dir / "model.safetensors"
save_file({"blocks.0.weight": torch.arange(4, dtype=torch.float32)}, str(checkpoint_file))
class Pipeline(nn.Module):
_supports_mmap_loading = True
def __init__(self):
super().__init__()
self.weights_sources = [
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
)
]
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
),
torch.device("cpu"),
)
model_map = backend._build_mmap_model_map(Pipeline())
assert model_map == {
"transformer.blocks.0.weight": (
"blocks.0.weight",
str(checkpoint_file),
)
}
def test_explicit_mmap_opt_in_rejects_missing_model_tensors(self, tmp_path, patched_offload_runtime, monkeypatch):
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
lambda: 1,
)
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
save_file(
{"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)},
str(checkpoint_dir / "model.safetensors"),
)
class Transformer(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)])
class Pipeline(nn.Module):
_supports_mmap_loading = True
def __init__(self):
super().__init__()
self.transformer = Transformer()
self.weights_sources = [
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
)
]
pipeline = Pipeline()
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
),
torch.device("cpu"),
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
with pytest.raises(RuntimeError, match="Explicit mmap loading left 1"):
backend._load_weights_via_mmap(pipeline, modules)
def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch):
"""DLO+AllGather should reject when TP world size > 1."""
import json
@@ -199,7 +199,7 @@ def _generate_bagel_image(
config_path = get_deploy_config_path("ci/bagel.yaml")
omni_kwargs: dict[str, Any] = {
"model": "ByteDance-Seed/BAGEL-7B-MoT",
"stage_configs_path": config_path,
"deploy_config": config_path,
"stage_init_timeout": 300,
}
if diffusion_quantization_config:
@@ -314,7 +314,7 @@ def test_bagel_text2img_mooncake_connector(run_level):
temp_config_file = _resolve_deploy_config(temp_config_file, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=temp_config_file,
deploy_config=temp_config_file,
stage_init_timeout=300,
) as runner:
generated_image = _generate_bagel_image(runner.omni)
@@ -251,7 +251,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=config_path,
deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_img2img(runner.omni, input_image)
if run_level == "advanced_model":
@@ -266,7 +266,7 @@ def test_bagel_text2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=config_path,
deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_text2img(runner.omni)
if run_level == "advanced_model":
+117 -50
View File
@@ -20,6 +20,11 @@ from PIL import Image
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("official", "omni"), required=True)
parser.add_argument(
"--pipeline-kind",
choices=("one_stage", "distilled", "two_stage"),
default="one_stage",
)
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model")
@@ -27,6 +32,8 @@ def _parse_args() -> argparse.Namespace:
parser.add_argument("--official-root", type=Path)
parser.add_argument("--checkpoint", type=Path)
parser.add_argument("--gemma-root", type=Path)
parser.add_argument("--spatial-upsampler", type=Path)
parser.add_argument("--distilled-lora", type=Path)
parser.add_argument("--enable-layerwise-offload", action="store_true")
return parser.parse_args()
@@ -91,9 +98,13 @@ def _configure_official_sdpa(pipeline: Any) -> None:
attention=attention,
masked_attention=attention,
)
owners_and_attributes = (
(pipeline.stage, "_transformer_builder"),
owners_and_attributes = [
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
]
owners_and_attributes.extend(
(stage, "_transformer_builder")
for name in ("stage", "stage_1", "stage_2")
if (stage := getattr(pipeline, name, None)) is not None
)
for owner, attribute in owners_and_attributes:
builder = getattr(owner, attribute)
@@ -104,6 +115,12 @@ def _configure_official_sdpa(pipeline: Any) -> None:
)
def _require_path(path: Path | None, description: str) -> str:
if path is None or not path.is_file():
raise ValueError(f"Official {description} is required and must exist: {path}")
return str(path)
@torch.inference_mode()
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.official_root is None or args.checkpoint is None or args.gemma_root is None:
@@ -113,16 +130,54 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
torch.backends.cuda.enable_cudnn_sdp(False)
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.types import OffloadMode
pipeline = TI2VidOneStagePipeline(
checkpoint_path=str(args.checkpoint),
gemma_root=str(args.gemma_root),
loras=(),
offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE,
)
if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
raise ValueError(
f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
)
checkpoint = _require_path(args.checkpoint, "checkpoint")
gemma_root = str(args.gemma_root)
if not Path(gemma_root).is_dir():
raise ValueError(f"Official Gemma root is required and must be a directory: {gemma_root}")
offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
pipeline: Any
if args.pipeline_kind == "one_stage":
pipeline = TI2VidOneStagePipeline(
checkpoint_path=checkpoint,
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
elif args.pipeline_kind == "distilled":
pipeline = DistilledPipeline(
distilled_checkpoint_path=checkpoint,
spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
else:
distilled_lora = [
LoraPathStrengthAndSDOps(
path=_require_path(args.distilled_lora, "distilled LoRA"),
strength=1.0,
sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
)
]
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=checkpoint,
distilled_lora=distilled_lora,
spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
_configure_official_sdpa(pipeline)
image_path = request.get("image")
images = (
@@ -140,34 +195,45 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
)
]
)
video, audio = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
if args.pipeline_kind == "distilled":
video, audio = pipeline(
prompt=request["prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
images=images,
)
else:
video, audio = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
_save_outputs(
args.output_dir,
@@ -179,6 +245,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
"checkpoint": str(args.checkpoint),
"pipeline_kind": args.pipeline_kind,
},
)
@@ -193,15 +260,9 @@ def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio")
audio_sample_rate = multimodal_output.get("audio_sample_rate")
if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput):
frames = frames.request_output
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio", audio)
audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
if isinstance(frames, OmniRequestOutput):
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if isinstance(frames, list) and len(frames) == 1:
frames = frames[0]
@@ -253,6 +314,10 @@ def _canonical_video(video: Any) -> torch.Tensor:
def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.model is None or args.model_class_name is None:
raise ValueError("Omni backend requires --model and --model-class-name")
if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
raise ValueError(
f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
)
attention_config = {"default": {"backend": "TORCH_SDPA"}}
@@ -273,7 +338,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
parallel_config=DiffusionParallelConfig(),
)
try:
model_class_name = get_model_class_name(omni)
detected_model_class_name = get_model_class_name(omni)
model_class_name = args.model_class_name
sampling_params = OmniDiffusionSamplingParams(
height=request["height"],
width=request["width"],
@@ -289,10 +355,9 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
}
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance)
prompt: dict[str, Any] = {
"prompt": request["prompt"],
"negative_prompt": request["negative_prompt"],
}
prompt: dict[str, Any] = {"prompt": request["prompt"]}
if request.get("negative_prompt"):
prompt["negative_prompt"] = request["negative_prompt"]
image_path = request.get("image")
if image_path is not None:
with Image.open(str(image_path)) as source_image:
@@ -312,6 +377,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"model": args.model,
"model_class_name": model_class_name,
"model_config_class_name": detected_model_class_name,
"pipeline_kind": args.pipeline_kind,
},
)
finally:
@@ -3,10 +3,9 @@
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
The comparison runs both runtimes through PyTorch SDPA and uses
``max_batch_size=4`` in the official reference to match Omni's fused guidance
batch. Video and audio guidance use the official non-HQ one-stage defaults;
only the generation shape and step count are reduced for CI runtime.
The original reduced one-stage guards remain unchanged. Four complementary
default-shape cases cover both model versions, both two-stage weight families,
and T2V/I2V without expanding every combination into a duplicate golden.
"""
from __future__ import annotations
@@ -17,8 +16,9 @@ import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from dataclasses import asdict, dataclass, replace
from pathlib import Path
from typing import Literal
import numpy as np
import pytest
@@ -53,74 +53,260 @@ NEGATIVE_PROMPT = (
# Both runtimes use PyTorch SDPA with the current Torch dispatch defaults.
ATTENTION_BACKEND = "torch_sdpa"
VIDEO_SSIM_MEAN_THRESHOLD = 0.95
VIDEO_SSIM_MIN_THRESHOLD = 0.90
VIDEO_PSNR_MEAN_THRESHOLD = 30.0
AUDIO_RELATIVE_L2_THRESHOLD = 0.2
AUDIO_COSINE_THRESHOLD = 0.95
@dataclass(frozen=True)
class LTXArtifact:
repo_id: str
filename: str
revision: str
env: str
repo_type: Literal["model", "dataset"] = "model"
@dataclass(frozen=True)
class LTXAccuracyThresholds:
video_ssim_mean: float
video_ssim_min: float
video_psnr_mean_db: float
# Zero disables gating for waveform-sensitive audio that is not expected
# to align with the official output.
audio_relative_l2: float
audio_cosine_similarity: float
STRICT_THRESHOLDS = LTXAccuracyThresholds(
video_ssim_mean=0.95,
video_ssim_min=0.90,
video_psnr_mean_db=30.0,
audio_relative_l2=0.20,
audio_cosine_similarity=0.95,
)
@dataclass(frozen=True)
class LTXAccuracyCase:
name: str
pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
model_id: str
model_revision: str
model_env: str
model_class_name: str
checkpoint_repo: str
checkpoint_filename: str
checkpoint_revision: str
checkpoint_env: str
stg_block: int
checkpoint: LTXArtifact
width: int
height: int
num_frames: int
num_inference_steps: int
seed: int
stg_block: int | None
thresholds: LTXAccuracyThresholds
prompt: str = PROMPT
image_repo: str | None = None
image_filename: str | None = None
image_revision: str | None = None
gemma_model_id: str | None = None
gemma_model_revision: str | None = None
gemma_model_env: str | None = None
image: LTXArtifact | None = None
spatial_upsampler: LTXArtifact | None = None
distilled_lora: LTXArtifact | None = None
# Official uses block streaming; Omni uses layerwise offload. Both preserve
# bf16 arithmetic while keeping full-resolution two-stage cases on one H100.
enable_layerwise_offload: bool = False
CASES = (
LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a"
LTX23_REVISION = "4229404625088d21c4f112eb640fb04a0900ee25"
LTX2_CHECKPOINT = LTXArtifact(
repo_id="Lightricks/LTX-2",
filename="ltx-2-19b-dev.safetensors",
revision=LTX2_REVISION,
env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
)
LTX2_DISTILLED_CHECKPOINT = replace(
LTX2_CHECKPOINT,
filename="ltx-2-19b-distilled.safetensors",
env="VLLM_TEST_LTX2_DISTILLED_OFFICIAL_CHECKPOINT",
)
LTX2_UPSAMPLER = replace(
LTX2_CHECKPOINT,
filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
env="VLLM_TEST_LTX2_UPSAMPLER",
)
LTX2_DISTILLED_LORA = replace(
LTX2_CHECKPOINT,
filename="ltx-2-19b-distilled-lora-384.safetensors",
env="VLLM_TEST_LTX2_DISTILLED_LORA",
)
LTX23_CHECKPOINT = LTXArtifact(
repo_id="Lightricks/LTX-2.3",
filename="ltx-2.3-22b-dev.safetensors",
revision=LTX23_REVISION,
env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
)
LTX23_DISTILLED_CHECKPOINT = replace(
LTX23_CHECKPOINT,
# The pinned Diffusers checkpoint contains the original merged distilled
# Transformer. Keep the reference Transformer on the same version; the
# spatial upsampler is versioned independently and remains on 1.1 below.
filename="ltx-2.3-22b-distilled.safetensors",
env="VLLM_TEST_LTX23_DISTILLED_OFFICIAL_CHECKPOINT",
)
LTX23_UPSAMPLER = replace(
LTX23_CHECKPOINT,
filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
env="VLLM_TEST_LTX23_UPSAMPLER",
)
LTX23_DISTILLED_LORA = replace(
LTX23_CHECKPOINT,
filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
env="VLLM_TEST_LTX23_DISTILLED_LORA",
)
I2V_IMAGE = LTXArtifact(
repo_id="huggingface/documentation-images",
filename="diffusers/svd/rocket.png",
revision="645d8364f0c7a101180b364811b5a11a362e4010",
env="VLLM_TEST_LTX_I2V_IMAGE",
repo_type="dataset",
)
LEGACY_CASES = (
LTXAccuracyCase(
name="ltx2",
pipeline_kind="one_stage",
model_id="Lightricks/LTX-2",
model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2",
checkpoint_filename="ltx-2-19b-dev.safetensors",
checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
checkpoint=LTX2_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=29,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3_i2v",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
image_repo="huggingface/documentation-images",
image_filename="diffusers/svd/rocket.png",
image_revision="645d8364f0c7a101180b364811b5a11a362e4010",
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
TWO_STAGE_CASES = (
LTXAccuracyCase(
name="ltx2_distilled_t2v",
pipeline_kind="distilled",
model_id="rootonchair/LTX-2-19b-distilled",
model_revision="388e2846f54aae51687498ffb6b27c7c2c9ce9e5",
model_env="VLLM_TEST_LTX2_DISTILLED_MODEL",
model_class_name="LTX2DistilledPipeline",
checkpoint=LTX2_DISTILLED_CHECKPOINT,
spatial_upsampler=LTX2_UPSAMPLER,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=8,
seed=10,
stg_block=None,
gemma_model_id="Lightricks/LTX-2",
gemma_model_revision=LTX2_REVISION,
gemma_model_env="VLLM_TEST_LTX2_MODEL",
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_distilled_i2v",
pipeline_kind="distilled",
model_id="diffusers/LTX-2.3-Distilled-Diffusers",
model_revision="432e0d3c2d1769aaa4d295f9243f7062bf6b47ee",
model_env="VLLM_TEST_LTX23_DISTILLED_MODEL",
model_class_name="LTX2DistilledPipeline",
checkpoint=LTX23_DISTILLED_CHECKPOINT,
spatial_upsampler=LTX23_UPSAMPLER,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=8,
seed=10,
stg_block=None,
gemma_model_id="diffusers/LTX-2.3-Diffusers",
gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
gemma_model_env="VLLM_TEST_LTX23_MODEL",
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
LTXAccuracyCase(
name="ltx2_two_stage_layer_fused_t2v",
pipeline_kind="two_stage",
model_id="Lightricks/LTX-2",
model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2TwoStagePipeline",
checkpoint=LTX2_CHECKPOINT,
spatial_upsampler=LTX2_UPSAMPLER,
distilled_lora=LTX2_DISTILLED_LORA,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=40,
seed=10,
stg_block=29,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_two_stage_layer_fused_i2v",
pipeline_kind="two_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2TwoStagePipeline",
checkpoint=LTX23_CHECKPOINT,
spatial_upsampler=LTX23_UPSAMPLER,
distilled_lora=LTX23_DISTILLED_LORA,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=30,
seed=10,
stg_block=28,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
start = time.perf_counter()
subprocess.run(command, env=env, timeout=timeout, check=True)
@@ -209,6 +395,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
"model_index.json",
"audio_vae/*",
"connectors/*",
"latent_upsampler/*",
"processor/*",
"scheduler/*",
"text_encoder/config.json",
@@ -223,74 +410,82 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
)
def _resolve_gemma_root(model: Path) -> Path:
def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path:
configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
if configured_root:
root = Path(configured_root)
assert root.is_dir(), f"Gemma root not found: {root}"
return root
if case.gemma_model_id is not None:
configured_model = os.environ.get(case.gemma_model_env or "")
if configured_model and Path(configured_model).is_dir():
return Path(configured_model)
return Path(
snapshot_download(
repo_id=case.gemma_model_id,
revision=case.gemma_model_revision,
allow_patterns=[
"processor/*",
"text_encoder/*",
"tokenizer/*",
],
)
)
return model
def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path:
configured_checkpoint = os.environ.get(case.checkpoint_env)
if configured_checkpoint:
checkpoint = Path(configured_checkpoint)
assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}"
return checkpoint
model_checkpoint = model / case.checkpoint_filename
if model_checkpoint.is_file():
return model_checkpoint
def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path:
configured_path = os.environ.get(artifact.env)
if configured_path:
path = Path(configured_path)
assert path.is_file(), f"Configured LTX artifact not found: {path}"
return path
if model is not None:
model_path = model / artifact.filename
if model_path.is_file():
return model_path
return Path(
hf_hub_download(
repo_id=case.checkpoint_repo,
filename=case.checkpoint_filename,
revision=case.checkpoint_revision,
repo_id=artifact.repo_id,
repo_type=None if artifact.repo_type == "model" else artifact.repo_type,
filename=artifact.filename,
revision=artifact.revision,
)
)
def _resolve_image(case: LTXAccuracyCase) -> Path | None:
if case.image_filename is None:
if case.image is None:
return None
if case.image_repo is None or case.image_revision is None:
raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.")
configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE")
if configured_image:
image = Path(configured_image)
assert image.is_file(), f"LTX I2V conditioning image not found: {image}"
return image
return Path(
hf_hub_download(
repo_id=case.image_repo,
repo_type="dataset",
filename=case.image_filename,
revision=case.image_revision,
)
)
return _resolve_artifact(case.image)
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
request: dict[str, object] = {
"pipeline_kind": case.pipeline_kind,
"prompt": case.prompt,
"negative_prompt": NEGATIVE_PROMPT,
"width": 512,
"height": 384,
"num_frames": 25,
"negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
"width": case.width,
"height": case.height,
"num_frames": case.num_frames,
"fps": 24,
"num_inference_steps": 20,
"seed": 42,
"video_cfg_scale": 3.0,
"audio_cfg_scale": 7.0,
"video_stg_scale": 1.0,
"audio_stg_scale": 1.0,
"video_modality_scale": 3.0,
"audio_modality_scale": 3.0,
"video_rescale_scale": 0.7,
"audio_rescale_scale": 0.7,
"video_stg_blocks": [case.stg_block],
"audio_stg_blocks": [case.stg_block],
"num_inference_steps": case.num_inference_steps,
"seed": case.seed,
}
if case.pipeline_kind != "distilled":
assert case.stg_block is not None
request.update(
video_cfg_scale=3.0,
audio_cfg_scale=7.0,
video_stg_scale=1.0,
audio_stg_scale=1.0,
video_modality_scale=3.0,
audio_modality_scale=3.0,
video_rescale_scale=0.7,
audio_rescale_scale=0.7,
video_stg_blocks=[case.stg_block],
audio_stg_blocks=[case.stg_block],
)
if image is not None:
request["image"] = str(image.resolve())
return request
@@ -299,20 +494,30 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
device = torch.device("cpu")
ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device)
psnr = PeakSignalNoiseRatio(data_range=1.0).to(device)
ssim_scores: list[float] = []
psnr_scores: list[float] = []
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
max_abs = 0.0
absolute_error_sum = 0.0
with torch.inference_mode():
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device)
ssim_scores.append(float(ssim(prediction_tensor, reference_tensor)))
psnr_scores.append(float(psnr(prediction_tensor, reference_tensor)))
ssim.reset()
psnr.reset()
difference = np.abs(reference_frame.astype(np.float64) - prediction_frame.astype(np.float64))
max_abs = max(max_abs, float(difference.max()))
absolute_error_sum += float(difference.sum())
return {
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
"max_abs": float(difference.max()),
"mean_abs": float(difference.mean()),
"max_abs": max_abs,
"mean_abs": absolute_error_sum / reference.size,
}
@@ -343,13 +548,19 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
"""Compare official and Omni raw AV outputs from the same E2E request."""
def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
"""Compare official and Omni raw AV outputs from the same single-H100 request."""
configured_artifact_root = os.environ.get("VLLM_TEST_LTX_ARTIFACT_ROOT")
if configured_artifact_root:
accuracy_artifact_root = Path(configured_artifact_root)
accuracy_artifact_root.mkdir(parents=True, exist_ok=True)
output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name)
official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official")
model = _resolve_model(case)
gemma_root = _resolve_gemma_root(model)
checkpoint = _resolve_checkpoint(case, model)
gemma_root = _resolve_gemma_root(case, model)
checkpoint = _resolve_artifact(case.checkpoint, model)
spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
image = _resolve_image(case)
request_path = output_root / "request.json"
request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
@@ -359,8 +570,13 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(runner),
"--request",
str(request_path),
"--pipeline-kind",
case.pipeline_kind,
]
if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}:
enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get(
"VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", ""
).lower() in {"1", "true", "yes", "on"}
if enable_layerwise_offload:
runner_args.append("--enable-layerwise-offload")
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
@@ -372,10 +588,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
)
official_output = output_root / "official"
_run(
_official_runner_prefix()
+ runner_args
+ [
official_args = _official_runner_prefix() + runner_args
official_args.extend(
[
"--backend",
"official",
"--output-dir",
@@ -386,12 +601,16 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(checkpoint),
"--gemma-root",
str(gemma_root),
],
env=env,
]
)
if spatial_upsampler is not None:
official_args.extend(["--spatial-upsampler", str(spatial_upsampler)])
if distilled_lora is not None:
official_args.extend(["--distilled-lora", str(distilled_lora)])
_run(official_args, env=env)
omni_output = output_root / "omni"
_run(
omni_args = (
[sys.executable]
+ runner_args
+ [
@@ -403,9 +622,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(model),
"--model-class-name",
case.model_class_name,
],
env=env,
]
)
_run(omni_args, env=env)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
@@ -423,18 +642,27 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
result = {
"case": case.name,
"task": "i2v" if image is not None else "t2v",
"pipeline_kind": case.pipeline_kind,
"model_class_name": case.model_class_name,
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"model_revision": case.model_revision,
"checkpoint_revision": case.checkpoint_revision,
"checkpoint_revision": case.checkpoint.revision,
"spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
"distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
"enable_layerwise_offload": enable_layerwise_offload,
"thresholds": asdict(case.thresholds),
"request": _request(case, image),
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD
assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD
assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD
assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
if case.thresholds.audio_relative_l2 > 0:
assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
if case.thresholds.audio_cosine_similarity > 0:
assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
@@ -17,7 +17,7 @@ from vllm_omni.platforms import current_omni_platform
# Override in CI if needed: AUDIOX_TEST_MODEL=<model-or-local-path>
models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models]
pytestmark = [
@@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml")
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG)
@@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
]
# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_cosy_deployment = get_deploy_config_path("cosyvoice3.yaml")
_cosy_model_path = str(_resolve_model_dir())
_OMNI_RUNNER_PARAMS = [
@@ -17,7 +17,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "baidu/ERNIE-Image"
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAM = (
MODEL,
None,
@@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None:
with OmniRunner(
MODEL,
stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
deploy_config=_get_deploy_config(async_chunk=async_chunk),
stage_init_timeout=600,
trust_remote_code=True,
) as omni_runner:
@@ -52,7 +52,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
def omni() -> Generator[Omni, None, None]:
with OmniRunner(
MODEL_NAME,
stage_configs_path=str(STAGE_CONFIG_PATH),
deploy_config=str(STAGE_CONFIG_PATH),
trust_remote_code=True,
) as runner:
yield runner.omni
@@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]:
with OmniRunner(
model,
seed=42,
stage_configs_path=None,
model_class_name="LongCatVideoAvatarPipeline",
additional_config={
"model_type": "avatar-v1.5",
@@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
@@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"]
_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models]
test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models]
def get_question(prompt_type: str = "text") -> str:
@@ -25,7 +25,7 @@ from vllm_omni import Omni
MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml")
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAM = (
MODEL_NAME,
STAGE_CONFIG,
@@ -20,13 +20,11 @@ from tests.helpers.stage_config import get_deploy_config_path
MODEL = "k2-fsa/OmniVoice"
DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
# The migrated test passes deploy_config through extra Omni kwargs.
# OmniRunner tuple: model, deploy config path, extra Omni kwargs.
_OMNI_RUNNER_PARAM = (
MODEL,
None,
DEPLOY_CONFIG,
{
"deploy_config": DEPLOY_CONFIG,
"trust_remote_code": True,
"log_stats": True,
},
@@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
"""Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
with OmniRunner(MODEL, stage_configs_path=None) as runner:
with OmniRunner(MODEL) as runner:
omni = runner.omni
baseline = _generate(omni)
img_1x = _generate(omni, lora_request, lora_scale=1.0)
@@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img(
def test_sensenova_u1_img2img(run_level):
"""Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
input_image = _load_input_image()
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
) as runner:
with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
if run_level == "full_model" or run_level == "advanced_model":
_validate_pixels(generated_image)
@@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit():
input_image = _load_input_image()
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
cache_backend="cache_dit",
) as runner:
_generate_sensenova_u1_img2img(runner.omni, input_image)
@@ -151,10 +151,7 @@ def _generate_sensenova_u1_image(
@hardware_test(res={"cuda": "H100"})
def test_sensenova_u1_text2img(run_level):
"""Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML)."""
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
) as runner:
with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
generated_image = _generate_sensenova_u1_image(runner.omni)
if run_level == "full_model" or run_level == "advanced_model":
_validate_pixels(generated_image)
@@ -168,7 +165,6 @@ def test_sensenova_u1_text2img_cache_dit():
"""Test SenseNova-U1 text2img with Cache-DiT enabled."""
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
stage_configs_path=None,
cache_backend="cache_dit",
) as runner:
_generate_sensenova_u1_image(runner.omni)
@@ -159,7 +159,7 @@ def test_soulxsinger_multistage_from_audio(
with OmniRunner(
model,
stage_configs_path=get_deploy_config_path(deploy_yaml),
deploy_config=get_deploy_config_path(deploy_yaml),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
@@ -212,7 +212,7 @@ def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) ->
with OmniRunner(
str(base_dir),
stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"),
async_chunk=False,
) as runner:
sampling = OmniDiffusionSamplingParams(
@@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput
MODEL = "stepfun-ai/Step-Audio-2-mini"
DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml")
TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})]
TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})]
pytestmark = [pytest.mark.slow, pytest.mark.tts]
@@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming():
async def _run():
async_omni = AsyncOmni(
model=MODEL,
stage_configs_path=STAGE_CONFIG,
deploy_config=STAGE_CONFIG,
stage_init_timeout=300,
enforce_eager=True,
)
+109
View File
@@ -0,0 +1,109 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Online-serving smoke tests for the public LTX-2 pipeline entries."""
import os
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers")
BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers")
PROMPT = "A serene lake at sunset with mountains in the background."
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
HSDP_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
def _server(
model: str,
model_class_name: str,
*extra_args: str,
) -> OmniServerParams:
return OmniServerParams(
model=model,
server_args=[
"--enforce-eager",
*extra_args,
"--model-class-name",
model_class_name,
],
)
def _cases():
return [
pytest.param(
_server(BASE_MODEL, "LTX2Pipeline"),
3,
id="one_stage",
marks=SINGLE_CARD_MARKS,
),
pytest.param(
_server(DISTILLED_MODEL, "LTX2DistilledPipeline"),
8,
id="distilled",
marks=SINGLE_CARD_MARKS,
),
pytest.param(
_server(
DISTILLED_MODEL,
"LTX2DistilledPipeline",
"--use-hsdp",
"--hsdp-shard-size",
"2",
),
8,
id="distilled_hsdp2",
marks=HSDP_MARKS,
),
pytest.param(
_server(
BASE_MODEL,
"LTX2TwoStagePipeline",
"--enable-layerwise-offload",
),
3,
id="two_stage",
marks=SINGLE_CARD_MARKS,
),
]
@pytest.mark.parametrize(
("omni_server", "num_inference_steps"),
_cases(),
indirect=["omni_server"],
)
def test_ltx_pipeline_entries(
omni_server: OmniServer,
num_inference_steps: int,
openai_client: OpenAIClientHandler,
subtests,
):
for task in ("t2v", "i2v"):
with subtests.test(task=task):
form_data = {
"prompt": PROMPT,
"model": omni_server.model,
"height": 128,
"width": 128,
"num_frames": 9,
"fps": 8,
"num_inference_steps": num_inference_steps,
"seed": 42,
}
request_config = {
"model": omni_server.model,
"form_data": form_data,
}
if task == "i2v":
request_config["image_reference"] = (
f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
)
openai_client.send_video_diffusion_request(request_config)
@@ -1,91 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Expansion tests for LTX-2 two-stage pipelines.
Coverage:
- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video)
"""
import os
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled")
PROMPT = "A serene lake at sunset with mountains in the background."
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
def _cases():
cases = []
cases.append(
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2DistilledPipeline",
],
),
False,
id="t2v_hsdp",
marks=PARALLEL_MARKS,
)
)
cases.append(
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2DistilledPipeline",
],
),
True,
id="i2v_hsdp",
marks=PARALLEL_MARKS,
)
)
return cases
@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"])
def test_ltx2_two_stage_hsdp(
omni_server: OmniServer,
is_i2v: bool,
openai_client: OpenAIClientHandler,
):
# Keep CI small while exercising the fixed 8+3-step distilled recipe.
# Height and width describe the final output; Stage 1 runs at half size.
form_data = {
"prompt": PROMPT,
"model": omni_server.model,
"height": 128,
"width": 128,
"num_frames": 9,
"fps": 8,
"num_inference_steps": 8,
"seed": 42,
}
request_config = {
"model": omni_server.model,
"form_data": form_data,
}
if is_i2v:
request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
openai_client.send_video_diffusion_request(request_config)
-139
View File
@@ -4,7 +4,6 @@ invariant to the specific attributes of vLLM config except in cases where we
explicitly patch values that differ from vLLM.
"""
import argparse
import inspect
import json
import os
@@ -20,7 +19,6 @@ from vllm.engine.arg_utils import EngineArgs
from vllm_omni.config.model import OmniModelConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_init_utils import build_engine_args_dict
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -149,19 +147,6 @@ def test_qwen3_tts_startup_task_type_is_validated():
)
def test_from_cli_args_picks_up_stage_configs_path():
"""from_cli_args should pick up stage_configs_path from namespace."""
ns = argparse.Namespace(
model="facebook/opt-125m",
stage_configs_path="/some/path.yaml",
custom_pipeline_args=None,
)
args = OmniEngineArgs.from_cli_args(ns)
assert args.stage_configs_path == "/some/path.yaml"
assert args.custom_pipeline_args is None
def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
"""Ensure Code2Wav mirrors stage max_model_len into nested HF overrides.
@@ -418,130 +403,6 @@ def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size
def test_stage_configs_path_field():
"""OmniEngineArgs with stage_configs_path should construct without error."""
args = OmniEngineArgs(stage_configs_path="/some/path.yaml")
assert args.stage_configs_path == "/some/path.yaml"
def test_strip_single_engine_args():
"""_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
kwargs = {
# Parent EngineArgs fields — stripped unless explicitly allowlisted
"compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
"tensor_parallel_size": 4,
"gpu_memory_utilization": 0.9,
"model": "some/model",
# Parent field that should be kept (allowlisted)
"worker_extension_cls": "some.Extension",
# OmniEngineArgs-only / non-engine fields — should pass through
"stage_configs_path": "/path/to/yaml",
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
"mode": "text-to-image",
"lora_path": "/some/lora",
}
filtered = AsyncOmniEngine._strip_single_engine_args(kwargs)
# Stripped — parent EngineArgs fields
assert "compilation_config" not in filtered
assert filtered["tensor_parallel_size"] == 4
assert "gpu_memory_utilization" not in filtered
assert "model" not in filtered
# Stripped — orchestrator-level OmniEngineArgs field
assert "stage_configs_path" not in filtered
# Kept
assert filtered["worker_extension_cls"] == "some.Extension"
assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"}
assert filtered["mode"] == "text-to-image"
assert filtered["lora_path"] == "/some/lora"
def test_strip_single_engine_args_model_does_not_trigger_warning(mocker):
"""model is always in kwargs (callers set it via from_cli_args/asdict),
so it should not cause the override warning by itself or appear in it."""
mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning")
# Typical caller kwargs: model is always present, no other parent
# EngineArgs fields are explicitly overridden.
AsyncOmniEngine._strip_single_engine_args(
{
"model": "some/model",
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
}
)
mock_warn.assert_not_called()
# When there *are* genuinely surprising overrides alongside model,
# the warning should mention them but not model. Keep-listed fields such as
# tensor_parallel_size are intentionally passed through and should not warn.
AsyncOmniEngine._strip_single_engine_args(
{
"model": "some/model",
"compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
"tensor_parallel_size": 4,
"custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
}
)
mock_warn.assert_called_once()
warned_args = mock_warn.call_args[0][-1] # the formatted arg list
assert "compilation_config" in warned_args
assert "tensor_parallel_size" not in warned_args
assert "model" not in warned_args
def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path):
"""New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs.
Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping
those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni
accuracy from ~78% to ~70%.
"""
from vllm_omni.entrypoints.utils import is_new_format_deploy_config
yaml_path = tmp_path / "minicpm_deploy.yaml"
yaml_path.write_text(
"""
pipeline: minicpmo_4_5
async_chunk: true
stages:
- stage_id: 0
devices: "0"
""".strip(),
encoding="utf-8",
)
path = str(yaml_path)
assert is_new_format_deploy_config(path)
kwargs = {
"stage_configs_path": path,
"trust_remote_code": True,
"interleave_mm_strings": True,
"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
"gpu_memory_utilization": 0.55,
}
# Mirror AsyncOmniEngine._resolve_stage_configs strip decision.
if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]):
base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs))
else:
base = dict(kwargs)
assert base.get("interleave_mm_strings") is True
assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}}
# Parent EngineArgs fields also remain available for deploy merge.
assert base.get("gpu_memory_utilization") == 0.55
legacy = tmp_path / "legacy.yaml"
legacy.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
assert not is_new_format_deploy_config(str(legacy))
legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy))
stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs)
assert "interleave_mm_strings" not in stripped
assert "media_io_kwargs" not in stripped
assert "gpu_memory_utilization" not in stripped
# For https://github.com/vllm-project/vllm-omni/issues/3293
def test_tensor_parallel_size_none_is_handled():
"""Ensure the tensor parallel size of None isn't forwarded."""
+1 -1
View File
@@ -68,7 +68,7 @@ async def test_abort():
with ExitStack() as after:
engine = AsyncOmni(
model=model,
stage_configs_path=stage_config,
deploy_config=stage_config,
)
after.callback(engine.shutdown)
+2 -2
View File
@@ -428,7 +428,7 @@ class TestSingleStageModeDetection:
assert engine.single_stage_mode is False
assert engine._single_stage_id_filter is None
def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture):
def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture):
duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
get_pipeline_config = mocker.patch(
"vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
@@ -441,7 +441,7 @@ class TestSingleStageModeDetection:
engine = self._make_engine_no_thread(
mocker,
stage_configs_path="/fake/duplex.yaml",
deploy_config="/fake/duplex.yaml",
)
get_pipeline_config.assert_called_once_with(
+2
View File
@@ -139,6 +139,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st
engine_extras={
"model": str(stage_model),
"attention_backend": "FLASHINFER",
"attention_config": {"use_trtllm_attention": False},
"hf_overrides": {"rope_scaling": {"factor": 2}},
"limit_mm_per_prompt": {"audio": 1},
"logits_processors": ["test.custom.LogitsProcessor"],
@@ -229,6 +230,7 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
assert thinker_args["omni_kv_config"] == {"need_send_cache": True}
assert thinker_args["worker_cls"] == "test.custom.Worker"
assert thinker_args["attention_backend"] == "FLASHINFER"
assert thinker_args["attention_config"] == {"use_trtllm_attention": False}
assert thinker_args["trust_remote_code"] is True
assert thinker_args["dtype"] == "float16"
assert thinker_args["distributed_executor_backend"] == "mp"
+1 -1
View File
@@ -364,7 +364,7 @@ async def test_omni_generate_request_id():
Text modality only; asserts caller-visible ids are preserved across the
multi-stage orchestrator path on H100.
"""
engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG)
try:
for req_id in _OMNI_REQ_IDS:
async for output in engine.generate(
@@ -478,12 +478,11 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
engine._strip_single_engine_args = lambda kwargs: kwargs
_, stage_configs = engine._resolve_stage_configs(
"dummy-model",
{
"stage_configs_path": "dummy.yaml",
"deploy_config": "dummy.yaml",
"additional_config": {"torchair_graph_config": {"enabled": True}},
},
trust_remote_code=False,
@@ -516,12 +515,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
engine._strip_single_engine_args = lambda kwargs: kwargs
_, stage_configs = engine._resolve_stage_configs(
"dummy-model",
{
"stage_configs_path": "dummy.yaml",
"deploy_config": "dummy.yaml",
"quantization_config": {"method": "bitsandbytes"},
},
trust_remote_code=False,
+1 -1
View File
@@ -193,7 +193,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
run_headless(args)
assert captured["args"][1] is None
assert len(captured["args"]) == 2
assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
+9 -7
View File
@@ -381,7 +381,6 @@ class TestLoadAndResolveStageConfigs:
config_path, stage_configs, _ = load_and_resolve_stage_configs(
model="black-forest-labs/FLUX.2-klein-4B",
stage_configs_path=None,
kwargs=kwargs,
trust_remote_code=False,
default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
@@ -390,9 +389,7 @@ class TestLoadAndResolveStageConfigs:
assert len(stage_configs) == 1
assert "dtype" in stage_configs[0]["engine_args"]
def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
self, tmp_path, mocker: MockerFixture
):
def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
deploy_path = tmp_path / "qwen3_multi.yaml"
deploy_path.write_text(
'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n',
@@ -413,18 +410,23 @@ class TestLoadAndResolveStageConfigs:
"vllm_omni.entrypoints.utils.load_stage_configs_from_model",
return_value=(returned_stage_configs, None),
)
cli_overrides = {
"interleave_mm_strings": True,
"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
"gpu_memory_utilization": 0.55,
}
config_path, stage_configs, _ = load_and_resolve_stage_configs(
model="dummy-model",
stage_configs_path=str(deploy_path),
kwargs={},
kwargs=cli_overrides,
trust_remote_code=True,
deploy_config_path=str(deploy_path),
)
load_stage_configs.assert_called_once_with(
"dummy-model",
trust_remote_code=True,
base_engine_args={},
base_engine_args=cli_overrides,
deploy_config_path=str(deploy_path),
stage_overrides=None,
strategy_config_path=None,
+4 -9
View File
@@ -2635,7 +2635,7 @@ class OmniRunner:
# affects the test runner wrapper.
init_timeout: int = 1800,
log_stats: bool = False,
stage_configs_path: str | None = None,
deploy_config: str | None = None,
**kwargs,
) -> None:
startup_t0 = time.perf_counter()
@@ -2653,7 +2653,7 @@ class OmniRunner:
stage_init_timeout=stage_init_timeout,
batch_timeout=batch_timeout,
init_timeout=init_timeout,
stage_configs_path=stage_configs_path,
deploy_config=deploy_config,
**kwargs,
)
startup_s = time.perf_counter() - startup_t0
@@ -3179,12 +3179,7 @@ def iter_omni_server(
print("OmniServer stopping...")
else:
if stage_config_path is not None:
from vllm_omni.entrypoints.utils import is_new_format_deploy_config
if is_new_format_deploy_config(stage_config_path):
server_args += ["--deploy-config", stage_config_path]
else:
server_args += ["--stage-configs-path", stage_config_path]
server_args += ["--deploy-config", stage_config_path]
with (
OmniServer(
@@ -3237,7 +3232,7 @@ def iter_omni_runner(
model = model_prefix + model
if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
model = resolve_tiny_model_path(model)
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner:
print("OmniRunner started successfully")
yield runner
print("OmniRunner stopping...")
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression tests for the dots.tts talker's stop-signal pairing and
per-request state lifecycle.
Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes
the methods under test read) — no weights, no GPU, no engine. Mirrors
tests/model_executor/models/voxcpm2/test_talker_state_eviction.py.
"""
from __future__ import annotations
import functools
from types import SimpleNamespace
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
torch = pytest.importorskip("torch")
_VOCAB_SIZE = 151672
@functools.lru_cache(maxsize=1)
def _dots_tts_talker_mod():
"""Defer talker import (pulls vLLM model_executor) until first use."""
from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import (
DotsTTSForConditionalGeneration,
_RequestState,
)
return DotsTTSForConditionalGeneration, _RequestState
def _make_bare_talker():
DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod()
talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration)
talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE)
talker._active_states = {}
talker._pending_requests = []
talker._results_queue = []
talker._audio_queue = []
talker._deferred_cleanup_ids = set()
talker._beta_trace = False
return talker
def _add_state(talker, req_id: str, *, is_stopping: bool = False):
_, _RequestState = _dots_tts_talker_mod()
state = _RequestState(request_id=req_id)
state.prefill_completed = True
state.is_stopping = is_stopping
talker._active_states[req_id] = state
return state
def _hidden(bsz: int) -> torch.Tensor:
return torch.zeros(bsz, 1536, dtype=torch.float32)
class TestComputeLogitsPairing:
"""compute_logits pairs _results_queue entries with batch rows
positionally, so every request must push exactly one entry per step —
including prefills, which push a (req_id, None) placeholder."""
def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None:
talker = _make_bare_talker()
_add_state(talker, "prefill-req")
decode_state = _add_state(talker, "decode-req", is_stopping=True)
decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]])
# Row 0 is the prefill (None placeholder), row 1 the stopping decode.
talker._results_queue = [
("prefill-req", None),
("decode-req", torch.tensor([[0.1, 0.9]])),
]
logits = talker.compute_logits(_hidden(2))
assert logits.shape == (2, _VOCAB_SIZE)
# Prefill row: forced continue.
assert logits[0, 0].item() == 1.0
assert logits[0, 1].item() == float("-inf")
# Stopping decode row: stop slot wins.
assert logits[1, 1].item() == 1.0
assert logits[1, 0].item() == 0.0
# Everything else stays masked.
assert logits[0, 2].item() == float("-inf")
assert logits[1, 2].item() == float("-inf")
# Queue drained.
assert talker._results_queue == []
def test_below_threshold_stop_logits_force_continue(self) -> None:
"""Raw (continue, stop) softmax must not reach the sampler: a
prob_stop in (0.5, 0.8] would flip greedy argmax even though
is_stopping stayed False (the 0.8-threshold contract)."""
talker = _make_bare_talker()
_add_state(talker, "req", is_stopping=False)
talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))]
logits = talker.compute_logits(_hidden(1))
assert logits[0, 0].item() == 1.0
assert logits[0, 1].item() == float("-inf")
def test_empty_queue_defaults_to_all_continue(self) -> None:
talker = _make_bare_talker()
logits = talker.compute_logits(_hidden(3))
assert torch.all(logits[:, 0] == 1.0)
assert torch.all(logits[:, 1] == float("-inf"))
class TestStateEviction:
"""on_requests_finished fires BEFORE forward(), so eviction must be
deferred to _flush_deferred_cleanup at the end of the step — the
finishing request's audio still drains through the current forward."""
def test_finished_request_survives_until_flush(self) -> None:
talker = _make_bare_talker()
_add_state(talker, "req-a")
_add_state(talker, "req-b")
talker.on_requests_finished(["req-a"])
assert "req-a" in talker._active_states # not evicted yet
talker._flush_deferred_cleanup()
assert "req-a" not in talker._active_states
assert "req-b" in talker._active_states
assert talker._deferred_cleanup_ids == set()
def test_unknown_request_id_is_ignored(self) -> None:
talker = _make_bare_talker()
talker.on_requests_finished(["never-seen"])
talker._flush_deferred_cleanup()
assert talker._active_states == {}
@@ -1326,7 +1326,9 @@ class TestRegistry:
assert "higgs_multimodal_qwen3" in OMNI_PIPELINES
def test_deploy_yaml_exists(self):
def test_deploy_yaml_uses_native_flashinfer(self):
from vllm_omni.config.stage_config import load_deploy_config
for name in (
"higgs_multimodal_qwen3.yaml",
"higgs_multimodal_qwen3_high_throughput.yaml",
@@ -1334,6 +1336,10 @@ class TestRegistry:
):
yaml_path = get_deploy_config_path(name)
assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
deploy_config = load_deploy_config(yaml_path)
stage_0 = next(stage for stage in deploy_config.stages if stage.stage_id == 0)
assert stage_0.engine_extras["attention_backend"] == "FLASHINFER"
assert stage_0.engine_extras["attention_config"] == {"use_trtllm_attention": False}
# ---- AC-3: Prompt Builder ----
+7 -4
View File
@@ -233,10 +233,13 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None:
}
)
assert get_extra_body_params("LTX2Pipeline") == expected
assert get_extra_output_params("LTX2Pipeline") == frozenset()
assert get_extra_body_params("LTX2DistilledPipeline") == expected
assert get_extra_output_params("LTX2DistilledPipeline") == frozenset()
for pipeline_name in (
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
):
assert get_extra_body_params(pipeline_name) == expected
assert get_extra_output_params(pipeline_name) == frozenset()
@pytest.mark.core_model
@@ -27,6 +27,7 @@ EXCLUDED_MODELS = [
"OvisImagePipeline",
"WanPipeline",
"WanVACEPipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
+1 -1
View File
@@ -147,7 +147,7 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
# (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time.
omni = Omni(
model,
stage_configs_path=yaml_path,
deploy_config=yaml_path,
stage_init_timeout=1800,
init_timeout=3600,
async_chunk=async_mode,
+2 -1
View File
@@ -53,7 +53,6 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
"model",
"omni",
"output_modalities",
"stage_configs_path",
"stage_id",
"tokenizer",
}
@@ -91,6 +90,7 @@ class _ModelEngineOverrides(TypedDict, total=False):
trust_remote_code: bool
dtype: Any
attention_backend: Any
attention_config: Any
moe_backend: str
hf_overrides: Any
limit_mm_per_prompt: dict[str, Any]
@@ -322,6 +322,7 @@ class OmniStageModelConfig:
trust_remote_code: bool = False
dtype: Any = "auto"
attention_backend: Any = None
attention_config: Any = None
moe_backend: str = "auto"
hf_overrides: Any = None
limit_mm_per_prompt: dict[str, Any] | None = None
+2
View File
@@ -47,6 +47,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import (
)
from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE
from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE
from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
@@ -144,6 +145,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE,
"wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
"voxcpm2": VOXCPM2_PIPELINE,
"dots_tts": DOTS_TTS_PIPELINE,
"cosyvoice3": COSYVOICE3_PIPELINE,
"audex_tts": AUDEX_TTS_PIPELINE,
"audex_tta": AUDEX_TTA_PIPELINE,
+1 -42
View File
@@ -4,7 +4,6 @@
from __future__ import annotations
import dataclasses
import functools
import re
import warnings
@@ -56,7 +55,7 @@ def build_stage_runtime_overrides(
``internal_keys`` defaults to the union of
``arg_utils.internal_blacklist_keys()`` and ``arg_utils.SHARED_FIELDS``
so that neither orchestrator-only fields nor shared-pipeline fields
(``model`` / ``stage_configs_path`` / ``log_stats`` / ``stage_id``) leak
(``model`` / ``log_stats`` / ``stage_id``) leak
into a stage's per-stage runtime overrides — the orchestrator sets those
uniformly for every stage, they are not per-stage knobs. Callers can
pass an explicit set for tests or specialized flows.
@@ -90,46 +89,6 @@ def build_stage_runtime_overrides(
return result
def strip_parent_engine_args(
kwargs: dict[str, Any],
*,
parent_fields: dict[str, dataclasses.Field],
keep_keys: set[str] | frozenset[str] = frozenset(),
strip_keys: set[str] | frozenset[str] = frozenset(),
no_warn_keys: set[str] | frozenset[str] = frozenset(),
) -> tuple[dict[str, Any], list[str]]:
"""Strip parent ``EngineArgs`` fields before merging into stage YAML."""
overridden: list[str] = []
result: dict[str, Any] = {}
for key, value in kwargs.items():
if key in strip_keys:
continue
if key not in parent_fields or key in keep_keys:
result[key] = value
continue
field_def = parent_fields[key]
if field_def.default is not dataclasses.MISSING:
default = field_def.default
elif field_def.default_factory is not dataclasses.MISSING:
default = field_def.default_factory()
else:
default = dataclasses.MISSING
if default is dataclasses.MISSING or value is None:
continue
if dataclasses.is_dataclass(default) and not isinstance(default, type):
default = asdict(default)
if value != default and key not in no_warn_keys:
overridden.append(key)
return result, sorted(overridden)
def _apply_diffusion_parallel_runtime_overrides(
engine_args: dict[str, Any],
runtime_overrides: dict[str, Any],
+59
View File
@@ -44,6 +44,65 @@ vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
--deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml
```
## Why Stage 0 Pins Native FlashInfer
The three Stage 0 profiles set `attention_backend: FLASHINFER` and also set:
```yaml
attention_config:
use_trtllm_attention: false
```
`FLASHINFER` names the overall vLLM attention backend. Starting with vLLM
0.25, that backend can internally route Hopper decode to the TRT-LLM XQA
kernel. The XQA eligibility check answers whether the kernel can run, not
whether it is the fastest kernel for the current workload. In particular, the
decode builder can select XQA statically when all of the following hold:
- the explicit `use_trtllm_attention` override is not `false`;
- the device is Hopper/SM90 and the XQA cubin is available; and
- the query-head count is divisible by the KV-head count.
Higgs has 32 query heads and 8 KV heads, so it passes the head-layout check.
The static gate does not include a model dtype, KV-cache dtype, context-length,
batch-shape, or measured-performance comparison. This matters because kernel
support is not a performance guarantee.
Higgs also has an unusual output shape that must not be confused with its
attention shape. One autoregressive audio step advances one token per request.
The resulting hidden row is then projected into eight codebook distributions
with shape `[1, 8, codebook_size]`. The eight codebooks therefore do not turn
the attention query into `q_len=8`; at concurrency one, XQA still receives a
single BF16 query row. At this small per-step shape, any XQA advantage for
other KV dtypes, batch sizes, or context lengths does not necessarily apply.
The difference is paid in every transformer layer and every generated audio
step, so even a small per-call disadvantage accumulates in TTFP and RTF.
The upstream XQA integration also showed that the result is dtype-sensitive:
its FP8 benchmark favored XQA, while its BF16 benchmark reported 21.16 ms mean
TPOT for XQA versus 19.47 ms for FA3. That is not a direct native-FlashInfer
comparison, but it demonstrates why SM90 compatibility alone is insufficient
to choose the fastest BF16 kernel. See the
[vLLM XQA benchmark](https://github.com/vllm-project/vllm/pull/43232#issuecomment-4501176335).
The end-to-end H800 A/B in
[issue #5584](https://github.com/vllm-project/vllm-omni/issues/5584#issuecomment-5188700218)
isolated the routing decision on the same vLLM 0.25 stack:
| Decode path | TTFP (ms) | RTF |
|---|---:|---:|
| XQA (automatic default) | 175.76 | 0.3405 |
| Native FlashInfer (`use_trtllm_attention: false`) | 166.39 | 0.3233 |
| vLLM 0.24 reference | 166.68 | 0.3209 |
This proves that XQA selection causes the observed local regression and that
native FlashInfer recovers the pre-v0.25 range. It does not by itself identify
one defective XQA instruction or memory transaction. Establishing that deeper
kernel-level cause requires an attention-only microbenchmark and GPU profiler
trace for the exact Higgs batch/context distribution. Until vLLM has a
workload-aware selector or XQA is faster for this BF16 shape, the model-level
pin is the deterministic choice.
## Notes
- Stage 1 remains `enforce_eager: true` in both profiles.
+28
View File
@@ -0,0 +1,28 @@
# dots.tts deployment config (single-stage AR talker).
#
# Conservative first-integration settings: eager execution, no CUDA
# Graphs, no batched side-path optimizations yet — see voxcpm2.yaml for
# the runtime knobs a mature deployment of this model family grows.
#
# Verified locally on RTX 5080 16GB.
async_chunk: false
dtype: bfloat16
stages:
- stage_id: 0
max_num_seqs: 4
gpu_memory_utilization: 0.8
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
async_scheduling: true
max_num_batched_tokens: 4096
max_model_len: 4096
devices: "0"
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 4096
repetition_penalty: 1.0
@@ -38,6 +38,11 @@ stages:
gpu_memory_utilization: 0.6
enforce_eager: true
attention_backend: FLASHINFER
# Eight codebooks share one attention row, so decode remains BF16 q_len=1
# per request. vLLM's SM90 XQA capability gate does not benchmark this
# shape; keep the faster native path (README_higgs_audio_v3.md, #5584).
attention_config:
use_trtllm_attention: false
trust_remote_code: true
# Stage-0 prefix caching is on by default. The talker's HS opt-out and
# deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py)
Loaded 100 of 146 files, more files were not shown because too many files have changed in this diff. Show more