[Diffusion] Document DLO mmap model opt-in

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/cuda/test-nightly.yml

@@ -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

README.md

@@ -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.

docker/Dockerfile.npu

@@ -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

docker/Dockerfile.npu.a3

@@ -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

docker/Dockerfile.npu.ci

@@ -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

docker/Dockerfile.npu.ci.a3

@@ -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

docs/.nav.yml

@@ -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

docs/community/contact_us.md

@@ -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.

docs/community/governance.md

@@ -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
 

docs/community/meetings.md

@@ -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/).

docs/community/meetups.md

@@ -1 +0,0 @@
-# Meetups

docs/community/volunteers.md

@@ -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!

docs/contributing/README.md

@@ -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
 

docs/design/feature/distributed_layerwise_offload.md

@@ -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.

docs/models/supported_models.md

@@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

docs/usage/faq.md

@@ -10,12 +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: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).
+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.

docs/user_guide/diffusion/cpu_offload.md

@@ -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.

docs/user_guide/quantization/fp8.md

@@ -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 | ✅︎ |

examples/offline_inference/covo_audio/end2end.py

@@ -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",

examples/offline_inference/dynin_omni/end2end.py

@@ -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)

examples/offline_inference/hunyuan_image3/end2end.py

@@ -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}")

examples/offline_inference/minicpmo/end2end.py

@@ -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,

examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py

@@ -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:

examples/offline_inference/qwen3_omni/end2end.py

@@ -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",

examples/offline_inference/step_audio2/end2end.py

@@ -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,

examples/offline_inference/text_to_image/text_to_image.py

@@ -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

examples/offline_inference/text_to_speech/README.md

@@ -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.

examples/offline_inference/text_to_speech/dots_tts/end2end.py

@@ -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()

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -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",

examples/offline_inference/text_to_speech/glm_tts/end2end.py

@@ -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,
     )

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

@@ -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",

mkdocs.yml

@@ -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

recipes/LTX/LTX-2.md

@@ -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)

recipes/rednote-hilab/dots.tts.md

@@ -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.

tests/config/test_config_factory.py

@@ -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``)."""

tests/config/test_omni_config.py

@@ -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",

tests/diffusion/layers/test_fused_qk_norm_rope.py

@@ -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)

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -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():

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -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(

tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py

@@ -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

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -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",

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -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"]

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -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,

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -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

tests/diffusion/quantization/test_quantization_fp8.py

@@ -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:

tests/distributed/omni_connectors/test_bagel_mooncake_connector.py

@@ -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)

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -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":

tests/e2e/accuracy/ltx/run_ltx_reference.py

@@ -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:

tests/e2e/accuracy/ltx/test_ltx_official_similarity.py

@@ -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

tests/e2e/offline_inference/test_audiox_model_expansion.py

@@ -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 = [

tests/e2e/offline_inference/test_bagel_expansion.py

@@ -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)
 
 

tests/e2e/offline_inference/test_cosyvoice3_expansion.py

@@ -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 = [

tests/e2e/offline_inference/test_ernie_image_expansion.py

@@ -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,

tests/e2e/offline_inference/test_glm_tts_expansion.py

@@ -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:

tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py

@@ -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

tests/e2e/offline_inference/test_longcat_video_avatar.py

@@ -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",

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

@@ -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"

tests/e2e/offline_inference/test_minicpmo_4_5.py

@@ -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:

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

@@ -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,

tests/e2e/offline_inference/test_omnivoice_expansion.py

@@ -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,
     },

tests/e2e/offline_inference/test_sensenova_u1_expansion.py

@@ -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)

tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py

@@ -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)

tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py

@@ -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)

tests/e2e/offline_inference/test_soulxsinger_expansion.py

@@ -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(

tests/e2e/offline_inference/test_step_audio2_expansion.py

@@ -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]
 

tests/e2e/offline_inference/test_voxtral_tts_expansion.py

@@ -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,
         )

tests/e2e/online_serving/test_ltx.py

@@ -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)

tests/e2e/online_serving/test_ltx2_expansion.py

@@ -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)

tests/engine/test_arg_utils.py

@@ -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."""

tests/engine/test_async_omni_engine_abort.py

@@ -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)
 

tests/engine/test_single_stage_mode.py

