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

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

docker/Dockerfile.npu

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

docker/Dockerfile.npu.a3

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

docker/Dockerfile.npu.ci

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

docker/Dockerfile.npu.ci.a3

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

docs/.nav.yml

@@ -71,7 +71,7 @@ nav:
         - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
         - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
         - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
-        - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
+        - VAE Parallel: user_guide/diffusion/parallelism/vae_parallel.md
       - Attention Backends: user_guide/diffusion/attention_backends.md
       - Regional Compilation: user_guide/diffusion/regional_compilation.md
       - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
@@ -123,7 +123,7 @@ nav:
           - Pipeline Parallel: design/feature/pipeline_parallel.md
           - Sequence Parallel: design/feature/sequence_parallel.md
           - Tensor Parallel: design/feature/tensor_parallel.md
-          - VAE Patch Parallelism: design/feature/vae_parallel.md
+          - VAE Parallel: design/feature/vae_parallel.md
         - Attention Optimization:
           - Skip-Softmax: design/feature/skip_softmax.md
         - Cache-DiT: design/feature/cache_dit.md
@@ -163,7 +163,8 @@ nav:
 - CLI Reference: cli
 - Community:
   - Governance: community/governance.md
-  - community/*
+  - Contact Us: community/contact_us.md
+  - Community Meetings: community/meetings.md
   - Slack: https://slack.vllm.ai
   - Blog: https://blog.vllm.ai
   - Forum: https://discuss.vllm.ai

docs/community/contact_us.md

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

docs/community/governance.md

@@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
 - [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
 
-## Diffusion Workload Division
+## Reviewer Routing
 
-This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
+For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).
 
-### Runtime-related
+`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.
-
-- 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

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

docs/community/meetups.md

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

docs/community/volunteers.md

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

docs/contributing/README.md

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

docs/design/architecture_overview.md

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

docs/design/feature/vae_parallel.md

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

docs/design/index.md

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

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

docs/usage/metrics.md

@@ -1,23 +1,23 @@
 # Production Metrics
 
-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.
+vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server.
 
 ```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,42 +25,66 @@ 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_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: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: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_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_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_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_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_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_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_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.
 
-## Cross-Stage Transfer Metrics (`vllm:omni_`)
+## 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:`)
 
 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_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_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
-| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `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_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
 
 ## vLLM Engine Metrics (`vllm:`)
 
@@ -83,9 +107,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:` request tracking + latency | With `--log-stats` |
-| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
+| `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:` 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` |
 
@@ -94,5 +118,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_` / `vllm_omni`.
+- 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.
+- 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 Parallelism](vae_parallelism.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
+| **[VAE Parallel](vae_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
 | **[HSDP](hsdp.md)**                                | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs              |
 | **[Expert Parallelism](expert_parallel.md)**       | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0)                                          |
 

docs/user_guide/diffusion/parallelism/vae_parallel.md

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

docs/user_guide/diffusion_features.md

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

examples/offline_inference/covo_audio/end2end.py

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

examples/offline_inference/dynin_omni/end2end.py

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

examples/offline_inference/hunyuan_image3/end2end.py

@@ -73,7 +73,6 @@ def parse_args():
     )
     parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
     parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
-    parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
     parser.add_argument(
         "--stream",
         action="store_true",
@@ -146,13 +145,7 @@ def main():
     else:
         bot_task = args.bot_task
 
-    if args.deploy_config is not None and args.stage_configs_path is not None:
+    deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
-        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,
@@ -168,10 +161,7 @@ def main():
 
     if additional_config is not None:
         omni_kwargs["additional_config"] = additional_config
-    if deploy_config is not None:
+    omni_kwargs["deploy_config"] = deploy_config
-        omni_kwargs["deploy_config"] = deploy_config
-    else:
-        omni_kwargs["stage_configs_path"] = stage_configs_path
 
     omni = Omni(**omni_kwargs)
 
@@ -272,10 +262,7 @@ def main():
     print(f"  Modality: {args.modality}")
     print(f"  Prompt task: {task}")
     print(f"  Bot task: {bot_task}")
-    if deploy_config is not None:
+    print(f"  Deploy config: {deploy_config}")
-        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,12 +452,6 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to a stage configs file (deprecated; prefer --deploy-config).",
-    )
     parser.add_argument(
         "--deploy-config",
         type=str,

examples/offline_inference/qwen3_omni/end2end.py

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

examples/offline_inference/step_audio2/end2end.py

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

examples/offline_inference/text_to_image/text_to_image.py

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

examples/offline_inference/text_to_speech/fish_speech/end2end.py

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

examples/offline_inference/text_to_speech/glm_tts/end2end.py

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

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

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

mkdocs.yml

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

tests/config/test_config_factory.py

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

tests/config/test_omni_config.py

@@ -582,8 +582,12 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
     assert cfg.data_parallel_size == 3
     assert cfg.cfg_parallel_size == 3
     assert cfg.mask_sp_padding is True
+    assert cfg.vae_parallel_mode == "tile"
     assert cfg.world_size == 72
 
+    auto_cfg = OmniStageDiffusionParallelConfig(vae_parallel_mode="auto")
+    assert auto_cfg.vae_parallel_mode == "auto"
+
 
 def test_parallel_config_derived_fields_are_not_init_inputs():
     with pytest.raises(ValidationError):

tests/diffusion/distributed/test_wan_spatial_shard.py

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

tests/diffusion/quantization/test_quantization_fp8.py

@@ -189,7 +189,7 @@ def _generate_single_stage_video(
 
 
 def _generate_bagel_image(
-    quantization_config: str | None = None,
+    diffusion_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",
-        "stage_configs_path": config_path,
+        "deploy_config": config_path,
         "stage_init_timeout": 300,
     }
-    if quantization_config:
+    if diffusion_quantization_config:
-        omni_kwargs["quantization_config"] = quantization_config
+        omni_kwargs["diffusion_quantization_config"] = diffusion_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 FP8 quantization_config generates a valid image.
+    """BAGEL with diffusion-stage FP8 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(quantization_config="fp8")
+    image, _ = _generate_bagel_image(diffusion_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(quantization_config=None)
+    image, _ = _generate_bagel_image(diffusion_quantization_config=None)
     image.save("test_bagel_bf16.png")
 
 

tests/diffusion/test_diffusion_output_formatter.py

@@ -95,6 +95,7 @@ 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",
-            stage_configs_path=temp_config_file,
+            deploy_config=temp_config_file,
             stage_init_timeout=300,
         ) as runner:
             generated_image = _generate_bagel_image(runner.omni)

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

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

tests/e2e/offline_inference/test_audiox_model_expansion.py

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

tests/e2e/offline_inference/test_bagel_expansion.py

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

tests/e2e/offline_inference/test_cosyvoice3_expansion.py

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

tests/e2e/offline_inference/test_ernie_image_expansion.py

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

tests/e2e/offline_inference/test_glm_tts_expansion.py

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

tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py

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

tests/e2e/offline_inference/test_longcat_video_avatar.py

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

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

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

tests/e2e/offline_inference/test_minicpmo_4_5.py

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

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

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

tests/e2e/offline_inference/test_omnivoice_expansion.py

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

tests/e2e/offline_inference/test_sensenova_u1_expansion.py

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

tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py

@@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img(
 def test_sensenova_u1_img2img(run_level):
     """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
     input_image = _load_input_image()
-    with OmniRunner(
+    with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
-        "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)
@@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit():
     input_image = _load_input_image()
     with OmniRunner(
         "SenseNova/SenseNova-U1-8B-MoT",
-        stage_configs_path=None,
         cache_backend="cache_dit",
     ) as runner:
         _generate_sensenova_u1_img2img(runner.omni, input_image)

tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py

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

tests/e2e/offline_inference/test_soulxsinger_expansion.py

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

tests/e2e/offline_inference/test_step_audio2_expansion.py

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

tests/e2e/offline_inference/test_voxtral_tts_expansion.py

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

tests/engine/test_arg_utils.py

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

tests/engine/test_async_omni_engine_abort.py

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

tests/engine/test_single_stage_mode.py

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

tests/entrypoints/test_async_omni.py

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

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -478,12 +478,11 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-    engine._strip_single_engine_args = lambda kwargs: kwargs
 
     _, stage_configs = engine._resolve_stage_configs(
         "dummy-model",
         {
-            "stage_configs_path": "dummy.yaml",
+            "deploy_config": "dummy.yaml",
             "additional_config": {"torchair_graph_config": {"enabled": True}},
         },
         trust_remote_code=False,
@@ -516,12 +515,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-    engine._strip_single_engine_args = lambda kwargs: kwargs
 
     _, stage_configs = engine._resolve_stage_configs(
         "dummy-model",
         {
-            "stage_configs_path": "dummy.yaml",
+            "deploy_config": "dummy.yaml",
             "quantization_config": {"method": "bitsandbytes"},
         },
         trust_remote_code=False,

tests/entrypoints/test_serve.py

@@ -76,6 +76,18 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
+def test_serve_parser_defaults_vae_parallel_mode_to_tile_and_accepts_auto() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    default_args = parser.parse_args(["serve", "fake-model", "--omni"])
+    auto_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "auto"])
+
+    assert default_args.vae_parallel_mode == "tile"
+    assert auto_args.vae_parallel_mode == "auto"
+
+
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",
@@ -193,7 +205,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
     with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
         run_headless(args)
 
-    assert captured["args"][1] is None
+    assert len(captured["args"]) == 2
     assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
     assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
     assert captured["strategy_config_path"] == "/tmp/strategy.yaml"

tests/entrypoints/test_utils.py

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

tests/helpers/runtime.py

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

tests/metrics/test_modality.py

@@ -42,6 +42,10 @@ _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,
 ]
 
 
@@ -55,6 +59,10 @@ 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:
@@ -156,6 +164,50 @@ 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.
 # ---------------------------------------------------------------------------
@@ -183,6 +235,18 @@ 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):
@@ -239,8 +303,88 @@ class TestObserveModalityAtFinalize:
         )
         assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
 
-    def test_non_audio_output_type_skipped(self):
+    def test_diffusion_path_full(self):
-        # Image / video / text output types are out of scope for this module.
+        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):
         stub = _StubModMetrics()
         for output_type in ("text", "image", "video"):
             observe_modality_at_finalize(
@@ -251,7 +395,7 @@ class TestObserveModalityAtFinalize:
                 stage_metrics=_Bag(stage_gen_time_ms=100.0),
                 engine_outputs=_Bag(),
             )
-        assert stub.calls == []
+        assert not any(c[0].startswith("observe_audio") or c[0].startswith("inc_audio") for c in stub.calls)
 
     def test_replica_id_none_skipped(self):
         stub = _StubModMetrics()

tools/baselines/omni_hash_baseline.py

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

vllm_omni/config/omni_config.py

@@ -53,7 +53,6 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
         "model",
         "omni",
         "output_modalities",
-        "stage_configs_path",
         "stage_id",
         "tokenizer",
     }
@@ -474,9 +473,10 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
         if self.ulysses_mode not in {"strict", "advanced_uaa"}:
             raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
-        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
-                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of "
+                "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 

vllm_omni/config/stage_config.py

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

vllm_omni/deploy/moss_tts_local.yaml

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

vllm_omni/diffusion/data.py

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

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

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

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

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

vllm_omni/diffusion/output_formatter.py

@@ -138,6 +138,7 @@ 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,9 +167,6 @@ 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.
@@ -231,7 +228,6 @@ 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
@@ -490,7 +486,6 @@ 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
@@ -582,7 +577,6 @@ 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,7 +9,6 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
-import dataclasses
 import json
 import queue
 import threading
@@ -24,7 +23,6 @@ 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
@@ -34,7 +32,6 @@ 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
@@ -69,7 +66,6 @@ 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,
 )
@@ -85,39 +81,6 @@ 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:
@@ -210,10 +173,8 @@ class AsyncOmniEngine:
                 self._omni_master_port,
             )
 
-        # Stage resolution pops deploy_config, so get pipeline-wide settings
+        # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand.
-        # beforehand. The stage CLI exposes the same deploy YAML through
+        deploy_config_path = kwargs.get("deploy_config")
-        # 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
@@ -1116,40 +1077,6 @@ 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.
 
@@ -1188,7 +1115,6 @@ 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)
@@ -1196,26 +1122,15 @@ 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 stage_configs_path/model factory."
+                "Ignoring it and resolving stages from deploy_config/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,
-            stage_configs_path,
+            kwargs,
-            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,

vllm_omni/entrypoints/async_omni.py

@@ -123,9 +123,8 @@ class AsyncOmni(EngineClient, OmniBase):
     Args:
         model: Model name or path to load.
         **kwargs: Additional keyword arguments.
-            - stage_configs_path: Optional path to YAML file containing stage
+            - deploy_config: Optional path to a deploy YAML. If None,
-              configurations. If None, configurations are resolved from model
+              configurations are resolved from the model pipeline factory.
-              pipeline factory.
             - log_stats: Whether to enable statistics logging.
             - stage_init_timeout: Timeout for per-stage initialization.
             - init_timeout: Total timeout for orchestrator startup.

vllm_omni/entrypoints/cli/serve.py

@@ -790,9 +790,10 @@ class OmniServeCommand(CLISubcommand):
             "--vae-parallel-mode",
             type=str,
             default="tile",
-            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
+            choices=["tile", "auto", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
-            "'tile' (default) uses patch/tile parallel decode; "
+            "'tile' (default) uses patch/tile decode; 'auto' selects the longer spatial-shard axis for "
+            "validated Wan tiled decodes on a full DiT group, then falls back to patch/tile decode; "
             "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
             "decoder feature maps along height/width and exchanges halo regions. The "
             "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "
@@ -945,8 +946,6 @@ def run_headless(args: TrackingNamespace) -> None:
 
     config_path, stage_configs, _ = load_and_resolve_stage_configs(
         model,
-        # The serve CLI no longer accepts legacy stage_args YAMLs.
-        None,
         args_dict,
         # store_true cannot express an explicit False: absent maps to None
         # ("not specified") so the deploy yaml's per-stage value applies.

vllm_omni/entrypoints/omni_base.py

@@ -568,6 +568,7 @@ class OmniBase(PDDisaggregationMixin):
             msg_id = id(result)
             consumed = self._consumed_metric_message_ids(req_id)
             if msg_id not in consumed:
+                metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs)
                 metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
                 consumed.add(msg_id)
 

vllm_omni/entrypoints/openai/api_server.py

@@ -679,7 +679,7 @@ async def build_async_omni_from_stage_config(
         EngineClient instance (AsyncOmni) ready for use
 
     Note:
-        Stage configurations are loaded from args.stage_configs_path if provided,
+        Stage configurations are loaded from ``args.deploy_config`` when provided,
         otherwise from the model's default configuration.
     """
 
@@ -1156,7 +1156,7 @@ async def omni_init_app_state(
     state.openai_serving_duplex = None
     if state.openai_serving_chat is not None and should_enable_duplex_endpoint(
         state.stage_configs,
-        config_path=getattr(args, "stage_configs_path", None) or getattr(args, "deploy_config", None),
+        config_path=getattr(args, "deploy_config", None),
     ):
         from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler
 

vllm_omni/entrypoints/utils.py

@@ -6,7 +6,6 @@ from dataclasses import fields, is_dataclass
 from pathlib import Path
 from typing import Any, get_args, get_origin
 
-import yaml
 from vllm.logger import init_logger
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
@@ -30,15 +29,6 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent
 logger = init_logger(__name__)
 
 
-def is_new_format_deploy_config(path: str | None) -> bool:
-    """Return True when *path* is a deploy YAML (``stages:`` without legacy ``stage_args:``)."""
-    if not path or not os.path.exists(path):
-        return False
-    with open(path, encoding="utf-8") as f:
-        peek = yaml.safe_load(f) or {}
-    return isinstance(peek, dict) and "stages" in peek and "stage_args" not in peek
-
-
 _DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = {
     "GlmImagePipeline": "glm_image",
 }
@@ -606,7 +596,6 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
 
 def load_and_resolve_stage_configs(
     model: str,
-    stage_configs_path: str | None,
     kwargs: dict | None,
     *,
     trust_remote_code: bool | None,
@@ -615,17 +604,15 @@ def load_and_resolve_stage_configs(
     stage_overrides: dict[str, dict[str, Any]] | None = None,
     strategy_config_path: str | None = None,
 ) -> tuple[str, list, str | None]:
-    """Load stage configurations from model or YAML file with fallback to defaults.
+    """Load stage configurations from a deploy YAML or model defaults.
 
     Args:
         model: Model name or path
-        stage_configs_path: Optional path to legacy YAML (stage_args format)
         kwargs: Engine arguments to merge with stage configs
         trust_remote_code: Whether to trust remote code while resolving the model config.
         default_stage_cfg_factory: Optional callable that takes no args and returns
             default stage config list when no configs are found
-        deploy_config_path: Optional path to deploy YAML (new format).
+        deploy_config_path: Optional path to a deploy YAML.
-            Mutually exclusive with ``stage_configs_path``.
         stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
             Keys are stage_id strings, values are dicts of overrides.
         strategy_config_path: Optional path to a composable-parallel
@@ -636,65 +623,21 @@ def load_and_resolve_stage_configs(
         the strategy-derived pipeline-wide load-balancer policy (``None`` when no
         strategy set one), returned for the engine to apply.
     """
-    if stage_configs_path is not None and deploy_config_path is not None:
+    config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model)
-        raise ValueError(
+    stage_configs, omni_lb_policy = load_stage_configs_from_model(
-            "--stage-configs-path and --deploy-config are mutually exclusive: "
+        model,
-            "they use different path resolution rules and loading paths. "
+        trust_remote_code=trust_remote_code,
-            "Use --deploy-config for new-format YAMLs (preferred); "
+        base_engine_args=kwargs,
-            "--stage-configs-path is kept only for the legacy `stage_args` format "
+        deploy_config_path=deploy_config_path,
-            "and will be removed in a future release."
+        stage_overrides=stage_overrides,
-        )
+        strategy_config_path=strategy_config_path,
-    if stage_configs_path is not None and deploy_config_path is None:
+    )
-        if not os.path.exists(stage_configs_path):
+    if not stage_configs:
-            raise FileNotFoundError(
+        if default_stage_cfg_factory is not None:
-                f"--stage-configs-path {stage_configs_path!r} does not exist. "
+            default_stage_cfg = default_stage_cfg_factory()
-                "Legacy `stage_configs/` yamls were replaced by `vllm_omni/deploy/<model>.yaml`; "
+            stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
-                "use --deploy-config. See docs/configuration/stage_configs.md."
-            )
-        if is_new_format_deploy_config(stage_configs_path):
-            deploy_config_path = stage_configs_path
-            stage_configs_path = None
         else:
-            logger.warning(
+            stage_configs = []
-                "--stage-configs-path is deprecated; migrate %r and use --deploy-config.",
-                stage_configs_path,
-            )
-
-    omni_lb_policy: str | None = None
-    if deploy_config_path is not None:
-        config_path = deploy_config_path
-        stage_configs, omni_lb_policy = load_stage_configs_from_model(
-            model,
-            trust_remote_code=trust_remote_code,
-            base_engine_args=kwargs,
-            deploy_config_path=deploy_config_path,
-            stage_overrides=stage_overrides,
-            strategy_config_path=strategy_config_path,
-        )
-        if not stage_configs:
-            if default_stage_cfg_factory is not None:
-                default_stage_cfg = default_stage_cfg_factory()
-                stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
-            else:
-                stage_configs = []
-    elif stage_configs_path is None:
-        config_path = resolve_model_config_path(model)
-        stage_configs, omni_lb_policy = load_stage_configs_from_model(
-            model,
-            trust_remote_code=trust_remote_code,
-            base_engine_args=kwargs,
-            stage_overrides=stage_overrides,
-            strategy_config_path=strategy_config_path,
-        )
-        if not stage_configs:
-            if default_stage_cfg_factory is not None:
-                default_stage_cfg = default_stage_cfg_factory()
-                stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
-            else:
-                stage_configs = []
-    else:
-        config_path = stage_configs_path
-        stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=kwargs)
 
     stage_configs = filter_stages(config_path, stage_configs, kwargs)
     logger.debug(f"stage_configs: {stage_configs}")

vllm_omni/metrics/definitions.py

@@ -134,6 +134,15 @@ AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK
 AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS
 
 
+# ============================================================================
+# Diffusion family (per-stage + per-replica diffusion timing breakdowns)
+# ============================================================================
+DIFFUSION_EXEC_S = METRIC_PREFIX + "diffusion_exec_s"
+DIFFUSION_EXEC_PER_STEP_S = METRIC_PREFIX + "diffusion_exec_per_step_s"
+DIFFUSION_PREPROCESS_S = METRIC_PREFIX + "diffusion_preprocess_s"
+DIFFUSION_POSTPROCESS_S = METRIC_PREFIX + "diffusion_postprocess_s"
+
+
 # ============================================================================
 # Cross-stage Transfer family (per-physical-hop TX/RX/in-flight timings)
 # ============================================================================

vllm_omni/metrics/modality.py

@@ -78,6 +78,35 @@ _audio_skipped_family = Counter(
 )
 
 
+# ----------------------------------------------------------------------------
+# Diffusion family
+# ----------------------------------------------------------------------------
+_diffusion_exec_family = Histogram(
+    defs.DIFFUSION_EXEC_S,
+    "DiT forward pass execution time per request in seconds.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_BUCKETS,
+)
+_diffusion_exec_per_step_family = Histogram(
+    defs.DIFFUSION_EXEC_PER_STEP_S,
+    "DiT forward pass execution time per denoising step in seconds.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+_diffusion_preprocess_family = Histogram(
+    defs.DIFFUSION_PREPROCESS_S,
+    "Diffusion input preprocessing time per request in seconds.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+_diffusion_postprocess_family = Histogram(
+    defs.DIFFUSION_POSTPROCESS_S,
+    "Diffusion output postprocessing (VAE decode) time per request in seconds.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+
+
 class OmniModalityMetrics:
     """Per-modality observe API. Stage/replica are passed at observe time
     because a single OmniModalityMetrics instance per pipeline serves all
@@ -139,6 +168,44 @@ class OmniModalityMetrics:
             reason=reason or "unknown",
         ).inc()
 
+    # ---- Diffusion --------------------------------------------------------
+
+    def observe_diffusion_exec(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats:
+            return
+        _diffusion_exec_family.labels(
+            model_name=self._model_name,
+            stage=stage,
+            replica=replica,
+        ).observe(seconds)
+
+    def observe_diffusion_exec_per_step(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats:
+            return
+        _diffusion_exec_per_step_family.labels(
+            model_name=self._model_name,
+            stage=stage,
+            replica=replica,
+        ).observe(seconds)
+
+    def observe_diffusion_preprocess(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats:
+            return
+        _diffusion_preprocess_family.labels(
+            model_name=self._model_name,
+            stage=stage,
+            replica=replica,
+        ).observe(seconds)
+
+    def observe_diffusion_postprocess(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats:
+            return
+        _diffusion_postprocess_family.labels(
+            model_name=self._model_name,
+            stage=stage,
+            replica=replica,
+        ).observe(seconds)
+
 
 def observe_modality_at_finalize(
     mod_metrics: OmniModalityMetrics,
@@ -161,43 +228,49 @@ def observe_modality_at_finalize(
     """
     if replica_id is None or stage_metrics is None or output_type is None:
         return
-    if output_type != "audio":
-        return
 
     stage_label = str(stage_id)
     replica_label = str(replica_id)
     gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0
     mm_out = extract_mm_output(engine_outputs)
 
-    sample_rate = defs.resolve_audio_sample_rate(mm_out)
+    if output_type == "audio":
-    # `stage_metrics.audio_generated_frames` is the legacy per-chunk
+        sample_rate = defs.resolve_audio_sample_rate(mm_out)
-    # accumulator field on StageRequestStats. No production path currently
+        n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
-    # fills it, so the fallback below is the live source — but we leave the
+        if n_frames == 0:
-    # field lookup in place in case the accumulator gets re-wired upstream.
+            n_frames = count_audio_frames(mm_out)
-    n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
+        mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
-    if n_frames == 0:
+        duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
-        n_frames = count_audio_frames(mm_out)
+        if duration_s > 0:
-    mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
+            mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
-    duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
+            mod_metrics.observe_audio_rtf(
-    if duration_s > 0:
+                stage_label,
-        mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
+                replica_label,
-        mod_metrics.observe_audio_rtf(
+                defs.compute_audio_rtf(gen_time_s, duration_s),
-            stage_label,
+            )
-            replica_label,
+        else:
-            defs.compute_audio_rtf(gen_time_s, duration_s),
+            mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
-        )
+
-    else:
+    dm = getattr(stage_metrics, "diffusion_metrics", None)
-        # Request completed (finish_reason ∈ {stop, length} — error paths
+    if dm:
-        # don't reach finalize) but no audio samples were produced. Covers
+        _key_map = {
-        # silent `return None` skips in the talker→code2wav stage
+            "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec,
-        # processors and the `parsed.append((0,0))` malformed-length path
+            "preprocess_time_s": mod_metrics.observe_diffusion_preprocess,
-        # in qwen3-tts code2wav. raise-paths surface via the upstream
+            "postprocess_time_s": mod_metrics.observe_diffusion_postprocess,
-        # vllm:request_success_total{finished_reason="error"} channel and
+        }
-        # never reach this branch.
+        for key, observe_fn in _key_map.items():
-        mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
+            val = dm.get(key)
-    # audio_underrun / continuity are emitted from the streaming path in
+            if val is not None:
-    # observe_audio_streaming_finalize; finalize is too late for the
+                observe_fn(stage_label, replica_label, float(val))
-    # per-chunk timeline they need.
+
+        exec_time = dm.get("diffusion_engine_exec_time_s")
+        num_steps = dm.get("num_inference_steps")
+        if exec_time is not None and num_steps and num_steps > 0:
+            mod_metrics.observe_diffusion_exec_per_step(
+                stage_label,
+                replica_label,
+                float(exec_time) / int(num_steps),
+            )
 
 
 def observe_audio_first_packet(

vllm_omni/metrics/stats.py

@@ -44,7 +44,7 @@ class StageRequestStats:
     final_output_type: str | None = None
     request_id: str | None = None
     postprocess_time_ms: float = 0.0
-    diffusion_metrics: dict[str, int] = None
+    diffusion_metrics: dict[str, float] = None
     audio_generated_frames: int = 0
     audio_sample_rate: int = 0
     audio_duration_s: float = 0.0
@@ -579,7 +579,7 @@ class OrchestratorAggregator:
         if final_output_type is not None:
             stats.final_output_type = final_output_type
         stats.diffusion_metrics = (
-            {k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
+            {k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
             if req_id in self.diffusion_metrics
             else None
         )
@@ -644,8 +644,19 @@ class OrchestratorAggregator:
         if isinstance(diffusion_metrics, list):
             diffusion_metrics = diffusion_metrics[0]
         if diffusion_metrics:
+            _MS_TO_S = {
+                "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s",
+                "preprocess_time_ms": "preprocess_time_s",
+                "postprocess_time_ms": "postprocess_time_s",
+                "diffusion_engine_total_time_ms": "diffusion_engine_total_time_s",
+            }
             for key, value in diffusion_metrics.items():
-                self.diffusion_metrics[req_id][key] += value
+                if value is None:
+                    continue
+                if key in _MS_TO_S:
+                    self.diffusion_metrics[req_id][_MS_TO_S[key]] = float(value) / 1000.0
+                else:
+                    self.diffusion_metrics[req_id][key] = value
 
     def on_forward(
         self,

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -730,7 +730,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
         has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
 
         # Run forward pass
-        clear_kv_metadata = self.speculative_config is None
+        defer_kv_connector_finalize = self.speculative_config is not None and (
+            get_pp_group().is_last_rank or self.broadcast_pp_output
+        )
         with (
             record_function_or_nullcontext("forward"),
             set_ascend_forward_context(
@@ -744,13 +746,12 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
                 model_instance=self.model,
                 skip_compiled=has_encoder_input,
                 has_sinks=self._has_sinks,
-                input_ids=input_ids,
                 eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
             ),
             self.maybe_get_kv_connector_output(
                 scheduler_output,
                 **(
-                    {"defer_finalize": not clear_kv_metadata}
+                    {"defer_finalize": defer_kv_connector_finalize}
                 ),
             ) as kv_connector_output,
         ):

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -381,7 +381,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
         has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
 
         # Run forward pass
-        clear_kv_metadata = self.speculative_config is None
+        defer_kv_connector_finalize = self.speculative_config is not None and (
+            get_pp_group().is_last_rank or self.broadcast_pp_output
+        )
         with (
             record_function_or_nullcontext("forward"),
             set_ascend_forward_context(
@@ -395,12 +397,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
                 model_instance=self.model,
                 skip_compiled=has_encoder_input,
                 has_sinks=self._has_sinks,
-                input_ids=input_ids,
                 eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
             ),
             self.maybe_get_kv_connector_output(
                 scheduler_output,
-                **({"defer_finalize": not clear_kv_metadata}),
+                **({"defer_finalize": defer_kv_connector_finalize}),
             ) as kv_connector_output,
         ):
             if self.cache_config.mamba_cache_mode == "align":
@@ -654,9 +655,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
         num_active_loras: int = 0,
         profile_seq_lens: int | None = None,
         profile_cpp: bool = False,
+        randomize_inputs: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # only support eager mode and piecewise graph now
-        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode()
         # If cudagraph_mode.decode_mode() == FULL and
         # cudagraph_mode.separate_routine(). This means that we are using
         # different graphs and/or modes for mixed prefill-decode batches vs.
@@ -917,7 +919,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
                 batch_descriptor=batch_desc,
                 model_instance=self.model,
                 has_sinks=self._has_sinks,
-                input_ids=input_ids,
                 eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
             ):
                 outputs = self.model(

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -125,9 +125,10 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         num_active_loras: int = 0,
         profile_seq_lens: int | None = None,
         profile_cpp: bool = False,
+        randomize_inputs: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # only support eager mode and piecewise graph now
-        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode()
         # If cudagraph_mode.decode_mode() == FULL and
         # cudagraph_mode.separate_routine(). This means that we are using
         # different graphs and/or modes for mixed prefill-decode batches vs.
@@ -367,19 +368,21 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
                 if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"):
                     return self.drafter.model.compute_logits(hidden_states[dummy_indices])
 
-            with set_ascend_forward_context(
+            with (
-                attn_metadata,
+                self.maybe_randomize_inputs(input_ids, inputs_embeds, randomize_inputs=randomize_inputs),
-                self.vllm_config,
+                set_ascend_forward_context(
-                num_tokens=num_tokens_padded,
+                    attn_metadata,
-                num_tokens_across_dp=num_tokens_across_dp,
+                    self.vllm_config,
-                in_profile_run=is_profile,
+                    num_tokens=num_tokens_padded,
-                num_actual_tokens=num_tokens_padded,
+                    num_tokens_across_dp=num_tokens_across_dp,
-                aclgraph_runtime_mode=cudagraph_runtime_mode,
+                    in_profile_run=is_profile,
-                batch_descriptor=batch_desc,
+                    num_actual_tokens=num_tokens_padded,
-                model_instance=self.model,
+                    aclgraph_runtime_mode=cudagraph_runtime_mode,
-                has_sinks=self._has_sinks,
+                    batch_descriptor=batch_desc,
-                input_ids=input_ids,
+                    model_instance=self.model,
-                eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
+                    has_sinks=self._has_sinks,
+                    eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
+                ),
             ):
                 # ---------------------------------------Omni-new----------------------------------------------
                 if getattr(self.model, "talker", None) is not None and self.has_talker_mtp: