Merge origin/main into codex-vae-spatial-decode-auto

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_ltx.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"
         mirror_hardwares: h100_2
 
       - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
@@ -477,16 +477,10 @@ 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,7 +16,6 @@ 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.

docs/.nav.yml

@@ -71,7 +71,7 @@ nav:
         - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
         - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
         - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
-        - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
+        - VAE Parallel: user_guide/diffusion/parallelism/vae_parallel.md
       - Attention Backends: user_guide/diffusion/attention_backends.md
       - Regional Compilation: user_guide/diffusion/regional_compilation.md
       - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
@@ -123,7 +123,7 @@ nav:
           - Pipeline Parallel: design/feature/pipeline_parallel.md
           - Sequence Parallel: design/feature/sequence_parallel.md
           - Tensor Parallel: design/feature/tensor_parallel.md
-          - VAE Patch Parallelism: design/feature/vae_parallel.md
+          - VAE Parallel: design/feature/vae_parallel.md
         - Attention Optimization:
           - Skip-Softmax: design/feature/skip_softmax.md
         - Cache-DiT: design/feature/cache_dit.md

docs/design/architecture_overview.md

@@ -398,7 +398,7 @@ the source for configuration, compatibility, and implementation details.
 * **Composable parallelism:** Diffusion stages can combine [CFG-Parallel](feature/cfg_parallel.md),
   [Expert Parallel](feature/expert_parallel.md), [HSDP](feature/hsdp.md),
   [Pipeline Parallel](feature/pipeline_parallel.md), [Sequence Parallel](feature/sequence_parallel.md),
-  [Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](feature/vae_parallel.md)
+  [Tensor Parallel](feature/tensor_parallel.md), and [VAE Parallel](feature/vae_parallel.md)
   according to the pipeline and hardware topology.
 * **Attention and cache acceleration:** [Skip-Softmax](feature/skip_softmax.md),
   [Cache-DiT](feature/cache_dit.md), and [TeaCache](feature/teacache.md)

docs/design/feature/vae_parallel.md

@@ -1,42 +1,57 @@
-# VAE Patch Parallelism
+# VAE Parallel
 
-This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
+This document defines the VAE parallelism strategy contract and describes how to
-We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
+add a strategy to a diffusion model. Tile/patch parallelism is the portable
+baseline. Spatial-shard parallelism is a model-specific decode strategy that can
+be selected explicitly or through a request-aware policy.
+
+We use **Qwen-Image** as the reference implementation for tile-parallel decode,
+**Wan2.2** for tile-parallel encode, and **Wan2.2** as the first spatial-shard
+decode implementation. Other VAE families can add spatial strategies later
+without changing the public mode meanings defined here.
 
 ---
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
+- [Strategy Selection Contract](#strategy-selection-contract)
-- [Encode Parallel Implementation](#encode-parallel-implementation)
+- [Tile/Patch Implementation (Decode)](#tilepatch-implementation-decode)
+- [Tile/Patch Encode Implementation](#tilepatch-encode-implementation)
 - [Testing](#testing)
 - [Reference Implementations](#reference-implementations)
+- [Spatially-Sharded Decode](#spatially-sharded-decode)
+- [Adding Another VAE Strategy](#adding-another-vae-strategy)
 - [Summary](#summary)
 
 ---
 
 ## Overview
 
-### What is Vae Patch parallel?
+### What is VAE parallelism?
 
-**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
+VAE parallelism distributes VAE **encoding** or **decoding** across multiple
+ranks. The implementation currently has two strategy families:
 
-+ Split into multiple spatial tiles
+| Strategy | Unit of work | Communication | Current scope |
+|----------|--------------|---------------|---------------|
+| Tile/patch | Independent overlapping tiles | Gather, stitch, and optionally broadcast | Encode and decode across supported distributed VAEs |
+| Spatial shard | One feature map sharded along a spatial axis | Per-layer halo exchange followed by output gather | Wan decode |
 
-+ Distributed across multiple ranks
+Both strategies preserve the model-specific receptive field and reconstruct the
+same logical output as single-rank execution. They differ in how work is divided:
+tile parallelism assigns complete overlapping tiles to ranks, while spatial
+sharding keeps one logical feature map partitioned throughout the decoder.
 
-+ Encoded/Decoded in parallel
+This approach can:
 
-+ Merged to reconstruct the final output
++ Distribute computation across multiple devices
++ Reduce peak memory usage per device
++ Accelerate encoding or decoding latency
 
-This approach:
+The actual benefit depends on shape, topology, communication cost, and VAE
-
+architecture. A strategy must not become the global default without parity and
-+ Distributes computation across multiple devices
+performance evidence on its supported models.
-
-+ Reduces peak memory usage per device
-
-+ Accelerates encoding/decoding latency
 
 ### When to Use Encode vs Decode Parallel
 
@@ -45,8 +60,16 @@ This approach:
 | **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
 | **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
 
-### Architecture
+### Architecture and ownership
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
+
+`DiffusionParallelConfig.vae_parallel_mode` carries the public strategy request.
+The shared registry forwards that value to the VAE, but each distributed VAE
+adapter owns its supported modes, request-shape selection, topology validation,
+and fallback behavior. Model-specific policy must not be added to the shared
+runner or registry.
+
+**DistributedVaeExecutor** is the model-agnostic core for tile-parallel VAE
+encoding and decoding.
 
 The executor is model-agnostic and accepts three function parameters:
 
@@ -76,6 +99,11 @@ This design separates:
 
 + Model-specific tiling and merging logic
 
+Spatial-shard implementations use the same configured VAE process group but do
+not use the tile task executor. They keep the feature map sharded across decoder
+layers and own their halo-exchange and final-gather contracts inside the model's
+distributed autoencoder implementation.
+
 #### Why split / exec / merge is necessary?
 
 The latent tensor cannot be arbitrarily partitioned.
@@ -92,7 +120,34 @@ Therefore:
 
 + Merge must perform blending to avoid seams
 
-## Step-by-Step Implementation (Decode)
+## Strategy Selection Contract
+
+`vae_parallel_mode="tile"` is the compatibility-preserving global default.
+Other values are stable strategy requests, but support remains model-specific.
+
+| Mode | Meaning | Current implementation | Unsupported or ineligible request |
+|------|---------|------------------------|-----------------------------------|
+| `tile` | Use the existing tile/patch executor | All distributed VAE adapters | Not applicable |
+| `auto` | Let the VAE adapter select a validated strategy for each request | Wan decode selects height or width spatial sharding | Wan uses tile when ineligible; unsupported VAE families reject or document their fallback |
+| `spatial_shard_height` | Force spatial-shard decode along height | Wan decode | Wan warns and uses tile when ineligible; unsupported VAE families reject |
+| `spatial_shard_width` | Force spatial-shard decode along width | Wan decode | Wan warns and uses tile when ineligible; unsupported VAE families reject |
+
+The shared configuration layer validates the mode name. It does not decide
+whether a model supports that mode. A VAE adapter adding a mode must:
+
+1. Keep `tile` as an exact feature-off path.
+2. Select a strategy per request rather than from stale process-global state.
+3. Validate input rank/shape and process-group membership before collectives.
+4. Reject or explicitly fall back from unsupported combinations; never silently
+   reinterpret one forced strategy as another.
+5. Preserve direct encode/decode behavior when Diffusers does not select a tiled
+   path.
+
+`auto` is a policy hook, not a universal heuristic. Each VAE family may select
+from only the strategies it has validated. Adding another family must not change
+Wan's selector or the meanings of the explicit modes.
+
+## Tile/Patch Implementation (Decode)
 
 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -251,7 +306,7 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```
 
-## Encode Parallel Implementation
+## Tile/Patch Encode Implementation
 
 For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
 
@@ -449,18 +504,24 @@ Complete examples in the codebase:
 
 ---
 
-## Spatially-Sharded Decode (Wan)
+## Spatially-Sharded Decode
 
-The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).
+The tile-parallel executor above assigns independent spatial **tiles** to ranks.
+A spatial-shard backend instead partitions one global decoder feature map along
+height or width. Wan is the first supported implementation.
+
+`DiffusionParallelConfig.vae_parallel_mode="tile"` remains the default. For Wan,
+opt into request-shape selection with `"auto"`, or force an axis with
+`"spatial_shard_height"` or `"spatial_shard_width"`.
 
 ### How it differs from tile parallel
 
-| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|-----------------------------------------------|
+|--------|--------------------------|----------------------|-----------------------------------------------|
-| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
+| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
-| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
-| Scope | Decode + encode | Decode only |
+| Scope | Decode + encode | Wan decode selection; encode remains tiled | Wan decode only |
 
 Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.
 
@@ -476,21 +537,65 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
   -> DistributedAutoencoderKLWan.tiled_decode dispatch
 ```
 
-`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode.
+`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it
+requires distributed decode to be enabled, a 5D latent, and
+`vae_patch_parallel_size == DiT group size`. The method is reached only after
+Diffusers has selected tiled decode for the request. In `auto` mode the longer
+axis that can cover every rank is selected, with width winning ties. If neither
+axis can cover the group, auto mode uses tile decode. This policy is cached only
+through normal kernel shape caches; it is reevaluated for every request.
 
 ### Notes
 
-- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
+- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely.
 - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
 
+## Adding Another VAE Strategy
+
+Add future strategies through the distributed autoencoder that owns the model
+semantics. Do not add model-name dispatch to the shared runner, registry, or
+`DistributedVaeExecutor`.
+
+For a new VAE family:
+
+1. **Declare capability**: document which modes, operations, shapes, dtypes, and
+   process-group layouts the adapter supports.
+2. **Implement model-local selection**: keep `auto` policy and explicit-mode
+   dispatch in the distributed autoencoder. The selector may use request shape,
+   but must be deterministic and side-effect free.
+3. **Preserve the baseline**: retain exact tile/direct fallbacks and state-dict
+   compatibility. Structural wrappers must not register duplicate parameters or
+   alter checkpoint names.
+4. **Define communication**: derive shards from the configured VAE group, keep
+   collectives symmetric, specify padding/halo ownership, and return the same
+   output placement as the tile path.
+5. **Bound request state**: make the selected axis and scratch buffers
+   request-scoped, release temporary buffers on success and failure, and support
+   alternating modes in one long-lived worker.
+6. **Validate independently**: compare each explicit strategy against the
+   single-rank reference, then test `auto`, tile fallback, unsupported topology,
+   and a mixed request sequence. Record per-rank peak memory and decode latency
+   before documenting a preferred mode.
+
+Promote a selector into shared infrastructure only after at least two VAE
+families use the same inputs, invariants, fallback behavior, and lifecycle. Until
+then, shared code owns transport and configuration while each model adapter owns
+policy.
+
 ---
 
 ## Summary
 
-Adding VAE Patch Parallel support to diffusion model:
+Adding VAE parallel support to a diffusion model:
 
-1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
+1. **Implement the distributed VAE** - Inherit from the base VAE class and
-2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
+   `DistributedVaeMixin`.
-3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
+2. **Keep tile as the baseline** - For tile decode, refactor `tiled_decode` into
-4. **Change VAE model in pipeline** - Use the distributed version
+   `tile_split`/`tile_exec`/`tile_merge`.
-5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`
+3. **Add encode parallelism when needed** - Implement
+   `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models.
+4. **Add model-specific strategies locally** - Own selectors, topology checks,
+   communication, and fallback in the distributed autoencoder.
+5. **Wire the pipeline** - Replace the pipeline VAE with its distributed version.
+6. **Test strategy and lifecycle parity** - Compare one rank with every supported
+   strategy, including feature-off, fallback, and mixed-request sequences.

docs/design/index.md

@@ -48,7 +48,7 @@ implementation contract; it is not, by itself, a general support claim.
 - [Pipeline Parallel](feature/pipeline_parallel.md)
 - [Sequence Parallel](feature/sequence_parallel.md)
 - [Tensor Parallel](feature/tensor_parallel.md)
-- [VAE Patch Parallelism](feature/vae_parallel.md)
+- [VAE Parallel](feature/vae_parallel.md)
 
 #### Attention optimization
 

docs/models/supported_models.md

@@ -42,8 +42,7 @@ 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` | ✅︎ | ✅︎ | | | — |
-| `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 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
-| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -59,7 +58,6 @@ 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/user_guide/diffusion/parallelism/overview.md

@@ -10,7 +10,7 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
 | **[Sequence Parallelism](sequence_parallel.md)**   | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
 | **[CFG-Parallel](cfg_parallel.md)**                | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation                         |
 | **[Pipeline Parallelism](pipeline_parallel.md)**   | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory             |
-| **[VAE Parallelism](vae_parallelism.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
+| **[VAE Parallel](vae_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
 | **[HSDP](hsdp.md)**                                | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs              |
 | **[Expert Parallelism](expert_parallel.md)**       | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0)                                          |
 

docs/user_guide/diffusion/parallelism/vae_parallel.md

@@ -1,4 +1,4 @@
-# VAE Parallelism Guide
+# VAE Parallel Guide
 
 
 ## Table of Content
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
-| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
+| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"`, `"auto"` (Wan request-shape selection), `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding is currently Wan-only. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
 
 Additional requirements:
 
@@ -128,7 +128,7 @@ Additional requirements:
 
 ## Spatially-Sharded Decode (Wan)
 
-The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`.
+The default `vae_parallel_mode="tile"` keeps patch/tile parallel decode unchanged. For Wan, set `vae_parallel_mode="auto"` to select spatial-shard decode only after a request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. Auto mode splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests and partial groups fall back to tile decode. Use `"spatial_shard_height"` or `"spatial_shard_width"` to force an axis.
 
 Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.
 
@@ -141,7 +141,7 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(
         tensor_parallel_size=2,
         vae_patch_parallel_size=2,               # must match the DiT group size
-        vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
+        vae_parallel_mode="auto",                # opt in to Wan request-shape selection
     ),
 )
 ```
@@ -152,14 +152,14 @@ Or from the CLI / serving entrypoint:
 vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
     --tensor-parallel-size 2 \
     --vae-patch-parallel-size 2 \
-    --vae-parallel-mode spatial_shard_width
+    --vae-parallel-mode auto
 ```
 
 **Constraints and behavior:**
 
-- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
+- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other VAE families retain their existing mode handling and may reject explicit Wan spatial modes.
-- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
+- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
-- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).
+- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
 
 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
 

docs/user_guide/diffusion_features.md

@@ -61,7 +61,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
 |--------|-------------|----------|
 | **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
 | **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
-| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
+| **[VAE Parallel](diffusion/parallelism/vae_parallel.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
 
 ### Extensions
 

docs/user_guide/quantization/fp8.md

@@ -80,14 +80,12 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 
 ## Model Type Support
 
-### Diffusion Models
+### Diffusion Model (Qwen-Image, Wan2.2)
 
 | 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/text_to_speech/README.md

@@ -23,7 +23,6 @@ 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 |
@@ -431,31 +430,6 @@ 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

@@ -1,125 +0,0 @@
-"""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()

recipes/LTX/LTX-2.md

@@ -4,20 +4,11 @@
 
 ## Pipelines
 
-| `--model-class-name` | Task | Required checkpoint repositories |
+| Model | Checkpoint | Task | `--model-class-name` | Batching |
-|---|---|---|
+|---|---|---|---|---|
-| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
+| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
-| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
+| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
-| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
-| `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
@@ -25,10 +16,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.
 
-`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
+`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
-then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
+or provide one initial image for I2V. As in the official LTX pipelines,
-merged distilled Transformer in both stages. Both entries support T2V and I2V;
+distilled mode is selected explicitly rather than inferred from the checkpoint
-select their class explicitly.
+name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
 
 ## API Migration
 
@@ -77,28 +68,24 @@ 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`.
 
-## Two-Stage Defaults
+## Distilled Two-Stage Defaults
 
-| Parameter | Ordinary | Full-distilled |
+| Parameter | Value |
-|---|---:|---:|
+|---|---:|
-| Final width × height | 1536 × 1024 | 1536 × 1024 |
+| Final width × height | 1536 × 1024 |
-| Stage 1 width × height | 768 × 512 | 768 × 512 |
+| Stage 1 width × height | 768 × 512 |
-| Frames / frame rate | 121 / 24 | 121 / 24 |
+| Frames / frame rate | 121 / 24 |
-| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 |
+| Stage 1 / Stage 2 denoise steps | 8 / 3 |
-| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only |
+| Guidance | Fixed positive-only |
 
-API dimensions are final dimensions and must be divisible by 64. All LTX
+The dimensions passed to the API are the final output dimensions and must be
-requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or
+divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
-LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so
+The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
-`num_inference_steps`, when supplied, must be `8`. Both entries reject custom
+`8` when supplied.
-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 one-stage checkpoint:
+Start either checkpoint:
 
 ```bash
 vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
@@ -108,19 +95,11 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
 vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
 ```
 
-Start a two-stage checkpoint with its explicit class:
+Start the distilled checkpoint with its two-stage entry:
 
 ```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
@@ -156,9 +135,8 @@ together with `input_reference`.
 
 ## Guidance
 
-One-stage and ordinary Stage 1 support independent video/audio CFG,
+LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
-spatio-temporal guidance (STG), cross-modality guidance, and rescaling.
+cross-modality guidance, and rescaling:
-Distilled stages and ordinary Stage 2 are fixed positive-only.
 
 | Parameter | Default | Effect | Alias |
 |---|---:|---|---|
@@ -194,9 +172,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 and ordinary Stage 1 have four Transformer passes
+The default one-stage recipe has four Transformer passes per denoise step:
-per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod`
+`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
-(cross-modality). The useful balanced configurations are therefore:
+configurations are therefore:
 
 | `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
 |---:|---:|---:|---|
@@ -217,8 +195,9 @@ 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. Positive-only distilled phases do not benefit
+configured parallel dimensions. Distilled pipelines use positive-only
-from `cfg_parallel_size > 1`.
+guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
+them.
 
 ### Python API
 
@@ -263,22 +242,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 two-stage. |
+| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
 | `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 and ordinary Stage 1, 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, 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 and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. |
+| `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_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` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. |
+| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
-| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. |
+| `audio_latents` | tensor, `None` | One-stage initial audio latents; 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 and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. |
+| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative 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. |
@@ -294,13 +273,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`.
 
 ### Recipe-Specific Request Capabilities
 
-| Override | One-stage | Ordinary two-stage | Distilled two-stage |
+| Override | One-stage | Distilled two-stage |
-|---|---|---|---|
+|---|---|---|
-| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only |
+| Guidance | Supported | Fixed positive-only |
-| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected |
+| Negative prompt/embeddings | Supported | Rejected |
-| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 |
+| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
-| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules |
+| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
-| Video/audio latents | Supported | Rejected | Rejected |
+| Video/audio latents | Supported | Rejected |
 
 These capability checks apply equally to direct `forward` keywords and
 values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
@@ -330,10 +309,6 @@ 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
 
@@ -348,15 +323,14 @@ 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`.
-- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution
+- Current two-stage support is limited to the distilled checkpoint; official
-  remains out of scope.
+  LTX two-stage/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

@@ -1,152 +0,0 @@
-# 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_omni_config.py

@@ -471,7 +471,6 @@ def test_sub_config_fields_match_structured_scopes():
         "trust_remote_code",
         "dtype",
         "attention_backend",
-        "attention_config",
         "moe_backend",
         "hf_overrides",
         "limit_mm_per_prompt",
@@ -583,8 +582,12 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
     assert cfg.data_parallel_size == 3
     assert cfg.cfg_parallel_size == 3
     assert cfg.mask_sp_padding is True
+    assert cfg.vae_parallel_mode == "tile"
     assert cfg.world_size == 72
 
+    auto_cfg = OmniStageDiffusionParallelConfig(vae_parallel_mode="auto")
+    assert auto_cfg.vae_parallel_mode == "auto"
+
 
 def test_parallel_config_derived_fields_are_not_init_inputs():
     with pytest.raises(ValidationError):

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
@@ -10,6 +11,7 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
+    select_auto_spatial_shard_split_dim,
 )
 from vllm_omni.platforms import current_omni_platform
 
@@ -17,6 +19,23 @@ from vllm_omni.platforms import current_omni_platform
 # the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
 
 
+@contextmanager
+def _active_spatial_shard(split_dim: str):
+    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
+        wan_spatial_shard.SpatialShardContext(
+            input_extent=2,
+            local_input_extent=2,
+            split_dim=split_dim,
+            rank=0,
+            world_size=3,
+        )
+    )
+    try:
+        yield
+    finally:
+        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_split_for_parallel_decode_pads_uneven_height():
@@ -332,7 +351,7 @@ def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.Mon
         lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim),
     )
     module = ShrinkingWanAttentionBlock()
-    wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width")
+    wan_spatial_shard._patch_attention_block(module, group=object())
     x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
     token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
         wan_spatial_shard.SpatialShardContext(
@@ -377,12 +396,14 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
     mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    mid = mid_rank_pad(x)
+    with _active_spatial_shard("height"):
+        mid = mid_rank_pad(x)
     assert mid.shape == (1, 1, 2, 3)
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
     last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    last = last_rank_pad(x)
+    with _active_spatial_shard("height"):
+        last = last_rank_pad(x)
     assert last.shape == (1, 1, 3, 3)
 
 
@@ -392,13 +413,15 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
     x = torch.ones((1, 1, 2, 2))
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
-    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    mid = mid_rank_pad(x)
+    with _active_spatial_shard("width"):
+        mid = mid_rank_pad(x)
     assert mid.shape == (1, 1, 2, 2)
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
-    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    last = last_rank_pad(x)
+    with _active_spatial_shard("width"):
+        last = last_rank_pad(x)
     assert last.shape == (1, 1, 2, 3)
 
 
@@ -421,6 +444,111 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
     assert vae._spatial_shard_decode_enabled(z) is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "z_shape,world_size,expected",
+    [
+        ((1, 16, 1, 60, 104), 4, "width"),
+        ((1, 16, 5, 60, 104), 4, "width"),
+        ((1, 16, 5, 104, 60), 4, "height"),
+        ((1, 16, 5, 2, 2), 4, None),
+        ((1, 16, 5, 64, 64), 4, "width"),
+        ((1, 16, 64, 64), 4, None),
+        ((1, 16, 5, 64, 64), 1, None),
+    ],
+)
+def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
+    assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
+        lambda group=None: 4,
+    )
+
+    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
+    vae.use_tiling = True
+    vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
+    vae.is_distributed_enabled = lambda: True
+    large = torch.zeros((1, 16, 5, 60, 104))
+    unshardable = torch.zeros((1, 16, 1, 2, 2))
+    unsupported = torch.zeros((1, 16, 60, 104))
+
+    assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
+    assert vae._spatial_shard_decode_enabled(large) is True
+    assert vae._spatial_shard_decode_enabled(unshardable) is False
+    assert vae._spatial_shard_decode_enabled(unsupported) is False
+
+    vae.distributed_executor.parallel_size = 2
+    assert vae._spatial_shard_decode_enabled(large) is False
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_set_parallel_size_prepares_spatial_wrappers_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
+    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
+    executor = SimpleNamespace(group=object(), parallel_size=1, parallel_mode="tile")
+
+    def set_parallel_size(parallel_size: int, mode: str = "tile") -> None:
+        executor.parallel_size = parallel_size
+        executor.parallel_mode = mode
+
+    executor.set_parallel_size = set_parallel_size
+    object.__setattr__(vae, "distributed_executor", executor)
+    installed = []
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
+        lambda group=None: 4,
+    )
+    monkeypatch.setattr(
+        wan_spatial_shard,
+        "install_wan_spatial_shard_decode",
+        lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
+    )
+
+    vae.set_parallel_size(2, mode="auto")
+    assert installed == []
+
+    vae.set_parallel_size(4, mode="auto")
+    assert installed == [(vae, vae.distributed_executor.group, "height")]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@torch.no_grad()
+def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
+    vae = DistributedAutoencoderKLWan(
+        base_dim=8,
+        decoder_base_dim=8,
+        z_dim=4,
+        dim_mult=[1, 2],
+        num_res_blocks=1,
+        attn_scales=[],
+        temperal_downsample=[False],
+        latents_mean=[0.0] * 4,
+        latents_std=[1.0] * 4,
+        scale_factor_temporal=1,
+        scale_factor_spatial=2,
+    ).eval()
+    latents = torch.randn(1, 4, 1, 4, 5)
+    parameter_names = {name for name, _ in vae.named_parameters()}
+    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
+    reference = vae.decode(latents, return_dict=False)[0]
+
+    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
+    assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
+    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
+
+    assert {name for name, _ in vae.named_parameters()} == parameter_names
+    vae.load_state_dict(state_dict, strict=True)
+    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
@@ -566,3 +694,99 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
     assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
         f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
     )
+
+
+def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        destroy_model_parallel,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = master_port
+    device = current_omni_platform.get_torch_device(rank)
+    current_omni_platform.set_device(device)
+    backend = current_omni_platform.dist_backend
+    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
+    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
+
+    try:
+        vae = DistributedAutoencoderKLWan.from_pretrained(
+            _SPATIAL_SHARD_MODEL,
+            subfolder=_SPATIAL_SHARD_SUBFOLDER,
+            torch_dtype=torch.float32,
+        ).to(device=device, dtype=torch.float32)
+        vae.eval()
+        generator = torch.Generator(device=device).manual_seed(1)
+        small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
+        landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
+        portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
+
+        with torch.inference_mode():
+            vae.use_tiling = False
+            vae.set_parallel_size(1, mode="tile")
+            small_reference = vae.decode(small, return_dict=False)[0]
+
+            vae.use_tiling = True
+            vae.set_parallel_size(2, mode="tile")
+            tile_before_install = vae.decode(small, return_dict=False)[0]
+
+            vae.set_parallel_size(2, mode="auto")
+            assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
+            small_sharded = vae.decode(small, return_dict=False)[0]
+            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
+            vae.decode(landscape, return_dict=False)
+
+            vae.set_parallel_size(2, mode="tile")
+            tile_after_install = vae.decode(small, return_dict=False)[0]
+
+            vae.use_tiling = False
+            vae.set_parallel_size(1, mode="tile")
+            portrait_reference = vae.decode(portrait, return_dict=False)[0]
+
+            vae.use_tiling = True
+            vae.set_parallel_size(2, mode="auto")
+            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
+            portrait_sharded = vae.decode(portrait, return_dict=False)[0]
+
+        assert all(
+            buffer is None
+            for module in vae.decoder.modules()
+            for name, buffer in module._buffers.items()
+            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
+        )
+
+        if rank == 0:
+            small_diff = (small_sharded.float() - small_reference.float()).abs()
+            portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
+            return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
+            return_dict["small_max_abs_diff"] = small_diff.max().item()
+            return_dict["small_mean_abs_diff"] = small_diff.mean().item()
+            return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
+            return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
+    finally:
+        destroy_model_parallel()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_auto_spatial_shard_switches_shape_paths_safely():
+    manager = mp.get_context("spawn").Manager()
+    return_dict = manager.dict()
+    mp.spawn(
+        _auto_spatial_shard_decode_worker,
+        args=(return_dict, "29502"),
+        nprocs=2,
+        join=True,
+    )
+
+    assert return_dict.get("tile_fallback_exact") is True
+    assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
+    assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
+    assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
+    assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE

tests/diffusion/layers/test_fused_qk_norm_rope.py

@@ -1,82 +0,0 @@
-# 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_matches_official_values_and_token_major_layout():
+def test_prepare_video_latents_samples_directly_in_packed_token_space():
     pipeline = _make_pipeline(LTX2Pipeline)
     pipeline.vae_spatial_compression_ratio = 8
     pipeline.vae_temporal_compression_ratio = 8
@@ -44,7 +44,6 @@ def test_prepare_video_latents_matches_official_values_and_token_major_layout():
     )
 
     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():
@@ -82,42 +81,6 @@ 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():
@@ -174,9 +137,7 @@ 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)
-    packed = latent_ops.pack_audio_latents(latents)
+    torch.testing.assert_close(padded, 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,19 +12,6 @@ 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."""
 
@@ -212,7 +199,6 @@ 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

@@ -1,483 +0,0 @@
-# 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,21 +13,13 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.models.ltx2 import (
+from vllm_omni.diffusion.models.ltx2 import ltx2_components
-    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 (
@@ -41,7 +33,6 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
     LTXGuidancePlan,
     LTXGuidanceSpec,
     LTXModalityGuidance,
-    _repeat_batch,
     build_perturbation_kwargs,
     combine_guided_x0,
 )
@@ -49,29 +40,18 @@ 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 (
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
-    LTX2DistilledPipeline,
-    LTX2TwoStagePipeline,
-)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -124,22 +104,10 @@ 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
 
@@ -155,72 +123,6 @@ 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
 
@@ -242,104 +144,16 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
     assert pipe.reports_stage_durations
 
 
-def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
+def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
-    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
@@ -516,15 +330,6 @@ 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(
         [
@@ -575,32 +380,6 @@ 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))
@@ -768,19 +547,11 @@ 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 not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
+    assert 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)
@@ -823,8 +594,7 @@ def test_denoise_executor_owns_progress_and_interrupt():
     torch.testing.assert_close(state.audio, torch.tensor(11.0))
 
 
-@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
+def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
-def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
     request_inputs = LTXRequestInputs(
         prompt="prompt",
         negative_prompt="",
@@ -901,15 +671,12 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
             return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
 
     prompt_context_sentinel = prompt_context
-    pipeline = object.__new__(pipeline_cls)
+    pipeline = object.__new__(LTX2DistilledPipeline)
     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)
@@ -919,18 +686,14 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
 
     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))
-    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), 4.0))
-    torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
 
 
-@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
+def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
-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__(pipeline_cls)
+    pipeline = object.__new__(LTX2DistilledPipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.vae_spatial_compression_ratio = 32
     pipeline.vae_temporal_compression_ratio = 8
@@ -987,18 +750,6 @@ 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"),
     [
@@ -1207,98 +958,6 @@ 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
@@ -1773,7 +1432,6 @@ 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:
@@ -1819,11 +1477,6 @@ 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",
@@ -1835,9 +1488,6 @@ class TestRegistryIntegration:
             "LTX23ImageToVideoPipeline",
             "LTX2TwoStagesPipeline",
             "LTX2ImageToVideoTwoStagesPipeline",
-            "LTX2ImageToVideoTwoStagePipeline",
-            "LTX23TwoStagePipeline",
-            "LTX23ImageToVideoTwoStagePipeline",
         }
         assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
 
@@ -1847,7 +1497,6 @@ class TestRegistryIntegration:
 
         expected = [
             "LTX2Pipeline",
-            "LTX2TwoStagePipeline",
             "LTX2DistilledPipeline",
             "LTX2T2VDMD2Pipeline",
             "LTX2I2VDMD2Pipeline",
@@ -1860,7 +1509,6 @@ class TestRegistryIntegration:
         "model_class_name",
         [
             "LTX2Pipeline",
-            "LTX2TwoStagePipeline",
             "LTX2DistilledPipeline",
             "LTX2T2VDMD2Pipeline",
             "LTX2I2VDMD2Pipeline",

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -1,20 +1,13 @@
 # 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 (
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
-    LTX2AudioVideoAttnProcessor,
+
-    LTX2VideoTransformer3DModel,
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-    _LTX2ParallelAttention,
-    _make_rms_norm,
-)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -73,83 +66,3 @@ 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/e2e/accuracy/ltx/run_ltx_reference.py

@@ -20,11 +20,6 @@ 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")
@@ -32,8 +27,6 @@ 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()
 
@@ -98,13 +91,9 @@ def _configure_official_sdpa(pipeline: Any) -> None:
         attention=attention,
         masked_attention=attention,
     )
-    owners_and_attributes = [
+    owners_and_attributes = (
+        (pipeline.stage, "_transformer_builder"),
         (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)
@@ -115,12 +104,6 @@ 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:
@@ -130,54 +113,16 @@ 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
 
-    if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
+    pipeline = TI2VidOneStagePipeline(
-        raise ValueError(
+        checkpoint_path=str(args.checkpoint),
-            f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
+        gemma_root=str(args.gemma_root),
-        )
+        loras=(),
-    checkpoint = _require_path(args.checkpoint, "checkpoint")
+        offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE,
-    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 = (
@@ -195,45 +140,34 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
             )
         ]
     )
-    if args.pipeline_kind == "distilled":
+    video, audio = pipeline(
-        video, audio = pipeline(
+        prompt=request["prompt"],
-            prompt=request["prompt"],
+        negative_prompt=request["negative_prompt"],
-            seed=request["seed"],
+        seed=request["seed"],
-            height=request["height"],
+        height=request["height"],
-            width=request["width"],
+        width=request["width"],
-            num_frames=request["num_frames"],
+        num_frames=request["num_frames"],
-            frame_rate=request["fps"],
+        frame_rate=request["fps"],
-            images=images,
+        num_inference_steps=request["num_inference_steps"],
-        )
+        video_guider_params=MultiModalGuiderParams(
-    else:
+            cfg_scale=request["video_cfg_scale"],
-        video, audio = pipeline(
+            stg_scale=request["video_stg_scale"],
-            prompt=request["prompt"],
+            rescale_scale=request["video_rescale_scale"],
-            negative_prompt=request["negative_prompt"],
+            modality_scale=request["video_modality_scale"],
-            seed=request["seed"],
+            skip_step=0,
-            height=request["height"],
+            stg_blocks=request["video_stg_blocks"],
-            width=request["width"],
+        ),
-            num_frames=request["num_frames"],
+        audio_guider_params=MultiModalGuiderParams(
-            frame_rate=request["fps"],
+            cfg_scale=request["audio_cfg_scale"],
-            num_inference_steps=request["num_inference_steps"],
+            stg_scale=request["audio_stg_scale"],
-            video_guider_params=MultiModalGuiderParams(
+            rescale_scale=request["audio_rescale_scale"],
-                cfg_scale=request["video_cfg_scale"],
+            modality_scale=request["audio_modality_scale"],
-                stg_scale=request["video_stg_scale"],
+            skip_step=0,
-                rescale_scale=request["video_rescale_scale"],
+            stg_blocks=request["audio_stg_blocks"],
-                modality_scale=request["video_modality_scale"],
+        ),
-                skip_step=0,
+        images=images,
-                stg_blocks=request["video_stg_blocks"],
+        max_batch_size=4,
-            ),
+    )
-            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,
@@ -245,7 +179,6 @@ 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,
         },
     )
 
@@ -260,9 +193,15 @@ 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 not frames.images:
+        if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput):
-            raise ValueError("No video frames found in OmniRequestOutput")
+            frames = frames.request_output
-        frames = frames.images
+            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 isinstance(frames, list) and len(frames) == 1:
         frames = frames[0]
@@ -314,10 +253,6 @@ 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"}}
 
@@ -338,8 +273,7 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
         parallel_config=DiffusionParallelConfig(),
     )
     try:
-        detected_model_class_name = get_model_class_name(omni)
+        model_class_name = get_model_class_name(omni)
-        model_class_name = args.model_class_name
         sampling_params = OmniDiffusionSamplingParams(
             height=request["height"],
             width=request["width"],
@@ -355,9 +289,10 @@ 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"]}
+        prompt: dict[str, Any] = {
-        if request.get("negative_prompt"):
+            "prompt": request["prompt"],
-            prompt["negative_prompt"] = request["negative_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:
@@ -377,8 +312,6 @@ 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,9 +3,10 @@
 
 """E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
 
-The original reduced one-stage guards remain unchanged. Four complementary
+The comparison runs both runtimes through PyTorch SDPA and uses
-default-shape cases cover both model versions, both two-stage weight families,
+``max_batch_size=4`` in the official reference to match Omni's fused guidance
-and T2V/I2V without expanding every combination into a duplicate golden.
+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.
 """
 
 from __future__ import annotations
@@ -16,9 +17,8 @@ import shutil
 import subprocess
 import sys
 import time
-from dataclasses import asdict, dataclass, replace
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal
 
 import numpy as np
 import pytest
@@ -53,260 +53,74 @@ 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
-@dataclass(frozen=True)
+VIDEO_PSNR_MEAN_THRESHOLD = 30.0
-class LTXArtifact:
+AUDIO_RELATIVE_L2_THRESHOLD = 0.2
-    repo_id: str
+AUDIO_COSINE_THRESHOLD = 0.95
-    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: LTXArtifact
+    checkpoint_repo: str
-    width: int
+    checkpoint_filename: str
-    height: int
+    checkpoint_revision: str
-    num_frames: int
+    checkpoint_env: str
-    num_inference_steps: int
+    stg_block: int
-    seed: int
-    stg_block: int | None
-    thresholds: LTXAccuracyThresholds
     prompt: str = PROMPT
-    gemma_model_id: str | None = None
+    image_repo: str | None = None
-    gemma_model_revision: str | None = None
+    image_filename: str | None = None
-    gemma_model_env: str | None = None
+    image_revision: 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
 
 
-LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a"
+CASES = (
-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=LTX2_REVISION,
+        model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
         model_env="VLLM_TEST_LTX2_MODEL",
         model_class_name="LTX2Pipeline",
-        checkpoint=LTX2_CHECKPOINT,
+        checkpoint_repo="Lightricks/LTX-2",
-        width=512,
+        checkpoint_filename="ltx-2-19b-dev.safetensors",
-        height=384,
+        checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
-        num_frames=25,
+        checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
-        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=LTX23_CHECKPOINT,
+        checkpoint_repo="Lightricks/LTX-2.3",
-        width=512,
+        checkpoint_filename="ltx-2.3-22b-dev.safetensors",
-        height=384,
+        checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
-        num_frames=25,
+        checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
-        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=LTX23_CHECKPOINT,
+        checkpoint_repo="Lightricks/LTX-2.3",
-        width=512,
+        checkpoint_filename="ltx-2.3-22b-dev.safetensors",
-        height=384,
+        checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
-        num_frames=25,
+        checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
-        num_inference_steps=20,
-        seed=42,
         stg_block=28,
-        thresholds=STRICT_THRESHOLDS,
+        image_repo="huggingface/documentation-images",
-        image=I2V_IMAGE,
+        image_filename="diffusers/svd/rocket.png",
+        image_revision="645d8364f0c7a101180b364811b5a11a362e4010",
     ),
 )
 
 
-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)
@@ -395,7 +209,6 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
                 "model_index.json",
                 "audio_vae/*",
                 "connectors/*",
-                "latent_upsampler/*",
                 "processor/*",
                 "scheduler/*",
                 "text_encoder/config.json",
@@ -410,82 +223,74 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
     )
 
 
-def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path:
+def _resolve_gemma_root(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_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path:
+def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path:
-    configured_path = os.environ.get(artifact.env)
+    configured_checkpoint = os.environ.get(case.checkpoint_env)
-    if configured_path:
+    if configured_checkpoint:
-        path = Path(configured_path)
+        checkpoint = Path(configured_checkpoint)
-        assert path.is_file(), f"Configured LTX artifact not found: {path}"
+        assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}"
-        return path
+        return checkpoint
-    if model is not None:
+    model_checkpoint = model / case.checkpoint_filename
-        model_path = model / artifact.filename
+    if model_checkpoint.is_file():
-        if model_path.is_file():
+        return model_checkpoint
-            return model_path
     return Path(
         hf_hub_download(
-            repo_id=artifact.repo_id,
+            repo_id=case.checkpoint_repo,
-            repo_type=None if artifact.repo_type == "model" else artifact.repo_type,
+            filename=case.checkpoint_filename,
-            filename=artifact.filename,
+            revision=case.checkpoint_revision,
-            revision=artifact.revision,
         )
     )
 
 
 def _resolve_image(case: LTXAccuracyCase) -> Path | None:
-    if case.image is None:
+    if case.image_filename is None:
         return None
-    return _resolve_artifact(case.image)
+    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,
+        )
+    )
 
 
 def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
     request: dict[str, object] = {
-        "pipeline_kind": case.pipeline_kind,
         "prompt": case.prompt,
-        "negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
+        "negative_prompt": NEGATIVE_PROMPT,
-        "width": case.width,
+        "width": 512,
-        "height": case.height,
+        "height": 384,
-        "num_frames": case.num_frames,
+        "num_frames": 25,
         "fps": 24,
-        "num_inference_steps": case.num_inference_steps,
+        "num_inference_steps": 20,
-        "seed": case.seed,
+        "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],
     }
-    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
@@ -494,30 +299,20 @@ 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] = []
-    max_abs = 0.0
+    for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
-    absolute_error_sum = 0.0
+        reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
-    with torch.inference_mode():
+        prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
-        for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
+        ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
-            reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device)
+        psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
-            prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device)
+    difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
-            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": max_abs,
+        "max_abs": float(difference.max()),
-        "mean_abs": absolute_error_sum / reference.size,
+        "mean_abs": float(difference.mean()),
     }
 
 
@@ -548,19 +343,13 @@ 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_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
+def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
-    """Compare official and Omni raw AV outputs from the same single-H100 request."""
+    """Compare official and Omni raw AV outputs from the same E2E 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(case, model)
+    gemma_root = _resolve_gemma_root(model)
-    checkpoint = _resolve_artifact(case.checkpoint, model)
+    checkpoint = _resolve_checkpoint(case, 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")
@@ -570,13 +359,8 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
         str(runner),
         "--request",
         str(request_path),
-        "--pipeline-kind",
-        case.pipeline_kind,
     ]
-    enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get(
+    if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}:
-        "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
@@ -588,9 +372,10 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
     )
 
     official_output = output_root / "official"
-    official_args = _official_runner_prefix() + runner_args
+    _run(
-    official_args.extend(
+        _official_runner_prefix()
-        [
+        + runner_args
+        + [
             "--backend",
             "official",
             "--output-dir",
@@ -601,16 +386,12 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
             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"
-    omni_args = (
+    _run(
         [sys.executable]
         + runner_args
         + [
@@ -622,9 +403,9 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
             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())
@@ -642,27 +423,18 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
     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"] >= case.thresholds.video_ssim_mean
+    assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD
-    assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
+    assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD
-    assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
+    assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD
-    if case.thresholds.audio_relative_l2 > 0:
+    assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD
-        assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
+    assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD
-    if case.thresholds.audio_cosine_similarity > 0:
-        assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity

tests/e2e/online_serving/test_ltx.py

@@ -1,109 +0,0 @@
-# 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

@@ -0,0 +1,91 @@
+# 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_stage_engine_args.py

@@ -139,7 +139,6 @@ 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"],
@@ -230,7 +229,6 @@ 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_serve.py

@@ -76,6 +76,18 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
+def test_serve_parser_defaults_vae_parallel_mode_to_tile_and_accepts_auto() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    default_args = parser.parse_args(["serve", "fake-model", "--omni"])
+    auto_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "auto"])
+
+    assert default_args.vae_parallel_mode == "tile"
+    assert auto_args.vae_parallel_mode == "auto"
+
+
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",

tests/model_executor/models/dots_tts/__init__.py

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

tests/model_executor/models/dots_tts/test_talker.py

@@ -1,137 +0,0 @@
-# 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,9 +1326,7 @@ class TestRegistry:
 
         assert "higgs_multimodal_qwen3" in OMNI_PIPELINES
 
-    def test_deploy_yaml_uses_native_flashinfer(self):
+    def test_deploy_yaml_exists(self):
-        from vllm_omni.config.stage_config import load_deploy_config
-
         for name in (
             "higgs_multimodal_qwen3.yaml",
             "higgs_multimodal_qwen3_high_throughput.yaml",
@@ -1336,10 +1334,6 @@ 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,13 +233,10 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None:
         }
     )
 
-    for pipeline_name in (
+    assert get_extra_body_params("LTX2Pipeline") == expected
-        "LTX2Pipeline",
+    assert get_extra_output_params("LTX2Pipeline") == frozenset()
-        "LTX2TwoStagePipeline",
+    assert get_extra_body_params("LTX2DistilledPipeline") == expected
-        "LTX2DistilledPipeline",
+    assert get_extra_output_params("LTX2DistilledPipeline") == frozenset()
-    ):
-        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,7 +27,6 @@ EXCLUDED_MODELS = [
     "OvisImagePipeline",
     "WanPipeline",
     "WanVACEPipeline",
-    "LTX2TwoStagePipeline",
     "LTX2DistilledPipeline",
     "LTX2T2VDMD2Pipeline",
     "LTX2I2VDMD2Pipeline",

vllm_omni/config/omni_config.py

@@ -90,7 +90,6 @@ 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,7 +321,6 @@ 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
@@ -475,9 +473,10 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
         if self.ulysses_mode not in {"strict", "advanced_uaa"}:
             raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
-        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
-                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of "
+                "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 

vllm_omni/config/pipeline_registry.py

@@ -47,7 +47,6 @@ 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
@@ -145,7 +144,6 @@ 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/deploy/README_higgs_audio_v3.md

@@ -44,65 +44,6 @@ 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

@@ -1,28 +0,0 @@
-# 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,11 +38,6 @@ 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,11 +37,6 @@ 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,11 +43,6 @@ 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/diffusion/cache/teacache/extractors.py

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

vllm_omni/diffusion/data.py

@@ -212,6 +212,9 @@ class DiffusionParallelConfig:
 
     - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
       of spatial tiles and the results are stitched on rank 0.
+    - "auto": Select the longer spatial-shard axis for validated Wan tiled
+      decodes on a full DiT-group topology; otherwise use patch/tile parallel
+      decode.
     - "spatial_shard_height": Spatially-sharded decode that splits decoder
       feature maps along height and exchanges halo rows around spatial
       convolutions.
@@ -251,8 +254,9 @@ class DiffusionParallelConfig:
         assert self.allgather_degree > 0, "AllGather degree must be > 0"
         assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
         assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
-        assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
+        assert self.vae_parallel_mode in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}, (
-            "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+            "vae_parallel_mode must be one of "
+            "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
             f"but got {self.vae_parallel_mode!r}."
         )
         if self.allgather_degree > 1:

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -23,6 +23,30 @@ from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 
 
+def select_auto_spatial_shard_split_dim(
+    z_shape: tuple[int, ...],
+    world_size: int,
+) -> str | None:
+    """Select a profitable Wan spatial-shard axis for one latent shape.
+
+    This selector runs only after Diffusers has already chosen ``tiled_decode``
+    for a spatially large request. Splitting the longer eligible axis reduces
+    halo traffic for rectangular requests and improves per-rank balance.
+    """
+    if len(z_shape) != 5 or world_size <= 1:
+        return None
+
+    _, _, _, height, width = z_shape
+    candidates = [
+        (extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size
+    ]
+    if not candidates:
+        return None
+    # Prefer width for square latents: it was modestly faster in the existing
+    # four-GPU Wan benchmark and keeps the choice deterministic.
+    return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
+
+
 class OmniAutoencoderKLWan(AutoencoderKLWan):
     def _execution_context(self):
         try:
@@ -56,6 +80,19 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         model.init_distributed()
         return model
 
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
+        super().set_parallel_size(parallel_size, mode=mode)
+        if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
+            return
+
+        executor = self.distributed_executor
+        world_size = dist.get_world_size(group=executor.group)
+        if parallel_size <= 1 or parallel_size != world_size:
+            return
+
+        split_dim = "width" if mode == "spatial_shard_width" else "height"
+        wan_spatial_shard.install_wan_spatial_shard_decode(self, executor.group, split_dim=split_dim)
+
     def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
         _, _, num_frames, height, width = z.shape
         sample_height = height * self.spatial_compression_ratio
@@ -270,20 +307,26 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         dec = torch.clamp(dec, min=-1.0, max=1.0)
         return dec
 
-    def _spatial_shard_decode_split_dim(self) -> str | None:
+    def _spatial_shard_decode_split_dim(
+        self,
+        z: torch.Tensor | None = None,
+        world_size: int | None = None,
+    ) -> str | None:
         mode = self.distributed_executor.parallel_mode
         if mode == "spatial_shard_width":
             return "width"
         if mode == "spatial_shard_height":
             return "height"
+        if mode == "auto" and z is not None and world_size is not None:
+            return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size)
         return None
 
     def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        split_dim = self._spatial_shard_decode_split_dim()
+        if self.distributed_executor.parallel_mode == "tile":
-        if split_dim is None:
             return False
         if z.ndim != 5:
-            logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
+            if self.distributed_executor.parallel_mode != "tile":
+                logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
             return False
         if not self.is_distributed_enabled():
             return False
@@ -292,23 +335,32 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         world_size = dist.get_world_size(group=group)
         requested_size = int(self.distributed_executor.parallel_size)
         if requested_size != world_size:
-            logger.warning(
+            if self.distributed_executor.parallel_mode == "auto":
-                "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
+                logger.debug(
-                "to match the DIT group size; falling back to tiled decode. "
+                    "Wan VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
-                "requested=%s dit_group=%s split_dim=%s",
+                    requested_size,
-                requested_size,
+                    world_size,
-                world_size,
+                )
-                split_dim,
+            elif self._spatial_shard_decode_split_dim() is not None:
-            )
+                logger.warning(
+                    "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
+                    "to match the DIT group size; falling back to tiled decode. "
+                    "requested=%s dit_group=%s split_dim=%s",
+                    requested_size,
+                    world_size,
+                    self._spatial_shard_decode_split_dim(),
+                )
             return False
-        return True
+        return self._spatial_shard_decode_split_dim(z, world_size) is not None
 
     def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
         if self._spatial_shard_decode_enabled(z):
-            split_dim = self._spatial_shard_decode_split_dim()
+            world_size = dist.get_world_size(group=self.distributed_executor.group)
+            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
+            assert split_dim is not None
             logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim)
             return wan_spatial_shard.spatial_shard_decode(
                 self,

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -52,6 +52,11 @@ _SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar(
 )
 
 
+def _active_spatial_shard_split_dim() -> str | None:
+    context = _SPATIAL_SHARD_CONTEXT.get()
+    return context.split_dim if context is not None else None
+
+
 def _spatial_dim(split_dim: str) -> int:
     if split_dim == "height":
         return -2
@@ -332,21 +337,27 @@ class WanDistZeroPad2d(nn.Module):
         padding: tuple[int, int, int, int],
         group: dist.ProcessGroup,
         *,
-        split_dim: str = "height",
+        defer_split_after_padding: bool = False,
-        split_padding: tuple[int, int] | None = None,
     ) -> None:
         super().__init__()
         self.padding = padding
-        self.split_dim = split_dim
+        self.defer_split_after_padding = defer_split_after_padding
-        default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1])
-        self.split_padding = split_padding or default_split_padding
         self.group = group
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        split_dim = _active_spatial_shard_split_dim()
+        if split_dim is None:
+            return F.pad(x, self.padding)
+
         rank, world_size = _rank_world(self.group)
         left, right, top, bottom = self.padding
+        if self.defer_split_after_padding:
+            if split_dim == "height":
+                bottom = 0
+            else:
+                right = 0
         if world_size > 1:
-            if self.split_dim == "height":
+            if split_dim == "height":
                 top = top if rank == 0 else 0
                 bottom = bottom if rank == world_size - 1 else 0
             else:
@@ -360,8 +371,7 @@ class WanDistConv2d(nn.Conv2d):
         self,
         source: nn.Conv2d,
         group: dist.ProcessGroup,
-        split_dim: str = "height",
+        deferred_padding: tuple[int, int, int, int] | None = None,
-        split_padding: tuple[int, int] | None = None,
     ):
         super().__init__(
             source.in_channels,
@@ -378,75 +388,117 @@ class WanDistConv2d(nn.Conv2d):
         )
         self.load_state_dict(source.state_dict())
         self.group = group