@@ -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(

tests/engine/test_stage_engine_args.py

@@ -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"

tests/entrypoints/test_async_omni.py

@@ -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(

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -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,

tests/entrypoints/test_serve.py

@@ -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"

tests/entrypoints/test_utils.py

@@ -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,

tests/helpers/runtime.py

@@ -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...")

tests/model_executor/models/dots_tts/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/model_executor/models/dots_tts/test_talker.py

@@ -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 == {}

tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py

@@ -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 ----

tests/model_extras/test_model_extras.py

@@ -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

tests/model_tests/diffusion/test_alignment.py

@@ -27,6 +27,7 @@ EXCLUDED_MODELS = [
     "OvisImagePipeline",
     "WanPipeline",
     "WanVACEPipeline",
+    "LTX2TwoStagePipeline",
     "LTX2DistilledPipeline",
     "LTX2T2VDMD2Pipeline",
     "LTX2I2VDMD2Pipeline",

tools/baselines/omni_hash_baseline.py

@@ -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,

vllm_omni/config/omni_config.py

@@ -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

vllm_omni/config/pipeline_registry.py

@@ -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,

vllm_omni/config/stage_config.py

@@ -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],

vllm_omni/deploy/README_higgs_audio_v3.md

@@ -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.

vllm_omni/deploy/dots_tts.yaml

@@ -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

vllm_omni/deploy/higgs_multimodal_qwen3.yaml

@@ -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)

vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml

@@ -37,6 +37,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)

vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml

@@ -43,6 +43,11 @@ stages:
     gpu_memory_utilization: 0.6
     enforce_eager: false
     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
     enable_prefix_caching: true
     async_scheduling: false

vllm_omni/deploy/moss_tts_local.yaml

@@ -31,12 +31,12 @@ stages:
     output_connectors:
       to_stage_1: shm
     default_sampling_params:
-      temperature: 1.0
-      top_p: 0.95
-      top_k: 50
+      temperature: 1.7
+      top_p: 0.8
+      top_k: 25
       max_tokens: 4096
       seed: 42
-      repetition_penalty: 1.1
+      repetition_penalty: 1.0
 
   - stage_id: 1
     max_num_seqs: 64

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1263,6 +1263,7 @@ def extract_minimax_h3_context(
         _BF16_DTYPE,
         _FORWARD_SUPPORTED_KWARGS,
         MINIMAX_H3_ADALN_MODALITY_NUM,
+        _build_rope_table,
         _modulate_scale_shift,
         _required_kwarg,
     )
