Fix LTX2 channels-last review scope

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

.buildkite/cuda/test-nightly.yml

@@ -345,7 +345,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json

.buildkite/npu/test-npu-nightly.yml

@@ -115,7 +115,7 @@ steps:
           DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export CACHE_DIT_VERSION=1.5.0
+          - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

.readthedocs.yml

@@ -11,15 +11,10 @@ build:
     post_checkout:
       - git fetch --unshallow || true
     pre_install:
-      # `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1
-      # rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2).
-      - python -m pip install --upgrade pip
       # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
       # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
-      # 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url`
-      # replaces PyPI, so an sdist fallback cannot reach its build deps. The
-      # quotes keep YAML from reading `:all:` as a mapping key.
-      - "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch"
+      # 1 GB memory cap and time out the mkdocs build.
+      - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
 
 mkdocs:
   configuration: mkdocs.yml

benchmarks/lingbot_video/dense_pipeline_parity.py

@@ -167,11 +167,9 @@ def main() -> None:
     ]
     native_cmd = [
         sys.executable,
-        "examples/offline_inference/text_to_video/text_to_video.py",
+        "examples/offline_inference/text_to_video/text_to_video_lingbot.py",
         "--model",
         args.model,
-        "--model-class-name",
-        "LingBotVideoPipeline",
         "--prompt",
         args.prompt,
         "--output",

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.27.0
+ARG VLLM_ASCEND_TAG=v0.26.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.27.0-a3
+ARG VLLM_ASCEND_TAG=v0.26.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.27.0
+ARG VLLM_ASCEND_TAG=v0.26.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.npu.ci.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.27.0-a3
+ARG VLLM_ASCEND_TAG=v0.26.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

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 Parallel: user_guide/diffusion/parallelism/vae_parallel.md
+        - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.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 Parallel: design/feature/vae_parallel.md
+          - VAE Patch Parallelism: design/feature/vae_parallel.md
         - Attention Optimization:
           - Skip-Softmax: design/feature/skip_softmax.md
         - Cache-DiT: design/feature/cache_dit.md
@@ -163,8 +163,7 @@ nav:
 - CLI Reference: cli
 - Community:
   - Governance: community/governance.md
-  - Contact Us: community/contact_us.md
-  - Community Meetings: community/meetings.md
+  - community/*
   - Slack: https://slack.vllm.ai
   - Blog: https://blog.vllm.ai
   - Forum: https://discuss.vllm.ai

docs/community/contact_us.md

@@ -1,9 +1,5 @@
 # Contact Us
 
-- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues).
-- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/).
-- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md).
-
-## Security Disclosures
-
-Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request.
+- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
+- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
+- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature

docs/community/governance.md

@@ -52,11 +52,31 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
 - [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
 
-## Reviewer Routing
+## Diffusion Workload Division
 
-For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).
+This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
 
-`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights.
+### Runtime-related
+
+- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
+
+### Optimization
+
+- Parallel: @RuixiangMa, @wtomin
+- Attention backend: @ZJY0516
+- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
+- Offload: @yuanheng-zhao
+- Intel GPU support: @xuechendi
+- Quantization: @lishunyang12, @david6666666
+- Kernel fusion / communication-computation: @ZJY0516
+
+### Docs & Test
+
+- Diffusion docs: @RuixiangMa, @wtomin
+
+## Meetings
+
+Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
 
 ## Committer Nomination Process
 

docs/community/meetings.md

@@ -1,8 +0,0 @@
-# Community Meetings
-
-vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth:
-
-- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0)
-- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0)
-
-Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/).

docs/community/meetups.md

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

docs/community/volunteers.md

@@ -0,0 +1,12 @@
+# Volunteers for Bugfix and CI
+
+We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
+
+If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
+
+| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
+|----------|----------|----------|----------|----------|----------|
+| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
+| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
+
+We kindly welcome more contributors to fix bugs and contribute new features!

docs/contributing/README.md

@@ -3,7 +3,7 @@
 Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
 
 !!! note
-    vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. See [Community Meetings](../community/meetings.md) for current schedules, access details, agendas, and past notes.
+    We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
 
 ## Getting Started
 
@@ -122,7 +122,7 @@ Repository skills guide the agent, but they do not replace contributor judgment,
 If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
 
 !!! important
-    Do not report suspected security vulnerabilities through a public issue, pull request, discussion, or Slack channel. Follow the [security disclosure instructions](../community/contact_us.md#security-disclosures) to arrange a private report.
+    If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
 
 ## Pull Requests & Code Reviews
 

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 Parallel](feature/vae_parallel.md)
+  [Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](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,57 +1,42 @@
-# VAE Parallel
+# VAE Patch Parallelism
 
-This document defines the VAE parallelism strategy contract and describes how to
-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.
+This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
+We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
 
 ---
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Strategy Selection Contract](#strategy-selection-contract)
-- [Tile/Patch Implementation (Decode)](#tilepatch-implementation-decode)
-- [Tile/Patch Encode Implementation](#tilepatch-encode-implementation)
+- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
+- [Encode Parallel Implementation](#encode-parallel-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 parallelism?
+### What is Vae Patch parallel?
 
-VAE parallelism distributes VAE **encoding** or **decoding** across multiple
-ranks. The implementation currently has two strategy families:
+**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
 
-| 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 |
++ Split into multiple spatial tiles
 
-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.
++ Distributed across multiple ranks
 
-This approach can:
++ Encoded/Decoded in parallel
 
-+ Distribute computation across multiple devices
-+ Reduce peak memory usage per device
-+ Accelerate encoding or decoding latency
++ Merged to reconstruct the final output
 
-The actual benefit depends on shape, topology, communication cost, and VAE
-architecture. A strategy must not become the global default without parity and
-performance evidence on its supported models.
+This approach:
+
++ Distributes computation across multiple devices
+
++ Reduces peak memory usage per device
+
++ Accelerates encoding/decoding latency
 
 ### When to Use Encode vs Decode Parallel
 
@@ -60,16 +45,8 @@ performance evidence on its supported models.
 | **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
 | **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
 
-### Architecture and ownership
-
-`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.
+### Architecture
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
 
 The executor is model-agnostic and accepts three function parameters:
 
@@ -99,11 +76,6 @@ 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.
@@ -120,34 +92,7 @@ Therefore:
 
 + Merge must perform blending to avoid seams
 
-## 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-by-Step Implementation (Decode)
 
 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -306,7 +251,7 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```
 
-## Tile/Patch Encode Implementation
+## Encode Parallel 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.
 
@@ -504,24 +449,18 @@ Complete examples in the codebase:
 
 ---
 
-## Spatially-Sharded Decode
+## Spatially-Sharded Decode (Wan)
 
-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"`.
+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"`).
 
 ### How it differs from tile parallel
 
-| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|----------------------|-----------------------------------------------|
-| 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 | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| 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 | Wan decode selection; encode remains tiled | Wan decode only |
+| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+|--------|--------------------------|-----------------------------------------------|
+| Unit of work | Independent overlapping tiles | 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) |
+| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Scope | Decode + encode | 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`.
 
@@ -537,65 +476,21 @@ 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`. 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.
+`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.
 
 ### Notes
 
-- 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.
+- 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.
 - 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 parallel support to a diffusion model:
+Adding VAE Patch Parallel support to diffusion model:
 
-1. **Implement the distributed VAE** - Inherit from the base VAE class and
-   `DistributedVaeMixin`.
-2. **Keep tile as the baseline** - For tile decode, refactor `tiled_decode` into
-   `tile_split`/`tile_exec`/`tile_merge`.
-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.
+1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
+2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
+3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
+4. **Change VAE model in pipeline** - Use the distributed version
+5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`

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 Parallel](feature/vae_parallel.md)
+- [VAE Patch Parallelism](feature/vae_parallel.md)
 
 #### Attention optimization
 

docs/design/metrics.md

@@ -1,6 +1,6 @@
 # Prometheus Metrics Design
 
-This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics.
+This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
 
 ## Objectives
 

docs/usage/faq.md

@@ -10,12 +10,12 @@ A: Refer to [GPU memory calculation and configuration](../configuration/gpu_memo
 
 > Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
 
-A: Search the current [issues](https://github.com/vllm-project/vllm-omni/issues) first. If no existing issue matches a bug or CI failure, open a [bug report](https://github.com/vllm-project/vllm-omni/issues/new?template=400-bug-report.yml) with a minimal reproducer, environment details, and relevant logs. For technical or usage questions, use the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or the [vLLM Forum](https://discuss.vllm.ai/).
+A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
 
 > Q: Does vLLM-Omni support AWQ or any other quantization?
 
-A: AWQ is available through the inherited vLLM quantization registry, and vLLM-Omni supports several additional methods. Actual support and validation vary by model, component, hardware, and checkpoint. See the [quantization overview](../user_guide/quantization/overview.md) for validated combinations and method-specific guides.
+A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).
 
 > Q: Does vLLM-Omni support multimodal streaming input and output?
 
-A: Yes, for specific models and serving paths. Current capabilities include [streaming video input](../serving/video_stream_api.md), [streaming diffusion output](../user_guide/diffusion/execution_modes.md#streaming-output), and experimental full-duplex real-time audio described in the [project overview](../README.md). Check the model and API documentation because input and output support is model-specific.
+A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!

docs/usage/metrics.md

@@ -1,23 +1,23 @@
 # Production Metrics
 
-vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server.
+vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
 
 ```bash
 vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
 curl http://localhost:8000/metrics
 ```
 
-**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm_omni:*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
+**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
 
 ## Metric Namespaces
 
 | Prefix | Source | Present when |
 |--------|--------|--------------|
-| `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
+| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
 | `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
 | `http_` / `process_` | Uvicorn / Python runtime | Always |
 
-## Pipeline-Level Metrics (`vllm_omni:`)
+## Pipeline-Level Metrics (`vllm:omni_`)
 
 Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
 
@@ -25,66 +25,42 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the
 
 | Metric | Type | Labels | Description |
 |--------|------|--------|-------------|
-| `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
-| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
-| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
+| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
+| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
 
 ### Latency
 
 | Metric | Type | Labels | Description |
 |--------|------|--------|-------------|
-| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
+| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
 
-## Audio Modality Metrics (`vllm_omni:`)
+## Audio Modality Metrics (`vllm:omni_`)
 
 Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
 
 | Metric | Type | Extra label | Description |
 |--------|------|-------------|-------------|
-| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
-| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
-| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
-| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
-| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
-| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
-| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
+| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
+| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
 
 The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
 
-## Diffusion Metrics (`vllm_omni:`)
-
-Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`.
-
-| Metric | Type | Description |
-|--------|------|-------------|
-| `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds |
-| `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) |
-| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds |
-| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing (VAE decode) time per request in seconds |
-
-`diffusion_exec_s` uses `SECONDS_BUCKETS` (50 ms → 300 s); the other three use `SECONDS_FAST_BUCKETS` (1 ms → 60 s). `diffusion_exec_per_step_s` is skipped when `num_inference_steps` is not set.
-
-### Example PromQL
-
-```promql
-# P99 DiT exec time by stage
-histogram_quantile(0.99, rate(vllm_omni:diffusion_exec_s_bucket[5m]))
-
-# Ratio of preprocess to total exec wall time
-rate(vllm_omni:diffusion_preprocess_s_sum[5m])
-  / rate(vllm_omni:diffusion_exec_s_sum[5m])
-```
-
-## Cross-Stage Transfer Metrics (`vllm_omni:`)
+## Cross-Stage Transfer Metrics (`vllm:omni_`)
 
 Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
 
 | Metric | Type | Description |
 |--------|------|-------------|
-| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
-| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
-| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
-| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
 
 ## vLLM Engine Metrics (`vllm:`)
 
@@ -107,9 +83,9 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
 
 | Metric group | Multi-stage LLM (Qwen3-Omni) |
 |---|---|
-| `vllm_omni:` request tracking + latency | With `--log-stats` |
-| `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage |
-| `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
+| `vllm:omni_` request tracking + latency | With `--log-stats` |
+| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
+| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
 | `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
 | `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
 
@@ -118,5 +94,5 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
 - All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
 - Counters use the `_total` suffix (auto-appended by `prometheus_client`).
 - Sizes use the `_bytes` suffix.
-- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm_omni:`.
-- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm_omni:audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
+- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.

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 Parallel](vae_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
+| **[VAE Parallelism](vae_parallelism.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_parallelism.md

@@ -1,4 +1,4 @@
-# VAE Parallel Guide
+# VAE Parallelism 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"`, `"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). |
+| `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). |
 
 Additional requirements:
 
@@ -128,7 +128,7 @@ Additional requirements:
 
 ## Spatially-Sharded Decode (Wan)
 
-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.
+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"`.
 
 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="auto",                # opt in to Wan request-shape selection
+        vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
     ),
 )
 ```
@@ -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 auto
+    --vae-parallel-mode spatial_shard_width
 ```
 
 **Constraints and behavior:**
 
-- 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**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
-- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
+- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` 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.
+- `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).
 
 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 Parallel](diffusion/parallelism/vae_parallel.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
+| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
 
 ### Extensions
 

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -12,7 +12,6 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
 | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
-| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
 | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -36,25 +35,6 @@ python text_to_video.py \
   --output t2v_out.mp4
 ```
 
-### LingBot-Video
-
-The shared runner recognizes both official LingBot checkpoint IDs and selects
-the native `LingBotVideoPipeline` automatically:
-
-```bash
-python text_to_video.py \
-  --model robbyant/lingbot-video-dense-1.3b \
-  --prompt "a robotic arm picks up a red block" \
-  --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
-  --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
-  --output lingbot_t2v.mp4
-```
-
-Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local
-paths without `lingbot` in the directory name, add
-`--model-class-name LingBotVideoPipeline`. Model-only options use
-`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`.
-
 ### LTX-2
 
 ```bash

examples/offline_inference/covo_audio/end2end.py

@@ -59,7 +59,7 @@ def main(args):
 
     omni = Omni(
         model=args.model_name,
-        deploy_config=args.deploy_config,
+        stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
         batch_timeout=args.batch_timeout,
@@ -167,10 +167,10 @@ def parse_args():
         help="Sampling rate for audio loading (default: 16000).",
     )
     parser.add_argument(
-        "--deploy-config",
+        "--stage-configs-path",
         type=str,
         default=None,
-        help="Path to a deploy config YAML file.",
+        help="Path to stage configs YAML file.",
     )
     parser.add_argument(
         "--log-stats",

examples/offline_inference/dynin_omni/end2end.py

@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
     parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
     parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
     parser.add_argument(
-        "--deploy-config",
+        "--stage-config-path",
         type=str,
         default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
         help="Path to deploy config yaml.",
@@ -1394,8 +1394,8 @@ def main() -> None:
 
     from vllm_omni.entrypoints.omni import Omni
 
-    deploy_config_path = str(Path(args.deploy_config).expanduser())
-    omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
+    stage_config_path = str(Path(args.stage_config_path).expanduser())
+    omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
     sampling_params_list = [
         SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
         for _ in range(omni.num_stages)

examples/offline_inference/hunyuan_image3/end2end.py

@@ -73,6 +73,7 @@ def parse_args():
     )
     parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
     parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
+    parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
     parser.add_argument(
         "--stream",
         action="store_true",
@@ -145,7 +146,13 @@ def main():
     else:
         bot_task = args.bot_task
 
-    deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
+    if args.deploy_config is not None and args.stage_configs_path is not None:
+        raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
+
+    deploy_config = args.deploy_config
+    stage_configs_path = args.stage_configs_path
+    if deploy_config is None and stage_configs_path is None:
+        deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
 
     omni_kwargs = {
         "model": args.model,
@@ -161,7 +168,10 @@ def main():
 
     if additional_config is not None:
         omni_kwargs["additional_config"] = additional_config
-    omni_kwargs["deploy_config"] = deploy_config
+    if deploy_config is not None:
+        omni_kwargs["deploy_config"] = deploy_config
+    else:
+        omni_kwargs["stage_configs_path"] = stage_configs_path
 
     omni = Omni(**omni_kwargs)
 
@@ -262,7 +272,10 @@ def main():
     print(f"  Modality: {args.modality}")
     print(f"  Prompt task: {task}")
     print(f"  Bot task: {bot_task}")
-    print(f"  Deploy config: {deploy_config}")
+    if deploy_config is not None:
+        print(f"  Deploy config: {deploy_config}")
+    else:
+        print(f"  Stage config: {stage_configs_path}")
     print(f"  Num stages: {omni.num_stages}")
     if args.modality in ("text2img", "img2img"):
         print(f"  Inference steps: {args.steps}")

examples/offline_inference/minicpmo/end2end.py

@@ -452,6 +452,12 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to a stage configs file (deprecated; prefer --deploy-config).",
+    )
     parser.add_argument(
         "--deploy-config",
         type=str,

examples/offline_inference/qwen3_omni/end2end.py

@@ -484,10 +484,10 @@ def parse_args():
         help="Path to a .txt file with one prompt per line (preferred).",
     )
     parser.add_argument(
-        "--deploy-config",
+        "--stage-configs-path",
         type=str,
         default=None,
-        help="Path to a deploy config YAML.",
+        help="Path to a stage configs file.",
     )
     parser.add_argument(
         "--video-path",

examples/offline_inference/step_audio2/end2end.py

@@ -232,8 +232,15 @@ def main(args):
         raise ValueError(f"Unknown query type: {args.query_type}")
 
     # Initialize vLLM-Omni with Step-Audio2
-    # Resolve an explicit deploy config or use the bundled task default.
-    if args.deploy_config:
+    # Resolve deploy config path. Keep explicit stage_configs_path for legacy
+    # custom configs, but use bundled deploy configs by default.
+    deploy_config_path = None
+    stage_config_path = None
+    if args.stage_configs_path and args.deploy_config:
+        raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
+    if args.stage_configs_path:
+        stage_config_path = args.stage_configs_path
+    elif args.deploy_config:
         deploy_config_path = args.deploy_config
     else:
         configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
@@ -255,7 +262,10 @@ def main(args):
         "ray_address": args.ray_address,
         "trust_remote_code": True,
     }
-    omni_kwargs["deploy_config"] = deploy_config_path
+    if stage_config_path is not None:
+        omni_kwargs["stage_configs_path"] = stage_config_path
+    else:
+        omni_kwargs["deploy_config"] = deploy_config_path
 
     omni_llm = Omni(**omni_kwargs)
 
@@ -408,6 +418,12 @@ def parse_args():
             "  - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
         ),
     )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
+    )
     parser.add_argument(
         "--deploy-config",
         type=str,

examples/offline_inference/text_to_image/text_to_image.py

@@ -92,6 +92,12 @@ def parse_args() -> argparse.Namespace:
         "meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
         "stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
     )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
+    )
     parser.add_argument(
         "--deploy-config",
         type=str,
@@ -469,6 +475,8 @@ def main():
         omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
     if args.enforce_eager is not None:
         omni_kwargs["enforce_eager"] = args.enforce_eager
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
     if args.deploy_config:
         omni_kwargs["deploy_config"] = args.deploy_config
     if use_nextstep:
@@ -507,8 +515,8 @@ def main():
     print(f"  Image size: {args.width}x{args.height}")
     if args.lora_path:
         print(f"  LoRA: scale={args.lora_scale}")
-    if args.deploy_config:
-        print(f"  deploy-config: {args.deploy_config}")
+    if args.stage_configs_path:
+        print(f"  stage-configs-path: {args.stage_configs_path}")
     print(f"{'=' * 60}\n")
 
     # Build LoRA request when --lora-path is set

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 logger = logging.getLogger(__name__)
 
 DEFAULT_MODEL = "fishaudio/s2-pro"
-DEFAULT_DEPLOY_CONFIG = os.path.join(
+DEFAULT_STAGE_CONFIG = os.path.join(
     os.path.dirname(__file__),
     "..",
     "..",
@@ -152,7 +152,10 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
-    omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
+    else:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
     return omni_kwargs
 
 
@@ -275,6 +278,12 @@ def parse_args():
         default=None,
         help="Deploy config YAML path (under vllm_omni/deploy/).",
     )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to stage configs YAML (legacy, prefer --deploy-config).",
+    )
     parser.add_argument(
         "--output-dir",
         default="output_audio",

examples/offline_inference/text_to_speech/glm_tts/end2end.py

@@ -145,7 +145,7 @@ def main(args):
 
     omni = Omni(
         model=args.model,
-        deploy_config=deploy_config_path,
+        stage_configs_path=deploy_config_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

@@ -490,10 +490,10 @@ def parse_args():
         help="Path to a .txt file with one prompt per line (preferred).",
     )
     parser.add_argument(
-        "--deploy-config",
+        "--stage-configs-path",
         type=str,
         default=None,
-        help="Path to a deploy config YAML.",
+        help="Path to a stage configs file.",
     )
     parser.add_argument(
         "--audio-path",

examples/offline_inference/text_to_video/text_to_video.md

@@ -8,7 +8,6 @@ A unified script for text-to-video generation. Supports multiple models with mod
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
-| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
 | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -58,31 +57,18 @@ python text_to_video.py \
 
 ### LingBot-Video
 
-The shared runner recognizes the official dense and MoE checkpoint IDs, selects
-`LingBotVideoPipeline`, and builds the canonical video request envelope:
+The dedicated LingBot example explicitly selects the video output modality:
 
 ```bash
-python text_to_video.py \
+python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
   --model robbyant/lingbot-video-dense-1.3b \
-  --prompt "a robotic arm picks up a red block" \
   --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
-  --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
   --output lingbot_t2v.mp4
 ```
 
 Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
 checkpoint. Requested frame counts are rounded upward to the causal
-VAE's `4n+1` grid. For a local checkpoint path whose name does not contain
-`lingbot`, pass `--model-class-name LingBotVideoPipeline`.
-
-LingBot-only options such as `batch_cfg` and `output_type` use the shared
-model-extra channel, for example:
-
-```bash
-python text_to_video.py \
-  --model robbyant/lingbot-video-dense-1.3b \
-  --extra-body '{"batch_cfg": true, "output_type": "np"}'
-```
+VAE's `4n+1` grid.
 
 ### LTX-2
 

examples/offline_inference/text_to_video/text_to_video.py

@@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
-from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
+from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
@@ -80,17 +80,6 @@ _MODEL_PRESETS = {
         "fps": 16,
         "output": "helios_output.mp4",
     },
-    "lingbot": {
-        "model_class_name": "LingBotVideoPipeline",
-        "height": 192,
-        "width": 320,
-        "num_frames": 9,
-        "num_inference_steps": 2,
-        "guidance_scale": 3.0,
-        "fps": 24,
-        "flow_shift": 3.0,
-        "output": "lingbot_video_output.mp4",
-    },
     "ltx2": {
         "height": 512,
         "width": 768,
@@ -121,51 +110,27 @@ _MODEL_PRESETS = {
 def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
     model_lower = model.lower()
     class_lower = (model_class_name or "").lower()
-    if "lingbot" in model_lower or "lingbotvideo" in class_lower:
-        return _MODEL_PRESETS["lingbot"]
-    if "ltx" in class_lower or "ltx" in model_lower:
-        if "distilled" in class_lower or "distilled" in model_lower:
-            return _MODEL_PRESETS["ltx2_distilled"]
-        if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
-            return _MODEL_PRESETS["ltx23"]
+    if "distilled" in class_lower or "distilled" in model_lower:
+        return _MODEL_PRESETS["ltx2_distilled"]
+    if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
+        return _MODEL_PRESETS["ltx23"]
+    if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
         return _MODEL_PRESETS["ltx2"]
-    if "vace" in model_lower or "vace" in class_lower:
+    if "vace" in model_lower:
         return _MODEL_PRESETS["vace"]
     # Edge must be matched before the generic cosmos branch (its "cosmos" substring would
     # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
     if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
         return _MODEL_PRESETS["cosmos3_edge"]
-    if "cosmos" in model_lower or "cosmos" in class_lower:
+    if "cosmos" in model_lower:
         return _MODEL_PRESETS["cosmos"]
-    if "hunyuan" in model_lower or "hunyuan" in class_lower:
+    if "hunyuan" in model_lower:
         return _MODEL_PRESETS["hunyuan"]
-    if "helios" in model_lower or "helios" in class_lower:
+    if "helios" in model_lower:
         return _MODEL_PRESETS["helios"]
     return _MODEL_PRESETS["wan"]
 
 
-def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
-    """Build the canonical request envelope for the shared T2V example."""
-    result: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["video"],
-    }
-    if negative_prompt is not None:
-        result["negative_prompt"] = negative_prompt
-    return result
-
-
-def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
-    """Normalize decoded frames according to the pipeline's output contract."""
-    if not tensor.is_floating_point():
-        return tensor
-    if source_range == "negative_one_to_one":
-        return tensor.clamp(-1, 1) * 0.5 + 0.5
-    if source_range == "zero_to_one":
-        return tensor.clamp(0, 1)
-    raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
-
-
 def parse_profiler_config(value: str) -> dict[str, Any]:
     try:
         config = json.loads(value)
@@ -195,7 +160,7 @@ def parse_extra_body(value: str) -> dict[str, Any]:
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Generate a video from a text prompt. "
-        "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
+        "Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
     )
     parser.add_argument(
         "--model",
@@ -443,7 +408,6 @@ def main():
     for key, default_val in preset.items():
         if getattr(args, key.replace("-", "_"), None) is None:
             setattr(args, key.replace("-", "_"), default_val)
-    model_class_name = args.model_class_name
 
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
     # Cache-dit config (Wan2.2 only)
@@ -516,7 +480,6 @@ def main():
     omni = Omni(**omni_kwargs)
     model_class_name = get_model_class_name(omni) or model_class_name
     declared_extra_body_params = get_extra_body_params(model_class_name)
-    output_tensor_range = get_output_tensor_range(model_class_name)
 
     if profiler_enabled:
         print("[Profiler] Starting profiling...")
@@ -550,11 +513,12 @@ def main():
             lora_path=lora_path,
         )
 
-    negative_prompt = args.negative_prompt
-    if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
+    prompt_dict = {"prompt": args.prompt}
+    if args.negative_prompt is not None:
+        prompt_dict["negative_prompt"] = args.negative_prompt
+    elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
         # Preserve the historical empty-prompt behavior for non-LTX examples.
-        negative_prompt = ""
-    prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)
+        prompt_dict["negative_prompt"] = ""
 
     extra_args = {}
     if lora_request:
@@ -666,7 +630,8 @@ def main():
                 frame_tensor = frame_tensor[0]
             if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
                 frame_tensor = frame_tensor.permute(1, 2, 0)
-            frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
+            if frame_tensor.is_floating_point():
+                frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
             return frame_tensor.float().numpy()
         if isinstance(frame, np.ndarray):
             frame_array = frame
@@ -718,7 +683,8 @@ def main():
                 video_tensor = video_tensor[0]
         elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
             video_tensor = video_tensor.permute(1, 2, 3, 0)
-        video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
+        if video_tensor.is_floating_point():
+            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
         video_array = video_tensor.float().numpy()
     elif isinstance(frames, np.ndarray):
         video_array = frames

examples/offline_inference/text_to_video/text_to_video_lingbot.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
+    parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
+    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
+    parser.add_argument("--negative-prompt", default=None)
+    parser.add_argument("--output", default="lingbot_video_output.mp4")
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num-frames", type=int, default=9)
+    parser.add_argument("--num-inference-steps", type=int, default=2)
+    parser.add_argument("--guidance-scale", type=float, default=3.0)
+    parser.add_argument("--flow-shift", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
+    parser.add_argument("--batch-cfg", action="store_true")
+    return parser.parse_args()
+
+
+def _peak_memory_mb(output: Any) -> float:
+    result = output[0] if isinstance(output, list) and output else output
+    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _extract_video(output: Any) -> Any:
+    if isinstance(output, list):
+        output = output[0] if output else None
+    if isinstance(output, OmniRequestOutput):
+        if output.final_output_type != "image":
+            raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
+        if not output.images:
+            raise ValueError("No video frames found in OmniRequestOutput.")
+        return output.images[0] if len(output.images) == 1 else output.images
+    return output
+
+
+def _to_frame_array(video: Any) -> np.ndarray:
+    if isinstance(video, torch.Tensor):
+        tensor = video.detach().cpu()
+        if tensor.dim() == 5:
+            tensor = tensor[0]
+        if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
+            if tensor.shape[0] in (3, 4):
+                tensor = tensor.permute(1, 2, 3, 0)
+            elif tensor.shape[1] in (3, 4):
+                tensor = tensor.permute(0, 2, 3, 1)
+        if tensor.is_floating_point():
+            if float(tensor.min()) < 0.0:
+                tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
+            else:
+                tensor = tensor.clamp(0, 1)
+        return tensor.float().numpy()
+    if isinstance(video, np.ndarray):
+        array = video[0] if video.ndim == 5 else video
+        if np.issubdtype(array.dtype, np.integer):
+            return array.astype(np.float32) / 255.0
+        return array.astype(np.float32, copy=False)
+    if isinstance(video, list):
+        return np.stack([_to_frame_array(frame) for frame in video], axis=0)
+    return np.asarray(video, dtype=np.float32)
+
+
+def main() -> None:
+    args = parse_args()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    omni = Omni(
+        model=args.model,
+        model_class_name="LingBotVideoPipeline",
+        flow_shift=args.flow_shift,
+        parallel_config=DiffusionParallelConfig(),
+    )
+
+    prompt: dict[str, Any] = {
+        "prompt": args.prompt,
+        "modalities": ["video"],
+        "num_frames": args.num_frames,
+    }
+    if args.negative_prompt is not None:
+        prompt["negative_prompt"] = args.negative_prompt
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        generator=generator,
+        output_type=args.output_type,
+        extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
+    )
+
+    start = time.perf_counter()
+    output = omni.generate(prompt, sampling_params)
+    elapsed = time.perf_counter() - start
+    print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
+    peak_mb = _peak_memory_mb(output)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
+
+    video = _to_frame_array(_extract_video(output))
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    from diffusers.utils import export_to_video
+
+    export_to_video(list(video), str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+    omni.close()
+
+
+if __name__ == "__main__":
+    main()

mkdocs.yml

@@ -73,8 +73,6 @@ plugins:
   - awesome-nav
   - redirects:
       redirect_maps:
-        community/meetups.md: community/meetings.md
-        community/volunteers.md: community/contact_us.md
         user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
         user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
         user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md

recipes/Robbyant/LingBot-Video.md

@@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends.
 - Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
 - MoE checkpoint: <https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b>
 - Upstream project: <https://github.com/Robbyant/lingbot-video>
-- Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
+- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
 - Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
 
 ## Hardware Support
@@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights.
 
 ```bash
 CUDA_VISIBLE_DEVICES=0 \
-python examples/offline_inference/text_to_video/text_to_video.py \
+python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
   --model robbyant/lingbot-video-dense-1.3b \
   --prompt "a robotic arm picks up a red block" \
   --output lingbot_t2v.mp4 \
@@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video.py \
 
 ```bash
 CUDA_VISIBLE_DEVICES=0 \
-python examples/offline_inference/text_to_video/text_to_video.py \
+python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
   --model robbyant/lingbot-video-moe-30b-a3b \
   --prompt "a robotic arm picks up a red block" \
   --output lingbot_moe_t2v.mp4 \

recipes/Tencent/HunyuanImage-3.0-Instruct.md

@@ -82,7 +82,7 @@ example and the DiT sections below.
 - Optional environment variables:
 
 ```bash
-export CACHE_DIT_VERSION=1.5.0
+export CACHE_DIT_VERSION=1.3.0
 ```
 
 HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`

requirements/common.txt

@@ -16,7 +16,7 @@ diffusers==0.38.0
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
-cache-dit==1.5.0
+cache-dit==1.3.0
 tqdm>=4.66.0
 torchsde>=0.2.6
 openai-whisper>=20250625
@@ -34,5 +34,3 @@ pydantic-settings>=2.0.0
 # base image does not ship it, so declare it explicitly. The loader also guards
 # this import so a missing gguf no longer breaks all diffusion model loading.
 gguf>=0.10.0
-# Required for running the Cosmos3 model with guardrails enabled (default behavior).
-cosmos-guardrail>=0.3.1

tests/config/test_config_factory.py

@@ -6,7 +6,7 @@ Unit tests for StageConfigFactory and related classes.
 
 import importlib
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from pathlib import Path
 from unittest.mock import patch
 
@@ -34,8 +34,9 @@ from vllm_omni.config.stage_config import (
     load_deploy_config,
     merge_pipeline_deploy,
     pipeline_cfg_resolver,
+    strip_parent_engine_args,
 )
-from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
+from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -474,7 +475,7 @@ class TestStageConfigFactory:
         cli_overrides = {
             "gpu_memory_utilization": 0.9,
             "model": "some_model",  # Internal
-            "deploy_config": "/path",  # Internal
+            "stage_configs_path": "/path",  # Internal
             "batch_timeout": 10,  # Internal
         }
 
@@ -482,7 +483,7 @@ class TestStageConfigFactory:
 
         assert overrides["gpu_memory_utilization"] == 0.9
         assert "model" not in overrides
-        assert "deploy_config" not in overrides
+        assert "stage_configs_path" not in overrides
         assert "batch_timeout" not in overrides
 
     def test_per_stage_override_excludes_internal_keys(self):
@@ -524,6 +525,49 @@ class TestStageResolutionHelpers:
         assert "model" not in overrides
         assert "parallel_config" not in overrides
 
+    def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
+        parent_fields = {f.name: f for f in fields(EngineArgs)}
+        filtered, overridden = strip_parent_engine_args(
+            {
+                "model": "some/model",
+                "stage_configs_path": "/tmp/stages.yaml",
+                "tensor_parallel_size": 4,
+                "worker_extension_cls": "some.Extension",
+                "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
+            },
+            parent_fields=parent_fields,
+            keep_keys={"worker_extension_cls"},
+            strip_keys={"stage_configs_path"},
+            no_warn_keys={"model"},
+        )
+
+        assert filtered == {
+            "worker_extension_cls": "some.Extension",
+            "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
+        }
+        assert overridden == ["tensor_parallel_size"]
+
+    def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
+        parent_fields = {f.name: f for f in fields(EngineArgs)}
+        filtered, overridden = strip_parent_engine_args(
+            {
+                "model": "some/model",
+                "stage_configs_path": "/tmp/stages.yaml",
+                "allowed_local_media_path": "/data/qwentts",
+                "allowed_media_domains": ["example.com"],
+            },
+            parent_fields=parent_fields,
+            keep_keys={"allowed_local_media_path", "allowed_media_domains"},
+            strip_keys={"stage_configs_path"},
+            no_warn_keys={"model"},
+        )
+
+        assert filtered == {
+            "allowed_local_media_path": "/data/qwentts",
+            "allowed_media_domains": ["example.com"],
+        }
+        assert overridden == []
+
 
 class TestPipelineDiscovery:
     """Tests for the central pipeline registry (``OMNI_PIPELINES``)."""

tests/config/test_omni_config.py

@@ -582,12 +582,8 @@ 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):
@@ -783,6 +779,26 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
 
 
+def test_from_pipeline_config_routes_ltx2_model_extra(tmp_path):
+    deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    extras:",
+                "      ltx2_use_channels_last_3d: true",
+            ]
+        )
+    )
+
+    stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert stage.diffusion_config.extras["ltx2_use_channels_last_3d"] is True
+
+
 def test_structured_diffusion_config_rejects_invalid_compile_granularity():
     with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
         omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,5 +1,4 @@
 import os
-from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
@@ -11,7 +10,6 @@ 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
 
@@ -19,23 +17,6 @@ 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():
@@ -351,7 +332,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())
+    wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width")
     x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
     token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
         wan_spatial_shard.SpatialShardContext(
@@ -396,14 +377,12 @@ 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())
-    with _active_spatial_shard("height"):
-        mid = mid_rank_pad(x)
+    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())
-    with _active_spatial_shard("height"):
-        last = last_rank_pad(x)
+    last = last_rank_pad(x)
     assert last.shape == (1, 1, 3, 3)
 
 
@@ -413,15 +392,13 @@ 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())
-    with _active_spatial_shard("width"):
-        mid = mid_rank_pad(x)
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="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())
-    with _active_spatial_shard("width"):
-        last = last_rank_pad(x)
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    last = last_rank_pad(x)
     assert last.shape == (1, 1, 2, 3)
 
 
@@ -444,111 +421,6 @@ 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):
@@ -694,99 +566,3 @@ 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/models/ltx2/test_ltx2_vae.py

@@ -11,6 +11,100 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+class TestLTXChannelsLast3D:
+    @pytest.mark.parametrize(
+        ("extras", "expected"),
+        [({}, False), ({"ltx2_use_channels_last_3d": True}, True)],
+    )
+    def test_channels_last_opt_in_is_ltx2_model_extra(self, extras, expected):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
+
+        assert _ltx2_use_channels_last_3d(SimpleNamespace(extras=extras)) is expected
+
+    def test_channels_last_opt_in_rejects_non_boolean_model_extra(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
+
+        with pytest.raises(TypeError, match="ltx2_use_channels_last_3d"):
+            _ltx2_use_channels_last_3d(SimpleNamespace(extras={"ltx2_use_channels_last_3d": "true"}))
+
+    @pytest.mark.parametrize(
+        ("causal", "kernel_size"),
+        [
+            (True, 3),
+            (False, 3),
+            (True, (3, 1, 1)),
+            (True, 1),
+        ],
+    )
+    def test_layout_preserving_causal_conv_matches_original(self, causal, kernel_size):
+        from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
+
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            _enable_channels_last_3d,
+            _weight_is_channels_last_3d,
+        )
+
+        torch.manual_seed(0)
+        conv = LTX2VideoCausalConv3d(4, 5, kernel_size).to(dtype=torch.float32).eval()
+        hidden_states = torch.randn(1, 4, 5, 6, 7, dtype=torch.float32)
+
+        # Pointwise Conv3d storage is valid in both contiguous formats because
+        # all spatial/temporal kernel dimensions are singleton.
+        if kernel_size != 1:
+            assert not _weight_is_channels_last_3d(conv)
+        with torch.no_grad():
+            expected = conv(hidden_states.clone(), causal=causal)
+
+        converted, patched = _enable_channels_last_3d(conv)
+
+        assert (converted, patched) == (1, 1)
+        assert _weight_is_channels_last_3d(conv)
+        with torch.no_grad():
+            actual = conv(hidden_states.clone(), causal=causal)
+
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4)
+
+    def test_temporal_pad_replicates_edges_in_channels_last_3d(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            _temporal_pad_channels_last_3d,
+        )
+
+        torch.manual_seed(0)
+        hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
+
+        actual = _temporal_pad_channels_last_3d(hidden_states, left=2, right=1)
+        expected = torch.cat(
+            [
+                hidden_states[:, :, :1].repeat(1, 1, 2, 1, 1),
+                hidden_states,
+                hidden_states[:, :, -1:],
+            ],
+            dim=2,
+        )
+
+        assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    def test_contiguous_weight_keeps_original_padding_path(self, monkeypatch):
+        from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
+
+        import vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 as ltx2_vae
+
+        conv = LTX2VideoCausalConv3d(4, 5, 3).to(dtype=torch.float32).eval()
+        hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
+
+        def fail_channels_last_pad(*_args, **_kwargs):
+            raise AssertionError("contiguous weights must retain the original padding path")
+
+        monkeypatch.setattr(ltx2_vae, "_temporal_pad_channels_last_3d", fail_channels_last_pad)
+
+        with torch.no_grad():
+            expected = conv(hidden_states, causal=True)
+            actual = ltx2_vae._layout_preserving_causal_conv_forward(conv, hidden_states, causal=True)
+
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
 class TestLTXOutputRank:
     @pytest.mark.parametrize(
         ("distributed_vae_state", "expected_decode_calls"),

tests/diffusion/quantization/test_quantization_fp8.py

@@ -189,7 +189,7 @@ def _generate_single_stage_video(
 
 
 def _generate_bagel_image(
-    diffusion_quantization_config: str | None = None,
+    quantization_config: str | None = None,
     num_inference_steps: int = 15,
 ) -> tuple[Any, float]:
     """Generate an image with BAGEL (multi-stage: LLM + Diffusion).
@@ -199,11 +199,11 @@ def _generate_bagel_image(
     config_path = get_deploy_config_path("ci/bagel.yaml")
     omni_kwargs: dict[str, Any] = {
         "model": "ByteDance-Seed/BAGEL-7B-MoT",
-        "deploy_config": config_path,
+        "stage_configs_path": config_path,
         "stage_init_timeout": 300,
     }
-    if diffusion_quantization_config:
-        omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config
+    if quantization_config:
+        omni_kwargs["quantization_config"] = quantization_config
 
     model_name = omni_kwargs.pop("model")
     with OmniRunner(model_name, **omni_kwargs) as runner:
@@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory():
 
 @hardware_test(res={"cuda": "H100"})
 def test_bagel_fp8_generates_image():
-    """BAGEL with diffusion-stage FP8 generates a valid image.
+    """BAGEL with FP8 quantization_config generates a valid image.
 
     FP8 should only apply to the diffusion stage (Stage-1), not the
     LLM stage (Stage-0). We verify this by checking:
       1. Image is generated successfully
       2. LLM stage finish_reason is 'stop' (not 'length' from garbled output)
     """
-    image, _ = _generate_bagel_image(diffusion_quantization_config="fp8")
+    image, _ = _generate_bagel_image(quantization_config="fp8")
     image.save("test_bagel_fp8.png")
 
 
 @hardware_test(res={"cuda": "H100"})
 def test_bagel_bf16_generates_image():
     """BAGEL without quantization generates a valid image (baseline)."""
-    image, _ = _generate_bagel_image(diffusion_quantization_config=None)
+    image, _ = _generate_bagel_image(quantization_config=None)
     image.save("test_bagel_bf16.png")
 
 

tests/diffusion/test_diffusion_output_formatter.py

@@ -95,7 +95,6 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
         "resolution": 512,
         "width": None,
         "height": None,
-        "num_inference_steps": 1,
     }
 
 

tests/distributed/omni_connectors/test_bagel_mooncake_connector.py

@@ -314,7 +314,7 @@ def test_bagel_text2img_mooncake_connector(run_level):
         temp_config_file = _resolve_deploy_config(temp_config_file, run_level)
         with OmniRunner(
             "ByteDance-Seed/BAGEL-7B-MoT",
-            deploy_config=temp_config_file,
+            stage_configs_path=temp_config_file,
             stage_init_timeout=300,
         ) as runner:
             generated_image = _generate_bagel_image(runner.omni)

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -251,7 +251,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
     config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
     with OmniRunner(
         "ByteDance-Seed/BAGEL-7B-MoT",
-        deploy_config=config_path,
+        stage_configs_path=config_path,
     ) as runner:
         generated_image = _generate_bagel_img2img(runner.omni, input_image)
         if run_level == "advanced_model":
@@ -266,7 +266,7 @@ def test_bagel_text2img_shared_memory_connector(run_level):
     config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
     with OmniRunner(
         "ByteDance-Seed/BAGEL-7B-MoT",
-        deploy_config=config_path,
+        stage_configs_path=config_path,
     ) as runner:
         generated_image = _generate_bagel_text2img(runner.omni)
         if run_level == "advanced_model":

tests/e2e/offline_inference/test_audiox_model_expansion.py

@@ -17,7 +17,7 @@ from vllm_omni.platforms import current_omni_platform
 # Override in CI if needed: AUDIOX_TEST_MODEL=<model-or-local-path>
 models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]
 
-# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
 _OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models]
 
 pytestmark = [

tests/e2e/offline_inference/test_bagel_expansion.py

@@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
 BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml")
 DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
 
-# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
+# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
 _OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG)
 
 

tests/e2e/offline_inference/test_cosyvoice3_expansion.py

@@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
     ]
 
 
-# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
+# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
 _cosy_deployment = get_deploy_config_path("cosyvoice3.yaml")
 _cosy_model_path = str(_resolve_model_dir())
 _OMNI_RUNNER_PARAMS = [

tests/e2e/offline_inference/test_ernie_image_expansion.py

@@ -17,7 +17,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 MODEL = "baidu/ERNIE-Image"
 
-# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
 _OMNI_RUNNER_PARAM = (
     MODEL,
     None,

tests/e2e/offline_inference/test_glm_tts_expansion.py

@@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None:
 
     with OmniRunner(
         MODEL,
-        deploy_config=_get_deploy_config(async_chunk=async_chunk),
+        stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
         stage_init_timeout=600,
         trust_remote_code=True,
     ) as omni_runner:

tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py

@@ -52,7 +52,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 def omni() -> Generator[Omni, None, None]:
     with OmniRunner(
         MODEL_NAME,
-        deploy_config=str(STAGE_CONFIG_PATH),
+        stage_configs_path=str(STAGE_CONFIG_PATH),
         trust_remote_code=True,
     ) as runner:
         yield runner.omni

tests/e2e/offline_inference/test_longcat_video_avatar.py

@@ -84,6 +84,7 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]:
     with OmniRunner(
         model,
         seed=42,
+        stage_configs_path=None,
         model_class_name="LongCatVideoAvatarPipeline",
         additional_config={
             "model_type": "avatar-v1.5",

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

@@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16
 MODEL_PATH = "bytedance-research/MammothModa2-Preview"
 T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
 
-_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
+_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
 
 # Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
 _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"

tests/e2e/offline_inference/test_minicpmo_4_5.py

@@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"]
 _CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
 
 
-test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models]
+test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models]
 
 
 def get_question(prompt_type: str = "text") -> str:

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

@@ -25,7 +25,7 @@ from vllm_omni import Omni
 MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
 STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml")
 
-# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
 _OMNI_RUNNER_PARAM = (
     MODEL_NAME,
     STAGE_CONFIG,

tests/e2e/offline_inference/test_omnivoice_expansion.py

@@ -20,11 +20,13 @@ from tests.helpers.stage_config import get_deploy_config_path
 MODEL = "k2-fsa/OmniVoice"
 DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
 
-# OmniRunner tuple: model, deploy config path, extra Omni kwargs.
+# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
+# The migrated test passes deploy_config through extra Omni kwargs.
 _OMNI_RUNNER_PARAM = (
     MODEL,
-    DEPLOY_CONFIG,
+    None,
     {
+        "deploy_config": DEPLOY_CONFIG,
         "trust_remote_code": True,
         "log_stats": True,
     },

tests/e2e/offline_inference/test_sensenova_u1_expansion.py

@@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
     """Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
     lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
 
-    with OmniRunner(MODEL) as runner:
+    with OmniRunner(MODEL, stage_configs_path=None) as runner:
         omni = runner.omni
         baseline = _generate(omni)
         img_1x = _generate(omni, lora_request, lora_scale=1.0)

tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py

@@ -164,7 +164,10 @@ def _generate_sensenova_u1_img2img(
 def test_sensenova_u1_img2img(run_level):
     """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
     input_image = _load_input_image()
-    with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
+    with OmniRunner(
+        "SenseNova/SenseNova-U1-8B-MoT",
+        stage_configs_path=None,
+    ) as runner:
         generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
         if run_level == "full_model" or run_level == "advanced_model":
             _validate_pixels(generated_image)
@@ -179,6 +182,7 @@ def test_sensenova_u1_img2img_cache_dit():
     input_image = _load_input_image()
     with OmniRunner(
         "SenseNova/SenseNova-U1-8B-MoT",
+        stage_configs_path=None,
         cache_backend="cache_dit",
     ) as runner:
         _generate_sensenova_u1_img2img(runner.omni, input_image)

tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py

@@ -151,7 +151,10 @@ def _generate_sensenova_u1_image(
 @hardware_test(res={"cuda": "H100"})
 def test_sensenova_u1_text2img(run_level):
     """Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML)."""
-    with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
+    with OmniRunner(
+        "SenseNova/SenseNova-U1-8B-MoT",
+        stage_configs_path=None,
+    ) as runner:
         generated_image = _generate_sensenova_u1_image(runner.omni)
         if run_level == "full_model" or run_level == "advanced_model":
             _validate_pixels(generated_image)
@@ -165,6 +168,7 @@ def test_sensenova_u1_text2img_cache_dit():
     """Test SenseNova-U1 text2img with Cache-DiT enabled."""
     with OmniRunner(
         "SenseNova/SenseNova-U1-8B-MoT",
+        stage_configs_path=None,
         cache_backend="cache_dit",
     ) as runner:
         _generate_sensenova_u1_image(runner.omni)

tests/e2e/offline_inference/test_soulxsinger_expansion.py

@@ -159,7 +159,7 @@ def test_soulxsinger_multistage_from_audio(
 
     with OmniRunner(
         model,
-        deploy_config=get_deploy_config_path(deploy_yaml),
+        stage_configs_path=get_deploy_config_path(deploy_yaml),
         async_chunk=False,
     ) as runner:
         sampling = OmniDiffusionSamplingParams(
@@ -212,7 +212,7 @@ def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) ->
 
     with OmniRunner(
         str(base_dir),
-        deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"),
+        stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
         async_chunk=False,
     ) as runner:
         sampling = OmniDiffusionSamplingParams(

tests/e2e/offline_inference/test_step_audio2_expansion.py

@@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput
 
 MODEL = "stepfun-ai/Step-Audio-2-mini"
 DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml")
-TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})]
+TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})]
 
 pytestmark = [pytest.mark.slow, pytest.mark.tts]
 

tests/e2e/offline_inference/test_voxtral_tts_expansion.py

@@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming():
     async def _run():
         async_omni = AsyncOmni(
             model=MODEL,
-            deploy_config=STAGE_CONFIG,
+            stage_configs_path=STAGE_CONFIG,
             stage_init_timeout=300,
             enforce_eager=True,
         )

tests/engine/test_arg_utils.py

@@ -4,6 +4,7 @@ invariant to the specific attributes of vLLM config except in cases where we
 explicitly patch values that differ from vLLM.
 """
 
+import argparse
 import inspect
 import json
 import os
@@ -19,6 +20,7 @@ from vllm.engine.arg_utils import EngineArgs
 
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -147,6 +149,19 @@ def test_qwen3_tts_startup_task_type_is_validated():
         )
 
 
+def test_from_cli_args_picks_up_stage_configs_path():
+    """from_cli_args should pick up stage_configs_path from namespace."""
+    ns = argparse.Namespace(
+        model="facebook/opt-125m",
+        stage_configs_path="/some/path.yaml",
+        custom_pipeline_args=None,
+    )
+
+    args = OmniEngineArgs.from_cli_args(ns)
+    assert args.stage_configs_path == "/some/path.yaml"
+    assert args.custom_pipeline_args is None
+
+
 def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
     """Ensure Code2Wav mirrors stage max_model_len into nested HF overrides.
 
@@ -403,6 +418,130 @@ def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
     assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size
 
 
+def test_stage_configs_path_field():
+    """OmniEngineArgs with stage_configs_path should construct without error."""
+    args = OmniEngineArgs(stage_configs_path="/some/path.yaml")
+    assert args.stage_configs_path == "/some/path.yaml"
+
+
+def test_strip_single_engine_args():
+    """_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
+    kwargs = {
+        # Parent EngineArgs fields — stripped unless explicitly allowlisted
+        "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
+        "tensor_parallel_size": 4,
+        "gpu_memory_utilization": 0.9,
+        "model": "some/model",
+        # Parent field that should be kept (allowlisted)
+        "worker_extension_cls": "some.Extension",
+        # OmniEngineArgs-only / non-engine fields — should pass through
+        "stage_configs_path": "/path/to/yaml",
+        "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
+        "mode": "text-to-image",
+        "lora_path": "/some/lora",
+    }
+
+    filtered = AsyncOmniEngine._strip_single_engine_args(kwargs)
+
+    # Stripped — parent EngineArgs fields
+    assert "compilation_config" not in filtered
+    assert filtered["tensor_parallel_size"] == 4
+    assert "gpu_memory_utilization" not in filtered
+    assert "model" not in filtered
+
+    # Stripped — orchestrator-level OmniEngineArgs field
+    assert "stage_configs_path" not in filtered
+
+    # Kept
+    assert filtered["worker_extension_cls"] == "some.Extension"
+    assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"}
+    assert filtered["mode"] == "text-to-image"
+    assert filtered["lora_path"] == "/some/lora"
+
+
+def test_strip_single_engine_args_model_does_not_trigger_warning(mocker):
+    """model is always in kwargs (callers set it via from_cli_args/asdict),
+    so it should not cause the override warning by itself or appear in it."""
+    mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning")
+
+    # Typical caller kwargs: model is always present, no other parent
+    # EngineArgs fields are explicitly overridden.
+    AsyncOmniEngine._strip_single_engine_args(
+        {
+            "model": "some/model",
+            "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
+        }
+    )
+    mock_warn.assert_not_called()
+
+    # When there *are* genuinely surprising overrides alongside model,
+    # the warning should mention them but not model. Keep-listed fields such as
+    # tensor_parallel_size are intentionally passed through and should not warn.
+    AsyncOmniEngine._strip_single_engine_args(
+        {
+            "model": "some/model",
+            "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
+            "tensor_parallel_size": 4,
+            "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
+        }
+    )
+    mock_warn.assert_called_once()
+    warned_args = mock_warn.call_args[0][-1]  # the formatted arg list
+    assert "compilation_config" in warned_args
+    assert "tensor_parallel_size" not in warned_args
+    assert "model" not in warned_args
+
+
+def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path):
+    """New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs.
+
+    Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping
+    those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni
+    accuracy from ~78% to ~70%.
+    """
+    from vllm_omni.entrypoints.utils import is_new_format_deploy_config
+
+    yaml_path = tmp_path / "minicpm_deploy.yaml"
+    yaml_path.write_text(
+        """
+pipeline: minicpmo_4_5
+async_chunk: true
+stages:
+  - stage_id: 0
+    devices: "0"
+""".strip(),
+        encoding="utf-8",
+    )
+    path = str(yaml_path)
+    assert is_new_format_deploy_config(path)
+
+    kwargs = {
+        "stage_configs_path": path,
+        "trust_remote_code": True,
+        "interleave_mm_strings": True,
+        "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
+        "gpu_memory_utilization": 0.55,
+    }
+    # Mirror AsyncOmniEngine._resolve_stage_configs strip decision.
+    if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]):
+        base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs))
+    else:
+        base = dict(kwargs)
+    assert base.get("interleave_mm_strings") is True
+    assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}}
+    # Parent EngineArgs fields also remain available for deploy merge.
+    assert base.get("gpu_memory_utilization") == 0.55
+
+    legacy = tmp_path / "legacy.yaml"
+    legacy.write_text("stage_args:\n  - stage_id: 0\n", encoding="utf-8")
+    assert not is_new_format_deploy_config(str(legacy))
+    legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy))
+    stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs)
+    assert "interleave_mm_strings" not in stripped
+    assert "media_io_kwargs" not in stripped
+    assert "gpu_memory_utilization" not in stripped
+
+
 # For https://github.com/vllm-project/vllm-omni/issues/3293
 def test_tensor_parallel_size_none_is_handled():
     """Ensure the tensor parallel size of None isn't forwarded."""

tests/engine/test_async_omni_engine_abort.py

@@ -68,7 +68,7 @@ async def test_abort():
     with ExitStack() as after:
         engine = AsyncOmni(
             model=model,
-            deploy_config=stage_config,
+            stage_configs_path=stage_config,
         )
         after.callback(engine.shutdown)
 

tests/engine/test_single_stage_mode.py

@@ -428,7 +428,7 @@ class TestSingleStageModeDetection:
         assert engine.single_stage_mode is False
         assert engine._single_stage_id_filter is None
 
-    def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture):
+    def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture):
         duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
         get_pipeline_config = mocker.patch(
             "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
@@ -441,7 +441,7 @@ class TestSingleStageModeDetection:
 
         engine = self._make_engine_no_thread(
             mocker,
-            deploy_config="/fake/duplex.yaml",
+            stage_configs_path="/fake/duplex.yaml",
         )
 
         get_pipeline_config.assert_called_once_with(

tests/entrypoints/openai_api/test_serving_speech.py

@@ -680,7 +680,7 @@ class TestSpeechAPI:
 
     @pytest.mark.asyncio
     async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture):
-        """Test public diffusion speech success and extra_params propagation."""
+        """Test that extra_params are correctly applied to sampling_params_list in diffusion mode."""
         # Mock the engine client
         mock_engine = mocker.MagicMock()
 
@@ -704,11 +704,7 @@ class TestSpeechAPI:
 
         req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"})
 
-        response = await server.create_speech(req)
-
-        assert response.status_code == 200
-        assert response.media_type == "audio/wav"
-        assert response.body == b"dummy"
+        await server._create_diffusion_speech(req)
 
         # Verify generate was called
         mock_engine.generate.assert_called_once()
@@ -741,19 +737,6 @@ class TestTTSMethods:
         yield server
         server.shutdown()
 
-    def test_diffusion_audio_encode_speed_skips_ar_adapter(self, mocker: MockerFixture):
-        """Pure-diffusion speech keeps generic speed handling outside AR adapters."""
-        server = OmniOpenAIServingSpeech.for_diffusion(
-            diffusion_engine=mocker.MagicMock(),
-            model_name="test-model",
-        )
-        resolve_adapter = mocker.patch.object(serving_speech_module, "resolve_adapter")
-
-        speed = server._audio_encode_speed(OpenAICreateSpeechRequest(input="Hello", speed=1.25))
-
-        assert speed == 1.25
-        resolve_adapter.assert_not_called()
-
     def test_is_tts_detection_no_stage(self, speech_server):
         """Test TTS model detection when no TTS stage exists."""
         # Fixture creates server with stage_configs = [] -> _is_tts should be False

tests/entrypoints/test_async_omni.py

@@ -364,7 +364,7 @@ async def test_omni_generate_request_id():
     Text modality only; asserts caller-visible ids are preserved across the
     multi-stage orchestrator path on H100.
     """
-    engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG)
+    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
     try:
         for req_id in _OMNI_REQ_IDS:
             async for output in engine.generate(

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -20,6 +20,7 @@ def test_default_stage_config_includes_cache_backend():
             "cache_backend": "cache_dit",
             "cache_config": '{"Fn_compute_blocks": 2}',
             "vae_use_slicing": True,
+            "extras": {"ltx2_use_channels_last_3d": True},
             "ulysses_degree": 2,
         }
     )[0]
@@ -29,6 +30,7 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["cache_backend"] == "cache_dit"
     assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
     assert engine_args["vae_use_slicing"] is True
+    assert engine_args["extras"]["ltx2_use_channels_last_3d"] is True
     assert engine_args["parallel_config"].ulysses_degree == 2
     assert engine_args["model_stage"] == "diffusion"
 
@@ -478,11 +480,12 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+    engine._strip_single_engine_args = lambda kwargs: kwargs
 
     _, stage_configs = engine._resolve_stage_configs(
         "dummy-model",
         {
-            "deploy_config": "dummy.yaml",
+            "stage_configs_path": "dummy.yaml",
             "additional_config": {"torchair_graph_config": {"enabled": True}},
         },
         trust_remote_code=False,
@@ -515,11 +518,12 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+    engine._strip_single_engine_args = lambda kwargs: kwargs
 
     _, stage_configs = engine._resolve_stage_configs(
         "dummy-model",
         {
-            "deploy_config": "dummy.yaml",
+            "stage_configs_path": "dummy.yaml",
             "quantization_config": {"method": "bitsandbytes"},
         },
         trust_remote_code=False,

tests/entrypoints/test_serve.py

@@ -76,18 +76,6 @@ 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",
@@ -205,7 +193,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
     with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
         run_headless(args)
 
-    assert len(captured["args"]) == 2
+    assert captured["args"][1] is None
     assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
     assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
     assert captured["strategy_config_path"] == "/tmp/strategy.yaml"

tests/entrypoints/test_utils.py

@@ -381,6 +381,7 @@ class TestLoadAndResolveStageConfigs:
 
         config_path, stage_configs, _ = load_and_resolve_stage_configs(
             model="black-forest-labs/FLUX.2-klein-4B",
+            stage_configs_path=None,
             kwargs=kwargs,
             trust_remote_code=False,
             default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
@@ -389,7 +390,9 @@ class TestLoadAndResolveStageConfigs:
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
 
-    def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
+    def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
+        self, tmp_path, mocker: MockerFixture
+    ):
         deploy_path = tmp_path / "qwen3_multi.yaml"
         deploy_path.write_text(
             'stages:\n  - stage_id: 0\n    devices: "0"\n  - stage_id: 1\n    devices: "1,2,3"\n    num_replicas: 3\n',
@@ -410,23 +413,18 @@ class TestLoadAndResolveStageConfigs:
             "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
             return_value=(returned_stage_configs, None),
         )
-        cli_overrides = {
-            "interleave_mm_strings": True,
-            "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
-            "gpu_memory_utilization": 0.55,
-        }
 
         config_path, stage_configs, _ = load_and_resolve_stage_configs(
             model="dummy-model",
-            kwargs=cli_overrides,
+            stage_configs_path=str(deploy_path),
+            kwargs={},
             trust_remote_code=True,
-            deploy_config_path=str(deploy_path),
         )
 
         load_stage_configs.assert_called_once_with(
             "dummy-model",
             trust_remote_code=True,
-            base_engine_args=cli_overrides,
+            base_engine_args={},
             deploy_config_path=str(deploy_path),
             stage_overrides=None,
             strategy_config_path=None,

tests/examples/offline_inference/test_text_to_video_prompts.py

@@ -1,91 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from examples.offline_inference.text_to_video.text_to_video import (
-    _MODEL_PRESETS,
-    _detect_preset,
-    _normalize_float_tensor,
-    build_text_to_video_prompt,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.mark.parametrize(
-    ("negative_prompt", "expected_negative_prompt"),
-    [(None, None), ("", ""), ("blurry", "blurry")],
-)
-def test_text_to_video_builds_canonical_prompt(
-    negative_prompt: str | None,
-    expected_negative_prompt: str | None,
-) -> None:
-    result = build_text_to_video_prompt("a robot waves", negative_prompt)
-
-    assert result["prompt"] == "a robot waves"
-    assert result["modalities"] == ["video"]
-    if expected_negative_prompt is None:
-        assert "negative_prompt" not in result
-    else:
-        assert result["negative_prompt"] == expected_negative_prompt
-
-
-@pytest.mark.parametrize(
-    ("model", "model_class_name", "preset_name"),
-    [
-        ("robbyant/lingbot-video-dense-1.3b", None, "lingbot"),
-        ("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"),
-        ("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"),
-        ("BestWishYsh/Helios-Distilled", None, "helios"),
-        ("/models/cosmos-edge", "Cosmos3EdgeVFMTransformer", "cosmos3_edge"),
-        ("/models/vace", "WanVACEPipeline", "vace"),
-    ],
-)
-def test_detect_preset_is_scoped_to_model_family(
-    model: str,
-    model_class_name: str | None,
-    preset_name: str,
-) -> None:
-    assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name]
-
-
-def test_lingbot_preset_matches_lingbot_defaults() -> None:
-    assert _MODEL_PRESETS["lingbot"] == {
-        "model_class_name": "LingBotVideoPipeline",
-        "height": 192,
-        "width": 320,
-        "num_frames": 9,
-        "num_inference_steps": 2,
-        "guidance_scale": 3.0,
-        "fps": 24,
-        "flow_shift": 3.0,
-        "output": "lingbot_video_output.mp4",
-    }
-
-
-def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None:
-    frames = torch.tensor([0.0, 0.25, 1.0, 1.5])
-
-    assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0]))
-
-
-def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None:
-    frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5])
-
-    assert torch.equal(
-        _normalize_float_tensor(frames, "negative_one_to_one"),
-        torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]),
-    )
-
-
-def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None:
-    frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])]
-
-    normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames]
-
-    assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0]))
-    assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0]))

tests/helpers/runtime.py

@@ -2635,7 +2635,7 @@ class OmniRunner:
         # affects the test runner wrapper.
         init_timeout: int = 1800,
         log_stats: bool = False,
-        deploy_config: str | None = None,
+        stage_configs_path: str | None = None,
         **kwargs,
     ) -> None:
         startup_t0 = time.perf_counter()
@@ -2653,7 +2653,7 @@ class OmniRunner:
             stage_init_timeout=stage_init_timeout,
             batch_timeout=batch_timeout,
             init_timeout=init_timeout,
-            deploy_config=deploy_config,
+            stage_configs_path=stage_configs_path,
             **kwargs,
         )
         startup_s = time.perf_counter() - startup_t0
@@ -3179,7 +3179,12 @@ def iter_omni_server(
                 print("OmniServer stopping...")
         else:
             if stage_config_path is not None:
-                server_args += ["--deploy-config", stage_config_path]
+                from vllm_omni.entrypoints.utils import is_new_format_deploy_config
+
+                if is_new_format_deploy_config(stage_config_path):
+                    server_args += ["--deploy-config", stage_config_path]
+                else:
+                    server_args += ["--stage-configs-path", stage_config_path]
 
             with (
                 OmniServer(
@@ -3232,7 +3237,7 @@ def iter_omni_runner(
         model = model_prefix + model
         if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
             model = resolve_tiny_model_path(model)
-        with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner:
+        with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
             print("OmniRunner started successfully")
             yield runner
             print("OmniRunner stopping...")

tests/metrics/test_modality.py

@@ -42,10 +42,6 @@ _EXPECTED_FAMILIES = [
     defs.AUDIO_UNDERRUN_S,
     defs.AUDIO_CONTINUITY_OK_METRIC,
     defs.AUDIO_SKIPPED_REQUESTS_METRIC,
-    defs.DIFFUSION_EXEC_S,
-    defs.DIFFUSION_EXEC_PER_STEP_S,
-    defs.DIFFUSION_PREPROCESS_S,
-    defs.DIFFUSION_POSTPROCESS_S,
 ]
 
 
@@ -59,10 +55,6 @@ class TestRegistration:
         mod.observe_audio_underrun("s", "r", 0.01)
         mod.inc_audio_continuity_ok("s", "r", 100)
         mod.inc_audio_skipped("s", "r", "malformed_codec")
-        mod.observe_diffusion_exec("s", "r", 0.5)
-        mod.observe_diffusion_exec_per_step("s", "r", 0.01)
-        mod.observe_diffusion_preprocess("s", "r", 0.01)
-        mod.observe_diffusion_postprocess("s", "r", 0.2)
 
         out = generate_latest(REGISTRY).decode()
         for name in _EXPECTED_FAMILIES:
@@ -164,50 +156,6 @@ class TestAudio:
         assert unknown == 1.0
 
 
-# ---------------------------------------------------------------------------
-# Diffusion observe API
-# ---------------------------------------------------------------------------
-
-
-class TestDiffusion:
-    def test_diffusion_exec_observed(self, mod: OmniModalityMetrics) -> None:
-        stage, replica = "dit_exec", "0"
-        mod.observe_diffusion_exec(stage, replica, 1.25)
-        out = generate_latest(REGISTRY).decode()
-        prefix = f'{defs.DIFFUSION_EXEC_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
-        assert _sample_value(out, prefix) == 1.25
-
-    def test_diffusion_preprocess_observed(self, mod: OmniModalityMetrics) -> None:
-        stage, replica = "dit_pre", "0"
-        mod.observe_diffusion_preprocess(stage, replica, 0.03)
-        out = generate_latest(REGISTRY).decode()
-        prefix = f'{defs.DIFFUSION_PREPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
-        assert _sample_value(out, prefix) == pytest.approx(0.03)
-
-    def test_diffusion_postprocess_observed(self, mod: OmniModalityMetrics) -> None:
-        stage, replica = "dit_post", "0"
-        mod.observe_diffusion_postprocess(stage, replica, 0.8)
-        out = generate_latest(REGISTRY).decode()
-        prefix = f'{defs.DIFFUSION_POSTPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
-        assert _sample_value(out, prefix) == 0.8
-
-    def test_diffusion_exec_per_step_observed(self, mod: OmniModalityMetrics) -> None:
-        stage, replica = "dit_per_step", "0"
-        mod.observe_diffusion_exec_per_step(stage, replica, 0.078)
-        out = generate_latest(REGISTRY).decode()
-        prefix = f'{defs.DIFFUSION_EXEC_PER_STEP_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
-        assert _sample_value(out, prefix) == pytest.approx(0.078)
-
-    def test_log_stats_false_skips(self) -> None:
-        silent = OmniModalityMetrics(model_name="silent-model", log_stats=False)
-        silent.observe_diffusion_exec("s", "r", 1.0)
-        silent.observe_diffusion_exec_per_step("s", "r", 0.02)
-        silent.observe_diffusion_preprocess("s", "r", 0.01)
-        silent.observe_diffusion_postprocess("s", "r", 0.5)
-        out = generate_latest(REGISTRY).decode()
-        assert 'model_name="silent-model"' not in out
-
-
 # ---------------------------------------------------------------------------
 # Stub-driven routing tests for the finalize / streaming helpers.
 # ---------------------------------------------------------------------------
@@ -235,18 +183,6 @@ class _StubModMetrics:
     def inc_audio_skipped(self, s, r, reason):
         self.calls.append(("inc_audio_skipped", s, r, reason))
 
-    def observe_diffusion_exec(self, s, r, v):
-        self.calls.append(("observe_diffusion_exec", s, r, v))
-
-    def observe_diffusion_exec_per_step(self, s, r, v):
-        self.calls.append(("observe_diffusion_exec_per_step", s, r, v))
-
-    def observe_diffusion_preprocess(self, s, r, v):
-        self.calls.append(("observe_diffusion_preprocess", s, r, v))
-
-    def observe_diffusion_postprocess(self, s, r, v):
-        self.calls.append(("observe_diffusion_postprocess", s, r, v))
-
 
 class _Bag:
     def __init__(self, **kw):
@@ -303,88 +239,8 @@ class TestObserveModalityAtFinalize:
         )
         assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
 
-    def test_diffusion_path_full(self):
-        stub = _StubModMetrics()
-        stage_metrics = _Bag(
-            stage_gen_time_ms=1000.0,
-            diffusion_metrics={
-                "preprocess_time_s": 0.05,
-                "diffusion_engine_exec_time_s": 1.2,
-                "postprocess_time_s": 0.3,
-                "image_num": 1,
-                "resolution": 1024,
-                "num_inference_steps": 50,
-            },
-        )
-        observe_modality_at_finalize(
-            stub,
-            output_type="image",
-            stage_id=2,
-            replica_id=0,
-            stage_metrics=stage_metrics,
-            engine_outputs=_Bag(),
-        )
-        assert ("observe_diffusion_preprocess", "2", "0", 0.05) in stub.calls
-        assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
-        assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls
-        assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls
-
-    def test_diffusion_per_step_skipped_without_num_inference_steps(self):
-        stub = _StubModMetrics()
-        stage_metrics = _Bag(
-            stage_gen_time_ms=1000.0,
-            diffusion_metrics={
-                "diffusion_engine_exec_time_s": 1.2,
-                "preprocess_time_s": 0.05,
-                "postprocess_time_s": 0.3,
-            },
-        )
-        observe_modality_at_finalize(
-            stub,
-            output_type="image",
-            stage_id=2,
-            replica_id=0,
-            stage_metrics=stage_metrics,
-            engine_outputs=_Bag(),
-        )
-        assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
-        assert not any(c[0] == "observe_diffusion_exec_per_step" for c in stub.calls)
-
-    def test_diffusion_metrics_none_skipped(self):
-        stub = _StubModMetrics()
-        observe_modality_at_finalize(
-            stub,
-            output_type="image",
-            stage_id=2,
-            replica_id=0,
-            stage_metrics=_Bag(stage_gen_time_ms=500.0, diffusion_metrics=None),
-            engine_outputs=_Bag(),
-        )
-        assert not any(c[0].startswith("observe_diffusion") for c in stub.calls)
-
-    def test_diffusion_zero_values_observed(self):
-        stub = _StubModMetrics()
-        stage_metrics = _Bag(
-            stage_gen_time_ms=500.0,
-            diffusion_metrics={
-                "preprocess_time_s": 0.0,
-                "diffusion_engine_exec_time_s": 0.0,
-                "postprocess_time_s": 0.0,
-            },
-        )
-        observe_modality_at_finalize(
-            stub,
-            output_type="image",
-            stage_id=2,
-            replica_id=0,
-            stage_metrics=stage_metrics,
-            engine_outputs=_Bag(),
-        )
-        assert ("observe_diffusion_preprocess", "2", "0", 0.0) in stub.calls
-        assert ("observe_diffusion_exec", "2", "0", 0.0) in stub.calls
-        assert ("observe_diffusion_postprocess", "2", "0", 0.0) in stub.calls
-
-    def test_non_audio_output_type_skips_audio_but_still_checks_diffusion(self):
+    def test_non_audio_output_type_skipped(self):
+        # Image / video / text output types are out of scope for this module.
         stub = _StubModMetrics()
         for output_type in ("text", "image", "video"):
             observe_modality_at_finalize(
@@ -395,7 +251,7 @@ class TestObserveModalityAtFinalize:
                 stage_metrics=_Bag(stage_gen_time_ms=100.0),
                 engine_outputs=_Bag(),
             )
-        assert not any(c[0].startswith("observe_audio") or c[0].startswith("inc_audio") for c in stub.calls)
+        assert stub.calls == []
 
     def test_replica_id_none_skipped(self):
         stub = _StubModMetrics()

tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py

@@ -1,24 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from functools import lru_cache
 from types import SimpleNamespace
 from unittest.mock import Mock
 
 import pytest
 import torch
 import torch.nn as nn
-from vllm.platforms import current_platform
 
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
     HiFTGenerator,
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
-    CFMGraphWrapper,
     HiFTGraphWrapper,
 )
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+pytestmark = [pytest.mark.core_model]
 
 
 class _F0Predictor(nn.Module):
@@ -145,248 +142,3 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch)
     wrapper._capture.assert_not_called()
     wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source)
     assert result is wrapper.decode_fn.return_value
-
-
-# ---------------------------------------------------------------------------
-# CFMGraphWrapper tests
-# ---------------------------------------------------------------------------
-
-
-class _MiniDiT(nn.Module):
-    """Minimal DiT-like module with blocks_forward_chunk for CFM graph testing.
-
-    Mimics the upstream cosyvoice2 DiT's cache semantics:
-    - CausalConv1d.forward_chunk: cat([cnn_cache, x], dim=time) when cache is not None
-    - Attention.forward_chunk: cat([k, k_cache], dim=seq) when att_cache is not None
-    """
-
-    def __init__(self, in_dim: int = 16, hidden: int = 8, depth: int = 2) -> None:
-        super().__init__()
-        self.in_proj = nn.Linear(in_dim, hidden)
-        self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)])
-        self.final_layer = nn.Linear(hidden, in_dim)
-
-    def blocks_forward_chunk(
-        self,
-        x: torch.Tensor,
-        t: torch.Tensor,
-        mask: torch.Tensor | None,
-        cnn_cache: torch.Tensor | None = None,
-        att_cache: torch.Tensor | None = None,
-        cnn_cache_buffer: torch.Tensor | None = None,
-        att_cache_buffer: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        x = x.transpose(1, 2)
-        x = self.in_proj(x)
-        for b_idx in range(len(self.blocks)):
-            # Simulate CausalConv1d: cnn_cache contributes to the first frames.
-            # Upstream: if cnn_cache[b_idx] is None, creates zeros internally (same as zeros).
-            cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None
-            if cnn_b is not None:
-                x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2)
-            # Simulate Attention: att_cache contributes a bias when non-empty.
-            # Upstream: if att_cache[b_idx] is None, skips cat entirely (different path).
-            att_b = att_cache[b_idx] if att_cache is not None else None
-            if att_b is not None and att_b.shape[3] > 0:
-                x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1)
-            x = self.blocks[b_idx](x)
-            x = x + t
-            cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous()
-            dt = x.shape[1]
-            att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1)
-        x = self.final_layer(x)
-        x = x.transpose(1, 2)
-        return x
-
-
-def _cfm_inputs(
-    batch_size: int, chunk_size: int, old_att_len: int, *, device: str = "cuda"
-) -> tuple[torch.Tensor, ...]:
-    depth = 2
-    hidden = 8
-    estimator_input = torch.randn(batch_size, 16, chunk_size, device=device)
-    time_emb = torch.randn(batch_size, 1, hidden, device=device)
-    cnn_cache = torch.randn(depth, batch_size, hidden, 2, device=device)
-    att_cache = torch.randn(depth, batch_size, 1, old_att_len, hidden, device=device)
-    cnn_out = torch.empty(depth, batch_size, hidden, 2, device=device)
-    att_out = torch.empty(depth, batch_size, 1, old_att_len + chunk_size, hidden, device=device)
-    return estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    pool = torch.cuda.graph_pool_handle()
-    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
-
-    torch.manual_seed(0)
-    estimator = _MiniDiT().eval().cuda()
-    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
-
-    with torch.inference_mode():
-        for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)):
-            inputs = _cfm_inputs(2, chunk_size, old_att_len)
-
-            eager_inputs = tuple(v.clone() for v in inputs)
-            with torch.no_grad():
-                eager_result = estimator.blocks_forward_chunk(
-                    eager_inputs[0],
-                    eager_inputs[1],
-                    None,
-                    eager_inputs[2],
-                    eager_inputs[3],
-                    eager_inputs[4],
-                    eager_inputs[5],
-                )
-
-            wrapper.replay(*inputs)
-            replay_inputs = tuple(v.clone() for v in inputs)
-            graph_result, graph_cnn, graph_att = wrapper.replay(*replay_inputs)
-
-            torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
-            torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
-            torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
-
-    wrapper._capture_graph.cache_clear()
-
-
-def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper:
-    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
-    wrapper = object.__new__(CFMGraphWrapper)
-    wrapper.max_graphs = max_graphs
-    wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
-    wrapper.device = torch.device("cuda")
-
-    @lru_cache(maxsize=max_graphs)
-    def _capture_graph(key: tuple):
-        return wrapper._do_capture(key)
-
-    wrapper._capture_graph = _capture_graph
-    wrapper._do_capture = Mock(return_value=None)
-    return wrapper
-
-
-def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
-    wrapper = _cfm_mock_wrapper(monkeypatch)
-    wrapper._capture_graph = Mock(return_value=None)
-
-    inputs = _cfm_inputs(2, 10, 0)
-    wrapper.replay(*inputs)
-
-    wrapper._capture_graph.assert_called_once()
-    wrapper.graph_fn.assert_called_once()
-
-
-def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
-    wrapper = _cfm_mock_wrapper(monkeypatch)
-    wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None)
-
-    inputs = _cfm_inputs(2, 10, 0)
-    result = wrapper.replay(*inputs)
-
-    wrapper.graph_fn.assert_called_once()
-    assert result[0] is wrapper.graph_fn.return_value
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Verify that LRU cache evicts oldest entries when full."""
-    pool = torch.cuda.graph_pool_handle()
-    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
-
-    torch.manual_seed(0)
-    estimator = _MiniDiT().eval().cuda()
-    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2)
-
-    with torch.inference_mode():
-        inputs_a = _cfm_inputs(2, 10, 0)
-        wrapper.replay(*inputs_a)
-        assert wrapper._capture_graph.cache_info().currsize == 1
-
-        inputs_b = _cfm_inputs(2, 12, 0)
-        wrapper.replay(*inputs_b)
-        assert wrapper._capture_graph.cache_info().currsize == 2
-
-        wrapper.replay(*inputs_a)
-        assert wrapper._capture_graph.cache_info().currsize == 2
-
-        inputs_c = _cfm_inputs(2, 14, 0)
-        wrapper.replay(*inputs_c)
-        assert wrapper._capture_graph.cache_info().currsize == 2
-
-        wrapper.replay(*inputs_b)
-        hits_before = wrapper._capture_graph.cache_info().hits
-        wrapper.replay(*inputs_b)
-        hits_after = wrapper._capture_graph.cache_info().hits
-        assert hits_after == hits_before + 1
-
-    wrapper._capture_graph.cache_clear()
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_cfm_none_cache_parity_between_graph_and_eager(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """Verify that None→zeros (graph path) matches [None]*depth (eager path).
-
-    setup_batch calls _decode_cfm with cnn_cache=None on every request's
-    prompt pass. The graph path replaces None with zeros; the eager path
-    passes [None]*depth. Both must produce identical outputs.
-    """
-    torch.accelerator.synchronize()
-    pool = torch.cuda.graph_pool_handle()
-    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
-
-    torch.manual_seed(0)
-    estimator = _MiniDiT().eval().cuda()
-    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
-
-    depth = len(estimator.blocks)
-    hidden = 8
-    batch_size = 2
-    chunk_size = 20
-
-    estimator_input = torch.randn(batch_size, 16, chunk_size, device="cuda")
-    time_emb = torch.randn(batch_size, 1, hidden, device="cuda")
-    cnn_out = torch.empty(depth, batch_size, hidden, 2, device="cuda")
-    att_out = torch.empty(depth, batch_size, 1, chunk_size, hidden, device="cuda")
-
-    # Eager path: cnn_cache=[None]*depth, att_cache=[None]*depth
-    eager_input = estimator_input.clone()
-    eager_time = time_emb.clone()
-    eager_cnn_out = torch.empty_like(cnn_out)
-    eager_att_out = torch.empty_like(att_out)
-    with torch.no_grad():
-        eager_result = estimator.blocks_forward_chunk(
-            eager_input,
-            eager_time,
-            None,
-            [None] * depth,
-            [None] * depth,
-            eager_cnn_out,
-            eager_att_out,
-        )
-
-    # Graph path: cnn_cache=zeros, att_cache=zero-length (simulating _estimator_step's None→zeros)
-    zero_cnn = torch.zeros_like(cnn_out)
-    zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
-    wrapper.replay(estimator_input, time_emb, zero_cnn, zero_att, cnn_out, att_out)
-    graph_input = estimator_input.clone()
-    graph_time = time_emb.clone()
-    graph_cnn_out = torch.empty_like(cnn_out)
-    graph_att_out = torch.empty_like(att_out)
-    graph_zero_cnn = torch.zeros_like(cnn_out)
-    graph_zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
-    graph_result, graph_cnn, graph_att = wrapper.replay(
-        graph_input,
-        graph_time,
-        graph_zero_cnn,
-        graph_zero_att,
-        graph_cnn_out,
-        graph_att_out,
-    )
-
-    torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
-    torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5)
-    torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5)

tests/model_extras/test_model_extras.py

@@ -17,7 +17,6 @@ from vllm_omni.model_extras import (
     build_x_to_text_prompt,
     get_extra_body_params,
     get_extra_output_params,
-    get_output_tensor_range,
     get_x_to_text_model_family,
     should_init_extra_args_for_non_diffusion_stages,
 )
@@ -135,8 +134,6 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
         }
     )
     assert get_extra_output_params("LingBotVideoPipeline") == frozenset()
-    assert get_output_tensor_range("LingBotVideoPipeline") == "zero_to_one"
-    assert get_output_tensor_range("WanPipeline") == "negative_one_to_one"
     assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
 
 

tests/model_tests/diffusion/diff_model_builders.py

@@ -13,13 +13,10 @@ def _shrink_dit_rope_config(
     target_head_dim: int = 32,
     target_heads: int = 4,
     default_axes_dims_rope: list[int] | None = None,
-    joint_attention_dim: int | None = None,
 ) -> dict:
     config["num_layers"] = num_layers
     if num_single_layers is not None:
         config["num_single_layers"] = num_single_layers
-    if joint_attention_dim is not None:
-        config["joint_attention_dim"] = joint_attention_dim
     axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope)
     factor = config["attention_head_dim"] / target_head_dim
     config["attention_head_dim"] = target_head_dim
@@ -53,55 +50,22 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict:
     return config
 
 
-def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict:
-    text_config = config["text_config"]
-    old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"]
-    text_config["num_hidden_layers"] = 2
-    text_config["intermediate_size"] = 64
-    text_config["hidden_size"] = hidden_size
-    text_config["num_attention_heads"] = 2
-    text_config["num_key_value_heads"] = 2
-    text_config["layer_types"] = text_config["layer_types"][:2]
-    factor = old_head_dim / (hidden_size / 2)
-    mrope_section = text_config["rope_scaling"]["mrope_section"]
-    text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section]
-    config["num_hidden_layers"] = 2
-    config["intermediate_size"] = 64
-    config["hidden_size"] = hidden_size
-    config["num_attention_heads"] = 2
-    config["num_key_value_heads"] = 2
-    config["vision_config"]["depth"] = 2
-    config["vision_config"]["intermediate_size"] = 64
-    config["vision_config"]["fullatt_block_indexes"] = [0, 1]
-    config["vision_config"]["out_hidden_size"] = hidden_size
-    return config
-
-
-def _shrink_qwen_transformer_config(config: dict, joint_attention_dim: int = 64) -> dict:
-    return _shrink_dit_rope_config(config, joint_attention_dim=joint_attention_dim)
-
-
 def tiny_qwen_image_builder() -> str:
+    def shrink_text_encoder(config: dict) -> dict:
+        config["num_hidden_layers"] = 2
+        config["intermediate_size"] = 64
+        config["text_config"]["num_hidden_layers"] = 2
+        config["text_config"]["intermediate_size"] = 64
+        config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2]
+        config["vision_config"]["depth"] = 2
+        config["vision_config"]["intermediate_size"] = 64
+        config["vision_config"]["fullatt_block_indexes"] = [0, 1]
+        return config
+
     return build_tiny_from_configs(
         "QwenImagePipeline",
         "Qwen/Qwen-Image",
-        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
-    )
-
-
-def tiny_qwen_image_edit_builder() -> str:
-    return build_tiny_from_configs(
-        "QwenImageEditPipeline",
-        "Qwen/Qwen-Image-Edit",
-        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
-    )
-
-
-def tiny_qwen_image_edit_plus_builder() -> str:
-    return build_tiny_from_configs(
-        "QwenImageEditPlusPipeline",
-        "Qwen/Qwen-Image-Edit-2511",
-        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
+        transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config},
     )
 
 

tests/model_tests/diffusion/model_settings.py

@@ -74,14 +74,4 @@ DIFFUSION_TEST_SETTINGS = {
         builder=diff_model_builders.tiny_flux2_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
     ),
-    "QwenImageEditPipeline": DiffusionModelTestOpts(
-        model="Qwen/Qwen-Image-Edit",
-        builder=diff_model_builders.tiny_qwen_image_edit_builder,
-        supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
-    ),
-    "QwenImageEditPlusPipeline": DiffusionModelTestOpts(
-        model="Qwen/Qwen-Image-Edit-2511",
-        builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
-        supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
-    ),
 }

tests/model_tests/diffusion/test_alignment.py

@@ -21,6 +21,8 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 # if you are adding a new model and see the tests below fail, please follow the pattern
 # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
 EXCLUDED_MODELS = [
+    "QwenImageEditPipeline",
+    "QwenImageEditPlusPipeline",
     "QwenImageLayeredPipeline",
     "GlmImagePipeline",
     "ZImagePipeline",

tools/baselines/omni_hash_baseline.py

@@ -147,7 +147,7 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
     # (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time.
     omni = Omni(
         model,
-        deploy_config=yaml_path,
+        stage_configs_path=yaml_path,
         stage_init_timeout=1800,
         init_timeout=3600,
         async_chunk=async_mode,

vllm_omni/config/omni_config.py

@@ -53,6 +53,7 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
         "model",
         "omni",
         "output_modalities",
+        "stage_configs_path",
         "stage_id",
         "tokenizer",
     }
@@ -473,10 +474,9 @@ 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", "auto", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
-                "vae_parallel_mode must be one of "
-                "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 

vllm_omni/config/stage_config.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 import functools
 import re
 import warnings
@@ -55,7 +56,7 @@ def build_stage_runtime_overrides(
     ``internal_keys`` defaults to the union of
     ``arg_utils.internal_blacklist_keys()`` and ``arg_utils.SHARED_FIELDS``
     so that neither orchestrator-only fields nor shared-pipeline fields
-    (``model`` / ``log_stats`` / ``stage_id``) leak
+    (``model`` / ``stage_configs_path`` / ``log_stats`` / ``stage_id``) leak
     into a stage's per-stage runtime overrides — the orchestrator sets those
     uniformly for every stage, they are not per-stage knobs. Callers can
     pass an explicit set for tests or specialized flows.
@@ -89,6 +90,46 @@ def build_stage_runtime_overrides(
     return result
 
 
+def strip_parent_engine_args(
+    kwargs: dict[str, Any],
+    *,
+    parent_fields: dict[str, dataclasses.Field],
+    keep_keys: set[str] | frozenset[str] = frozenset(),
+    strip_keys: set[str] | frozenset[str] = frozenset(),
+    no_warn_keys: set[str] | frozenset[str] = frozenset(),
+) -> tuple[dict[str, Any], list[str]]:
+    """Strip parent ``EngineArgs`` fields before merging into stage YAML."""
+    overridden: list[str] = []
+    result: dict[str, Any] = {}
+
+    for key, value in kwargs.items():
+        if key in strip_keys:
+            continue
+
+        if key not in parent_fields or key in keep_keys:
+            result[key] = value
+            continue
+
+        field_def = parent_fields[key]
+        if field_def.default is not dataclasses.MISSING:
+            default = field_def.default
+        elif field_def.default_factory is not dataclasses.MISSING:
+            default = field_def.default_factory()
+        else:
+            default = dataclasses.MISSING
+
+        if default is dataclasses.MISSING or value is None:
+            continue
+
+        if dataclasses.is_dataclass(default) and not isinstance(default, type):
+            default = asdict(default)
+
+        if value != default and key not in no_warn_keys:
+            overridden.append(key)
+
+    return result, sorted(overridden)
+
+
 def _apply_diffusion_parallel_runtime_overrides(
     engine_args: dict[str, Any],
     runtime_overrides: dict[str, Any],

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -19,8 +19,6 @@ connectors:
       # A future initial_codec_chunk_frames shape will use lazy capture unless
       # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
       enable_hift_graph: true
-      enable_cfm_graph: true
-      cfm_max_graphs: 32
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_2gpu.yaml

@@ -22,8 +22,6 @@ connectors:
       # A future initial_codec_chunk_frames shape will use lazy capture unless
       # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
       enable_hift_graph: true
-      enable_cfm_graph: true
-      cfm_max_graphs: 32
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_3gpu.yaml

@@ -22,8 +22,6 @@ connectors:
       # A future initial_codec_chunk_frames shape will use lazy capture unless
       # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
       enable_hift_graph: true
-      enable_cfm_graph: true
-      cfm_max_graphs: 32
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_8x4090.yaml

@@ -33,8 +33,6 @@ connectors:
       # A future initial_codec_chunk_frames shape will use lazy capture unless
       # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
       enable_hift_graph: true
-      enable_cfm_graph: true
-      cfm_max_graphs: 32
 
 stages:
   - stage_id: 0

vllm_omni/deploy/moss_tts_local.yaml

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

vllm_omni/diffusion/data.py

@@ -212,9 +212,6 @@ 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.
@@ -254,9 +251,8 @@ 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", "auto", "spatial_shard_height", "spatial_shard_width"}, (
-            "vae_parallel_mode must be one of "
-            "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
+        assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
+            "vae_parallel_mode must be one of {'tile', '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_ltx2.py

@@ -1,11 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import MethodType
 from typing import Any
 
 import torch
 from diffusers import AutoencoderKLLTX2Video
+from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
 from diffusers.models.autoencoders.vae import DecoderOutput
+from torch import nn
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
@@ -19,6 +22,89 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
 logger = init_logger(__name__)
 
 
+# Adapted from https://github.com/sgl-project/sglang/pull/27431.
+def _weight_is_channels_last_3d(module: LTX2VideoCausalConv3d) -> bool:
+    weight = module.conv.weight
+    return (
+        weight.dim() == 5
+        and hasattr(torch, "channels_last_3d")
+        and weight.is_contiguous(memory_format=torch.channels_last_3d)
+    )
+
+
+def _temporal_pad_channels_last_3d(
+    hidden_states: torch.Tensor,
+    *,
+    left: int,
+    right: int,
+) -> torch.Tensor:
+    """Replicate temporal edge frames directly into CL3D storage."""
+    batch, channels, frames, height, width = hidden_states.shape
+    padded = torch.empty(
+        (batch, channels, frames + left + right, height, width),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+        memory_format=torch.channels_last_3d,
+    )
+    padded[:, :, left : left + frames].copy_(hidden_states)
+    if left:
+        padded[:, :, :left].copy_(hidden_states[:, :, :1])
+    if right:
+        padded[:, :, left + frames :].copy_(hidden_states[:, :, -1:])
+    return padded
+
+
+def _layout_preserving_causal_conv_forward(
+    module: LTX2VideoCausalConv3d,
+    hidden_states: torch.Tensor,
+    causal: bool = True,
+) -> torch.Tensor:
+    """Diffusers LTX causal-conv math with a CL3D-preserving pad path."""
+    time_kernel_size = module.kernel_size[0]
+    use_channels_last_pad = hidden_states.dim() == 5 and _weight_is_channels_last_3d(module)
+
+    if use_channels_last_pad:
+        if causal:
+            left, right = time_kernel_size - 1, 0
+        else:
+            left = right = (time_kernel_size - 1) // 2
+        if left == 0 and right == 0:
+            hidden_states = hidden_states.contiguous(memory_format=torch.channels_last_3d)
+        else:
+            hidden_states = _temporal_pad_channels_last_3d(
+                hidden_states,
+                left=left,
+                right=right,
+            )
+    elif causal:
+        pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, time_kernel_size - 1, 1, 1))
+        hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
+    else:
+        pad_size = (time_kernel_size - 1) // 2
+        pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, pad_size, 1, 1))
+        pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, pad_size, 1, 1))
+        hidden_states = torch.concatenate([pad_left, hidden_states, pad_right], dim=2)
+
+    return module.conv(hidden_states)
+
+
+def _enable_channels_last_3d(model: nn.Module) -> tuple[int, int]:
+    """Convert LTX Conv3d weights and install layout-preserving causal pads."""
+    if not hasattr(torch, "channels_last_3d"):
+        return 0, 0
+
+    converted_convs = 0
+    patched_causal_convs = 0
+    for module in model.modules():
+        if isinstance(module, LTX2VideoCausalConv3d):
+            module.forward = MethodType(_layout_preserving_causal_conv_forward, module)
+            patched_causal_convs += 1
+        if isinstance(module, nn.Conv3d):
+            module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
+            converted_convs += 1
+    return converted_convs, patched_causal_convs
+
+
 class LTX2VaeExecutor(DistributedVaeExecutor):
     def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype):
         tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
@@ -80,6 +166,16 @@ class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMi
         model.init_distributed()
         return model
 
+    def enable_channels_last_3d(self) -> None:
+        """Opt into the layout-preserving CL3D path for this video VAE."""
+        converted_convs, patched_causal_convs = _enable_channels_last_3d(self)
+        logger.info(
+            "LTX2 video VAE: converted %d Conv3d weights to channels_last_3d and "
+            "installed %d layout-preserving causal conv paths",
+            converted_convs,
+            patched_causal_convs,
+        )
+
     def init_distributed(self):
         self.distributed_executor = LTX2VaeExecutor()
 

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -23,30 +23,6 @@ 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:
@@ -80,19 +56,6 @@ 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
@@ -307,26 +270,20 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         dec = torch.clamp(dec, min=-1.0, max=1.0)
         return dec
 
-    def _spatial_shard_decode_split_dim(
-        self,
-        z: torch.Tensor | None = None,
-        world_size: int | None = None,
-    ) -> str | None:
+    def _spatial_shard_decode_split_dim(self) -> 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:
-        if self.distributed_executor.parallel_mode == "tile":
+        split_dim = self._spatial_shard_decode_split_dim()
+        if split_dim is None:
             return False
         if z.ndim != 5:
-            if self.distributed_executor.parallel_mode != "tile":
-                logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
+            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
@@ -335,32 +292,23 @@ 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:
-            if self.distributed_executor.parallel_mode == "auto":
-                logger.debug(
-                    "Wan VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
-                    requested_size,
-                    world_size,
-                )
-            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(),
-                )
+            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,
+                split_dim,
+            )
             return False
-        return self._spatial_shard_decode_split_dim(z, world_size) is not None
+        return True
 
     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):
-            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
+            split_dim = self._spatial_shard_decode_split_dim()
             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,11 +52,6 @@ _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
@@ -337,27 +332,21 @@ class WanDistZeroPad2d(nn.Module):
         padding: tuple[int, int, int, int],
         group: dist.ProcessGroup,
         *,
-        defer_split_after_padding: bool = False,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ) -> None:
         super().__init__()
         self.padding = padding
-        self.defer_split_after_padding = defer_split_after_padding
+        self.split_dim = split_dim
+        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 split_dim == "height":
+            if self.split_dim == "height":
                 top = top if rank == 0 else 0
                 bottom = bottom if rank == world_size - 1 else 0
             else:
@@ -371,7 +360,8 @@ class WanDistConv2d(nn.Conv2d):
         self,
         source: nn.Conv2d,
         group: dist.ProcessGroup,
-        deferred_padding: tuple[int, int, int, int] | None = None,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ):
         super().__init__(
             source.in_channels,
@@ -388,117 +378,75 @@ class WanDistConv2d(nn.Conv2d):
         )
         self.load_state_dict(source.state_dict())
         self.group = group
-        self._direct_padding = source.padding
-        self._direct_padding_mode = source.padding_mode
-        self._direct_reversed_padding = source._reversed_padding_repeated_twice
-        self._deferred_padding = deferred_padding
-        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
-        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
-            raise ValueError(f"Wan spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
-        self._spatial_pad_h = int(source_padding[-2])
-        self._spatial_pad_w = int(source_padding[-1])
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        self.split_dim = split_dim
+        self.split_tensor_dim = _spatial_dim(split_dim)
+        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
+        self.halo_size = (kernel_extent - 1) // 2
+        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
+        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
+        if split_dim == "height":
+            if split_padding is None:
+                split_padding = (pad_h, pad_h)
+            self._non_split_padding = (pad_w, pad_w, 0, 0)
+            self.kernel_extent = self.kernel_size[-2]
+            self.stride_extent = self.stride[-2]
+        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:
-        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 = F.pad(x, self._non_split_padding)
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.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[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,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
         out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
-        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)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
         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=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = 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,
@@ -515,98 +463,88 @@ 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)
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        if split_dim == "height":
+            self.split_pad_left = int(self._source_padding[2])
+            self.split_pad_right = int(self._source_padding[3])
+            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:
-        split_dim = _active_spatial_shard_split_dim()
-        padding = list(self._source_padding)
+        padding = list(self._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=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.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[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,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
         out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
-        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)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
         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=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
         return trim_params
 
 
@@ -653,15 +591,12 @@ def _trim_local_conv_output(
     return out
 
 
-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> 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)
@@ -695,6 +630,7 @@ def _replace_child(
     name: str,
     child: nn.Module,
     group: dist.ProcessGroup,
+    split_dim: str,
 ) -> None:
     if child.__class__.__name__ == "WanCausalConv3d":
         setattr(
@@ -703,39 +639,51 @@ 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(
-                padding,
+                module_padding,
                 group,
-                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
+                split_dim=split_dim,
+                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
             ),
         )
         return
     if isinstance(child, nn.Conv2d):
-        deferred_padding = None
+        split_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):
-                deferred_padding = prev.padding
+                split_padding = prev.split_padding
             elif isinstance(prev, nn.ZeroPad2d):
-                deferred_padding = tuple(int(value) for value in prev.padding)
+                pad = 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,
-                deferred_padding=deferred_padding,
+                split_dim=split_dim,
+                split_padding=split_padding,
             ),
         )
 
@@ -743,20 +691,21 @@ 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)
+        _patch_attention_block(module, group, split_dim)
         inside_attention = True
 
     for name, child in list(module.named_children()):
         if child.__class__.__name__ == "WanCausalConv3d":
-            _replace_child(module, name, child, group)
+            _replace_child(module, name, child, group, split_dim)
             continue
         if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(module, name, child, group)
+            _replace_child(module, name, child, group, split_dim)
             continue
-        _patch_decoder_modules(child, group, inside_attention=inside_attention)
+        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
 
 
 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -767,25 +716,16 @@ 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 context-aware variants and attaching a dedicated
-    spatial forward. Outside that forward the replacements execute their exact
-    local PyTorch behavior, so a later auto-selected tile request remains safe.
-    The same installed decoder can alternate height and width sharding.
+    convolutions/padding for halo-exchanging variants and wrapping
+    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
+    instance and is applied only once (subsequent calls are no-ops). A given
+    instance is bound to a single ``split_dim``; switching between
+    ``"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
@@ -793,12 +733,18 @@ 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)
+    _patch_decoder_modules(decoder, group, split_dim)
     upsample_count = _decoder_upsample_count(decoder)
     orig_forward = decoder.forward
 
@@ -808,8 +754,6 @@ 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]
@@ -837,9 +781,10 @@ 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._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
+    decoder.forward = MethodType(_forward, decoder)
     vae._vllm_omni_wan_spatial_shard_installed = True
-    logger.info("Installed Wan VAE dynamic-axis spatial-shard decode.")
+    vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
 
 
 def spatial_shard_decode(
@@ -868,12 +813,11 @@ def spatial_shard_decode(
             decoded_chunks = []
             for i in range(z.shape[2]):
                 vae._conv_idx = [0]
-                chunk = vae.decoder._vllm_omni_spatial_shard_forward(
+                chunk = vae.decoder(
                     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)
@@ -887,10 +831,6 @@ 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/models/cosmos3/guardrails.py

@@ -65,10 +65,11 @@ except ImportError:
         # ``model_config["guardrails"]`` is False.
         def __init__(self, *args, **kwargs):
             raise ValueError(
-                f"The guardrails are enabled but not functional for {self.__class__}. This is most likely due to a"
-                f"missing cosmos-guardrail package. Please install cosmos-guardrail to be compliant with the "
+                f"You have disabled the safety checker for {self.__class__}. This is in violation of the "
                 "[NVIDIA Open Model License Agreement]"
-                "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license)."
+                "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). "
+                f"Please ensure that you are compliant with the license agreement."
+                f"Please install cosmos-guardrail package to enable safety checks."
             )
 
 

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -57,6 +57,17 @@ _LTX_COMPONENT_SUBFOLDERS = (
 )
 logger = logging.getLogger(__name__)
 
+_LTX2_CHANNELS_LAST_3D_EXTRA = "ltx2_use_channels_last_3d"
+
+
+def _ltx2_use_channels_last_3d(od_config: Any) -> bool:
+    """Read the LTX2-only channels-last opt-in from model extras."""
+    extras = getattr(od_config, "extras", {}) or {}
+    enabled = extras.get(_LTX2_CHANNELS_LAST_3D_EXTRA, False)
+    if not isinstance(enabled, bool):
+        raise TypeError(f"{_LTX2_CHANNELS_LAST_3D_EXTRA} must be a bool, got {type(enabled)!r}")
+    return enabled
+
 
 @dataclass(frozen=True)
 class LTXComponentProfile:
@@ -365,6 +376,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         local_files_only=local_files_only,
         dtype=dtype,
     )
+    if _ltx2_use_channels_last_3d(od_config):
+        pipeline.vae.enable_channels_last_3d()
     pipeline.audio_vae = _load_component(
         AutoencoderKLLTX2Audio,
         model,

vllm_omni/diffusion/output_formatter.py

@@ -138,7 +138,6 @@ def format_diffusion_outputs(
         ),
         "width": request.sampling_params.width,
         "height": request.sampling_params.height,
-        "num_inference_steps": request.sampling_params.num_inference_steps,
     }
 
     # Only the primary payload determines the response type. Some image models

vllm_omni/engine/arg_utils.py

@@ -167,6 +167,9 @@ class OmniEngineArgs(EngineArgs):
             Required when single-stage mode is active.
         omni_master_port: TCP port for the OmniMasterServer registration
             socket.  Required when single-stage mode is active.
+        stage_configs_path: Optional path to a JSON/YAML file containing
+            stage configurations for the multi-stage pipeline. If None,
+            stage configs are resolved from the model's default configuration.
         output_modalities: Optional list of output modality names to enable
             (e.g. ["text", "audio"]). If None, all modalities supported by
             the model are used.
@@ -228,6 +231,7 @@ class OmniEngineArgs(EngineArgs):
     omni_dp_size_local: int = 1
     omni_lb_policy: str = "random"
     omni_heartbeat_timeout: float = 30.0
+    stage_configs_path: str | None = None
     output_modalities: list[str] | None = None
     log_stats: bool = False
     custom_pipeline_args: dict[str, Any] | None = None
@@ -486,6 +490,7 @@ class OrchestratorArgs:
     ray_address: str | None = None
 
     # === Config Files ===
+    stage_configs_path: str | None = None
     deploy_config: str | None = None
     stage_overrides: str | None = None  # raw JSON string; parsed downstream
     # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays
@@ -577,6 +582,7 @@ SHARED_FIELDS: frozenset[str] = frozenset(
         "model",  # orch: detect model_type; engine: load weights
         "stage_id",  # orch: route (headless); engine: identity
         "log_stats",  # both want the flag
+        "stage_configs_path",  # orch: load legacy YAML; engine: may reference for validation
         "async_chunk",  # orch: read from CLI, redistribute; engine: per-stage flag
         "tokenizer",  # orch: detect model type; engine: tokenization
     }

vllm_omni/engine/async_omni_engine.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import dataclasses
 import json
 import queue
 import threading
@@ -23,6 +24,7 @@ import janus
 import torch
 from omegaconf import OmegaConf
 from vllm import envs as vllm_envs
+from vllm.engine.arg_utils import EngineArgs
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
@@ -32,6 +34,7 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot
 from vllm_omni.config.stage_config import (
     DuplexSessionRuntimeConfig,
     load_deploy_config,
+    strip_parent_engine_args,
 )
 from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
 from vllm_omni.diffusion.diffusion_engine import supports_audio_output
@@ -66,6 +69,7 @@ from vllm_omni.engine.stage_runtime import (
 )
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import (
+    is_new_format_deploy_config,
     load_and_resolve_stage_configs,
     parse_stage_overrides,
 )
@@ -81,6 +85,39 @@ if TYPE_CHECKING:
 
 _STARTUP_POLL_INTERVAL_S = 1.0
 _REQUEST_QUEUE_MAXSIZE = 256
+# ============================================================================
+# Parent-EngineArgs field-routing contracts (consumed by
+# AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set).
+# ============================================================================
+
+# Fields that must survive the "equal to default → strip" filter because
+# diffusion stages need them even when equal to vllm's default value
+# (e.g. colocate worker setup relies on worker_extension_cls being forwarded).
+_PARENT_ARGS_KEEP: frozenset[str] = frozenset(
+    {
+        "worker_extension_cls",
+        "allowed_local_media_path",
+        "allowed_media_domains",
+        # Legacy stage-config YAMLs may intentionally leave parallel or
+        # distributed knobs unspecified at the stage level and rely on
+        # top-level CLI values to fill them in during the per-stage merge.
+        # Keep these fields so stages that omit them can inherit CLI values,
+        # while stages with explicit YAML values still win because the legacy
+        # stage-config loader prefers stage-local engine args.
+        "tensor_parallel_size",
+    }
+)
+
+# Omni orchestrator-level fields consumed by ``_resolve_stage_configs`` that
+# must never leak into per-stage EngineArgs (``stage_configs_path`` would
+# trigger the ``create_model_config`` guard).
+_PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"})
+
+
+# Fields always populated by callers (via ``from_cli_args`` / ``asdict``) so
+# their presence as an override is never a surprise — suppress the
+# "override ignored" warning for these.
+_PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"})
 
 
 class AsyncOmniEngine:
@@ -173,8 +210,10 @@ class AsyncOmniEngine:
                 self._omni_master_port,
             )
 
-        # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand.
-        deploy_config_path = kwargs.get("deploy_config")
+        # Stage resolution pops deploy_config, so get pipeline-wide settings
+        # beforehand. The stage CLI exposes the same deploy YAML through
+        # stage_configs_path.
+        deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path")
         # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not
         # specified" so stage-config resolution can defer to the deploy yaml's
         # per-stage value (see ``with_trust_remote_code_override``). The
@@ -1036,6 +1075,7 @@ class AsyncOmniEngine:
             "streaming_output": kwargs.get("diffusion_streaming_output", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
             "extras": {
+                **(kwargs.get("extras") or {}),
                 "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
                 "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
             },
@@ -1077,6 +1117,40 @@ class AsyncOmniEngine:
         default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
         return default_stage_cfg
 
+    @staticmethod
+    def _strip_single_engine_args(kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Remove parent ``EngineArgs`` fields from *kwargs*.
+
+        When ``stage_configs_path`` is set, per-stage engine args are defined
+        in the YAML.  Top-level single-engine fields (``compilation_config``,
+        ``tensor_parallel_size``, …) must not leak into per-stage configs via
+        the ``base_engine_args`` merge in ``load_stage_configs_from_yaml`` —
+        they can cause type errors (e.g. ``compilation_config`` as a JSON
+        string rejected by ``VllmConfig``) or silently override YAML values.
+
+        Logs a warning for any parent field whose value differs from the
+        dataclass default, so users know their explicit overrides are ignored.
+        See the module-level ``_PARENT_ARGS_*`` constants for the routing
+        contracts this method enforces.
+        """
+        parent_fields: dict[str, dataclasses.Field] = {f.name: f for f in dataclasses.fields(EngineArgs)}
+        result, overridden = strip_parent_engine_args(
+            kwargs,
+            parent_fields=parent_fields,
+            keep_keys=_PARENT_ARGS_KEEP,
+            strip_keys=_PARENT_ARGS_STRIP,
+            no_warn_keys=_PARENT_ARGS_NO_WARN,
+        )
+
+        if overridden:
+            logger.warning(
+                "stage_configs_path is set — the following top-level engine "
+                "args are ignored (per-stage YAML takes precedence): %s",
+                ", ".join(sorted(overridden)),
+            )
+
+        return result
+
     def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None:
         """Apply a strategy-derived ``omni_lb_policy`` to the engine.
 
@@ -1115,6 +1189,7 @@ class AsyncOmniEngine:
     ) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
 
+        stage_configs_path = kwargs.get("stage_configs_path", None)
         deploy_config_path = kwargs.pop("deploy_config", None)
         strategy_config_path = kwargs.pop("strategy_config", None)
         stage_overrides_json = kwargs.pop("stage_overrides", None)
@@ -1122,15 +1197,26 @@ class AsyncOmniEngine:
         if explicit_stage_configs is not None:
             logger.warning(
                 "`stage_configs` is not part of the public API. "
-                "Ignoring it and resolving stages from deploy_config/model factory."
+                "Ignoring it and resolving stages from stage_configs_path/model factory."
             )
 
+        # Legacy ``stage_args`` YAMLs own per-stage EngineArgs completely, so
+        # strip parent CLI fields. New-format deploy YAMLs (also accepted via
+        # ``--stage-configs-path``) merge CLI overrides the same way as
+        # ``--deploy-config`` — do not strip, or flags like
+        # ``interleave_mm_strings`` / ``media_io_kwargs`` are silently dropped.
+        if stage_configs_path is not None and not is_new_format_deploy_config(stage_configs_path):
+            base_kwargs = self._strip_single_engine_args(kwargs)
+        else:
+            base_kwargs = kwargs
+
         # Parse --stage-overrides JSON string if provided
         stage_overrides = parse_stage_overrides(stage_overrides_json)
 
         config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
             model,
-            kwargs,
+            stage_configs_path,
+            base_kwargs,
             trust_remote_code=trust_remote_code,
             default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
             deploy_config_path=deploy_config_path,