-        self.split_dim = split_dim
+        self._direct_padding = source.padding
-        self.split_tensor_dim = _spatial_dim(split_dim)
+        self._direct_padding_mode = source.padding_mode
-        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
+        self._direct_reversed_padding = source._reversed_padding_repeated_twice
-        self.halo_size = (kernel_extent - 1) // 2
+        self._deferred_padding = deferred_padding
-        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
+        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
-        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
+        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
-        if split_dim == "height":
+            raise ValueError(f"Wan spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
-            if split_padding is None:
+        self._spatial_pad_h = int(source_padding[-2])
-                split_padding = (pad_h, pad_h)
+        self._spatial_pad_w = int(source_padding[-1])
-            self._non_split_padding = (pad_w, pad_w, 0, 0)
+        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-            self.kernel_extent = self.kernel_size[-2]
+        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-            self.stride_extent = self.stride[-2]
+        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
-        else:
-            if split_padding is None:
-                split_padding = (pad_w, pad_w)
-            self._non_split_padding = (0, 0, pad_h, pad_h)
-            self.kernel_extent = self.kernel_size[-1]
-            self.stride_extent = self.stride[-1]
-        self.split_pad_left, self.split_pad_right = split_padding
-        self._halo_recv_top_buf: torch.Tensor | None = None
-        self._halo_recv_bottom_buf: torch.Tensor | None = None
-        self._trim_cache: dict[int, tuple[int, int, int]] = {}
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.pad(x, self._non_split_padding)
+        split_dim = _active_spatial_shard_split_dim()
+        if split_dim is None:
+            if self._direct_padding_mode != "zeros":
+                x = F.pad(x, self._direct_reversed_padding, mode=self._direct_padding_mode)
+                padding: str | tuple[int, int] = (0, 0)
+            else:
+                padding = self._direct_padding
+            return F.conv2d(x, self.weight, self.bias, self.stride, padding, self.dilation, self.groups)
+
+        split_tensor_dim = _spatial_dim(split_dim)
+        if split_dim == "height":
+            kernel_extent = self.kernel_size[-2]
+            stride_extent = self.stride[-2]
+            non_split_padding = (self._spatial_pad_w, self._spatial_pad_w, 0, 0)
+            default_split_padding = (self._spatial_pad_h, self._spatial_pad_h)
+            deferred_split_padding = (
+                (self._deferred_padding[2], self._deferred_padding[3]) if self._deferred_padding is not None else None
+            )
+        else:
+            kernel_extent = self.kernel_size[-1]
+            stride_extent = self.stride[-1]
+            non_split_padding = (0, 0, self._spatial_pad_h, self._spatial_pad_h)
+            default_split_padding = (self._spatial_pad_w, self._spatial_pad_w)
+            deferred_split_padding = (
+                (self._deferred_padding[0], self._deferred_padding[1]) if self._deferred_padding is not None else None
+            )
+        split_pad_left, split_pad_right = deferred_split_padding or default_split_padding
+        halo_size = (kernel_extent - 1) // 2
+
+        x = F.pad(x, non_split_padding)
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=self.halo_size,
+            halo_size=halo_size,
-            split_dim=self.split_dim,
+            split_dim=split_dim,
             recv_top_buf=self._halo_recv_top_buf,
             recv_bottom_buf=self._halo_recv_bottom_buf,
         )
-        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
+        shift, start, upper_bound = self._get_trim_params(
+            x.shape[split_tensor_dim],
+            split_dim=split_dim,
+            halo_size=halo_size,
+            split_pad_left=split_pad_left,
+            split_pad_right=split_pad_right,
+            kernel_extent=kernel_extent,
+            stride_extent=stride_extent,
+        )
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                self.split_tensor_dim,
+                split_tensor_dim,
                 shift,
-                x_padded.shape[self.split_tensor_dim] - shift,
+                x_padded.shape[split_tensor_dim] - shift,
             )
         out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=split_dim)
 
-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+    def _get_trim_params(
-        trim_params = self._trim_cache.get(local_extent)
+        self,
+        local_extent: int,
+        *,
+        split_dim: str,
+        halo_size: int,
+        split_pad_left: int,
+        split_pad_right: int,
+        kernel_extent: int,
+        stride_extent: int,
+    ) -> tuple[int, int, int]:
+        cache_key = (split_dim, local_extent)
+        trim_params = self._trim_cache.get(cache_key)
         if trim_params is None:
             rank, world_size = _rank_world(self.group)
             trim_params = _compute_conv_trim_params(
                 local_extent=local_extent,
                 rank=rank,
                 world_size=world_size,
-                halo_size=self.halo_size,
+                halo_size=halo_size,
-                pad_before=self.split_pad_left,
+                pad_before=split_pad_left,
-                pad_after=self.split_pad_right,
+                pad_after=split_pad_right,
-                kernel_extent=self.kernel_extent,
+                kernel_extent=kernel_extent,
-                stride_extent=self.stride_extent,
+                stride_extent=stride_extent,
             )
-            self._trim_cache[local_extent] = trim_params
+            self._trim_cache[cache_key] = trim_params
         return trim_params
 
 
 class WanDistCausalConv3d(nn.Conv3d):
+    # Cross-PR compatibility contract for the lossless Wan data-movement
+    # installer: it may fuse the no-context direct path, but must leave this
+    # wrapper in control whenever a spatial-shard context is active.
+    _vllm_omni_dynamic_spatial_shard_conv = True
+
     def __init__(
         self,
         source: nn.Conv3d,
         group: dist.ProcessGroup,
-        split_dim: str = "height",
     ):
         super().__init__(
             source.in_channels,
@@ -463,88 +515,98 @@ class WanDistCausalConv3d(nn.Conv3d):
         )
         self.load_state_dict(source.state_dict())
         self.group = group
-        self.split_dim = split_dim
-        self.split_tensor_dim = _spatial_dim(split_dim)
         source_padding = getattr(source, "_padding", None)
         if source_padding is None:
             p_t, p_h, p_w = source.padding
             source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0)
         self._source_padding = tuple(source_padding)
-        if split_dim == "height":
+        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-            self.split_pad_left = int(self._source_padding[2])
+        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-            self.split_pad_right = int(self._source_padding[3])
+        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
-            self.kernel_extent = self.kernel_size[-2]
-            self.stride_extent = self.stride[-2]
-            padding = (
-                self._source_padding[0],
-                self._source_padding[1],
-                0,
-                0,
-                self._source_padding[4],
-                self._source_padding[5],
-            )
-        else:
-            self.split_pad_left = int(self._source_padding[0])
-            self.split_pad_right = int(self._source_padding[1])
-            self.kernel_extent = self.kernel_size[-1]
-            self.stride_extent = self.stride[-1]
-            padding = (
-                0,
-                0,
-                self._source_padding[2],
-                self._source_padding[3],
-                self._source_padding[4],
-                self._source_padding[5],
-            )
-        self.halo_size = (self.kernel_extent - 1) // 2
-        self._padding = padding if self.halo_size > 0 else self._source_padding
-        self._halo_recv_top_buf: torch.Tensor | None = None
-        self._halo_recv_bottom_buf: torch.Tensor | None = None
-        self._trim_cache: dict[int, tuple[int, int, int]] = {}
 
     def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        padding = list(self._padding)
+        split_dim = _active_spatial_shard_split_dim()
+        padding = list(self._source_padding)
         if cache_x is not None and padding[4] > 0:
             cache_x = cache_x.to(x.device)
             x = torch.cat([cache_x, x], dim=2)
             padding[4] -= cache_x.shape[2]
 
+        if split_dim is None:
+            return super().forward(F.pad(x, padding))
+
+        split_tensor_dim = _spatial_dim(split_dim)
+        if split_dim == "height":
+            split_pad_left = int(padding[2])
+            split_pad_right = int(padding[3])
+            kernel_extent = self.kernel_size[-2]
+            stride_extent = self.stride[-2]
+            padding[2] = 0
+            padding[3] = 0
+        else:
+            split_pad_left = int(padding[0])
+            split_pad_right = int(padding[1])
+            kernel_extent = self.kernel_size[-1]
+            stride_extent = self.stride[-1]
+            padding[0] = 0
+            padding[1] = 0
+        halo_size = (kernel_extent - 1) // 2
+
         x = F.pad(x, padding)
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=self.halo_size,
+            halo_size=halo_size,
-            split_dim=self.split_dim,
+            split_dim=split_dim,
             recv_top_buf=self._halo_recv_top_buf,
             recv_bottom_buf=self._halo_recv_bottom_buf,
         )
-        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
+        shift, start, upper_bound = self._get_trim_params(
+            x.shape[split_tensor_dim],
+            split_dim=split_dim,
+            halo_size=halo_size,
+            split_pad_left=split_pad_left,
+            split_pad_right=split_pad_right,
+            kernel_extent=kernel_extent,
+            stride_extent=stride_extent,
+        )
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                self.split_tensor_dim,
+                split_tensor_dim,
                 shift,
-                x_padded.shape[self.split_tensor_dim] - shift,
+                x_padded.shape[split_tensor_dim] - shift,
             )
         out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=split_dim)
 
-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+    def _get_trim_params(
-        trim_params = self._trim_cache.get(local_extent)
+        self,
+        local_extent: int,
+        *,
+        split_dim: str,
+        halo_size: int,
+        split_pad_left: int,
+        split_pad_right: int,
+        kernel_extent: int,
+        stride_extent: int,
+    ) -> tuple[int, int, int]:
+        cache_key = (split_dim, local_extent)
+        trim_params = self._trim_cache.get(cache_key)
         if trim_params is None:
             rank, world_size = _rank_world(self.group)
             trim_params = _compute_conv_trim_params(
                 local_extent=local_extent,
                 rank=rank,
                 world_size=world_size,
-                halo_size=self.halo_size,
+                halo_size=halo_size,
-                pad_before=self.split_pad_left,
+                pad_before=split_pad_left,
-                pad_after=self.split_pad_right,
+                pad_after=split_pad_right,
-                kernel_extent=self.kernel_extent,
+                kernel_extent=kernel_extent,
-                stride_extent=self.stride_extent,
+                stride_extent=stride_extent,
             )
-            self._trim_cache[local_extent] = trim_params
+            self._trim_cache[cache_key] = trim_params
         return trim_params
 
 
@@ -591,12 +653,15 @@ def _trim_local_conv_output(
     return out
 
 
-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
     if getattr(module, "_vllm_omni_spatial_shard_attention", False):
         return
     orig_forward = module.forward
 
     def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
+        split_dim = _active_spatial_shard_split_dim()
+        if split_dim is None:
+            return orig_forward(x, *args, **kwargs)
         _, world_size = _rank_world(group)
         if world_size <= 1:
             return orig_forward(x, *args, **kwargs)
@@ -630,7 +695,6 @@ def _replace_child(
     name: str,
     child: nn.Module,
     group: dist.ProcessGroup,
-    split_dim: str,
 ) -> None:
     if child.__class__.__name__ == "WanCausalConv3d":
         setattr(
@@ -639,51 +703,39 @@ def _replace_child(
             WanDistCausalConv3d(
                 child,
                 group,
-                split_dim=split_dim,
             ),
         )
         return
     if isinstance(child, nn.ZeroPad2d):
         padding = tuple(int(p) for p in child.padding)
-        module_padding = padding
-        if parent.__class__.__name__ == "Sequential":
-            # Let the following WanDistConv2d account for global after-edge
-            # padding; this module handles the non-split dimension and before edge.
-            if split_dim == "height":
-                module_padding = (padding[0], padding[1], padding[2], 0)
-            else:
-                module_padding = (padding[0], 0, padding[2], padding[3])
         setattr(
             parent,
             name,
             WanDistZeroPad2d(
-                module_padding,
+                padding,
                 group,
-                split_dim=split_dim,
+                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
-                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
             ),
         )
         return
     if isinstance(child, nn.Conv2d):
-        split_padding = None
+        deferred_padding = None
         if name == "1" and parent.__class__.__name__ == "Sequential":
             # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
             # stride-2 conv with padding=0.  Only the last rank should see the
             # bottom/right global padding, which is approximated by split_padding.
             prev = getattr(parent, "0", None)
             if isinstance(prev, WanDistZeroPad2d):
-                split_padding = prev.split_padding
+                deferred_padding = prev.padding
             elif isinstance(prev, nn.ZeroPad2d):
-                pad = prev.padding
+                deferred_padding = tuple(int(value) for value in prev.padding)
-                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
         setattr(
             parent,
             name,
             WanDistConv2d(
                 child,
                 group,
-                split_dim=split_dim,
+                deferred_padding=deferred_padding,
-                split_padding=split_padding,
             ),
         )
 
@@ -691,21 +743,20 @@ def _replace_child(
 def _patch_decoder_modules(
     module: nn.Module,
     group: dist.ProcessGroup,
-    split_dim: str,
     inside_attention: bool = False,
 ) -> None:
     if module.__class__.__name__ == "WanAttentionBlock":
-        _patch_attention_block(module, group, split_dim)
+        _patch_attention_block(module, group)
         inside_attention = True
 
     for name, child in list(module.named_children()):
         if child.__class__.__name__ == "WanCausalConv3d":
-            _replace_child(module, name, child, group, split_dim)
+            _replace_child(module, name, child, group)
             continue
         if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(module, name, child, group, split_dim)
+            _replace_child(module, name, child, group)
             continue
-        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
+        _patch_decoder_modules(child, group, inside_attention=inside_attention)
 
 
 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -716,16 +767,25 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
     return count
 
 
+def _clear_wan_spatial_shard_runtime_buffers(vae: Any) -> None:
+    decoder = getattr(vae, "decoder", None)
+    if decoder is None:
+        return
+    for module in decoder.modules():
+        if "_halo_recv_top_buf" in module._buffers:
+            module._halo_recv_top_buf = None
+        if "_halo_recv_bottom_buf" in module._buffers:
+            module._halo_recv_bottom_buf = None
+
+
 def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
     """Patch ``vae.decoder`` once for spatially-sharded decode.
 
     This mutates the already-loaded decoder in place by swapping its spatial
-    convolutions/padding for halo-exchanging variants and wrapping
+    convolutions/padding for context-aware variants and attaching a dedicated
-    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
+    spatial forward. Outside that forward the replacements execute their exact
-    instance and is applied only once (subsequent calls are no-ops). A given
+    local PyTorch behavior, so a later auto-selected tile request remains safe.
-    instance is bound to a single ``split_dim``; switching between
+    The same installed decoder can alternate height and width sharding.
-    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
-    otherwise.
 
     Only group-relative rank 0 assembles the final decoded frame, mirroring the
     distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
@@ -733,18 +793,12 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
     """
     _spatial_dim(split_dim)
     if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
-        installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height")
-        if installed_split_dim != split_dim:
-            raise ValueError(
-                "Wan spatial-shard VAE decoder was already patched for "
-                f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split."
-            )
         return
     decoder = getattr(vae, "decoder", None)
     if decoder is None:
         raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
 
-    _patch_decoder_modules(decoder, group, split_dim)
+    _patch_decoder_modules(decoder, group)
     upsample_count = _decoder_upsample_count(decoder)
     orig_forward = decoder.forward
 
@@ -754,6 +808,8 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
         feat_cache: list[torch.Tensor] | None = None,
         feat_idx: list[int] | None = None,
         first_chunk: bool = False,
+        *,
+        split_dim: str,
     ) -> torch.Tensor:
         if feat_idx is None:
             feat_idx = [0]
@@ -781,10 +837,9 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
             _SPATIAL_SHARD_CONTEXT.reset(token)
         return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0)
 
-    decoder.forward = MethodType(_forward, decoder)
+    decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
     vae._vllm_omni_wan_spatial_shard_installed = True
-    vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    logger.info("Installed Wan VAE dynamic-axis spatial-shard decode.")
-    logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
 
 
 def spatial_shard_decode(
@@ -813,11 +868,12 @@ def spatial_shard_decode(
             decoded_chunks = []
             for i in range(z.shape[2]):
                 vae._conv_idx = [0]
-                chunk = vae.decoder(
+                chunk = vae.decoder._vllm_omni_spatial_shard_forward(
                     x[:, :, i : i + 1, :, :],
                     feat_cache=vae._feat_map,
                     feat_idx=vae._conv_idx,
                     first_chunk=(i == 0),
+                    split_dim=split_dim,
                 )
                 if produce_output:
                     decoded_chunks.append(chunk)
@@ -831,6 +887,10 @@ def spatial_shard_decode(
                 out = z.new_zeros(0)
     finally:
         vae.clear_cache()
+        # Halo buffers are request-shape scratch space allocated outside the
+        # tagged weight pool. Drop live references so sleep/offload never has
+        # to preserve or discard them as model state.
+        _clear_wan_spatial_shard_runtime_buffers(vae)
 
     if not return_dict:
         return (out,)

vllm_omni/diffusion/layers/fused_qk_norm_rope.py

@@ -1,299 +0,0 @@
-# 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/models/ltx2/__init__.py

@@ -12,16 +12,12 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     LTX2Pipeline,
     LTX2T2VDMD2Pipeline,
 )
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
-    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

@@ -1,156 +0,0 @@
-# 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,8 +9,7 @@ import inspect
 import json
 import logging
 import os
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -20,8 +19,6 @@ 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
@@ -35,7 +32,6 @@ 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,
@@ -74,9 +70,6 @@ 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(
@@ -108,96 +101,16 @@ 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:
@@ -479,34 +392,16 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         )
 
     if "latent_upsampler" in profile.resident_modules:
-        upsampler_config = os.path.join(model, "latent_upsampler", "config.json")
+        # BlurDownsample constructs an integer kernel that must be initialized
-        if os.path.isfile(upsampler_config) or not local_files_only:
+        # on CPU; component placement is handled uniformly after construction.
-            try:
+        with torch.device("cpu"):
-                pipeline.latent_upsampler = _load_component(
+            pipeline.latent_upsampler = _load_component(
-                    LTX2LatentUpsamplerModel,
+                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,
-                profile.artifact_repo_id,
+                "latent_upsampler",
-                profile.latent_upsampler_filename,
+                local_files_only=local_files_only,
+                dtype=dtype,
             )
-            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)
@@ -517,16 +412,6 @@ 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,14 +62,13 @@ 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 AV latents and the context used to produce them."""
+    """Denoised, unpacked AV latents and the context used to produce them."""
 
     forward_context: LTXForwardContext
     video: torch.Tensor
@@ -285,8 +284,7 @@ def prepare_rope_coords_stage(
     latents: torch.Tensor,
     audio_latents: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    transformer = pipeline.transformer
+    video_coords = pipeline.transformer.rope.prepare_video_coords(
-    video_coords = transformer.rope.prepare_video_coords(
         latents.shape[0],
         forward_ctx.latent_num_frames,
         forward_ctx.latent_height,
@@ -294,7 +292,7 @@ def prepare_rope_coords_stage(
         latents.device,
         fps=forward_ctx.request_inputs.frame_rate,
     )
-    audio_coords = transformer.audio_rope.prepare_audio_coords(
+    audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords(
         audio_latents.shape[0],
         forward_ctx.padded_audio_num_frames,
         audio_latents.device,
@@ -333,7 +331,6 @@ 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,
@@ -445,25 +442,11 @@ 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)

vllm_omni/diffusion/models/ltx2/ltx2_guidance.py

@@ -126,10 +126,7 @@ def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.
 
 def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
     """Convert x0 back to velocity using the official fp32 arithmetic order."""
-    # Official LTX converts the scalar scheduler sigma to a host value before
+    sigma = sigma.to(torch.float32)
-    # division. A CUDA scalar tensor selects different kernel arithmetic and
-    # can cross bf16 rounding boundaries late in the denoise trajectory.
-    sigma = sigma.to(torch.float32).item()
     return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype)
 
 
@@ -151,7 +148,6 @@ def combine_guided_x0(
     uncond_perturbed: torch.Tensor | float,
     uncond_modality: torch.Tensor | float,
     guidance: LTXModalityGuidance,
-    rescale_token_count: int | None = None,
 ) -> torch.Tensor:
     dtype = cond.dtype
     cond = cond.float()
@@ -167,18 +163,9 @@ def combine_guided_x0(
         # uses its full-tensor standard deviation. In Omni, dimension 0 may
         # instead contain independently scheduled requests. Reduce each item
         # separately so its result cannot depend on unrelated co-batched work.
-        stats_cond = cond
+        reduce_dims = tuple(range(1, pred.ndim))
-        stats_pred = pred
+        cond_std = cond.std(dim=reduce_dims, keepdim=True)
-        if rescale_token_count is not None:
+        pred_std = pred.std(dim=reduce_dims, keepdim=True)
-            if cond.ndim < 2 or not 0 < rescale_token_count <= cond.shape[1]:
-                raise ValueError(
-                    f"Guidance rescale token count must be in [1, {cond.shape[1]}], got {rescale_token_count}."
-                )
-            stats_cond = cond[:, :rescale_token_count]
-            stats_pred = pred[:, :rescale_token_count]
-        reduce_dims = tuple(range(1, stats_pred.ndim))
-        cond_std = stats_cond.std(dim=reduce_dims, keepdim=True)
-        pred_std = stats_pred.std(dim=reduce_dims, keepdim=True)
         eps = torch.finfo(pred.dtype).eps
         factor = cond_std / pred_std.clamp_min(eps)
         factor = torch.where(pred_std > eps, factor, torch.ones_like(factor))
@@ -188,8 +175,6 @@ def combine_guided_x0(
 
 
 def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor:
-    if repeats == 1:
-        return tensor
     return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1))
 
 
@@ -293,11 +278,6 @@ class LTXGuidanceExecutor:
         if model_pass_count > 1:
             denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, model_pass_count)
             denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, model_pass_count)
-            if denoise_ctx.audio_attention_mask is not None:
-                denoise_ctx.audio_attention_mask = _repeat_batch(
-                    denoise_ctx.audio_attention_mask,
-                    model_pass_count,
-                )
         return denoise_ctx
 
     @staticmethod
@@ -342,7 +322,6 @@ class LTXGuidanceExecutor:
         guidance: LTXModalityGuidance,
         *,
         model_sigma: torch.Tensor | None = None,
-        rescale_token_count: int | None = None,
     ) -> torch.Tensor:
         model_sigma = sigma if model_sigma is None else model_sigma
         # Official guidance reduces contiguous BSC tensors. The fused
@@ -355,7 +334,6 @@ class LTXGuidanceExecutor:
             uncond_perturbed=x0.get("ptb", 0.0),
             uncond_modality=x0.get("mod", 0.0),
             guidance=guidance,
-            rescale_token_count=rescale_token_count,
         )
         return velocity_from_x0(sample, guided, sigma)
 
@@ -456,7 +434,6 @@ class LTXGuidanceExecutor:
                 audio_splits,
                 forward_ctx.audio_scheduler.sigmas[index],
                 plan.spec.audio,
-                rescale_token_count=forward_ctx.original_audio_num_frames,
             ),
         )
 
@@ -532,7 +509,6 @@ class LTXGuidanceExecutor:
                 audio_splits,
                 forward_ctx.audio_scheduler.sigmas[index],
                 plan.spec.audio,
-                rescale_token_count=forward_ctx.original_audio_num_frames,
             ),
         )
 

vllm_omni/diffusion/models/ltx2/ltx2_latents.py

@@ -117,9 +117,7 @@ def create_noised_state(
         device=latents.device,
         dtype=latents.dtype,
     )
-    if isinstance(noise_scale, torch.Tensor):
+    return noise_scale * noise + (1 - noise_scale) * latents
-        noise_scale = noise_scale.float()
-    return torch.lerp(latents.float(), noise.float(), noise_scale).to(latents.dtype)
 
 
 def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor:
@@ -137,16 +135,6 @@ def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
     return latents[:, :num_frames]
 
 
-def clear_audio_padding(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
-    """Keep Omni's SP-only audio padding outside the sampler state."""
-    if not 0 < num_frames <= latents.shape[1]:
-        raise ValueError(f"Audio frame count must be in [1, {latents.shape[1]}], got {num_frames}.")
-    if num_frames == latents.shape[1]:
-        return latents
-    padding = latents.new_zeros(latents.shape[0], latents.shape[1] - num_frames, latents.shape[2])
-    return torch.cat([latents[:, :num_frames], padding], dim=1)
-
-
 def get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int:
     if sp_size > 1:
         audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size
@@ -217,12 +205,7 @@ def prepare_video_latents(
             f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
             f" size of {batch_size}. Make sure the batch size matches the length of the generators."
         )
-    latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+    return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-    # Official LTX samples the 5D video latent first and then patchifies it,
-    # leaving token-major storage behind the [B, tokens, channels] view.  The
-    # logical values are unchanged, but preserving that layout is important:
-    # CUDA selects a different bf16 GEMM path for a contiguous token tensor.
-    return latents.transpose(1, 2).contiguous().transpose(1, 2)
 
 
 def prepare_audio_latents(
@@ -243,16 +226,6 @@ def prepare_audio_latents(
     sp_size = getattr(pipeline.od_config.parallel_config, "sequence_parallel_size", 1) or 1
     padded_latent_length = get_sp_padded_audio_latent_length(original_latent_length, int(sp_size))
 
-    def pad_logical_latents(logical_latents: torch.Tensor) -> torch.Tensor:
-        if padded_latent_length == original_latent_length:
-            return logical_latents
-        padding = logical_latents.new_zeros(
-            logical_latents.shape[0],
-            padded_latent_length - original_latent_length,
-            logical_latents.shape[2],
-        )
-        return torch.cat([logical_latents, padding], dim=1)
-
     if latents is not None:
         if latents.ndim == 4:
             latents = pack_audio_latents(latents)
@@ -264,24 +237,29 @@ def prepare_audio_latents(
                 pipeline.audio_vae.latents_mean,
                 pipeline.audio_vae.latents_std,
             )
+        latents = create_noised_state(latents, noise_scale, generator)
 
         if latents.shape[1] not in {original_latent_length, padded_latent_length}:
             raise ValueError(
                 "Provided `audio_latents` has incompatible audio frame count "
                 f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}."
             )
-        # Padding is an Omni implementation detail, not part of the official
+        if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length:
-        # audio state. Draw request RNG only for logical tokens so later phases
+            padding = torch.zeros(
-        # remain seed-invariant across SP degrees, then append deterministic 0s.
+                latents.shape[0],
-        latents = create_noised_state(latents[:, :original_latent_length], noise_scale, generator)
+                padded_latent_length - original_latent_length,
-        latents = pad_logical_latents(latents)
+                latents.shape[2],
+                dtype=latents.dtype,
+                device=latents.device,
+            )
+            latents = torch.cat([latents, padding], dim=1)
         return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
 
-    shape = (batch_size, original_latent_length, num_channels_latents * latent_mel_bins)
+    shape = (batch_size, padded_latent_length, num_channels_latents * latent_mel_bins)
     if isinstance(generator, list) and len(generator) != batch_size:
         raise ValueError(
             f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
             f" size of {batch_size}. Make sure the batch size matches the length of the generators."
         )
-    latents = pad_logical_latents(randn_tensor(shape, generator=generator, device=device, dtype=dtype))
+    latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
     return latents, original_latent_length, padded_latent_length

vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py

@@ -1,430 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Construction and execution of a fixed LTX refinement-phase adapter.
-
-Checkpoint key parsing remains in :mod:`ltx2_adapter_parser`; this module
-resolves the declared artifact and makes the existing Transformer linears
-adapter-capable while retaining their quantization and tensor-parallel paths.
-"""
-
-from __future__ import annotations
-
-from collections import defaultdict
-from collections.abc import Iterable
-from dataclasses import dataclass
-from typing import Any
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import vllm.envs as envs
-from vllm.distributed import (
-    split_tensor_along_last_dim,
-    tensor_model_parallel_all_reduce,
-)
-from vllm.logger import init_logger
-from vllm.model_executor.layers.batch_invariant import linear_batch_invariant
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-    UnquantizedLinearMethod,
-)
-from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
-from vllm.platforms import current_platform
-
-from .ltx2_adapter_parser import AdapterManifest, AdapterTarget, iter_adapter_tensors, parse_ltx_adapter
-from .ltx2_components import resolve_ltx_artifact
-from .ltx2_recipes import LTX_DISTILLED_ADAPTER_SLOT
-
-logger = init_logger(__name__)
-
-
-@dataclass(frozen=True)
-class _AdapterLayout:
-    """Rank-local storage and execution layout for one LoRA pair."""
-
-    kind: str
-    input_size: int
-    output_size: int
-    a_input_start: int
-    a_input_size: int
-    b_output_start: int
-    b_output_size: int
-    output_start: int
-    reduce_results: bool = False
-
-
-def _module_device(module: nn.Module) -> torch.device:
-    for tensor in (*module.parameters(), *module.buffers()):
-        if tensor.device.type != "meta":
-            return tensor.device
-    raise RuntimeError(f"Cannot materialize an LTX phase adapter for meta module {module._get_name()}.")
-
-
-def _build_layout(layer: nn.Module, target: AdapterTarget) -> _AdapterLayout:
-    if isinstance(layer, QKVParallelLinear):
-        if target.slice not in {"q", "k", "v"}:
-            raise ValueError(
-                f"LTX packed QKV target {target.module!r} requires a q/k/v adapter slice, got {target.slice!r}."
-            )
-        q_local = layer.num_heads * layer.head_size
-        k_local = layer.num_kv_heads * layer.head_size
-        v_local = layer.num_kv_heads * layer.v_head_size
-        if target.slice == "q":
-            output_size = layer.total_num_heads * layer.head_size
-            b_output_start = layer.tp_rank * q_local
-            b_output_size = q_local
-            output_start = 0
-        elif target.slice == "k":
-            output_size = layer.total_num_kv_heads * layer.head_size
-            shard_rank = layer.tp_rank // layer.num_kv_head_replicas
-            b_output_start = shard_rank * k_local
-            b_output_size = k_local
-            output_start = q_local
-        else:
-            output_size = layer.total_num_kv_heads * layer.v_head_size
-            shard_rank = layer.tp_rank // layer.num_kv_head_replicas
-            b_output_start = shard_rank * v_local
-            b_output_size = v_local
-            output_start = q_local + k_local
-        return _AdapterLayout(
-            kind="qkv",
-            input_size=layer.input_size,
-            output_size=output_size,
-            a_input_start=0,
-            a_input_size=layer.input_size,
-            b_output_start=b_output_start,
-            b_output_size=b_output_size,
-            output_start=output_start,
-        )
-
-    if isinstance(layer, ColumnParallelLinear):
-        if layer.gather_output:
-            raise ValueError(
-                f"LTX phase adapter does not support gathered ColumnParallelLinear target {target.module!r}."
-            )
-        return _AdapterLayout(
-            kind="column",
-            input_size=layer.input_size,
-            output_size=layer.output_size,
-            a_input_start=0,
-            a_input_size=layer.input_size,
-            b_output_start=layer.tp_rank * layer.output_size_per_partition,
-            b_output_size=layer.output_size_per_partition,
-            output_start=0,
-        )
-
-    if isinstance(layer, RowParallelLinear):
-        return _AdapterLayout(
-            kind="row",
-            input_size=layer.input_size,
-            output_size=layer.output_size,
-            a_input_start=layer.tp_rank * layer.input_size_per_partition,
-            a_input_size=layer.input_size_per_partition,
-            b_output_start=0,
-            b_output_size=layer.output_size,
-            output_start=0,
-            reduce_results=layer.reduce_results,
-        )
-
-    if isinstance(layer, nn.Linear):
-        return _AdapterLayout(
-            kind="replicated",
-            input_size=layer.in_features,
-            output_size=layer.out_features,
-            a_input_start=0,
-            a_input_size=layer.in_features,
-            b_output_start=0,
-            b_output_size=layer.out_features,
-            output_start=0,
-        )
-
-    raise TypeError(f"LTX phase adapter target {target.module!r} has unsupported layer type {type(layer).__name__}.")
-
-
-class _AdapterPiece(nn.Module):
-    """One rank-local A/B pair stored as non-persistent buffers."""
-
-    def __init__(self, target: AdapterTarget, layout: _AdapterLayout) -> None:
-        super().__init__()
-        self.target = target
-        self.layout = layout
-        # The fixed graph is installed before base checkpoint loading.  The
-        # empty buffers are replaced only after the base layer's quant/TP layout
-        # has settled and before offload/compile snapshots are taken.
-        self.register_buffer("lora_a", torch.empty(0), persistent=False)
-        self.register_buffer("lora_b", torch.empty(0), persistent=False)
-
-    def load(self, lora_a: torch.Tensor, lora_b: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> None:
-        if tuple(lora_a.shape) != self.target.a_shape or tuple(lora_b.shape) != self.target.b_shape:
-            raise ValueError(
-                f"LTX adapter source shape changed for {self.target.source_module!r}: "
-                f"A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}."
-            )
-        if lora_a.ndim != 2 or lora_b.ndim != 2:
-            raise ValueError(f"LTX adapter target {self.target.source_module!r} must use rank-2 A/B tensors.")
-        rank = self.target.a_shape[0]
-        if lora_a.shape[1] != self.layout.input_size or lora_b.shape[0] != self.layout.output_size:
-            raise ValueError(
-                f"LTX adapter shape does not match {self.target.module!r}: "
-                f"expected A=({rank}, {self.layout.input_size}), "
-                f"B=({self.layout.output_size}, {rank}); "
-                f"got A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}."
-            )
-
-        local_a = lora_a.narrow(1, self.layout.a_input_start, self.layout.a_input_size)
-        local_b = lora_b.narrow(0, self.layout.b_output_start, self.layout.b_output_size)
-        self.lora_a = local_a.to(device=device, dtype=dtype)
-        self.lora_b = local_b.to(device=device, dtype=dtype)
-
-    def delta(self, input_: torch.Tensor) -> torch.Tensor:
-        if self.lora_a.numel() == 0 or self.lora_b.numel() == 0:
-            raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.")
-        if self.layout.kind == "row" and input_.shape[-1] != self.layout.a_input_size:
-            input_ = input_.narrow(-1, self.layout.a_input_start, self.layout.a_input_size)
-        input_ = input_.to(dtype=self.lora_a.dtype)
-        delta = F.linear(F.linear(input_, self.lora_a), self.lora_b)
-        if self.layout.kind == "row" and self.layout.reduce_results:
-            delta = tensor_model_parallel_all_reduce(delta)
-        return delta
-
-    def weight_delta(self) -> torch.Tensor:
-        """Materialize the rank-local BF16 weight delta used by official fusion."""
-        if self.lora_a.numel() == 0 or self.lora_b.numel() == 0:
-            raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.")
-        return torch.matmul(self.lora_b, self.lora_a)
-
-
-def _apply_unquantized_gemm_with_weight(
-    layer: nn.Module,
-    input_: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor | None,
-) -> torch.Tensor:
-    if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
-        return linear_batch_invariant(input_, weight, bias)
-    return dispatch_unquantized_gemm()(layer, input_, weight, bias)
-
-
-def _forward_with_weight(
-    layer: nn.Module,
-    input_: torch.Tensor,
-    weight: torch.Tensor,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
-    """Run an unquantized base layer with an ephemeral fused weight."""
-    if isinstance(layer, nn.Linear):
-        return F.linear(input_, weight, layer.bias)
-
-    quant_method = getattr(layer, "quant_method", None)
-    if not isinstance(quant_method, UnquantizedLinearMethod):
-        raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.")
-
-    if isinstance(layer, RowParallelLinear):
-        if layer.input_is_parallel:
-            input_parallel = input_
-        else:
-            input_parallel = split_tensor_along_last_dim(input_, num_partitions=layer.tp_size)[
-                layer.tp_rank
-            ].contiguous()
-        bias = None if layer.tp_rank > 0 or layer.skip_bias_add else layer.bias
-        output_parallel = _apply_unquantized_gemm_with_weight(layer, input_parallel, weight, bias)
-        if layer.reduce_results and layer.tp_size > 1:
-            output = tensor_model_parallel_all_reduce(output_parallel)
-        else:
-            output = output_parallel
-    elif isinstance(layer, ColumnParallelLinear):
-        bias = layer.bias if not layer.skip_bias_add else None
-        output = _apply_unquantized_gemm_with_weight(layer, input_, weight, bias)
-    else:
-        raise TypeError(f"Unsupported LTX layer-fused linear type {type(layer).__name__}.")
-
-    if not layer.return_bias:
-        return output
-    return output, layer.bias if layer.skip_bias_add else None
-
-
-class _PhaseAdapterLinear(nn.Module):
-    """Delegate the base linear unchanged and add an enabled adapter slot."""
-
-    def __init__(
-        self,
-        base_layer: nn.Module,
-        targets: Iterable[AdapterTarget],
-        *,
-        adapter_dtype: torch.dtype,
-        layer_fused: bool = False,
-    ) -> None:
-        super().__init__()
-        self.base_layer = base_layer
-        self.adapter_dtype = adapter_dtype
-        if layer_fused and not (
-            isinstance(base_layer, nn.Linear)
-            or isinstance(getattr(base_layer, "quant_method", None), UnquantizedLinearMethod)
-        ):
-            raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.")
-        self.layer_fused = layer_fused
-        pieces = [_AdapterPiece(target, _build_layout(base_layer, target)) for target in targets]
-        self.pieces = nn.ModuleList(pieces)
-        self._pieces_by_target = {piece.target: piece for piece in pieces}
-        self._enabled = False
-
-    def load_target(self, target: AdapterTarget, lora_a: torch.Tensor, lora_b: torch.Tensor) -> None:
-        try:
-            piece = self._pieces_by_target[target]
-        except KeyError as exc:
-            raise ValueError(f"Adapter target {target.module!r} is not installed on this layer.") from exc
-        piece.load(lora_a, lora_b, device=_module_device(self.base_layer), dtype=self.adapter_dtype)
-
-    def set_enabled(self, enabled: bool) -> None:
-        self._enabled = enabled
-
-    def _add_adapter_delta(self, input_: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
-        if not self._enabled:
-            return output
-        if output.requires_grad:
-            output = output.clone()
-        for piece in self.pieces:
-            delta = piece.delta(input_).to(dtype=output.dtype)
-            start = piece.layout.output_start
-            output[..., start : start + piece.layout.b_output_size].add_(delta)
-        return output
-
-    @torch.no_grad()
-    def _fused_weight(self) -> torch.Tensor:
-        base_weight = self.base_layer.weight
-        fused_weight = base_weight.clone()
-        for piece in self.pieces:
-            start = piece.layout.output_start
-            base_slice = base_weight.narrow(0, start, piece.layout.b_output_size)
-            delta = piece.weight_delta()
-            if delta.shape != base_slice.shape:
-                raise ValueError(
-                    f"LTX layer-fused delta shape {tuple(delta.shape)} does not match base weight slice "
-                    f"{tuple(base_slice.shape)} for {piece.target.module!r}."
-                )
-            # Match official fusion exactly: materialize B@A in the adapter
-            # dtype, then add the base tensor in-place before the main GEMM.
-            delta.add_(base_slice)
-            fused_weight.narrow(0, start, piece.layout.b_output_size).copy_(delta)
-        return fused_weight
-
-    def forward(self, input_: torch.Tensor):
-        if not self._enabled:
-            return self.base_layer(input_)
-        if self.layer_fused:
-            return _forward_with_weight(self.base_layer, input_, self._fused_weight())
-        output = self.base_layer(input_)
-        if isinstance(output, tuple):
-            return (self._add_adapter_delta(input_, output[0]), *output[1:])
-        return self._add_adapter_delta(input_, output)
-
-
-class LTXPhaseAdapterRuntime:
-    """Install and operate one fixed phase adapter on an existing transformer."""
-
-    def __init__(
-        self,
-        transformer: nn.Module,
-        manifest: AdapterManifest,
-        *,
-        dtype: torch.dtype,
-        layer_fused: bool = False,
-    ) -> None:
-        self.transformer = transformer
-        self.manifest = manifest
-        self._wrappers: dict[str, _PhaseAdapterLinear] = {}
-        self._target_names: tuple[str, ...] = ()
-        self._materialized = False
-        grouped: dict[str, list[AdapterTarget]] = defaultdict(list)
-        for target in manifest.targets:
-            grouped[target.module].append(target)
-
-        original_layers = {name: self.transformer.get_submodule(name) for name in grouped}
-        for name, targets in grouped.items():
-            wrapper = _PhaseAdapterLinear(
-                original_layers[name],
-                targets,
-                adapter_dtype=dtype,
-                layer_fused=layer_fused,
-            )
-            parent_name, _, child_name = name.rpartition(".")
-            parent = self.transformer.get_submodule(parent_name) if parent_name else self.transformer
-            if child_name not in parent._modules:
-                raise ValueError(f"LTX phase adapter target {name!r} is no longer a direct module child.")
-            parent._modules[child_name] = wrapper
-            self._wrappers[name] = wrapper
-
-        self._target_names = tuple(sorted(self._wrappers, key=len, reverse=True))
-        # LTX's custom loader receives the original checkpoint names.  Keep the
-        # adaptation local to that loader rather than changing AutoWeightsLoader.
-        self.transformer._phase_adapter_parameter_name = self.base_parameter_name
-        logger.info(
-            "Installed %d LTX %s phase-adapter linears",
-            len(self._wrappers),
-            "layer-fused" if layer_fused else "dynamic",
-        )
-
-    def base_parameter_name(self, name: str) -> str:
-        for target_name in self._target_names:
-            prefix = f"{target_name}."
-            if name.startswith(prefix):
-                return f"{target_name}.base_layer.{name[len(prefix) :]}"
-        return name
-
-    def finalize(self) -> None:
-        """Stream rank-local A/B buffers after base weight processing completes."""
-        if self._materialized:
-            return
-        for target, lora_a, lora_b in iter_adapter_tensors(self.manifest):
-            self._wrappers[target.module].load_target(target, lora_a, lora_b)
-        self._materialized = True
-        logger.info("Materialized %d LTX phase-adapter tensor pairs", len(self.manifest.targets))
-
-    def activate(self, adapter_slot: str | None) -> None:
-        if adapter_slot is not None:
-            if adapter_slot != LTX_DISTILLED_ADAPTER_SLOT:
-                raise ValueError(f"Unknown LTX phase adapter slot {adapter_slot!r}.")
-            if not self._materialized:
-                raise RuntimeError("LTX phase adapter data must be materialized before activation.")
-        for wrapper in self._wrappers.values():
-            wrapper.set_enabled(adapter_slot is not None)
-
-
-def build_ltx_phase_adapter(pipeline: Any) -> LTXPhaseAdapterRuntime | None:
-    """Build the fixed adapter required by a multi-phase LTX recipe."""
-    adapter_slots = {phase.adapter_slot for phase in pipeline.pipeline_recipe.phases if phase.adapter_slot is not None}
-    if not adapter_slots:
-        return None
-    if adapter_slots != {LTX_DISTILLED_ADAPTER_SLOT}:
-        raise ValueError(f"Unsupported LTX phase adapter slots: {sorted(adapter_slots)}.")
-
-    profile = pipeline.component_profile
-    if profile.distilled_lora_filename is None or profile.artifact_repo_id is None:
-        raise ValueError(f"{profile.name} does not declare the required distilled adapter artifact.")
-    if getattr(pipeline.od_config, "lora_path", None) is not None:
-        raise ValueError(
-            f"{pipeline.__class__.__name__} reserves LoRA execution for its phase adapter; "
-            "request or static LoRA composition is not supported yet."
-        )
-
-    quantization_config = getattr(pipeline.od_config, "quantization_config", None)
-    layer_fused = quantization_config is None
-    adapter_path = resolve_ltx_artifact(
-        pipeline.od_config.model,
-        profile.artifact_repo_id,
-        profile.distilled_lora_filename,
-    )
-
-    if layer_fused and pipeline.od_config.dtype is not torch.bfloat16:
-        raise ValueError("LTX layer-fused phase LoRA requires a bfloat16 pipeline dtype.")
-
-    manifest = parse_ltx_adapter(pipeline.transformer, adapter_path)
-    runtime = LTXPhaseAdapterRuntime(
-        pipeline.transformer,
-        manifest,
-        dtype=pipeline.od_config.dtype,
-        layer_fused=layer_fused,
-    )
-    return runtime

vllm_omni/diffusion/models/ltx2/ltx2_recipes.py

@@ -24,7 +24,6 @@ LTX_DEFAULT_NEGATIVE_PROMPT = (
 
 LTX_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0)
 LTX_STAGE_2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0)
-LTX_DISTILLED_ADAPTER_SLOT = "ltx_distilled"
 
 
 @dataclass(frozen=True)
@@ -37,7 +36,6 @@ class LTXPhaseRecipe:
     sigmas: tuple[float, ...] | None = None
     noise_scale: float = 0.0
     input_transform: Literal["initial", "spatial_upsample"] = "initial"
-    adapter_slot: str | None = None
     allow_guidance_override: bool = True
     use_official_sigma_schedule: bool = True
 
@@ -159,59 +157,11 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe(
     fixed_num_inference_steps=True,
 )
 
-# LTX-2.3 full-distilled uses the same official fixed 8 + 3 sigma schedules
-# and request contract. Its distinct component profile selects the 22B merged
-# distilled Transformer, BWE vocoder, and matching spatial upsampler; no
-# distilled LoRA is loaded for either phase.
-LTX23_DISTILLED_TWO_STAGE_RECIPE = LTX2_DISTILLED_TWO_STAGE_RECIPE
-
-
-def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipelineRecipe:
-    return LTXPipelineRecipe(
-        height=one_stage_recipe.height * 2,
-        width=one_stage_recipe.width * 2,
-        num_frames=one_stage_recipe.num_frames,
-        frame_rate=one_stage_recipe.frame_rate,
-        num_inference_steps=one_stage_recipe.num_inference_steps,
-        negative_prompt=one_stage_recipe.negative_prompt,
-        phases=(
-            LTXPhaseRecipe(
-                name="generate_lowres",
-                guidance=one_stage_recipe.request_guidance,
-                spatial_downscale=2,
-                noise_scale=1.0,
-            ),
-            LTXPhaseRecipe(
-                name="refine",
-                guidance=LTXGuidanceSpec.positive_only(),
-                sigmas=LTX_STAGE_2_DISTILLED_SIGMAS,
-                noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0],
-                input_transform="spatial_upsample",
-                adapter_slot=LTX_DISTILLED_ADAPTER_SLOT,
-                allow_guidance_override=False,
-                use_official_sigma_schedule=False,
-            ),
-        ),
-        video_output_phase=1,
-        # The official second stage refines video only and deliberately
-        # discards its audio result. Decode the full-context Stage-1 audio.
-        audio_output_phase=0,
-        allow_request_sigmas=False,
-        allow_request_latents=False,
-    )
-
-
-LTX2_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX2_ONE_STAGE_RECIPE)
-LTX23_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX23_ONE_STAGE_RECIPE)
-
 
 _PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = {
     ("one_stage", "2"): LTX2_ONE_STAGE_RECIPE,
     ("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE,
-    ("two_stage", "2"): LTX2_TWO_STAGE_RECIPE,
-    ("two_stage", "2.3"): LTX23_TWO_STAGE_RECIPE,
     ("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE,
-    ("distilled_two_stage", "2.3"): LTX23_DISTILLED_TWO_STAGE_RECIPE,
     ("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE,
     ("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE,
 }

vllm_omni/diffusion/models/ltx2/ltx2_request.py

@@ -5,9 +5,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
 from dataclasses import dataclass, replace
-from typing import Any, Literal
+from typing import Any
 
 import torch
 
@@ -50,34 +49,6 @@ class LTXRequestInputs:
         return self.guidance.video.cfg_scale
 
 
-LTXCheckpointKind = Literal["distilled", "regular"]
-
-
-def validate_ltx_checkpoint(
-    scheduler_config: Mapping[str, Any],
-    *,
-    expected_kind: LTXCheckpointKind | None,
-    pipeline_name: str,
-) -> None:
-    """Ensure scheduler metadata matches the checkpoint required by a profile."""
-    if expected_kind is None:
-        return
-
-    if expected_kind == "distilled":
-        if scheduler_config.get("use_dynamic_shifting", True) or scheduler_config.get("shift_terminal") is not None:
-            raise ValueError(
-                f"{pipeline_name} requires a merged distilled LTX2 checkpoint; "
-                "the loaded scheduler metadata describes a non-distilled checkpoint."
-            )
-        return
-
-    if not scheduler_config.get("use_dynamic_shifting", False) or scheduler_config.get("shift_terminal") is None:
-        raise ValueError(
-            f"{pipeline_name} requires a regular non-distilled LTX checkpoint; "
-            "the loaded scheduler metadata describes a distilled checkpoint."
-        )
-
-
 def validate_pipeline_request(
     request_inputs: LTXRequestInputs,
     *,

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -47,7 +47,6 @@ from .ltx2_guidance import (
     LTXGuidanceExecutor,
     LTXGuidancePlan,
 )
-from .ltx2_phase_adapter import LTXPhaseAdapterRuntime, build_ltx_phase_adapter
 from .ltx2_recipes import (
     LTXPhaseRecipe,
     LTXPipelineRecipe,
@@ -138,21 +137,9 @@ class LTXRuntime(
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         del prefix
-        parallel_config = getattr(od_config, "parallel_config", None)
-        if getattr(parallel_config, "ulysses_mode", "strict") == "advanced_uaa":
-            raise ValueError(
-                f"{self.__class__.__name__} does not support ulysses_mode='advanced_uaa'. "
-                "Use the default ulysses_mode='strict' for LTX sequence parallelism."
-            )
         self.model_version = detect_ltx_model_version(od_config.model)
         self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version)
         self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version)
-        if getattr(od_config, "cache_backend", "none") == "cache_dit" and len(self.pipeline_recipe.phases) > 1:
-            raise ValueError(
-                f"{self.__class__.__name__} does not support cache_backend='cache_dit'. "
-                "Cache-DiT currently supports only one-stage LTX recipes; multi-stage recipes require "
-                "phase-aware cache enable and refresh."
-            )
         self._dit_modules = list(self.component_profile.dit_modules)
         self._encoder_modules = list(self.component_profile.encoder_modules)
         self._vae_modules = list(self.component_profile.vae_modules)
@@ -163,7 +150,6 @@ class LTXRuntime(
         super().__init__()
         self._guidance_plan = LTXGuidancePlan.build(self.pipeline_recipe.request_guidance)
         initialize_pipeline_components(self, od_config)
-        self._phase_adapter: LTXPhaseAdapterRuntime | None = build_ltx_phase_adapter(self)
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )
@@ -219,8 +205,6 @@ class LTXRuntime(
             max_sequence_length=max_sequence_length,
         )
         image = self._resolve_request_image(req, image, request_inputs)
-        if image is not None and not self.support_image_input:
-            raise ValueError(f"{self.__class__.__name__} does not support `image` input.")
         request_sigmas = self._resolve_request_sigmas(req, sigmas)
         validate_pipeline_request(
             request_inputs,
@@ -230,13 +214,6 @@ class LTXRuntime(
             pipeline_name=self.__class__.__name__,
             request_sigmas=request_sigmas,
         )
-        phase_adapter = getattr(self, "_phase_adapter", None)
-        if phase_adapter is not None and any(
-            getattr(sampling, "lora_request", None) is not None for sampling in req.sampling_params_list
-        ):
-            raise ValueError(
-                f"{self.__class__.__name__} cannot compose a request LoRA with its internal phase adapter."
-            )
         return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image)
 
     def _run_recipe(
@@ -281,20 +258,8 @@ class LTXRuntime(
         return self.decode_phase(output_phase)
 
     def _enter_phase(self, phase: LTXPhaseRecipe) -> None:
+        """Hook for a future phase-weight strategy."""
         self._active_phase_name = phase.name
-        phase_adapter = getattr(self, "_phase_adapter", None)
-        if phase_adapter is None:
-            if phase.adapter_slot is not None:
-                raise RuntimeError(f"LTX phase {phase.name!r} requires adapter slot {phase.adapter_slot!r}.")
-            return
-        phase_adapter.activate(phase.adapter_slot)
-
-    def eval(self):
-        result = super().eval()
-        phase_adapter = getattr(self, "_phase_adapter", None)
-        if phase_adapter is not None:
-            phase_adapter.finalize()
-        return result
 
     def prepare_latents(
         self,
@@ -727,7 +692,6 @@ class LTXRuntime(
             noise_pred_audio,
             timestep,
         )
-        audio = latent_ops.clear_audio_padding(audio, forward_ctx.original_audio_num_frames)
         return latent_ops.LTXAVState(video=video, audio=audio)
 
     def _unpack_and_denormalize_stage(

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -355,19 +355,6 @@ def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor:
     return attention_mask
 
 
-class _LTX2ParallelAttention(Attention):
-    """Preserve LTX SP collectives while applying model-specific padding masks."""
-
-    def _run_local_attention(self, query, key, value, attn_metadata):
-        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
-        if attention_mask is not None:
-            # LTX masks describe K/V padding. Keep this semantic local to LTX:
-            # Flash varlen applies a 2D mask to both Q and K, while SP uses a
-            # broadcast 4D key mask that only SDPA consumes correctly.
-            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
-        return super()._run_local_attention(query, key, value, attn_metadata)
-
-
 class LTX2AudioVideoAttnProcessor:
     r"""
     Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model.
@@ -409,10 +396,11 @@ class LTX2AudioVideoAttnProcessor:
             return None
 
         if self._is_sp_enabled():
-            # Keep the full key mask replicated while strict Ulysses reshards
+            # In SP, Ulysses expects a 2D padding mask that matches query length.
-            # Q/K/V. The broadcast Q dimension also works when cross-attention
+            # For cross-attention, encoder sequence length != query length, so drop the mask.
-            # has different Q and K lengths, without changing generic Ulysses.
+            if encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]:
-            return to_ltx_padding_mask(attention_mask)[:, None, None, :]
+                return None
+            return to_ltx_padding_mask(attention_mask)
 
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
@@ -708,7 +696,7 @@ class LTX2Attention(torch.nn.Module):
                 torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(),
             ]
         )
-        self.attn = _LTX2ParallelAttention(
+        self.attn = Attention(
             num_heads=self.query_num_heads,
             head_size=dim_head,
             num_kv_heads=self.kv_num_heads,
@@ -721,15 +709,7 @@ class LTX2Attention(torch.nn.Module):
         # LTX-2.3: per-head gated attention
         # leave unquantized for this linear
         if apply_gated_attention:
-            self.to_gate_logits = ColumnParallelLinear(
+            self.to_gate_logits = nn.Linear(query_dim, self.query_num_heads, bias=True)
-                query_dim,
-                heads,
-                bias=True,
-                gather_output=False,
-                return_bias=False,
-                quant_config=None,
-                prefix=f"{prefix}.to_gate_logits" if prefix else "to_gate_logits",
-            )
         else:
             self.to_gate_logits = None
 
@@ -1518,14 +1498,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     _hsdp_shard_conditions = [is_transformer_block_module]
     _sp_plan: dict[str, Any] | None = None
-    stacked_params_mapping = (
-        (".attn1.to_qkv", ".attn1.to_q", "q"),
-        (".attn1.to_qkv", ".attn1.to_k", "k"),
-        (".attn1.to_qkv", ".attn1.to_v", "v"),
-        (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"),
-        (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"),
-        (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"),
-    )
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
     }
@@ -1856,7 +1828,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
         audio_sigma: torch.Tensor | None = None,
         encoder_attention_mask: torch.Tensor | None = None,
         audio_encoder_attention_mask: torch.Tensor | None = None,
-        audio_attention_mask: torch.Tensor | None = None,
         num_frames: int | None = None,
         height: int | None = None,
         width: int | None = None,
@@ -1890,8 +1861,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`.
             audio_encoder_attention_mask (`torch.Tensor`, *optional*):
                 Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling.
-            audio_attention_mask (`torch.Tensor`, *optional*):
-                Optional audio-token key padding mask shared by audio self-attention and audio-to-video attention.
             num_frames (`int`, *optional*):
                 The number of latent video frames. Used if calculating the video coordinates for RoPE.
             height (`int`, *optional*):
@@ -2075,8 +2044,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 "ca_audio_rotary_emb": audio_cross_attn_rotary_emb,
                 "encoder_attention_mask": encoder_attention_mask,
                 "audio_encoder_attention_mask": audio_encoder_attention_mask,
-                "audio_self_attention_mask": audio_attention_mask,
-                "a2v_cross_attention_mask": audio_attention_mask,
                 "video_self_attention_perturbation_mask": perturbation_mask_for("video_self_attention", block_idx),
                 "audio_self_attention_perturbation_mask": perturbation_mask_for("audio_self_attention", block_idx),
                 "a2v_cross_attention_perturbation_mask": perturbation_mask_for("a2v_cross_attention", block_idx),
@@ -2130,8 +2097,16 @@ class LTX2VideoTransformer3DModel(nn.Module):
         Returns:
             Set of parameter names that were successfully loaded.
         """
+        stacked_params_mapping = [
+            (".attn1.to_qkv", ".attn1.to_q", "q"),
+            (".attn1.to_qkv", ".attn1.to_k", "k"),
+            (".attn1.to_qkv", ".attn1.to_v", "v"),
+            (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"),
+            (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"),
+            (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"),
+        ]
+
         params_dict = dict(self.named_parameters())
-        phase_adapter_parameter_name = getattr(self, "_phase_adapter_parameter_name", None)
         tp_size = get_tensor_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
         loaded_params: set[str] = set()
@@ -2152,12 +2127,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
             return weight
 
         for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in self.stacked_params_mapping:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
                 packed_name = name.replace(weight_name, param_name)
-                if callable(phase_adapter_parameter_name):
-                    packed_name = phase_adapter_parameter_name(packed_name)
                 if packed_name not in params_dict:
                     continue
                 name = packed_name
@@ -2166,8 +2139,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 weight_loader(param, loaded_weight, shard_id)
                 break
             else:
-                if callable(phase_adapter_parameter_name):
-                    name = phase_adapter_parameter_name(name)
                 if name not in params_dict:
                     logger.warning(
                         "Skipping transformer weight %s -- not found in model "

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py

@@ -1,48 +1,27 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Two-stage entry points for the LTX model family."""
+"""Distilled two-stage entry points for the LTX model family."""
 
-from typing import ClassVar
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-from .ltx2_components import (
+from .ltx2_components import LTX2_DISTILLED_COMPONENT_PROFILE
-    LTX2_DISTILLED_COMPONENT_PROFILE,
-    LTX2_TWO_STAGE_COMPONENT_PROFILE,
-)
 from .ltx2_components import (
     get_ltx2_post_process_func as get_ltx2_post_process_func,  # noqa: F401
 )
-from .ltx2_conditioning import LTXI2VConditioningMixin
+from .ltx2_recipes import LTX2_DISTILLED_TWO_STAGE_RECIPE
-from .ltx2_recipes import (
+from .pipeline_ltx2 import LTX2Pipeline
-    LTX2_DISTILLED_TWO_STAGE_RECIPE,
-    LTX2_TWO_STAGE_RECIPE,
-)
-from .ltx2_runtime import LTXRuntime
 
 
-class _LTX2TwoStageBase(LTXI2VConditioningMixin, LTXRuntime):
+class LTX2DistilledPipeline(LTX2Pipeline):
-    """Static discovery contract shared by the two-stage public entries."""
+    """Unified LTX-2 distilled two-stage T2V/I2V entry."""
-
-    _dit_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.dit_modules)
-    _encoder_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.encoder_modules)
-    _vae_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.vae_modules)
-    _resident_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.resident_modules)
-    supports_request_batch = False
-    support_image_input = True
-    unified_text_image_entry = True
-
-
-class LTX2TwoStagePipeline(_LTX2TwoStageBase):
-    """Regular checkpoint with low-resolution generation and LoRA refinement."""
-
-    pipeline_kind = "two_stage"
-    component_profile = LTX2_TWO_STAGE_COMPONENT_PROFILE
-    pipeline_recipe = LTX2_TWO_STAGE_RECIPE
-
-
-class LTX2DistilledPipeline(_LTX2TwoStageBase):
-    """Unified LTX-2/LTX-2.3 full-distilled two-stage T2V/I2V entry."""
 
     pipeline_kind = "distilled_two_stage"
     component_profile = LTX2_DISTILLED_COMPONENT_PROFILE
     pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE
+    supports_request_batch = False
+    support_image_input = True
+    unified_text_image_entry = True
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -33,7 +33,6 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
-from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
@@ -238,15 +237,6 @@ class MiniMaxH3Rope(nn.Module):
         return torch.cat((half, half), dim=-1)  # [S, 96]
 
 
-def _build_rope_table(freqs: torch.Tensor) -> torch.Tensor:
-    """Materialize H3's packed ``[cos(freqs[:48]), sin(freqs[:48])]`` table."""
-    half = freqs.shape[-1] // 2
-    return torch.cat(
-        (torch.cos(freqs[..., :half]), torch.sin(freqs[..., :half])),
-        dim=-1,
-    ).to(_BF16_DTYPE)
-
-
 class MiniMaxH3TimeEmbedder(nn.Module):
     def __init__(
         self,
@@ -463,7 +453,7 @@ class MiniMaxH3Attention(nn.Module):
         self,
         x: torch.Tensor,
         *,
-        rope_table: torch.Tensor | None,
+        rope_freqs: torch.Tensor | None,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int | None = None,
@@ -489,18 +479,11 @@ class MiniMaxH3Attention(nn.Module):
         q = q.view(total, self.num_heads, self.head_dim)
         k = k.view(total, self.num_kv_heads, self.head_dim)
         v = v.view(total, self.num_kv_heads, self.head_dim)
-        if rope_table is None:
+        q = self.q_norm(q)
-            q = self.q_norm(q)
+        k = self.k_norm(k)
-            k = self.k_norm(k)
+        if rope_freqs is not None:
-        else:
+            q = self._apply_rope(q, rope_freqs)
-            q, k = fused_qk_norm_rope(
+            k = self._apply_rope(k, rope_freqs)
-                q,
-                k,
-                self.q_norm.weight,
-                self.k_norm.weight,
-                rope_table,
-                self.q_norm.variance_epsilon,
-            )
 
         # The packed layout uses a second document for alignment padding.
         # Local/Ulysses backends unpad it, while Ring keeps aligned rows for
@@ -645,7 +628,7 @@ class MiniMaxH3TokenRefinerBlock(nn.Module):
     ) -> torch.Tensor:
         x = x + self.attn(
             self.norm1(x),
-            rope_table=None,
+            rope_freqs=None,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
         )
@@ -724,7 +707,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         *,
         t_emb: torch.Tensor,
         combined_indices: torch.Tensor,
-        rope_table: torch.Tensor,
+        rope_freqs: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int,
@@ -752,7 +735,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE)
         h = self.attn(
             h,
-            rope_table=rope_table,
+            rope_freqs=rope_freqs,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
             packed_total=packed_total,
@@ -834,10 +817,10 @@ class MiniMaxH3SPPrepare(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        rope_table: torch.Tensor,
+        rope_freqs: torch.Tensor,
         combined_indices: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        return hidden_states, rope_table, combined_indices
+        return hidden_states, rope_freqs, combined_indices
 
 
 class MiniMaxH3SPGather(nn.Module):
@@ -1170,8 +1153,8 @@ class MiniMaxH3DiTModel(nn.Module):
         if inverse_indices.shape[0] != seq_len:
             raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
         device = x.device
-        # Compute the packed RoPE table once for all transformer blocks.
+        # Compute RoPE frequencies over the full packed sequence.
-        rope_table = _build_rope_table(self.rope(img_position_ids).to(device))
+        rope_freqs = self.rope(img_position_ids).to(device)
 
         decoder_input, t_emb = self._embed(
             x=x,
@@ -1192,7 +1175,7 @@ class MiniMaxH3DiTModel(nn.Module):
 
         hidden = decoder_input
         cu_seqlens = cu_seqlens.to(device)
-        block_rope = rope_table
+        block_rope = rope_freqs
         block_combined = combined_indices
 
         hidden, block_rope, block_combined = self.sp_prepare(
@@ -1205,7 +1188,7 @@ class MiniMaxH3DiTModel(nn.Module):
                 hidden,
                 t_emb=t_emb,
                 combined_indices=block_combined,
-                rope_table=block_rope,
+                rope_freqs=block_rope,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,

vllm_omni/diffusion/registry.py

@@ -71,11 +71,6 @@ _DIFFUSION_MODELS = {
         "pipeline_ltx2",
         "LTX2Pipeline",
     ),
-    "LTX2TwoStagePipeline": (
-        "ltx2",
-        "pipeline_ltx2_two_stage",
-        "LTX2TwoStagePipeline",
-    ),
     "LTX2DistilledPipeline": (
         "ltx2",
         "pipeline_ltx2_two_stage",
@@ -521,7 +516,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanPipeline": "get_wan22_post_process_func",
     "WanVACEPipeline": "get_wan22_vace_post_process_func",
     "LTX2Pipeline": "get_ltx2_post_process_func",
-    "LTX2TwoStagePipeline": "get_ltx2_post_process_func",
     "LTX2DistilledPipeline": "get_ltx2_post_process_func",
     "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
     "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",

vllm_omni/entrypoints/cli/serve.py

@@ -790,9 +790,10 @@ class OmniServeCommand(CLISubcommand):
             "--vae-parallel-mode",
             type=str,
             default="tile",
-            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
+            choices=["tile", "auto", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
-            "'tile' (default) uses patch/tile parallel decode; "
+            "'tile' (default) uses patch/tile decode; 'auto' selects the longer spatial-shard axis for "
+            "validated Wan tiled decodes on a full DiT group, then falls back to patch/tile decode; "
             "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
             "decoder feature maps along height/width and exchanges halo regions. The "
             "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "

vllm_omni/model_executor/models/dots_tts/__init__.py

@@ -1,8 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# NOTE: Do not import model classes in this file. Importing any
-# submodule in this package triggers __init__.py execution, and
-# both the model registry and pipeline registry import submodules
-# directly — heavy imports here would be loaded as a side effect
-# even though nothing depends on these re-exports.

vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py

@@ -1,457 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copyright 2026 The vLLM-Omni team.
-# Copyright (c) rednote-hilab. All rights reserved.
-# Adapted from:
-# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone
-#
-# Merged from two upstream files (preserved as-is, inference behaviour
-# unchanged; only internal imports were folded inline):
-#   - modules/backbone/layers.py  -> Dropout, Mlp, rotate_half,
-#                                    apply_rotary_pos_emb, RotaryEmbedding,
-#                                    MultiHeadAttention
-#                                    (Conv1d / ConvTranspose1d are NOT
-#                                    re-vendored here — they already live in
-#                                    dots_tts_vocoder.py)
-#   - modules/backbone/dit.py     -> modulate, TimestepEmbedder, FinalLayer,
-#                                    DiTBlock, DiT
-#
-# Per-section upstream attributions are preserved at each section header.
-# Nothing in this file diverges from upstream behaviour.
-
-from __future__ import annotations
-
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-
-# ============================================================================
-# Building blocks (adapted from modules/backbone/layers.py)
-# ----------------------------------------------------------------------------
-# Conv1d / ConvTranspose1d from layers.py are already vendored in
-# dots_tts_vocoder.py (they are used by the BigVGAN-style AudioVAE).  The
-# pieces below are the ones consumed by the DiT flow-matching head.
-# ============================================================================
-
-
-class Dropout(nn.Module):
-    def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **_kwargs):
-        super().__init__()
-        if p < 0.0 or p > 1.0:
-            raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
-        self.p = p
-        self.inplace = inplace
-        self.force_drop = force_drop
-
-    def forward(self, x, **_kwargs):
-        return F.dropout(
-            x,
-            p=self.p,
-            training=True if self.force_drop else self.training,
-            inplace=self.inplace,
-        )
-
-
-class Mlp(nn.Module):
-    def __init__(
-        self,
-        hidden_size,
-        ffn_hidden_size=4096,
-        act_layer=nn.GELU,
-        dropout=0.0,
-        **_kwargs,
-    ):
-        super().__init__()
-        self.fc1 = nn.Linear(hidden_size, ffn_hidden_size)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(ffn_hidden_size, hidden_size)
-        self.drop = Dropout(dropout)
-
-    def forward(self, x, _mask=None):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        return self.drop(x)
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-
-
-@torch.autocast(enabled=False, device_type="cuda")
-def apply_rotary_pos_emb(pos, t):
-    if pos.dim() == 3:
-        pos = pos.unsqueeze(1)
-    return t * pos.cos() + rotate_half(t) * pos.sin()
-
-
-class RotaryEmbedding(nn.Module):
-    def __init__(self, dim, theta=50000):
-        super().__init__()
-        self.register_buffer(
-            "inv_freq",
-            1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)),
-            persistent=False,
-        )
-        self._theta = float(theta)
-
-    def _apply(self, fn):
-        inv_freq = self.inv_freq
-        super()._apply(fn)
-        self.inv_freq = inv_freq.to(device=self.inv_freq.device, dtype=torch.float32)
-        return self
-
-    @torch.autocast(enabled=False, device_type="cuda")
-    def forward(self, t):
-        inv_freq = self.inv_freq
-        if inv_freq.device != t.device:
-            raise RuntimeError(f"RotaryEmbedding buffer device mismatch: inv_freq={inv_freq.device} input={t.device}.")
-        t = t.to(dtype=inv_freq.dtype)
-        if t.dim() == 1:
-            freqs = torch.einsum("i , j -> i j", t, inv_freq)
-        else:
-            freqs = torch.einsum("bi, j -> bij", t, inv_freq)
-        return torch.cat((freqs, freqs), dim=-1)
-
-
-class MultiHeadAttention(nn.Module):
-    """Multi-head attention"""
-
-    def __init__(
-        self,
-        hidden_size: int,
-        num_heads: int = 8,
-        qkv_bias: bool = False,
-        qk_norm: bool = False,
-        attn_drop: float = 0.0,
-        dropout: float = 0.0,
-        norm_layer: str = "LayerNorm",
-        rotary_bias: bool = False,
-        rotary_theta: float | None = 50000,
-        **_kwargs,
-    ):
-        super().__init__()
-        assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads"
-        self.num_heads = num_heads
-        self.head_dim = hidden_size // num_heads
-        self.scale = self.head_dim**-0.5
-        self.rotary_bias = rotary_bias
-
-        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
-        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
-        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
-
-        norm_layer = getattr(nn, norm_layer)
-        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
-        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
-
-        self.attn_drop = Dropout(attn_drop)
-        self.o_proj = nn.Linear(hidden_size, hidden_size)
-        self.o_dropout = Dropout(dropout)
-
-        if self.rotary_bias:
-            self.rotary = RotaryEmbedding(self.head_dim, theta=rotary_theta)
-
-    def forward(self, q, k=None, v=None, mask=None, pos_ids=None, **_kwargs):
-        k = k or q
-        v = v or q
-        B, L, _ = q.shape
-        _, S, _ = v.shape
-        if mask is not None:
-            if mask.ndim == 2:  # [B, L]
-                assert L == S
-                mask = rearrange(mask, "b j -> b 1 1 j")
-                mask = mask.expand(-1, self.num_heads, L, -1)
-            elif mask.ndim == 3:  # [B, L, S]
-                assert mask.size(1) == L and mask.size(2) == S
-                mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
-
-        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
-        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
-        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
-        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
-        q, k = self.q_norm(q), self.k_norm(k)
-
-        # Apply rotary
-        if self.rotary_bias:
-            if L == S:
-                if pos_ids is None:
-                    rotary_emb = self.rotary(torch.arange(L, device=q.device))
-                else:
-                    rotary_emb = self.rotary(pos_ids)
-                q, k = (apply_rotary_pos_emb(rotary_emb, tensor) for tensor in (q, k))
-            else:
-                q_rotary_emb = self.rotary(torch.arange(L, device=q.device))
-                k_rotary_emb = self.rotary(torch.arange(S, device=k.device))
-                q = apply_rotary_pos_emb(q_rotary_emb, q)
-                k = apply_rotary_pos_emb(k_rotary_emb, k)
-
-        attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device)
-
-        if mask is not None:
-            attn_bias.masked_fill_(mask.logical_not(), float("-inf"))
-
-        out = F.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=attn_bias,
-            dropout_p=self.attn_drop.p if self.training else 0.0,
-        )
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.o_dropout(self.o_proj(out))
-
-    def decode_step(self, x, *, cache, positions: torch.Tensor):
-        if x.size(1) <= 0:
-            raise ValueError("MultiHeadAttention.decode_step expects a non-empty input.")
-        if positions.ndim != 1 or positions.size(0) != x.size(1):
-            raise ValueError("MultiHeadAttention.decode_step positions must match the decode block length.")
-
-        q = self.q_proj(x)
-        k = self.k_proj(x)
-        v = self.v_proj(x)
-
-        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
-        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
-        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
-        q, k = self.q_norm(q), self.k_norm(k)
-        block_len = q.size(2)
-
-        if self.rotary_bias:
-            rotary_emb = self.rotary(positions)
-            q = apply_rotary_pos_emb(rotary_emb, q)
-            k = apply_rotary_pos_emb(rotary_emb, k)
-
-        cached_k, cached_v = cache
-        cached_k.index_copy_(2, positions, k)
-        cached_v.index_copy_(2, positions, v)
-
-        cache_capacity = cached_k.size(2)
-        key_positions = torch.arange(
-            cache_capacity,
-            device=x.device,
-            dtype=torch.long,
-        ).unsqueeze(0)
-        query_positions = positions.unsqueeze(1)
-        causal_mask = key_positions <= query_positions
-        valid_mask = key_positions <= positions[-1]
-        attn_bias = torch.zeros(
-            q.size(0),
-            self.num_heads,
-            block_len,
-            cache_capacity,
-            dtype=q.dtype,
-            device=q.device,
-        )
-        attn_bias.masked_fill_(
-            (causal_mask & valid_mask).unsqueeze(0).unsqueeze(0).logical_not(),
-            float("-inf"),
-        )
-
-        out = F.scaled_dot_product_attention(
-            q,
-            cached_k,
-            cached_v,
-            attn_mask=attn_bias,
-            dropout_p=self.attn_drop.p if self.training else 0.0,
-        )
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.o_dropout(self.o_proj(out)), cache
-
-
-# ============================================================================
-# DiT flow-matching head (adapted from modules/backbone/dit.py)
-# ----------------------------------------------------------------------------
-# DiT.__init__ takes a ``transformer_config`` object exposing ``.to_dict()``,
-# ``.hidden_size`` and ``.num_layers``.  Upstream constructs it from a pydantic
-# config class higher up the stack; vLLM-Omni callers can pass anything that
-# satisfies the same minimal interface (e.g. a stdlib @dataclass with a
-# ``to_dict`` method or SimpleNamespace).  See dots_tts_talker.py for the
-# concrete soar-checkpoint config factory.
-# ============================================================================
-
-
-def modulate(x, shift, scale, **_kwargs):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-class TimestepEmbedder(nn.Module):
-    def __init__(self, hidden_size, frequency_embedding_size=256):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=True),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        half = dim // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-            device=t.device
-        )
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        return self.mlp(t_freq)
-
-
-class FinalLayer(nn.Module):
-    def __init__(self, hidden_size, output_size):
-        super().__init__()
-        self.adaLN_modulation = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(hidden_size, 2 * hidden_size, bias=True),
-        )
-        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-5)
-        self.linear = nn.Linear(hidden_size, output_size, bias=True)
-
-    def forward(self, x, c, **_kwargs):
-        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
-        x = modulate(self.norm(x), shift, scale)
-        return self.linear(x)
-
-
-class DiTBlock(nn.Module):
-    def __init__(
-        self,
-        attention: nn.Module,
-        ffn: nn.Module,
-        hidden_size: int = 1024,
-        modulation: bool = False,
-        eps: float = 1e-5,
-        **_kwargs,
-    ):
-        super().__init__()
-        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
-        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
-        self.attn = attention
-        self.ffn = ffn
-        self.modulation = modulation
-        if modulation:
-            self.adaLN_modulation = nn.Sequential(
-                nn.SiLU(),
-                nn.Linear(hidden_size, 6 * hidden_size, bias=True),
-            )
-
-    def forward(self, x, condition=None, mask=None, **kwargs):
-        if condition is None:
-            assert not self.modulation, "Without global condition, must set modulation to False"
-        else:
-            assert self.modulation, "With global condition, must set modulation to True"
-            shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(condition).chunk(
-                6, dim=1
-            )
-
-        if condition is not None:
-            pack_indices = kwargs.get("pack_indices")
-            if pack_indices is not None:
-                gate_attn = gate_attn[pack_indices]
-                gate_ffn = gate_ffn[pack_indices]
-            else:
-                gate_attn = gate_attn.unsqueeze(1)
-                gate_ffn = gate_ffn.unsqueeze(1)
-
-        if condition is not None:
-            x = x + gate_attn * self.attn(
-                modulate(self.norm1(x), shift_attn, scale_attn, **kwargs),
-                mask=mask,
-                **kwargs,
-            )
-        else:
-            x = x + self.attn(self.norm1(x), mask=mask, **kwargs)
-
-        if condition is not None:
-            x = x + gate_ffn * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn, **kwargs))
-        else:
-            x = x + self.ffn(self.norm2(x), mask=mask)
-        return x
-
-
-class DiT(nn.Module):
-    def __init__(
-        self,
-        in_dim,
-        out_dim,
-        transformer_config,
-        *,
-        mode: str = "flow_matching",
-    ):
-        super().__init__()
-        if mode not in {"flow_matching", "meanflow"}:
-            raise ValueError(f"DiT mode must be 'flow_matching' or 'meanflow', got {mode!r}.")
-
-        transformer_kwargs = transformer_config.to_dict()
-        model_dim = transformer_config.hidden_size
-        self.mode = mode
-        self.num_layers = transformer_config.num_layers
-
-        self.input_layer = nn.Linear(in_dim, model_dim)
-        self.time_embedder = TimestepEmbedder(model_dim)
-        if mode == "meanflow":
-            self.duration_embedder = TimestepEmbedder(model_dim)
-
-        self.blocks = nn.ModuleList()
-        for i in range(self.num_layers):
-            attn_block = MultiHeadAttention(**transformer_kwargs, name=f"layer_{i}")
-            ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **transformer_kwargs)
-            self.blocks.append(DiTBlock(attention=attn_block, ffn=ffn_block, **transformer_kwargs))
-
-        self.output_layer = FinalLayer(model_dim, out_dim)
-        self.initialize_weights()
-
-    def initialize_weights(self):
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-
-        self.apply(_basic_init)
-
-        nn.init.normal_(self.time_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.time_embedder.mlp[2].weight, std=0.02)
-
-        for block in self.blocks:
-            if hasattr(block, "adaLN_modulation"):
-                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-
-        nn.init.constant_(self.output_layer.adaLN_modulation[-1].weight, 0)
-        nn.init.constant_(self.output_layer.adaLN_modulation[-1].bias, 0)
-        nn.init.constant_(self.output_layer.linear.weight, 0)
-        nn.init.constant_(self.output_layer.linear.bias, 0)
-
-    def forward(
-        self,
-        x,
-        timesteps,
-        duration: torch.Tensor | None = None,
-        mask=None,
-        attn_mask=None,
-        g_cond: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        t = self.time_embedder(timesteps)
-        c = t
-        duration_embedder = getattr(self, "duration_embedder", None)
-        if duration_embedder is not None and duration is not None:
-            c = c + duration_embedder(duration)
-        if g_cond is not None:
-            c = c + g_cond
-
-        x = self.input_layer(x)
-        for block in self.blocks:
-            x = block(x, c, mask=attn_mask, **kwargs)
-        return self.output_layer(x, c, **kwargs)

vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py

@@ -1,355 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copyright 2026 The vLLM-Omni team.
-# Copyright (c) rednote-hilab. All rights reserved.
-# Adapted from:
-# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone
-#
-# Single upstream source (preserved as-is, inference behaviour unchanged;
-# only the internal `from dots_tts.modules.backbone.layers import ...` line
-# was folded to point at our sibling vendor files):
-#   - modules/backbone/semantic_encoder.py -> SemanticEncoderDecodeState,
-#                                             TransformerEncoderLayer,
-#                                             SuperviseEncoder,
-#                                             VAESemanticEncoder
-#
-# Building blocks (Conv1d / Mlp / MultiHeadAttention) are imported from the
-# sibling vendor files where they already live:
-#   - Conv1d              <- dots_tts_vocoder.py
-#   - Mlp                 <- dots_tts_dit.py
-#   - MultiHeadAttention  <- dots_tts_dit.py
-#
-# Nothing else in this file diverges from upstream behaviour.
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-
-from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import (
-    Mlp,
-    MultiHeadAttention,
-)
-from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import Conv1d
-
-
-@dataclass
-class SemanticEncoderDecodeState:
-    conv_tail: torch.Tensor
-    layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...]
-    seq_len: int
-
-
-class TransformerEncoderLayer(nn.Module):
-    def __init__(
-        self,
-        hidden_size,
-        num_heads=16,
-        ffn_hidden_size=4096,
-        attn_dropout=0.0,
-        ffn_dropout=0.0,
-        norm_layer="LayerNorm",
-        **kwargs,
-    ):
-        super().__init__()
-        self.attn = MultiHeadAttention(
-            hidden_size,
-            num_heads,
-            attn_drop=attn_dropout,
-            norm_layer=norm_layer,
-            **kwargs,
-        )
-        norm_cls = getattr(nn, norm_layer)
-        self.attn_norm = norm_cls(hidden_size)
-        self.ffn = Mlp(hidden_size, ffn_hidden_size, dropout=ffn_dropout, act_layer=nn.SiLU)
-        self.ffn_norm = norm_cls(hidden_size)
-        self.hidden_size = hidden_size
-
-    def _build_causal_mask(self, T: int, device):  # noqa: N803
-        return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device))
-
-    def _build_padding_mask(self, x_lens, max_len: int, device):
-        B = x_lens.size(0)
-        positions = torch.arange(max_len, device=device).unsqueeze(0).expand(B, -1)
-        return positions < x_lens.unsqueeze(1)
-
-    def _fuse_attn_mask(self, causal_mask, padding_mask):
-        if causal_mask is None and padding_mask is None:
-            return None
-        if causal_mask is None:
-            row = padding_mask.unsqueeze(2)
-            col = padding_mask.unsqueeze(1)
-            return row & col
-        if padding_mask is None:
-            return causal_mask.unsqueeze(0)
-
-        _B, _T = padding_mask.shape
-        causal = causal_mask.unsqueeze(0)
-        row = padding_mask.unsqueeze(2)
-        col = padding_mask.unsqueeze(1)
-        pad_2d = row & col
-        return causal & pad_2d
-
-    def forward(
-        self,
-        x,
-        x_lens=None,
-        causal=True,
-    ):
-        _B, T, C = x.shape
-        assert self.hidden_size == C
-        device = x.device
-
-        causal_mask = self._build_causal_mask(T, device) if causal else None
-        if x_lens is not None:
-            padding_mask = self._build_padding_mask(x_lens, T, device)
-        else:
-            padding_mask = None
-        fused_mask = self._fuse_attn_mask(causal_mask, padding_mask)
-
-        h = self.attn_norm(x)
-        h = self.attn(
-            q=h,
-            mask=fused_mask,
-        )
-        x = x + h
-
-        h = self.ffn_norm(x)
-        h = self.ffn(h)
-        return x + h
-
-    def decode_step(
-        self,
-        x,
-        *,
-        cache: tuple[torch.Tensor, torch.Tensor],
-        positions: torch.Tensor,
-    ):
-        if x.size(1) <= 0:
-            raise ValueError("TransformerEncoderLayer.decode_step expects a non-empty input.")
-
-        h = self.attn_norm(x)
-        h, cache = self.attn.decode_step(h, cache=cache, positions=positions)
-        x = x + h
-
-        h = self.ffn_norm(x)
-        h = self.ffn(h)
-        return x + h, cache
-
-
-class SuperviseEncoder(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.hidden_size = config.get("hidden_size", 1024)
-        self.layers = nn.ModuleList(
-            [
-                TransformerEncoderLayer(
-                    hidden_size=self.hidden_size,
-                    num_heads=config.get("num_heads", 16),
-                    ffn_hidden_size=config.get("ffn_hidden_size", 4096),
-                    norm_layer=config.get("norm_layer", "LayerNorm"),
-                )
-                for _ in range(config.get("num_layers", 6))
-            ]
-        )
-        self.causal = config.get("causal", False)
-
-    def forward(self, x, x_lens=None):
-        batch_size, seq_len, _ = x.shape
-        if x_lens is None:
-            x_lens = torch.full((batch_size,), seq_len, device=x.device, dtype=torch.long)
-        for layer in self.layers:
-            x = layer(x, x_lens=x_lens, causal=self.causal)
-        return x
-
-    def init_decode_state(
-        self,
-        *,
-        batch_size: int,
-        max_seq_len: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ):
-        layer_caches = []
-        for layer in self.layers:
-            cache_shape = (
-                batch_size,
-                layer.attn.num_heads,
-                max_seq_len,
-                layer.attn.head_dim,
-            )
-            layer_caches.append(
-                (
-                    torch.zeros(cache_shape, dtype=dtype, device=device),
-                    torch.zeros(cache_shape, dtype=dtype, device=device),
-                )
-            )
-        return tuple(layer_caches)
-
-    def reset_decode_state(
-        self,
-        layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...],
-    ) -> None:
-        if len(layer_caches) != len(self.layers):
-            raise ValueError("Layer cache count does not match encoder depth.")
-        for key_cache, value_cache in layer_caches:
-            key_cache.zero_()
-            value_cache.zero_()
-
-    def decode_step(self, x, *, layer_caches, positions: torch.Tensor):
-        if len(layer_caches) != len(self.layers):
-            raise ValueError("Layer cache count does not match encoder depth.")
-
-        for layer, cache in zip(self.layers, layer_caches, strict=True):
-            x, _ = layer.decode_step(x, cache=cache, positions=positions)
-        return x
-
-
-class VAESemanticEncoder(nn.Module):
-    def __init__(self, in_dim, out_dim, config):
-        super().__init__()
-        in_ds_rate = 2
-        self.patch_size = int(config.patch_size)
-        self.in_ds_rate = in_ds_rate
-        self.ds_proj = Conv1d(in_dim, in_dim, kernel_size=in_ds_rate, stride=in_ds_rate, causal=True)
-        self.in_proj = nn.Linear(in_dim, config.PatchEncoder.hidden_size)
-        self.encoder = SuperviseEncoder(config.PatchEncoder)
-        self.out_ds_rate = self.patch_size // in_ds_rate
-        self.out_proj = nn.Linear(config.PatchEncoder.hidden_size * self.out_ds_rate, out_dim)
-
-    def forward(self, x, x_lens=None):
-        x = self._downsample(x)
-        x = self.in_proj(x)
-        z = self.encoder(x, x_lens=x_lens)
-        return self._project_embeddings(z)
-
-    def init_decode_state(
-        self,
-        *,
-        max_audio_patch_count: int,
-        batch_size: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> SemanticEncoderDecodeState:
-        return SemanticEncoderDecodeState(
-            conv_tail=torch.zeros(
-                (batch_size, self.ds_proj.in_channels, self.ds_proj.left_padding),
-                dtype=dtype,
-                device=device,
-            ),
-            layer_caches=self.encoder.init_decode_state(
-                batch_size=batch_size,
-                max_seq_len=max_audio_patch_count * self.out_ds_rate,
-                device=device,
-                dtype=dtype,
-            ),
-            seq_len=0,
-        )
-
-    def reset_decode_state(self, state: SemanticEncoderDecodeState) -> None:
-        state.conv_tail.zero_()
-        self.encoder.reset_decode_state(state.layer_caches)
-        state.seq_len = 0
-
-    def prefill(
-        self,
-        x,
-        state: SemanticEncoderDecodeState,
-    ) -> tuple[torch.Tensor, SemanticEncoderDecodeState]:
-        if x.ndim != 3:
-            raise ValueError(f"VAESemanticEncoder.prefill expects rank-3 input, got {tuple(x.shape)}.")
-        if x.size(1) % self.patch_size != 0:
-            raise ValueError(f"Prompt latent length {x.size(1)} must be divisible by patch_size={self.patch_size}.")
-
-        if x.size(1) == 0:
-            return (
-                x.new_zeros((x.size(0), 0, self.out_proj.out_features)),
-                state,
-            )
-        if state.conv_tail.size(0) != x.size(0):
-            raise ValueError("VAESemanticEncoder.prefill batch size does not match decode state.")
-
-        step_inputs = self.in_proj(self._downsample(x))
-        expected_token_count = (x.size(1) // self.patch_size) * self.out_ds_rate
-        if step_inputs.size(1) != expected_token_count:
-            raise RuntimeError(
-                "Patch encoder prefill produced an unexpected token count: "
-                f"expected={expected_token_count} actual={step_inputs.size(1)}."
-            )
-
-        current_seq_len = state.seq_len
-        next_seq_len = current_seq_len + step_inputs.size(1)
-        cache_capacity = state.layer_caches[0][0].size(2)
-        if next_seq_len > cache_capacity:
-            raise ValueError(
-                "Patch encoder prefill exceeds decode-state capacity: "
-                f"required={next_seq_len} capacity={cache_capacity}."
-            )
-
-        positions = torch.arange(step_inputs.size(1), device=x.device, dtype=torch.long) + current_seq_len
-        encoded = self.encoder.decode_step(
-            step_inputs,
-            layer_caches=state.layer_caches,
-            positions=positions,
-        )
-        embedding = self._project_embeddings(encoded)
-        raw = x.transpose(1, 2)
-        state.conv_tail.copy_(raw[..., -self.ds_proj.left_padding :])
-        state.seq_len = next_seq_len
-        return embedding, state
-
-    def decode_patch(
-        self,
-        latent_patch,
-        conv_tail: torch.Tensor,
-        layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...],
-        positions: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if latent_patch.ndim != 3:
-            raise ValueError(f"VAESemanticEncoder.decode_patch expects rank-3 input, got {tuple(latent_patch.shape)}.")
-        if latent_patch.size(1) != self.patch_size:
-            raise ValueError(f"decode_patch expects patch length {self.patch_size}, got {latent_patch.size(1)}.")
-        if positions.ndim != 1 or positions.size(0) != self.out_ds_rate:
-            raise ValueError("decode_patch positions must be a rank-1 tensor matching out_ds_rate.")
-
-        step_inputs, conv_tail = self._downsample_step(
-            latent_patch,
-            conv_tail=conv_tail,
-        )
-        if step_inputs.size(1) != self.out_ds_rate:
-            raise RuntimeError(f"Downsample step produced {step_inputs.size(1)} tokens, expected {self.out_ds_rate}.")
-
-        encoded = self.encoder.decode_step(
-            step_inputs,
-            layer_caches=layer_caches,
-            positions=positions,
-        )
-        embedding = self._project_embeddings(encoded)
-        return embedding, conv_tail
-
-    def _downsample(self, x):
-        return self.ds_proj(x.transpose(1, 2)).transpose(1, 2)
-
-    def _project_embeddings(self, z):
-        if self.out_ds_rate > 1:
-            z = rearrange(z, "b (s d) h -> b s (d h)", d=self.out_ds_rate)
-        return self.out_proj(z)
-
-    def _downsample_step(self, latent_patch, *, conv_tail):
-        raw = latent_patch.transpose(1, 2)
-        conv_input = torch.cat([conv_tail, raw], dim=-1)
-
-        projected = F.conv1d(
-            conv_input,
-            self.ds_proj.weight,
-            self.ds_proj.bias,
-            stride=self.ds_proj.stride[0],
-            padding=0,
-            dilation=self.ds_proj.dilation[0],
-            groups=self.ds_proj.groups,
-        ).transpose(1, 2)
-        new_conv_tail = raw[..., -self.ds_proj.left_padding :]
-        return self.in_proj(projected), new_conv_tail

vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py

@@ -1,70 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""dots.tts prompt builder for vllm-omni's Omni engine + serving layer.
-
-Mirrors upstream ``build_generation_schedule`` (rednote-hilab/dots.tts @
-a393d2e data/pipelines/tokenizing.py:169) for the
-no-prompt-audio + audio-placeholder path — text-only zero-shot TTS.
-
-V1 returns real Qwen2 token IDs (not placeholders ``[1] * N`` like
-voxcpm2's builder).  Reason: our talker's preprocess() runs
-``self.model.embed_tokens(input_ids)`` directly, so real IDs flow
-through naturally.  voxcpm2 uses placeholders because its prefill mixes
-reference-audio embeddings (no token IDs); dots.tts V1 zero-shot is
-purely text + audio_gen_start, all real IDs.
-
-Voice clone (reference audio + speaker x-vector) is intentionally not
-wired here — V1 is zero-shot only.  Add reference_audio /
-voice_profile / prompt_text branches in a follow-up (mirror voxcpm2
-``build_voxcpm2_prompt``).
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-from vllm.inputs import tokens_input
-
-# Upstream rednote-hilab/dots.tts @ a393d2e
-# data/pipelines/tts_pipeline.py:17-18
-_TTS_TEXT_PREFIX = "[文本]"
-_TTS_AUDIO_PREFIX = "[文本对应语音]"
-_AUDIO_GEN_START_TOKEN = "<|audio_gen_start|>"
-
-
-def build_dots_tts_prompt(
-    tokenizer: Any,
-    text: str,
-) -> dict[str, Any]:
-    """Build a dots.tts prefill prompt dict for ``Omni.generate()``.
-
-    Returns ``{"prompt_token_ids": [...], "additional_information": {...}}``
-    where ``prompt_token_ids`` is the real Qwen2 token sequence:
-
-        [text_tokens of "[文本]<text>[文本对应语音]", audio_gen_start_id]
-
-    talker's preprocess() embeds these into ``hidden_states`` for the AR
-    base LM.  Decode then drives via ``state.curr_embed_for_next`` (DiT
-    → patch_encoder AR loopback).
-    """
-    if not isinstance(text, str) or not text.strip():
-        raise ValueError(f"text must be a non-empty string, got {text!r}")
-
-    prompt_text = f"{_TTS_TEXT_PREFIX}{text}{_TTS_AUDIO_PREFIX}"
-    text_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
-
-    audio_gen_start_id = tokenizer.convert_tokens_to_ids(_AUDIO_GEN_START_TOKEN)
-    if not isinstance(audio_gen_start_id, int) or audio_gen_start_id == tokenizer.unk_token_id:
-        raise ValueError(
-            f"Tokenizer does not know {_AUDIO_GEN_START_TOKEN!r} "
-            f"(got id={audio_gen_start_id!r}).  Did you load the "
-            f"dots.tts tokenizer (with added_tokens.json)?"
-        )
-    prompt_token_ids = list(text_ids) + [audio_gen_start_id]
-
-    prompt = tokens_input(prompt_token_ids=prompt_token_ids)
-    prompt["additional_information"] = {
-        # Reserved for future voice clone (ref_audio / speaker x-vector).
-        # V1 zero-shot leaves this empty; preprocess() ignores it.
-    }
-    return prompt

vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py

@@ -1,732 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copyright 2026 The vLLM-Omni team.
-# Copyright (c) rednote-hilab. All rights reserved.
-# Portions Copyright 3D-Speaker
-# (https://github.com/alibaba-damo-academy/3D-Speaker), Apache-2.0.
-# Adapted from:
-# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/speaker
-# (whose CAM++ implementation is itself adapted from 3D-Speaker).
-#
-# Merged from four upstream files (preserved as-is, inference behaviour
-# unchanged; only internal cross-module imports were folded inline):
-#   - modules/speaker/encoder.py         -> SpeakerXVectorFeatures
-#   - modules/speaker/fbank.py           -> fbank constants + extract_speaker_fbank
-#   - modules/speaker/campplus.py        -> FCM, CAMPPlus
-#   - modules/speaker/campplus_layers.py -> get_nonlinear, statistics_pooling,
-#                                           StatsPool, TDNNLayer, CAMLayer,
-#                                           CAMDenseTDNNLayer, CAMDenseTDNNBlock,
-#                                           TransitLayer, DenseLayer, BasicResBlock
-#
-# Plus two tiny helpers folded from utils/audio.py (used by
-# extract_speaker_fbank):
-#   - high_quality_resample, extract_fbank
-#
-# Nothing in this file diverges from upstream behaviour.
-
-from __future__ import annotations
-
-import math
-import random
-from collections import OrderedDict
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as cp
-import torchaudio
-import torchaudio.compliance.kaldi as Kaldi
-import torchaudio.functional as AF
-from torch.nn.utils.rnn import pad_sequence
-
-# ============================================================================
-# Constants (from modules/speaker/fbank.py)
-# ============================================================================
-
-_SPEAKER_FBANK_SAMPLE_RATE = 16000
-_SPEAKER_FBANK_N_MELS = 80
-_SPEAKER_FBANK_MEAN_NORM = True
-_SPEAKER_FBANK_DITHER = 0.0
-
-
-# ============================================================================
-# Audio helpers (folded from utils/audio.py — used by extract_speaker_fbank)
-# ============================================================================
-
-
-def high_quality_resample(x, orig_sr, target_sr):
-    return AF.resample(
-        x,
-        orig_freq=orig_sr,
-        new_freq=target_sr,
-        lowpass_filter_width=64,
-        rolloff=0.95,
-        resampling_method="sinc_interp_kaiser",
-    )
-
-
-def extract_fbank(
-    waveform: torch.Tensor,
-    *,
-    sample_rate: int,
-    n_mels: int,
-    dither: float = 0.0,
-    mean_norm: bool = False,
-) -> torch.Tensor:
-    if waveform.ndim == 1:
-        feature_input = waveform.unsqueeze(0)
-    elif waveform.ndim == 2:
-        feature_input = waveform if waveform.size(0) == 1 else waveform[0:1, :]
-    else:
-        raise ValueError(f"FBank expects a 1D or 2D waveform, got shape {tuple(waveform.shape)}.")
-    features = Kaldi.fbank(
-        feature_input,
-        num_mel_bins=n_mels,
-        sample_frequency=sample_rate,
-        dither=dither,
-    )
-    if mean_norm:
-        features = features - features.mean(dim=0, keepdim=True)
-    return features
-
-
-# ============================================================================
-# Speaker fbank entry (from modules/speaker/fbank.py)
-# ============================================================================
-
-
-def extract_speaker_fbank(
-    waveform: torch.Tensor,
-    *,
-    sample_rate: int,
-) -> torch.Tensor:
-    feature_input = waveform
-    if sample_rate != _SPEAKER_FBANK_SAMPLE_RATE:
-        feature_input = high_quality_resample(
-            waveform,
-            orig_sr=sample_rate,
-            target_sr=_SPEAKER_FBANK_SAMPLE_RATE,
-        )
-    return extract_fbank(
-        feature_input,
-        sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
-        n_mels=_SPEAKER_FBANK_N_MELS,
-        dither=_SPEAKER_FBANK_DITHER,
-        mean_norm=_SPEAKER_FBANK_MEAN_NORM,
-    )
-
-
-# ============================================================================
-# CAM++ building blocks (from modules/speaker/campplus_layers.py)
-# ----------------------------------------------------------------------------
-# 3D-Speaker, Apache-2.0.
-# ============================================================================
-
-
-def get_nonlinear(config_str, channels):
-    nonlinear = nn.Sequential()
-    for name in config_str.split("-"):
-        if name == "relu":
-            nonlinear.add_module("relu", nn.ReLU(inplace=True))
-        elif name == "prelu":
-            nonlinear.add_module("prelu", nn.PReLU(channels))
-        elif name == "batchnorm":
-            nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels))
-        elif name == "batchnorm_":
-            nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False))
-        else:
-            raise ValueError(f"Unexpected module ({name}).")
-    return nonlinear
-
-
-def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, _eps=1e-2):
-    mean = x.mean(dim=dim)
-    std = x.std(dim=dim, unbiased=unbiased)
-    stats = torch.cat([mean, std], dim=-1)
-    if keepdim:
-        stats = stats.unsqueeze(dim=dim)
-    return stats
-
-
-class StatsPool(nn.Module):
-    def forward(self, x):
-        return statistics_pooling(x)
-
-
-class TDNNLayer(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        bias=False,
-        config_str="batchnorm-relu",
-    ):
-        super().__init__()
-        if padding < 0:
-            assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
-            padding = (kernel_size - 1) // 2 * dilation
-        self.linear = nn.Conv1d(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            bias=bias,
-        )
-        self.nonlinear = get_nonlinear(config_str, out_channels)
-
-    def forward(self, x):
-        x = self.linear(x)
-        return self.nonlinear(x)
-
-
-class CAMLayer(nn.Module):
-    def __init__(
-        self,
-        bn_channels,
-        out_channels,
-        kernel_size,
-        stride,
-        padding,
-        dilation,
-        bias,
-        reduction=2,
-    ):
-        super().__init__()
-        self.linear_local = nn.Conv1d(
-            bn_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            bias=bias,
-        )
-        self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1)
-        self.relu = nn.ReLU(inplace=True)
-        self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1)
-        self.sigmoid = nn.Sigmoid()
-
-    def forward(self, x):
-        y = self.linear_local(x)
-        context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
-        context = self.relu(self.linear1(context))
-        m = self.sigmoid(self.linear2(context))
-        return y * m
-
-    def seg_pooling(self, x, seg_len=100, stype="avg"):
-        if stype == "avg":
-            seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
-        elif stype == "max":
-            seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
-        else:
-            raise ValueError("Wrong segment pooling type.")
-        shape = seg.shape
-        seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
-        return seg[..., : x.shape[-1]]
-
-
-class CAMDenseTDNNLayer(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        bn_channels,
-        kernel_size,
-        stride=1,
-        dilation=1,
-        bias=False,
-        config_str="batchnorm-relu",
-        memory_efficient=False,
-    ):
-        super().__init__()
-        assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
-        padding = (kernel_size - 1) // 2 * dilation
-        self.memory_efficient = memory_efficient
-        self.nonlinear1 = get_nonlinear(config_str, in_channels)
-        self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False)
-        self.nonlinear2 = get_nonlinear(config_str, bn_channels)
-        self.cam_layer = CAMLayer(
-            bn_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            bias=bias,
-        )
-
-    def bn_function(self, x):
-        return self.linear1(self.nonlinear1(x))
-
-    def forward(self, x):
-        if self.training and self.memory_efficient:
-            x = cp.checkpoint(self.bn_function, x)
-        else:
-            x = self.bn_function(x)
-        return self.cam_layer(self.nonlinear2(x))
-
-
-class CAMDenseTDNNBlock(nn.ModuleList):
-    def __init__(
-        self,
-        num_layers,
-        in_channels,
-        out_channels,
-        bn_channels,
-        kernel_size,
-        stride=1,
-        dilation=1,
-        bias=False,
-        config_str="batchnorm-relu",
-        memory_efficient=False,
-    ):
-        super().__init__()
-        for i in range(num_layers):
-            layer = CAMDenseTDNNLayer(
-                in_channels=in_channels + i * out_channels,
-                out_channels=out_channels,
-                bn_channels=bn_channels,
-                kernel_size=kernel_size,
-                stride=stride,
-                dilation=dilation,
-                bias=bias,
-                config_str=config_str,
-                memory_efficient=memory_efficient,
-            )
-            self.add_module(f"tdnnd{i + 1}", layer)
-
-    def forward(self, x):
-        for layer in self:
-            x = torch.cat([x, layer(x)], dim=1)
-        return x
-
-
-class TransitLayer(nn.Module):
-    def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
-        super().__init__()
-        self.nonlinear = get_nonlinear(config_str, in_channels)
-        self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
-
-    def forward(self, x):
-        x = self.nonlinear(x)
-        return self.linear(x)
-
-
-class DenseLayer(nn.Module):
-    def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
-        super().__init__()
-        self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
-        self.nonlinear = get_nonlinear(config_str, out_channels)
-
-    def forward(self, x):
-        if len(x.shape) == 2:
-            x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
-        else:
-            x = self.linear(x)
-        return self.nonlinear(x)
-
-
-class BasicResBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, in_planes, planes, stride=1):
-        super().__init__()
-        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes)
-
-        self.shortcut = nn.Sequential()
-        if stride != 1 or in_planes != self.expansion * planes:
-            self.shortcut = nn.Sequential(
-                nn.Conv2d(
-                    in_planes,
-                    self.expansion * planes,
-                    kernel_size=1,
-                    stride=(stride, 1),
-                    bias=False,
-                ),
-                nn.BatchNorm2d(self.expansion * planes),
-            )
-
-    def forward(self, x):
-        out = F.relu(self.bn1(self.conv1(x)))
-        out = self.bn2(self.conv2(out))
-        out += self.shortcut(x)
-        return F.relu(out)
-
-
-# ============================================================================
-# CAM++ model (from modules/speaker/campplus.py)
-# ----------------------------------------------------------------------------
-# 3D-Speaker, Apache-2.0.
-# ============================================================================
-
-
-class FCM(nn.Module):
-    def __init__(
-        self,
-        block=BasicResBlock,
-        num_blocks=(2, 2),
-        m_channels=32,
-        feat_dim=_SPEAKER_FBANK_N_MELS,
-    ):
-        super().__init__()
-        self.in_planes = m_channels
-        self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(m_channels)
-
-        self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
-        self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2)
-
-        self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(m_channels)
-        self.out_channels = m_channels * (feat_dim // 8)
-
-    def _make_layer(self, block, planes, num_blocks, stride):
-        strides = [stride] + [1] * (num_blocks - 1)
-        layers = []
-        for stride in strides:
-            layers.append(block(self.in_planes, planes, stride))
-            self.in_planes = planes * block.expansion
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = x.unsqueeze(1)
-        out = F.relu(self.bn1(self.conv1(x)))
-        out = self.layer1(out)
-        out = self.layer2(out)
-        out = F.relu(self.bn2(self.conv2(out)))
-
-        shape = out.shape
-        return out.reshape(shape[0], shape[1] * shape[2], shape[3])
-
-
-class CAMPPlus(nn.Module):
-    _TDNN_KERNEL_SIZE = 5
-    _TDNN_STRIDE = 2
-    _TDNN_PADDING = 2
-
-    def __init__(
-        self,
-        feat_dim=_SPEAKER_FBANK_N_MELS,
-        embedding_size=512,
-        growth_rate=32,
-        bn_size=4,
-        init_channels=128,
-        config_str="batchnorm-relu",
-        memory_efficient=True,
-    ):
-        super().__init__()
-
-        self.head = FCM(feat_dim=feat_dim)
-        channels = self.head.out_channels
-
-        self.xvector = nn.Sequential(
-            OrderedDict(
-                [
-                    (
-                        "tdnn",
-                        TDNNLayer(
-                            channels,
-                            init_channels,
-                            self._TDNN_KERNEL_SIZE,
-                            stride=self._TDNN_STRIDE,
-                            dilation=1,
-                            padding=-1,
-                            config_str=config_str,
-                        ),
-                    ),
-                ]
-            )
-        )
-        channels = init_channels
-        for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2), strict=True)):
-            block = CAMDenseTDNNBlock(
-                num_layers=num_layers,
-                in_channels=channels,
-                out_channels=growth_rate,
-                bn_channels=bn_size * growth_rate,
-                kernel_size=kernel_size,
-                dilation=dilation,
-                config_str=config_str,
-                memory_efficient=memory_efficient,
-            )
-            self.xvector.add_module(f"block{i + 1}", block)
-            channels = channels + num_layers * growth_rate
-            self.xvector.add_module(
-                f"transit{i + 1}",
-                TransitLayer(channels, channels // 2, bias=False, config_str=config_str),
-            )
-            channels //= 2
-
-        self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
-
-        self.xvector.add_module("stats", StatsPool())
-        self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_"))
-
-        for m in self.modules():
-            if isinstance(m, (nn.Conv1d, nn.Linear)):
-                nn.init.kaiming_normal_(m.weight.data)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-
-    @staticmethod
-    def _conv_output_lengths(lengths, kernel_size, stride=1, padding=0, dilation=1):
-        return (
-            torch.div(
-                lengths + 2 * padding - dilation * (kernel_size - 1) - 1,
-                stride,
-                rounding_mode="floor",
-            )
-            + 1
-        )
-
-    @staticmethod
-    def _make_length_mask(lengths, max_len, device):
-        lengths = lengths.to(device=device, dtype=torch.long).clamp(min=0, max=max_len)
-        return torch.arange(max_len, device=device).unsqueeze(0) < lengths.unsqueeze(1)
-
-    def _masked_stats_pooling(self, x, lengths, unbiased=True, eps=1e-2):
-        lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1, max=x.size(-1))
-        mask = self._make_length_mask(lengths, x.size(-1), x.device).unsqueeze(1)
-        mask = mask.to(dtype=x.dtype)
-
-        denom = lengths.to(dtype=x.dtype).view(-1, 1).clamp_min(1.0)
-        mean = (x * mask).sum(dim=-1) / denom
-
-        centered = (x - mean.unsqueeze(-1)) * mask
-        var_denom = (lengths - 1).clamp_min(1).to(dtype=x.dtype).view(-1, 1) if unbiased else denom
-        var = centered.pow(2).sum(dim=-1) / var_denom
-        std = torch.sqrt(var.clamp_min(eps))
-        return torch.cat([mean, std], dim=1)
-
-    def forward(self, x, lengths=None):
-        x = x.permute(0, 2, 1)  # (B,T,F) => (B,F,T)
-        x = self.head(x)
-        if lengths is not None:
-            lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1)
-
-        for name, module in self.xvector.named_children():
-            if name == "stats":
-                x = self._masked_stats_pooling(x, lengths) if lengths is not None else module(x)
-                continue
-
-            x = module(x)
-            if name == "tdnn" and lengths is not None:
-                lengths = self._conv_output_lengths(
-                    lengths,
-                    kernel_size=self._TDNN_KERNEL_SIZE,
-                    stride=self._TDNN_STRIDE,
-                    padding=self._TDNN_PADDING,
-                )
-
-        return x
-
-
-# ============================================================================
-# SpeakerXVectorFeatures wrapper (from modules/speaker/encoder.py)
-# ============================================================================
-
-
-class SpeakerXVectorFeatures(nn.Module):
-    """Speaker embedding extractor based on 3D-Speaker CAM++."""
-
-    def __init__(
-        self,
-        sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
-        campplus_embedding_size=512,
-        max_audio_seconds=10.0,
-    ):
-        super().__init__()
-
-        self.sample_rate = sample_rate
-        self.max_audio_seconds = float(max_audio_seconds)
-        self.model = CAMPPlus(
-            feat_dim=_SPEAKER_FBANK_N_MELS,
-            embedding_size=campplus_embedding_size,
-        )
-        self.resample = None
-        if self.sample_rate != _SPEAKER_FBANK_SAMPLE_RATE:
-            self.resample = torchaudio.transforms.Resample(
-                orig_freq=sample_rate,
-                new_freq=_SPEAKER_FBANK_SAMPLE_RATE,
-            )
-
-        for param in self.model.parameters():
-            param.requires_grad = False
-
-    @staticmethod
-    def _normalize_lengths(lengths, batch_size, max_length, device, *, min_length):
-        if lengths is None:
-            return torch.full(
-                (batch_size,),
-                max_length,
-                device=device,
-                dtype=torch.long,
-            )
-        return lengths.to(device=device, dtype=torch.long).clamp(
-            min=min_length,
-            max=max_length,
-        )
-
-    def _crop_audio(self, audio, audio_lengths=None):
-        original_lengths = self._normalize_lengths(
-            audio_lengths,
-            audio.size(0),
-            audio.size(-1),
-            audio.device,
-            min_length=0,
-        )
-        if self.max_audio_seconds <= 0:
-            return audio, original_lengths, original_lengths, torch.zeros_like(original_lengths)
-
-        max_input_length = round(self.sample_rate * self.max_audio_seconds)
-        cropped_audio = []
-        cropped_lengths = []
-        starts = []
-
-        for index, total_length_tensor in enumerate(original_lengths):
-            total_length = int(total_length_tensor.item())
-            cropped_length = min(total_length, max_input_length)
-            start = random.randint(0, total_length - cropped_length) if total_length > cropped_length else 0
-            cropped_audio.append(audio[index, start : start + cropped_length])
-            cropped_lengths.append(cropped_length)
-            starts.append(start)
-
-        return (
-            pad_sequence(
-                cropped_audio,
-                batch_first=True,
-                padding_value=0.0,
-            ),
-            original_lengths,
-            torch.tensor(
-                cropped_lengths,
-                device=audio.device,
-                dtype=torch.long,
-            ),
-            torch.tensor(starts, device=audio.device, dtype=torch.long),
-        )
-
-    def _crop_fbank(
-        self,
-        fbank,
-        fbank_lengths,
-        original_audio_lengths,
-        cropped_audio_lengths,
-        starts,
-    ):
-        original_fbank_lengths = self._normalize_lengths(
-            fbank_lengths,
-            fbank.size(0),
-            fbank.size(1),
-            fbank.device,
-            min_length=1,
-        )
-        cropped_fbank = []
-        cropped_fbank_lengths = []
-
-        for index, total_feat_length_tensor in enumerate(original_fbank_lengths):
-            total_audio_length = int(original_audio_lengths[index].item())
-            total_feat_length = int(total_feat_length_tensor.item())
-            start_audio = int(starts[index].item())
-            end_audio = start_audio + int(cropped_audio_lengths[index].item())
-
-            if total_audio_length > 0:
-                start_feat = math.floor(start_audio * total_feat_length / total_audio_length)
-                end_feat = math.ceil(end_audio * total_feat_length / total_audio_length)
-            else:
-                start_feat = 0
-                end_feat = 1
-
-            start_feat = min(start_feat, total_feat_length - 1)
-            end_feat = min(max(end_feat, start_feat + 1), total_feat_length)
-            cropped_fbank.append(fbank[index, start_feat:end_feat])
-            cropped_fbank_lengths.append(end_feat - start_feat)
-
-        return pad_sequence(
-            cropped_fbank,
-            batch_first=True,
-            padding_value=0.0,
-        ), torch.tensor(
-            cropped_fbank_lengths,
-            device=fbank.device,
-            dtype=torch.long,
-        )
-
-    def _extract_fbank_batch(self, audio, audio_lengths):
-        if self.resample is not None:
-            audio = self.resample(audio)
-            audio_lengths = torch.ceil(audio_lengths.float() * (_SPEAKER_FBANK_SAMPLE_RATE / self.sample_rate)).long()
-
-        audio_cpu = audio.detach().cpu()
-        features = []
-
-        for index, valid_length_tensor in enumerate(audio_lengths):
-            valid_length = int(valid_length_tensor.item())
-            waveform = audio_cpu[index, :valid_length]
-            if waveform.numel() == 0:
-                waveform = audio_cpu.new_zeros(1)
-            features.append(
-                extract_speaker_fbank(
-                    waveform,
-                    sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
-                )
-            )
-
-        fbank_lengths = torch.tensor(
-            [feature.size(0) for feature in features],
-            device=audio.device,
-            dtype=torch.long,
-        )
-        fbank = pad_sequence(
-            features,
-            batch_first=True,
-            padding_value=0.0,
-        ).to(device=audio.device, dtype=audio.dtype)
-        return fbank, fbank_lengths
-
-    @torch.no_grad()
-    @torch.autocast(enabled=False, device_type="cuda")
-    def forward(self, audio, audio_lengths=None, fbank=None, fbank_lengths=None, **_kwargs):
-        self.model.eval()
-        audio = audio.float()
-        if audio.dim() == 3:
-            if audio.size(1) != 1:
-                raise ValueError(f"Speaker encoder expects mono audio, got shape {tuple(audio.shape)}.")
-            audio = audio[:, 0]
-        elif audio.dim() != 2:
-            raise ValueError(f"Speaker encoder expects a 2D or 3D audio tensor, got shape {tuple(audio.shape)}.")
-
-        audio, original_audio_lengths, cropped_audio_lengths, starts = self._crop_audio(
-            audio,
-            audio_lengths=audio_lengths,
-        )
-
-        if fbank is None:
-            fbank, fbank_lengths = self._extract_fbank_batch(
-                audio,
-                cropped_audio_lengths,
-            )
-        else:
-            if not isinstance(fbank, torch.Tensor):
-                raise TypeError("Speaker encoder expects `fbank` to be a torch.Tensor.")
-            if fbank.dim() != 3 or fbank.size(0) != audio.size(0):
-                raise ValueError(
-                    f"Speaker encoder expects `fbank` with shape (B, T, F) and matching batch size, got {tuple(fbank.shape)}."  # noqa: E501
-                )
-            fbank, fbank_lengths = self._crop_fbank(
-                fbank.to(device=audio.device, dtype=torch.float32),
-                fbank_lengths,
-                original_audio_lengths,
-                cropped_audio_lengths,
-                starts,
-            )
-
-        return self.model(fbank, lengths=fbank_lengths)

vllm_omni/model_executor/models/dots_tts/pipeline.py

@@ -1,46 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""dots.tts pipeline topology (frozen).
-
-Single-stage AR TTS: text → 48 kHz waveform in one pass.
-Architecture: Qwen2.5-1.5B base LM (vLLM-native) → DiTAR / MeanFlow head
-→ AudioVAE.decode → BigVGAN → 48 kHz mono.
-
-Speaker conditioning via CAM++ x-vector injected in preprocess().
-
-Modeled after voxcpm2's single-stage pipeline (same "vLLM-native base LM
-+ side-path computation" pattern), but simplified: no residual LM, no FSQ.
-"""
-
-from vllm_omni.config.stage_config import (
-    PipelineConfig,
-    StageExecutionType,
-    StagePipelineConfig,
-)
-
-DOTS_TTS_PIPELINE = PipelineConfig(
-    model_type="dots_tts",
-    model_arch="DotsTTSForConditionalGeneration",
-    default_deploy_config_name="dots_tts.yaml",
-    stages=(
-        StagePipelineConfig(
-            stage_id=0,
-            model_stage="latent_generator",
-            execution_type=StageExecutionType.LLM_AR,
-            input_sources=(),
-            final_output=True,
-            final_output_type="audio",
-            owns_tokenizer=True,
-            engine_output_type="audio",
-            sampling_constraints={
-                "detokenize": False,
-                # Stop signal slot in compute_logits (voxcpm2 convention —
-                # talker emits logits[i, 1] = 1.0 when prob_stop > 0.8, vLLM
-                # matches against this list to terminate the request).  Not
-                # Qwen2.5 EOS (151643) — that ID would never come out of the
-                # 2-slot continue/stop softmax our compute_logits builds.
-                "stop_token_ids": [1],
-            },
-        ),
-    ),
-)

vllm_omni/model_executor/models/registry.py

@@ -301,12 +301,6 @@ _OMNI_MODELS = {
         "voxcpm2_talker",
         "VoxCPM2TalkerForConditionalGeneration",
     ),
-    ## dots.tts
-    "DotsTTSForConditionalGeneration": (
-        "dots_tts",
-        "dots_tts_talker",
-        "DotsTTSForConditionalGeneration",
-    ),
     ## Voxtral TTS
     "VoxtralTTSForConditionalGeneration": (
         "voxtral_tts",

vllm_omni/model_extras/registry.py

@@ -205,7 +205,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         }
         for model_class_name in (
             "LTX2Pipeline",
-            "LTX2TwoStagePipeline",
             "LTX2DistilledPipeline",
         )
     },

vllm_omni/transformers_utils/configs/__init__.py

@@ -20,7 +20,6 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech",
     "GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts",
     "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2",
-    "DotsTTSConfig": "vllm_omni.transformers_utils.configs.dots_tts",
     "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts",
     "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3",
     "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice",
@@ -43,7 +42,6 @@ __all__ = [
     "FishSpeechFastARConfig",
     "GLMTTSConfig",
     "VoxCPM2Config",
-    "DotsTTSConfig",
     "VoxtralTTSConfig",
     "CosyVoice3Config",
     "OmniVoiceConfig",
@@ -72,7 +70,6 @@ def __dir__():
 # Eagerly import all config modules so their AutoConfig.register() side-effects
 # run as soon as `vllm_omni.transformers_utils.configs` is imported.
 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3  # noqa: F401, E402
-from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3  # noqa: F401, E402

vllm_omni/transformers_utils/configs/dots_tts.py

@@ -1,102 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""HF config for dots.tts.
-
-dots.tts splits its config across two files in the HF checkpoint:
-  * ``config.json``      — top-level dots.tts fields (latent_dim / patch_size
-                            / PatchEncoder / DiT / vocoder / etc.).
-  * ``llm_config.json``  — full Qwen2 LM config (vocab_size / hidden_size /
-                            num_hidden_layers / rope_theta / etc.).
-
-``AutoConfig.from_pretrained`` only consumes ``config.json``: its dots.tts
-sub-blocks (``vocoder`` / ``DiT`` / ``PatchEncoder``) arrive through
-``**kwargs`` and land as attributes, which the talker's config factories
-consume directly.  The Qwen2 LM fields live in ``llm_config.json``, which
-AutoConfig never reads — this class hoists them to top-level attributes
-with defaults matching the released checkpoints (soar / base / mf all ship
-the same Qwen2.5-1.5B backbone).  vLLM's ``Qwen2Model.__init__`` then reads
-``config.hidden_size`` / ``config.vocab_size`` / etc. through
-``get_text_config()``, the same pattern voxcpm2 follows.
-
-Future work: read ``llm_config.json`` during ``from_pretrained`` so a
-checkpoint with a different LM shape works without editing this file.
-
-Upstream reference:
-  https://huggingface.co/rednote-hilab/dots.tts-soar/blob/main/llm_config.json
-"""
-
-from __future__ import annotations
-
-from transformers import AutoConfig
-from transformers.configuration_utils import PretrainedConfig
-
-
-class DotsTTSConfig(PretrainedConfig):
-    """Configuration for dots.tts integration (Qwen2.5-1.5B-Base AR backbone)."""
-
-    model_type = "dots_tts"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        # -- top-level dots.tts params --
-        lm_config: dict | None = None,
-        # -- LM defaults (overridden by lm_config dict if present) --
-        # All values match dots.tts-soar's llm_config.json.
-        vocab_size: int = 151672,
-        hidden_size: int = 1536,
-        intermediate_size: int = 8960,
-        num_hidden_layers: int = 28,
-        num_attention_heads: int = 12,
-        num_key_value_heads: int = 2,
-        max_position_embeddings: int = 131072,
-        rms_norm_eps: float = 1e-6,
-        rope_theta: float = 1_000_000.0,
-        tie_word_embeddings: bool = True,
-        hidden_act: str = "silu",
-        bos_token_id: int = 151643,
-        eos_token_id: int = 151643,
-        **kwargs,
-    ):
-        super().__init__(
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
-        # The dots.tts sub-blocks (vocoder / DiT / PatchEncoder / latent_dim
-        # / patch_size / ...) arrive via **kwargs and are stored as
-        # attributes by PretrainedConfig; the talker's config factories
-        # read them from there.
-        self.lm_config = lm_config or {}
-
-        # -- Hoist LM parameters to top-level for vLLM's Qwen2Model --
-        lm = self.lm_config
-        self.vocab_size = lm.get("vocab_size", vocab_size)
-        self.hidden_size = lm.get("hidden_size", hidden_size)
-        self.intermediate_size = lm.get("intermediate_size", intermediate_size)
-        self.num_hidden_layers = lm.get("num_hidden_layers", num_hidden_layers)
-        self.num_attention_heads = lm.get("num_attention_heads", num_attention_heads)
-        self.num_key_value_heads = lm.get("num_key_value_heads", num_key_value_heads)
-        self.max_position_embeddings = lm.get("max_position_embeddings", max_position_embeddings)
-        self.rms_norm_eps = lm.get("rms_norm_eps", rms_norm_eps)
-        self.rope_theta = lm.get("rope_theta", rope_theta)
-        self.hidden_act = lm.get("hidden_act", hidden_act)
-        self.head_dim = self.hidden_size // self.num_attention_heads
-
-        # vLLM's Qwen2Model reads ``max_window_layers`` and checks against
-        # ``num_hidden_layers`` only when ``use_sliding_window`` is true.
-        # soar disables sliding window — set both to the same value so the
-        # assertion (if executed) stays vacuous.
-        self.use_sliding_window = lm.get("use_sliding_window", False)
-        self.sliding_window = lm.get("sliding_window", None)
-        self.max_window_layers = lm.get("max_window_layers", self.num_hidden_layers)
-
-    def get_text_config(self, **kwargs):
-        """Return self — LM fields are hoisted to top-level attributes."""
-        return self
-
-
-AutoConfig.register("dots_tts", DotsTTSConfig)
-
-__all__ = ["DotsTTSConfig"]