@@ -1318,7 +1319,7 @@ def extract_minimax_h3_context(
         raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
     device = x.device
 
-    rope_freqs = module.rope(img_position_ids).to(device)
+    rope_table = _build_rope_table(module.rope(img_position_ids).to(device))
 
     decoder_input, t_emb = module._embed(
         x=x,
@@ -1351,7 +1352,7 @@ def extract_minimax_h3_context(
     def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
         hidden, block_rope, block_combined = module.sp_prepare(
             decoder_input,
-            rope_freqs,
+            rope_table,
             combined_indices,
         )
         for block in module.blocks:
@@ -1359,7 +1360,7 @@ def extract_minimax_h3_context(
                 hidden,
                 t_emb=t_emb,
                 combined_indices=block_combined,
-                rope_freqs=block_rope,
+                rope_table=block_rope,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -18,6 +18,7 @@ import zmq
 from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue
 from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.executor.multiproc_executor import set_multiprocessing_worker_envs
 
 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
@@ -294,6 +295,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         logger.info("Starting server...")
 
         num_gpus = cast(int, od_config.num_gpus)
+        # Without this, every worker inherits one Torch thread per core, so an
+        # N-GPU run oversubscribes the host by N x core_count. Honours a
+        # user-provided OMP_NUM_THREADS.
+        set_multiprocessing_worker_envs()
         mp.set_start_method("spawn", force=True)
         processes = []
 

vllm_omni/diffusion/layers/fused_qk_norm_rope.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Q/K RMSNorm followed by packed non-interleaved RoPE.
+
+The public contract is shared by diffusion attention implementations:
+
+* ``q`` and ``k`` are ``[tokens, heads, head_dim]``;
+* norm weights are one-dimensional ``[head_dim]`` tensors;
+* ``rope_table`` is ``[tokens, rotary_dim]`` and stores
+  ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``.
+
+The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K
+or rotary-product intermediates. Unsupported inputs use the eager reference.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch.library import Library
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
+
+_FUSED_HEAD_DIM = 128
+_FUSED_ROTARY_DIM = 96
+_HEADS_PER_PROGRAM = 8
+
+
+def _apply_rope_table(
+    x: torch.Tensor,
+    rope_table: torch.Tensor,
+    rotary_dim: int,
+) -> torch.Tensor:
+    half = rotary_dim // 2
+    cos = rope_table[..., :half].to(x.dtype).unsqueeze(1)
+    sin = rope_table[..., half:].to(x.dtype).unsqueeze(1)
+    first = x[..., :half]
+    second = x[..., half:rotary_dim]
+    return torch.cat(
+        (
+            first * cos - second * sin,
+            second * cos + first * sin,
+            x[..., rotary_dim:],
+        ),
+        dim=-1,
+    )
+
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _rms_norm_rope_kernel(
+        x_ptr,
+        weight_ptr,
+        rope_table_ptr,
+        out_ptr,
+        x_stride_t,
+        x_stride_h,
+        x_stride_d,
+        rope_stride_t,
+        out_stride_t,
+        out_stride_h,
+        out_stride_d,
+        num_heads: tl.constexpr,
+        head_dim: tl.constexpr,
+        rotary_half: tl.constexpr,
+        eps: tl.constexpr,
+        heads_per_program: tl.constexpr,
+    ):
+        token = tl.program_id(0)
+        head_group = tl.program_id(1)
+        heads = head_group * heads_per_program + tl.arange(0, heads_per_program)
+        dims = tl.arange(0, head_dim)
+        mask = heads[:, None] < num_heads
+        offsets = token * x_stride_t + heads[:, None] * x_stride_h + dims[None, :] * x_stride_d
+
+        x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        weight = tl.load(weight_ptr + dims).to(tl.float32)
+        inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps)
+        normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16)
+
+        rotary_dim = rotary_half * 2
+        pair_dims = tl.where(
+            dims < rotary_half,
+            dims + rotary_half,
+            tl.where(dims < rotary_dim, dims - rotary_half, dims),
+        )
+        pair_offsets = token * x_stride_t + heads[:, None] * x_stride_h + pair_dims[None, :] * x_stride_d
+        pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32)
+        pair_weight = tl.load(weight_ptr + pair_dims).to(tl.float32)
+        pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16)
+
+        freq_dims = tl.where(
+            dims < rotary_half,
+            dims,
+            tl.where(dims < rotary_dim, dims - rotary_half, 0),
+        )
+        table_offsets = token * rope_stride_t + freq_dims
+        cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32)
+        sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32)
+        first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin
+        second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin
+        output = tl.where(
+            dims < rotary_dim,
+            tl.where(dims < rotary_half, first, second),
+            normalized.to(tl.float32),
+        )
+
+        out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d
+        tl.store(out_ptr + out_offsets, output, mask=mask)
+
+
+def _eager_qk_norm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q_norm = F.rms_norm(q, (head_dim,), q_weight, eps)
+    k_norm = F.rms_norm(k, (head_dim,), k_weight, eps)
+    return (
+        _apply_rope_table(q_norm, rope_table, rotary_dim),
+        _apply_rope_table(k_norm, rope_table, rotary_dim),
+    )
+
+
+def _fused_cuda_supported(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_dim: int,
+    rotary_dim: int,
+) -> bool:
+    return (
+        HAS_TRITON
+        and current_platform.is_cuda()
+        and q.is_cuda
+        and k.is_cuda
+        and q.dtype == torch.bfloat16
+        and k.dtype == torch.bfloat16
+        and head_dim == _FUSED_HEAD_DIM
+        and rotary_dim == _FUSED_ROTARY_DIM
+    )
+
+
+def _launch_fused_rms_norm_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    tokens, heads, head_dim = x.shape
+    rotary_half = rope_table.shape[-1] // 2
+    out = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    if tokens == 0:
+        return out
+    grid = (tokens, triton.cdiv(heads, _HEADS_PER_PROGRAM))
+    _rms_norm_rope_kernel[grid](
+        x,
+        weight,
+        rope_table,
+        out,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        rope_table.stride(0),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        num_heads=heads,
+        head_dim=head_dim,
+        rotary_half=rotary_half,
+        eps=eps,
+        heads_per_program=_HEADS_PER_PROGRAM,
+        num_warps=8,
+    )
+    return out
+
+
+def _fused_qk_norm_rope_impl(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not _fused_cuda_supported(q, k, head_dim, rotary_dim):
+        return _eager_qk_norm_rope(
+            q,
+            k,
+            q_weight,
+            k_weight,
+            rope_table,
+            eps,
+            head_dim,
+            rotary_dim,
+        )
+    return (
+        _launch_fused_rms_norm_rope(q, q_weight, rope_table, eps),
+        _launch_fused_rms_norm_rope(k, k_weight, rope_table, eps),
+    )
+
+
+def _fused_qk_norm_rope_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim
+    return torch.empty_like(q), torch.empty_like(k)
+
+
+_OMNI_OP_LIB = Library("vllm_omni", "FRAGMENT")
+if not hasattr(torch.ops.vllm_omni, "fused_qk_norm_rope"):
+    direct_register_custom_op(
+        op_name="fused_qk_norm_rope",
+        op_func=_fused_qk_norm_rope_impl,
+        fake_impl=_fused_qk_norm_rope_fake,
+        mutates_args=[],
+        target_lib=_OMNI_OP_LIB,
+    )
+
+
+def fused_qk_norm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    *,
+    head_dim: int | None = None,
+    rotary_dim: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply Q/K RMSNorm and packed non-interleaved RoPE."""
+    if q.ndim != 3 or k.ndim != 3:
+        raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}")
+    if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]:
+        raise ValueError(f"q and k shapes are incompatible: {q.shape} and {k.shape}")
+    if q.dtype != k.dtype or q.device != k.device:
+        raise ValueError("q and k must have the same dtype and device")
+    if q.dtype not in (torch.bfloat16, torch.float16, torch.float32):
+        raise TypeError(f"Fused QK RMSNorm/RoPE requires floating inputs, got {q.dtype}")
+
+    head_dim = q.shape[-1] if head_dim is None else head_dim
+    rotary_dim = rope_table.shape[-1] if rotary_dim is None else rotary_dim
+    if q.shape[-1] != head_dim:
+        raise ValueError(f"Expected q/k head_dim={head_dim}, got {q.shape[-1]}")
+    if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2:
+        raise ValueError(f"rotary_dim must be even and in [2, {head_dim}], got {rotary_dim}")
+    if q_weight.shape != (head_dim,) or k_weight.shape != (head_dim,):
+        raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}")
+    if q_weight.device != q.device or k_weight.device != q.device:
+        raise ValueError("Q/K norm weights must be on the activation device")
+    if rope_table.device != q.device or rope_table.dtype != q.dtype:
+        raise ValueError("rope_table must have the same dtype and device as q/k")
+    if rope_table.shape != (q.shape[0], rotary_dim):
+        raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}")
+
+    q_weight = q_weight.contiguous()
+    k_weight = k_weight.contiguous()
+    rope_table = rope_table.contiguous()
+    if not _fused_cuda_supported(q, k, head_dim, rotary_dim):
+        return _fused_qk_norm_rope_impl(
+            q,
+            k,
+            q_weight,
+            k_weight,
+            rope_table,
+            eps,
+            head_dim,
+            rotary_dim,
+        )
+    return torch.ops.vllm_omni.fused_qk_norm_rope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        rope_table,
+        eps,
+        head_dim,
+        rotary_dim,
+    )
+
+
+__all__ = ["fused_qk_norm_rope"]

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -363,22 +363,39 @@ class DiffusersPipelineLoader:
             else:
                 model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
 
-                # Skip load_weights only for DLO+AllGather when the model
-                # supports mmap loading and no online quantization is active.
-                # This condition MUST match the gate in
-                # DistributedLayerwiseOffloadBackend.enable() so that the
-                # loader skips ⟺ enable() uses mmap.
-                from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading
+                # Skip load_weights when DLO can keep checkpoint tensors as
+                # mmap-backed host masters.  Storage selection is independent
+                # of whether DLO reconstructs weights with AllGather.  The
+                # no-AllGather path can use mmap only when the checkpoint is
+                # already the rank-local runtime layout (currently TP=1 and no
+                # HSDP); TP/HSDP still require the regular loader callbacks.
+                # The loader and backend share should_select_dlo_mmap() so
+                # skipping load_weights always selects mmap in enable().
+                from vllm_omni.diffusion.offloader.offload_plan import (
+                    should_select_dlo_mmap,
+                    supports_mmap_loading,
+                )
 
                 _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
                 _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
                 _has_online_quant = self._has_online_quant(model)
                 _supports_mmap = supports_mmap_loading(model)
+                _tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1))
+                _use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False))
 
-                _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant
+                _skip_load = _dist_offload and should_select_dlo_mmap(
+                    mmap_supported=_supports_mmap,
+                    dlo_use_allgather=_use_ag,
+                    has_online_quant=_has_online_quant,
+                    tensor_parallel_size=_tp_size,
+                    use_hsdp=_use_hsdp,
+                )
 
                 if _skip_load:
-                    logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())")
+                    logger.info(
+                        "DLO mmap storage active (%s): skipping load_weights (will map checkpoint weights in enable())",
+                        "AllGather" if _use_ag else "rank-local",
+                    )
                 else:
                     if _dist_offload and _use_ag and _has_online_quant:
                         raise ValueError(
@@ -441,11 +458,11 @@ class DiffusersPipelineLoader:
         """Whether any layer uses an online-quant method that defers weight
         materialization onto the ``meta`` device (upstream vLLM
         ``uses_meta_device=True``, e.g. online FP8)."""
-        for module in model.modules():
-            quant_method = getattr(module, "quant_method", None)
-            if getattr(quant_method, "uses_meta_device", False):
-                return True
-        return False
+        from vllm_omni.diffusion.offloader.offload_plan import (
+            has_online_quantization,
+        )
+
+        return has_online_quantization(model)
 
     def _apply_skip_softmax_calibration(self, model: nn.Module) -> None:
         from vllm_omni.diffusion.attention.backends.trtllm_calibration import (

vllm_omni/diffusion/models/ltx2/__init__.py

@@ -12,12 +12,16 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     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,
+)
 
 __all__ = [
     "LTX2Pipeline",
     "LTX2T2VDMD2Pipeline",
     "LTX2I2VDMD2Pipeline",
+    "LTX2TwoStagePipeline",
     "LTX2DistilledPipeline",
     "get_ltx2_post_process_func",
     "load_transformer_config",

vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Parse official LTX distilled adapters without materializing their tensors.
+
+The official LTX adapters are raw safetensors files rather than PEFT
+checkpoints.  This module is deliberately LTX-specific: it translates the
+official names into the locally constructed Diffusers transformer names and
+describes the resulting adapter targets.  Applying those targets is the job
+of :mod:`ltx2_phase_adapter`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from safetensors import safe_open
+
+_IGNORED_OFFICIAL_LORA_MODULES = {"text_embedding_projection.aggregate_embed"}
+_OFFICIAL_TO_DIFFUSERS_PREFIXES = (
+    ("audio_prompt_adaln_single.", "audio_prompt_adaln."),
+    ("prompt_adaln_single.", "prompt_adaln."),
+    ("audio_adaln_single.", "audio_time_embed."),
+    ("adaln_single.", "time_embed."),
+    ("audio_patchify_proj", "audio_proj_in"),
+    ("patchify_proj", "proj_in"),
+    ("av_ca_video_scale_shift_adaln_single.", "av_cross_attn_video_scale_shift."),
+    ("av_ca_audio_scale_shift_adaln_single.", "av_cross_attn_audio_scale_shift."),
+    ("av_ca_a2v_gate_adaln_single.", "av_cross_attn_video_a2v_gate."),
+    ("av_ca_v2a_gate_adaln_single.", "av_cross_attn_audio_v2a_gate."),
+)
+
+
+@dataclass(frozen=True)
+class AdapterTarget:
+    """One A/B pair mapped onto a local transformer linear layer."""
+
+    # Keep the original Diffusers module name for diagnostics after LTX packs
+    # Q/K/V into the runtime module path.
+    source_module: str
+    module: str
+    slice: str | None
+    a_key: str
+    b_key: str
+    a_shape: tuple[int, ...]
+    b_shape: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class AdapterManifest:
+    """Tensor-free description of the fixed refinement adapter."""
+
+    path: str
+    targets: tuple[AdapterTarget, ...]
+
+
+def _to_diffusers_module_name(name: str) -> str:
+    for official_prefix, diffusers_prefix in _OFFICIAL_TO_DIFFUSERS_PREFIXES:
+        if name.startswith(official_prefix):
+            return diffusers_prefix + name[len(official_prefix) :]
+    return name
+
+
+def _resolve_lora_target(transformer: nn.Module, module_name: str) -> tuple[str, str | int | None] | None:
+    modules = dict(transformer.named_modules())
+    if module_name in modules:
+        return module_name, None
+    for packed_suffix, sub_suffix, shard_id in getattr(transformer, "stacked_params_mapping", ()):
+        if sub_suffix in module_name:
+            packed_name = module_name.replace(sub_suffix, packed_suffix)
+            if packed_name in modules:
+                return packed_name, shard_id
+    return None
+
+
+def parse_ltx_adapter(transformer: nn.Module, path: str) -> AdapterManifest:
+    """Translate one official raw LTX adapter into local Transformer targets."""
+    path = str(Path(path))
+    targets: list[AdapterTarget] = []
+    unresolved: list[str] = []
+    seen: set[tuple[str, str | int | None]] = set()
+
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        tensor_names = set(handle.keys())
+        official_name_set = {
+            tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0]
+            for tensor_name in tensor_names
+            if tensor_name.endswith(".lora_A.weight")
+        }
+        official_b_names = {
+            tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0]
+            for tensor_name in tensor_names
+            if tensor_name.endswith(".lora_B.weight")
+        }
+        missing_a = sorted(official_b_names - official_name_set - _IGNORED_OFFICIAL_LORA_MODULES)
+        if missing_a:
+            raise ValueError(f"Missing paired LoRA A tensors for {missing_a} in {path}.")
+
+        official_names = sorted(official_name_set)
+        for official_name in official_names:
+            if official_name in _IGNORED_OFFICIAL_LORA_MODULES:
+                continue
+            key_prefix = f"diffusion_model.{official_name}"
+            key_a = f"{key_prefix}.lora_A.weight"
+            key_b = f"{key_prefix}.lora_B.weight"
+            if key_b not in tensor_names:
+                raise ValueError(f"Missing paired LoRA tensor {key_b!r} in {path}.")
+
+            a_shape = tuple(handle.get_slice(key_a).get_shape())
+            b_shape = tuple(handle.get_slice(key_b).get_shape())
+            if len(a_shape) != 2 or len(b_shape) != 2 or b_shape[1] != a_shape[0]:
+                raise ValueError(f"Invalid LoRA pair for {official_name}: A={a_shape}, B={b_shape}.")
+
+            source_module = _to_diffusers_module_name(official_name)
+            resolved = _resolve_lora_target(transformer, source_module)
+            if resolved is None:
+                unresolved.append(official_name)
+                continue
+            module, shard_id = resolved
+            target_id = (module, shard_id)
+            if target_id in seen:
+                raise ValueError(f"Official LTX LoRA contains duplicate target {module!r}, slice {shard_id!r}.")
+            seen.add(target_id)
+            targets.append(
+                AdapterTarget(
+                    source_module=source_module,
+                    module=module,
+                    slice=shard_id,
+                    a_key=key_a,
+                    b_key=key_b,
+                    a_shape=a_shape,
+                    b_shape=b_shape,
+                )
+            )
+
+    if unresolved:
+        raise ValueError(f"Official LTX LoRA contains unmapped modules: {unresolved}.")
+    if not targets:
+        raise ValueError(f"No Transformer LoRA tensors were found in {path}.")
+    return AdapterManifest(path=path, targets=tuple(targets))
+
+
+def iter_adapter_tensors(manifest: AdapterManifest) -> Iterator[tuple[AdapterTarget, torch.Tensor, torch.Tensor]]:
+    """Stream adapter pairs in manifest order without a full checkpoint copy."""
+
+    with safe_open(manifest.path, framework="pt", device="cpu") as handle:
+        for target in manifest.targets:
+            yield (
+                target,
+                handle.get_tensor(target.a_key),
+                handle.get_tensor(target.b_key),
+            )

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -9,7 +9,8 @@ import inspect
 import json
 import logging
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -19,6 +20,8 @@ from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
 from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
 from diffusers.video_processor import VideoProcessor
 from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+from safetensors.torch import load_file
 from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
@@ -32,6 +35,7 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
+from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint
 from .ltx2_transformer import (
     LTX2VideoTransformer3DModel,
     apply_interleaved_rotary_emb,
@@ -70,6 +74,9 @@ class LTXComponentProfile:
     video_vae_cls: type = AutoencoderKLLTX2Video
     vocoder_cls: type = LTX2Vocoder
     vocoder_fallback_cls: type | None = None
+    artifact_repo_id: str | None = None
+    latent_upsampler_filename: str | None = None
+    distilled_lora_filename: str | None = None
 
 
 LTX2_COMPONENT_PROFILE = LTXComponentProfile(
@@ -101,16 +108,96 @@ LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
     video_vae_cls=DistributedAutoencoderKLLTX2Video,
 )
 
+LTX23_DISTILLED_COMPONENT_PROFILE = replace(
+    LTX23_COMPONENT_PROFILE,
+    name="ltx2_3_distilled",
+    resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
+    artifact_repo_id="Lightricks/LTX-2.3",
+    latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
+)
+
+LTX2_TWO_STAGE_COMPONENT_PROFILE = replace(
+    LTX2_COMPONENT_PROFILE,
+    name="ltx2_two_stage",
+    resident_modules=(*LTX2_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
+    artifact_repo_id="Lightricks/LTX-2",
+    latent_upsampler_filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
+    distilled_lora_filename="ltx-2-19b-distilled-lora-384.safetensors",
+)
+
+LTX23_TWO_STAGE_COMPONENT_PROFILE = replace(
+    LTX23_COMPONENT_PROFILE,
+    name="ltx2_3_two_stage",
+    resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
+    artifact_repo_id="Lightricks/LTX-2.3",
+    latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
+    distilled_lora_filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
+)
+
 
 _COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = {
     ("one_stage", "2"): LTX2_COMPONENT_PROFILE,
     ("one_stage", "2.3"): LTX23_COMPONENT_PROFILE,
+    ("two_stage", "2"): LTX2_TWO_STAGE_COMPONENT_PROFILE,
+    ("two_stage", "2.3"): LTX23_TWO_STAGE_COMPONENT_PROFILE,
     ("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE,
+    ("distilled_two_stage", "2.3"): LTX23_DISTILLED_COMPONENT_PROFILE,
     ("dmd2", "2"): LTX2_COMPONENT_PROFILE,
     ("dmd2", "2.3"): LTX23_COMPONENT_PROFILE,
 }
 
 
+def resolve_ltx_artifact(
+    model: str,
+    repo_id: str,
+    filename: str,
+) -> str:
+    """Resolve an official LTX sidecar from the model root or its Hub repository."""
+    candidate = Path(model) / filename
+    if candidate.is_file():
+        return str(candidate)
+
+    try:
+        return hf_hub_download(repo_id=repo_id, filename=filename)
+    except Exception as exc:
+        raise FileNotFoundError(
+            f"Unable to resolve LTX artifact {filename!r}. Searched {candidate}; "
+            f"place the file in the model root or make {repo_id} available."
+        ) from exc
+
+
+def _load_ltx_latent_upsampler_single_file(path: str, dtype: torch.dtype) -> LTX2LatentUpsamplerModel:
+    """Load an official single-file upsampler into the Diffusers module."""
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        metadata = handle.metadata() or {}
+    raw_config = json.loads(metadata.get("config", "{}"))
+    config = {
+        "in_channels": raw_config.get("in_channels", 128),
+        "mid_channels": raw_config.get("mid_channels", 1024),
+        "num_blocks_per_stage": raw_config.get("num_blocks_per_stage", 4),
+        "dims": raw_config.get("dims", 3),
+        "spatial_upsample": raw_config.get("spatial_upsample", True),
+        "temporal_upsample": raw_config.get("temporal_upsample", False),
+        "rational_spatial_scale": raw_config.get("spatial_scale", 2.0),
+        "use_rational_resampler": raw_config.get("rational_resampler", False),
+    }
+    with torch.device("cpu"):
+        upsampler = LTX2LatentUpsamplerModel(**config).to(dtype=dtype)
+
+    state_dict = load_file(path, device="cpu")
+    if "upsampler.0.weight" in state_dict and hasattr(upsampler.upsampler, "conv"):
+        state_dict["upsampler.conv.weight"] = state_dict.pop("upsampler.0.weight")
+        state_dict["upsampler.conv.bias"] = state_dict.pop("upsampler.0.bias")
+    missing, unexpected = upsampler.load_state_dict(state_dict, strict=False)
+    unresolved_missing = set(missing) - {"upsampler.blur_down.kernel"}
+    if unresolved_missing or unexpected:
+        raise ValueError(
+            f"Invalid LTX latent upsampler {path}: missing={sorted(unresolved_missing)}, "
+            f"unexpected={sorted(unexpected)}."
+        )
+    return upsampler
+
+
 def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile:
     """Resolve component construction independently from execution recipes."""
     try:
@@ -392,16 +479,34 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         )
 
     if "latent_upsampler" in profile.resident_modules:
-        # BlurDownsample constructs an integer kernel that must be initialized
-        # on CPU; component placement is handled uniformly after construction.
-        with torch.device("cpu"):
-            pipeline.latent_upsampler = _load_component(
-                LTX2LatentUpsamplerModel,
+        upsampler_config = os.path.join(model, "latent_upsampler", "config.json")
+        if os.path.isfile(upsampler_config) or not local_files_only:
+            try:
+                pipeline.latent_upsampler = _load_component(
+                    LTX2LatentUpsamplerModel,
+                    model,
+                    "latent_upsampler",
+                    local_files_only=local_files_only,
+                    dtype=dtype,
+                )
+            except (OSError, ValueError):
+                if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
+                    raise
+                upsampler_path = resolve_ltx_artifact(
+                    model,
+                    profile.artifact_repo_id,
+                    profile.latent_upsampler_filename,
+                )
+                pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
+        else:
+            if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
+                raise FileNotFoundError(f"LTX latent upsampler component not found under {model}.")
+            upsampler_path = resolve_ltx_artifact(
                 model,
-                "latent_upsampler",
-                local_files_only=local_files_only,
-                dtype=dtype,
+                profile.artifact_repo_id,
+                profile.latent_upsampler_filename,
             )
+            pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
 
     transformer_config = load_transformer_config(model, "transformer", local_files_only)
     quant_config = getattr(od_config, "quantization_config", None)
@@ -412,6 +517,16 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         subfolder="scheduler",
         local_files_only=local_files_only,
     )
+    expected_checkpoint_kind: LTXCheckpointKind | None = None
+    if pipeline.pipeline_kind == "two_stage":
+        expected_checkpoint_kind = "regular"
+    elif pipeline.pipeline_kind == "distilled_two_stage":
+        expected_checkpoint_kind = "distilled"
+    validate_ltx_checkpoint(
+        pipeline.scheduler.config,
+        expected_kind=expected_checkpoint_kind,
+        pipeline_name=type(pipeline).__name__,
+    )
 
     pipeline.vae_spatial_compression_ratio = pipeline.vae.spatial_compression_ratio
     pipeline.vae_temporal_compression_ratio = pipeline.vae.temporal_compression_ratio

vllm_omni/diffusion/models/ltx2/ltx2_denoise.py

@@ -62,13 +62,14 @@ class LTXDenoiseContext:
     audio_latents: torch.Tensor
     video_coords: torch.Tensor
     audio_coords: torch.Tensor
+    audio_attention_mask: torch.Tensor | None = None
     conditioning_mask: torch.Tensor | None = None
     conditioning_mask_for_model: torch.Tensor | None = None
 
 
 @dataclass
 class LTXPhaseResult:
-    """Denoised, unpacked AV latents and the context used to produce them."""
+    """Denoised AV latents and the context used to produce them."""
 
     forward_context: LTXForwardContext
     video: torch.Tensor
@@ -284,7 +285,8 @@ def prepare_rope_coords_stage(
     latents: torch.Tensor,
     audio_latents: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    video_coords = pipeline.transformer.rope.prepare_video_coords(
+    transformer = pipeline.transformer
+    video_coords = transformer.rope.prepare_video_coords(
         latents.shape[0],
         forward_ctx.latent_num_frames,
         forward_ctx.latent_height,
@@ -292,7 +294,7 @@ def prepare_rope_coords_stage(
         latents.device,
         fps=forward_ctx.request_inputs.frame_rate,
     )
-    audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords(
+    audio_coords = transformer.audio_rope.prepare_audio_coords(
         audio_latents.shape[0],
         forward_ctx.padded_audio_num_frames,
         audio_latents.device,
@@ -331,6 +333,7 @@ def build_transformer_kwargs(
         # with a learned register, making all output context tokens valid.
         "encoder_attention_mask": None,
         "audio_encoder_attention_mask": None,
+        "audio_attention_mask": denoise_ctx.audio_attention_mask,
         "num_frames": forward_ctx.latent_num_frames,
         "height": forward_ctx.latent_height,
         "width": forward_ctx.latent_width,
@@ -442,11 +445,25 @@ class LTXPhaseExecutor:
             video_audio_step_adapter=video_audio_step_adapter,
         )
         video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents)
+        ring_degree = getattr(pipeline.od_config.parallel_config, "ring_degree", 1) or 1
+        if padded_audio_num_frames > original_audio_num_frames and ring_degree > 1:
+            raise ValueError(
+                "LTX audio padding requires an attention mask, which Ring sequence parallelism does not support. "
+                "Use Ulysses-only SP or choose a request whose audio latent length is divisible by the SP size."
+            )
         denoise_ctx = LTXDenoiseContext(
             latents=latents,
             audio_latents=audio_latents,
             video_coords=video_coords,
             audio_coords=audio_coords,
+            audio_attention_mask=(
+                torch.arange(padded_audio_num_frames, device=audio_latents.device)
+                .lt(original_audio_num_frames)
+                .unsqueeze(0)
+                .expand(audio_latents.shape[0], -1)
+                if padded_audio_num_frames > original_audio_num_frames
+                else None
+            ),
             conditioning_mask=conditioning_mask,
         )
         denoise_ctx = pipeline._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx)