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

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
[Refactor] Remove internal stage_configs_path plumbing (#5741 )
2026-08-13 23:49:11 +08:00 · 2026-08-13 21:29:03 +08:00 · 2026-08-13 20:01:01 +08:00 · 2026-08-13 17:04:06 +08:00 · 2026-08-13 08:56:26 +00:00 · 2026-08-13 06:59:13 +00:00
89 changed files with 1341 additions and 1044 deletions
@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.26.0
+ARG VLLM_ASCEND_TAG=v0.27.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}

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

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

 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -1,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
@@ -43,7 +43,6 @@ nav:
    - Overview: features/README.md
    - Runtime and Stage Execution:
      - Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
-      - Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
      - Sleep Mode: features/sleep_mode.md
    - Quantization:
      - Overview: user_guide/quantization/overview.md
@@ -72,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
@@ -124,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
@@ -164,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
@@ -1,5 +1,9 @@
 # Contact Us

- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature
+- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues).
+- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/).
+- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md).
+
+## Security Disclosures
+
+Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request.
@@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
 - [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support

-## Diffusion Workload Division
+## Reviewer Routing

-This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
+For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).

-### Runtime-related
-
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
-
-### Optimization
-
- Parallel: @RuixiangMa, @wtomin
- Attention backend: @ZJY0516
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
- Offload: @yuanheng-zhao
- Intel GPU support: @xuechendi
- Quantization: @lishunyang12, @david6666666
- Kernel fusion / communication-computation: @ZJY0516
-
-### Docs & Test
-
- Diffusion docs: @RuixiangMa, @wtomin
-
-## Meetings
-
-Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
+`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights.

 ## Committer Nomination Process

@@ -0,0 +1,8 @@
+# Community Meetings
+
+vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth:
+
+- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0)
+- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0)
+
+Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/).
@@ -1 +0,0 @@
-# Meetups
@@ -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!
@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m

 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance

+## Multi-Stage Recipes
+
+- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
+
 ## Optimization Features

 - **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
@@ -0,0 +1,171 @@
+# Prefill-Decode (PD) Disaggregation
+
+PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
+stages so prompt processing and token generation can run on different workers.
+
+This is documented as a stage-config recipe instead of a bundled YAML because the
+deployment-specific values usually change per environment:
+
+- GPU placement
+- `tensor_parallel_size`
+- connector backend and connector ports
+- connector IPs or bootstrap addresses
+
+Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
+and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
+changes below.
+
+## Requirements
+
+- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
+- A KV connector supported by vLLM, such as `MooncakeConnector`
+- Matching `tensor_parallel_size` on the prefill and decode thinker stages
+
+## 1. Split the thinker into prefill and decode stages
+
+Replace the original thinker stage with two stages:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_prefill_only: true
+    runtime:
+      devices: "0"
+    engine_args:
+      max_num_seqs: 16
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_producer"
+        kv_rank: 0
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25201
+    final_output: false
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm
+    is_decode_only: true
+    runtime:
+      devices: "1"
+    engine_args:
+      max_num_seqs: 64
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_consumer"
+        kv_rank: 1
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25202
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+```
+
+Notes:
+
+- `is_prefill_only: true` marks the thinker stage that only saves KV.
+- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
+- `kv_transfer_config` is required on both stages.
+- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
+  prefill side only processes the prompt and exports KV.
+
+## 2. Shift the downstream stages by one index
+
+After inserting the extra thinker stage, renumber the remaining stages:
+
+```yaml
+  - stage_id: 2
+    runtime:
+      devices: "2"
+    engine_input_source: [1]
+    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
+
+  - stage_id: 3
+    runtime:
+      devices: "2"
+    engine_args:
+      max_num_seqs: 1
+    engine_input_source: [2]
+    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
+```
+
+Compared with the default Qwen3-Omni config:
+
+- the talker becomes stage `2` instead of stage `1`
+- the code2wav stage becomes stage `3` instead of stage `2`
+- the talker now reads from decode stage `1`
+
+## 3. Add runtime edges for the four-stage pipeline
+
+```yaml
+runtime:
+  enabled: true
+  edges:
+    - from: 0
+      to: 1
+    - from: 1
+      to: 2
+    - from: 2
+      to: 3
+```
+
+## 4. Launch with your custom config
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --deploy-config /path/to/qwen3_omni_pd.yaml
+```
+
+## Operational Notes
+
+- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
+  pair. Keep prefill and decode at the same `tensor_parallel_size`.
+- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
+  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
+- Choose connector ports and addresses that match your deployment. The values
+  shown above are examples only.
@@ -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

@@ -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)
@@ -1,42 +1,57 @@
-# VAE Patch Parallelism
+# VAE Parallel

-This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
-We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
+This document defines the VAE parallelism strategy contract and describes how to
+add a strategy to a diffusion model. Tile/patch parallelism is the portable
+baseline. Spatial-shard parallelism is a model-specific decode strategy that can
+be selected explicitly or through a request-aware policy.
+
+We use **Qwen-Image** as the reference implementation for tile-parallel decode,
+**Wan2.2** for tile-parallel encode, and **Wan2.2** as the first spatial-shard
+decode implementation. Other VAE families can add spatial strategies later
+without changing the public mode meanings defined here.

 ---

 ## Table of Contents

 - [Overview](#overview)
- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
- [Encode Parallel Implementation](#encode-parallel-implementation)
+- [Strategy Selection Contract](#strategy-selection-contract)
+- [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:
-
-+ Distributes computation across multiple devices
-
-+ Reduces peak memory usage per device
-
-+ Accelerates encoding/decoding latency
+The actual benefit depends on shape, topology, communication cost, and VAE
+architecture. A strategy must not become the global default without parity and
+performance evidence on its supported models.

 ### 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
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
+### Architecture and ownership
+
+`DiffusionParallelConfig.vae_parallel_mode` carries the public strategy request.
+The shared registry forwards that value to the VAE, but each distributed VAE
+adapter owns its supported modes, request-shape selection, topology validation,
+and fallback behavior. Model-specific policy must not be added to the shared
+runner or registry.
+
+**DistributedVaeExecutor** is the model-agnostic core for tile-parallel VAE
+encoding and decoding.

 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"`) |
-|--------|--------------------------|-----------------------------------------------|
-| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
-| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
-| Scope | Decode + encode | Decode only |
+| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+|--------|--------------------------|----------------------|-----------------------------------------------|
+| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Scope | Decode + encode | Wan decode selection; encode remains tiled | Wan decode only |

 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`
-2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
-3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
-4. **Change VAE model in pipeline** - Use the distributed version
-5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`
+1. **Implement the distributed VAE** - Inherit from the base VAE class and
+   `DistributedVaeMixin`.
+2. **Keep tile as the baseline** - For tile decode, refactor `tiled_decode` into
+   `tile_split`/`tile_exec`/`tile_merge`.
+3. **Add encode parallelism when needed** - Implement
+   `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models.
+4. **Add model-specific strategies locally** - Own selectors, topology checks,
+   communication, and fallback in the distributed autoencoder.
+5. **Wire the pipeline** - Replace the pipeline VAE with its distributed version.
+6. **Test strategy and lifecycle parity** - Compare one rank with every supported
+   strategy, including feature-off, fallback, and mixed-request sequences.
@@ -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

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

@@ -20,19 +20,13 @@ an implementation contract has a user-facing workflow.
 | Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
 | Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |

-The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
-describes the experimental Qwen3-Omni runtime topology and its current
-integration boundaries. It is not yet a supported launch recipe.
-
 Some runtime designs are deliberately not promoted as standalone User Guide
 features yet:

- Generic, OmniConnector-based
-  [Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
-  yet a standalone User Guide feature. Its
+- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
  [OmniConnector implementations](../design/index.md#communication) describe
-  transport contracts, while practical connector configuration remains under
-  [Pipeline and deploy configurations](../configuration/stage_configs.md).
+  topology and transport contracts. Practical deployment configuration remains
+  under [Pipeline and deploy configurations](../configuration/stage_configs.md).
 - [Async Chunk](../design/feature/async_chunk.md) and
  [Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
  are model- and pipeline-dependent. Use the selected model's deploy
@@ -1,140 +0,0 @@
-# Prefill-Decode (PD) Disaggregation (experimental)
-
-!!! warning "Experimental design reference"
-    PD disaggregation is only partially integrated on the current `main` branch.
-    The runtime contains PD detection and routing scaffolding, but there is no
-    supported deploy overlay or end-to-end validated launch recipe yet. Treat this
-    page as a description of the current design, not as production guidance.
-
-Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
-stages: a prefill worker that processes the prompt and produces KV cache, and a
-decode worker that imports that cache and generates tokens. The Talker and
-Code2Wav stages remain downstream; PD does not disaggregate the entire
-Qwen3-Omni pipeline.
-
-## Current design
-
-### Runtime topology
-
-```mermaid
-flowchart LR
-    client["Client request"] --> orchestrator["Omni Orchestrator"]
-    orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
-    prefill -->|"KV blocks"| connector[("vLLM KV connector")]
-    prefill -->|"KV routing metadata"| orchestrator
-    orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
-    connector -->|"remote KV load"| decode
-    decode -->|"generated text"| client
-```
-
-The orchestrator owns the request lifecycle, while the vLLM KV connector owns
-the bulk KV transfer. The original prompt is submitted to both Thinker stages:
-the decode worker still needs the request tokens and multimodal features to
-construct its request, while the imported KV cache prevents it from repeating
-the completed prefill computation.
-
-For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
-
-| Stage | Role | Output or handoff |
-| --- | --- | --- |
-| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
-| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
-| 2 | Talker | Converts combined Thinker conditioning into codec codes |
-| 3 | Code2Wav | Converts codec codes into audio |
-
-Text remains a final output of the decode stage, and audio remains a final
-output of Code2Wav.
-
-### Request lifecycle
-
-1. At startup, PD detection looks for one stage marked `is_prefill_only` and
-   one downstream stage marked `is_decode_only`. The decode stage must list the
-   prefill stage as an input source.
-2. The orchestrator clones the Thinker sampling parameters for prefill, forces
-   `max_tokens=1`, clears stop conditions, and sets producer-side
-   `kv_transfer_params`. This makes prefill finish by length after exporting KV.
-3. When prefill finishes, the orchestrator captures its KV routing metadata and
-   Qwen multimodal output under the original request ID.
-4. The orchestrator submits the original prompt to decode with consumer-side
-   `kv_transfer_params`. Those parameters identify the transfer, prefill engine,
-   bootstrap endpoint, and, when required by the connector, remote request.
-5. The decode engine imports the KV cache and continues generation. The normal
-   Omni stage path then carries Thinker conditioning through Talker and
-   Code2Wav.
-
-### Control and data paths
-
-PD uses two different connector layers. They solve different handoffs and are
-not interchangeable.
-
-| Path | Payload | Owner |
-| --- | --- | --- |
-| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
-| PD data path | Thinker attention KV blocks | vLLM KV connector |
-| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
-
-### Detection and validation
-
-The current runtime accepts one logical prefill/decode pair per pipeline. Each
-stage may use tensor parallelism, but both sides of the pair must satisfy these
-invariants:
-
- both stages define `kv_transfer_config`
- prefill uses `kv_producer` or `kv_both`
- decode uses `kv_consumer` or `kv_both`
- both stages use the same KV connector
- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
- both stages use the same `tensor_parallel_size`
-
-These checks happen before request processing so an inconsistent PD topology
-fails during engine initialization.
-
-## What is implemented today
-
-| Area | Current state |
-| --- | --- |
-| Pair detection and topology validation | Implemented in the shared PD helper |
-| Prefill sampling preparation | Implemented for sync and async entrypoints |
-| Prefill-to-decode routing | Implemented in the current Orchestrator |
-| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
-| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
-| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
-| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
-| Performance | No non-regression or throughput/latency result is established for the current path |
-
-Because these pieces are still converging, the previous legacy `stage_args` YAML
-and `vllm serve` command have been removed. The serve CLI now accepts the
-deploy-based `stages:` schema, so retaining the legacy recipe would imply a
-launch path that current `main` rejects.
-
-## Configuration contract under development
-
-The intended deploy integration starts from the
-[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
-and expands its frozen three-stage pipeline into the four stages shown above.
-A complete integration must:
-
- split the Thinker definition without duplicating the Talker or Code2Wav model
-  contracts
- remap downstream stage IDs, input sources, and stage connectors
- inject producer and consumer `kv_transfer_config` values into the two Thinker
-  stages
- allocate separate device sets for prefill and decode
- preserve both final outputs: decode text and Code2Wav audio
- provide connector-specific routing metadata without hard-coding one connector
-  contract
-
-The required number of GPUs depends on tensor-parallel sizes and whether Talker
-and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
-for decode, and one for the two downstream audio stages.
-
-## Implementation map
-
- PD detection, validation, and sampling-parameter preparation:
-  [`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
- startup detection and connector metadata:
-  [`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
- request-scoped prefill-to-decode routing:
-  [`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
- Qwen3 downstream conditioning helpers:
-  [Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
@@ -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.
@@ -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_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
-| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+| `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
+| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
+| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |

 ### Latency

 | Metric | Type | Labels | Description |
 |--------|------|--------|-------------|
-| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
+| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |

-## Audio Modality Metrics (`vllm:omni_`)
+## Audio Modality Metrics (`vllm_omni:`)

 Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.

 | Metric | Type | Extra label | Description |
 |--------|------|-------------|-------------|
-| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
-| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
-| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
-| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
-| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
-| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
-| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
+| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
+| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |

 The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.

-## 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_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
-| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
-| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |

 ## vLLM Engine Metrics (`vllm:`)

@@ -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_` audio modality | With `--log-stats`, if pipeline has a talker stage |
-| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
+| `vllm_omni:` request tracking + latency | With `--log-stats` |
+| `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage |
+| `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
 | `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
 | `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |

@@ -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`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
+- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm_omni:`.
+- 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.
@@ -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)                                          |

@@ -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.
- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).
+- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other VAE families retain their existing mode handling and may reject explicit Wan spatial modes.
+- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
+- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.

 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:

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

@@ -59,7 +59,7 @@ def main(args):

    omni = Omni(
        model=args.model_name,
-        stage_configs_path=args.stage_configs_path,
+        deploy_config=args.deploy_config,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
        batch_timeout=args.batch_timeout,
@@ -167,10 +167,10 @@ def parse_args():
        help="Sampling rate for audio loading (default: 16000).",
    )
    parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
        type=str,
        default=None,
-        help="Path to stage configs YAML file.",
+        help="Path to a deploy config YAML file.",
    )
    parser.add_argument(
        "--log-stats",
@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
    parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
    parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
    parser.add_argument(
-        "--stage-config-path",
+        "--deploy-config",
        type=str,
        default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
        help="Path to deploy config yaml.",
@@ -1394,8 +1394,8 @@ def main() -> None:

    from vllm_omni.entrypoints.omni import Omni

-    stage_config_path = str(Path(args.stage_config_path).expanduser())
-    omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
+    deploy_config_path = str(Path(args.deploy_config).expanduser())
+    omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
    sampling_params_list = [
        SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
        for _ in range(omni.num_stages)
@@ -73,7 +73,6 @@ def parse_args():
    )
    parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
    parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
-    parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
    parser.add_argument(
        "--stream",
        action="store_true",
@@ -146,13 +145,7 @@ def main():
    else:
        bot_task = args.bot_task

-    if args.deploy_config is not None and args.stage_configs_path is not None:
-        raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
-
-    deploy_config = args.deploy_config
-    stage_configs_path = args.stage_configs_path
-    if deploy_config is None and stage_configs_path is None:
-        deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
+    deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]

    omni_kwargs = {
        "model": args.model,
@@ -168,10 +161,7 @@ def main():

    if additional_config is not None:
        omni_kwargs["additional_config"] = additional_config
-    if deploy_config is not None:
-        omni_kwargs["deploy_config"] = deploy_config
-    else:
-        omni_kwargs["stage_configs_path"] = stage_configs_path
+    omni_kwargs["deploy_config"] = deploy_config

    omni = Omni(**omni_kwargs)

@@ -272,10 +262,7 @@ def main():
    print(f"  Modality: {args.modality}")
    print(f"  Prompt task: {task}")
    print(f"  Bot task: {bot_task}")
-    if deploy_config is not None:
-        print(f"  Deploy config: {deploy_config}")
-    else:
-        print(f"  Stage config: {stage_configs_path}")
+    print(f"  Deploy config: {deploy_config}")
    print(f"  Num stages: {omni.num_stages}")
    if args.modality in ("text2img", "img2img"):
        print(f"  Inference steps: {args.steps}")
@@ -452,12 +452,6 @@ def parse_args():
        default=None,
        help="Path to a .txt file with one prompt per line (preferred).",
    )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to a stage configs file (deprecated; prefer --deploy-config).",
-    )
    parser.add_argument(
        "--deploy-config",
        type=str,
@@ -484,10 +484,10 @@ def parse_args():
        help="Path to a .txt file with one prompt per line (preferred).",
    )
    parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
        type=str,
        default=None,
-        help="Path to a stage configs file.",
+        help="Path to a deploy config YAML.",
    )
    parser.add_argument(
        "--video-path",
@@ -232,15 +232,8 @@ def main(args):
        raise ValueError(f"Unknown query type: {args.query_type}")

    # Initialize vLLM-Omni with Step-Audio2
-    # Resolve deploy config path. Keep explicit stage_configs_path for legacy
-    # custom configs, but use bundled deploy configs by default.
-    deploy_config_path = None
-    stage_config_path = None
-    if args.stage_configs_path and args.deploy_config:
-        raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
-    if args.stage_configs_path:
-        stage_config_path = args.stage_configs_path
-    elif args.deploy_config:
+    # Resolve an explicit deploy config or use the bundled task default.
+    if args.deploy_config:
        deploy_config_path = args.deploy_config
    else:
        configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
@@ -262,10 +255,7 @@ def main(args):
        "ray_address": args.ray_address,
        "trust_remote_code": True,
    }
-    if stage_config_path is not None:
-        omni_kwargs["stage_configs_path"] = stage_config_path
-    else:
-        omni_kwargs["deploy_config"] = deploy_config_path
+    omni_kwargs["deploy_config"] = deploy_config_path

    omni_llm = Omni(**omni_kwargs)

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

    # Build LoRA request when --lora-path is set
@@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 logger = logging.getLogger(__name__)

 DEFAULT_MODEL = "fishaudio/s2-pro"
-DEFAULT_STAGE_CONFIG = os.path.join(
+DEFAULT_DEPLOY_CONFIG = os.path.join(
    os.path.dirname(__file__),
    "..",
    "..",
@@ -152,10 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
    )
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
-    else:
-        omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
+    omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
    return omni_kwargs


@@ -278,12 +275,6 @@ def parse_args():
        default=None,
        help="Deploy config YAML path (under vllm_omni/deploy/).",
    )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to stage configs YAML (legacy, prefer --deploy-config).",
-    )
    parser.add_argument(
        "--output-dir",
        default="output_audio",
@@ -145,7 +145,7 @@ def main(args):

    omni = Omni(
        model=args.model,
-        stage_configs_path=deploy_config_path,
+        deploy_config=deploy_config_path,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
    )
@@ -490,10 +490,10 @@ def parse_args():
        help="Path to a .txt file with one prompt per line (preferred).",
    )
    parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
        type=str,
        default=None,
-        help="Path to a stage configs file.",
+        help="Path to a deploy config YAML.",
    )
    parser.add_argument(
        "--audio-path",
@@ -73,11 +73,12 @@ 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
        user_guide/quantization/index.md: user_guide/quantization/overview.md
-        configuration/pd_disaggregation.md: features/pd_disaggregation.md
        design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
        design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
  - glightbox
@@ -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``)."""
@@ -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):
@@ -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 = mid_rank_pad(x)
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
+    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 = last_rank_pad(x)
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
+    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
@@ -199,7 +199,7 @@ def _generate_bagel_image(
    config_path = get_deploy_config_path("ci/bagel.yaml")
    omni_kwargs: dict[str, Any] = {
        "model": "ByteDance-Seed/BAGEL-7B-MoT",
-        "stage_configs_path": config_path,
+        "deploy_config": config_path,
        "stage_init_timeout": 300,
    }
    if diffusion_quantization_config:
@@ -95,6 +95,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
        "resolution": 512,
        "width": None,
        "height": None,
+        "num_inference_steps": 1,
    }


@@ -314,7 +314,7 @@ def test_bagel_text2img_mooncake_connector(run_level):
        temp_config_file = _resolve_deploy_config(temp_config_file, run_level)
        with OmniRunner(
            "ByteDance-Seed/BAGEL-7B-MoT",
-            stage_configs_path=temp_config_file,
+            deploy_config=temp_config_file,
            stage_init_timeout=300,
        ) as runner:
            generated_image = _generate_bagel_image(runner.omni)
@@ -251,7 +251,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
    config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
    with OmniRunner(
        "ByteDance-Seed/BAGEL-7B-MoT",
-        stage_configs_path=config_path,
+        deploy_config=config_path,
    ) as runner:
        generated_image = _generate_bagel_img2img(runner.omni, input_image)
        if run_level == "advanced_model":
@@ -266,7 +266,7 @@ def test_bagel_text2img_shared_memory_connector(run_level):
    config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
    with OmniRunner(
        "ByteDance-Seed/BAGEL-7B-MoT",
-        stage_configs_path=config_path,
+        deploy_config=config_path,
    ) as runner:
        generated_image = _generate_bagel_text2img(runner.omni)
        if run_level == "advanced_model":
@@ -17,7 +17,7 @@ from vllm_omni.platforms import current_omni_platform
 # Override in CI if needed: AUDIOX_TEST_MODEL=<model-or-local-path>
 models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]

-# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
 _OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models]

 pytestmark = [
@@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
 BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml")
 DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"

-# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
+# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
 _OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG)


@@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
    ]


-# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
+# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
 _cosy_deployment = get_deploy_config_path("cosyvoice3.yaml")
 _cosy_model_path = str(_resolve_model_dir())
 _OMNI_RUNNER_PARAMS = [
@@ -17,7 +17,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 MODEL = "baidu/ERNIE-Image"

-# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
 _OMNI_RUNNER_PARAM = (
    MODEL,
    None,
@@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None:

    with OmniRunner(
        MODEL,
-        stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
+        deploy_config=_get_deploy_config(async_chunk=async_chunk),
        stage_init_timeout=600,
        trust_remote_code=True,
    ) as omni_runner:
@@ -52,7 +52,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 def omni() -> Generator[Omni, None, None]:
    with OmniRunner(
        MODEL_NAME,
-        stage_configs_path=str(STAGE_CONFIG_PATH),
+        deploy_config=str(STAGE_CONFIG_PATH),
        trust_remote_code=True,
    ) as runner:
        yield runner.omni
@@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]:
    with OmniRunner(
        model,
        seed=42,
-        stage_configs_path=None,
        model_class_name="LongCatVideoAvatarPipeline",
        additional_config={
            "model_type": "avatar-v1.5",
@@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16
 MODEL_PATH = "bytedance-research/MammothModa2-Preview"
 T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")

-_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
+_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)

 # Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
 _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
@@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"]
 _CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")


-test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models]
+test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models]


 def get_question(prompt_type: str = "text") -> str:
@@ -25,7 +25,7 @@ from vllm_omni import Omni
 MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
 STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml")

-# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
 _OMNI_RUNNER_PARAM = (
    MODEL_NAME,
    STAGE_CONFIG,
@@ -20,13 +20,11 @@ from tests.helpers.stage_config import get_deploy_config_path
 MODEL = "k2-fsa/OmniVoice"
 DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")

-# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
-# The migrated test passes deploy_config through extra Omni kwargs.
+# OmniRunner tuple: model, deploy config path, extra Omni kwargs.
 _OMNI_RUNNER_PARAM = (
    MODEL,
-    None,
+    DEPLOY_CONFIG,
    {
-        "deploy_config": DEPLOY_CONFIG,
        "trust_remote_code": True,
        "log_stats": True,
    },
@@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
    """Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
    lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")

-    with OmniRunner(MODEL, stage_configs_path=None) as runner:
+    with OmniRunner(MODEL) as runner:
        omni = runner.omni
        baseline = _generate(omni)
        img_1x = _generate(omni, lora_request, lora_scale=1.0)
@@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img(
 def test_sensenova_u1_img2img(run_level):
    """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
    input_image = _load_input_image()
-    with OmniRunner(
-        "SenseNova/SenseNova-U1-8B-MoT",
-        stage_configs_path=None,
-    ) as runner:
+    with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
        generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
        if run_level == "full_model" or run_level == "advanced_model":
            _validate_pixels(generated_image)
@@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit():
    input_image = _load_input_image()
    with OmniRunner(
        "SenseNova/SenseNova-U1-8B-MoT",
-        stage_configs_path=None,
        cache_backend="cache_dit",
    ) as runner:
        _generate_sensenova_u1_img2img(runner.omni, input_image)
@@ -151,10 +151,7 @@ def _generate_sensenova_u1_image(
@hardware_test(res={"cuda": "H100"})
 def test_sensenova_u1_text2img(run_level):
    """Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML)."""
-    with OmniRunner(
-        "SenseNova/SenseNova-U1-8B-MoT",
-        stage_configs_path=None,
-    ) as runner:
+    with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
        generated_image = _generate_sensenova_u1_image(runner.omni)
        if run_level == "full_model" or run_level == "advanced_model":
            _validate_pixels(generated_image)
@@ -168,7 +165,6 @@ def test_sensenova_u1_text2img_cache_dit():
    """Test SenseNova-U1 text2img with Cache-DiT enabled."""
    with OmniRunner(
        "SenseNova/SenseNova-U1-8B-MoT",
-        stage_configs_path=None,
        cache_backend="cache_dit",
    ) as runner:
        _generate_sensenova_u1_image(runner.omni)
@@ -159,7 +159,7 @@ def test_soulxsinger_multistage_from_audio(

    with OmniRunner(
        model,
-        stage_configs_path=get_deploy_config_path(deploy_yaml),
+        deploy_config=get_deploy_config_path(deploy_yaml),
        async_chunk=False,
    ) as runner:
        sampling = OmniDiffusionSamplingParams(
@@ -212,7 +212,7 @@ def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) ->

    with OmniRunner(
        str(base_dir),
-        stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
+        deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"),
        async_chunk=False,
    ) as runner:
        sampling = OmniDiffusionSamplingParams(
@@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput

 MODEL = "stepfun-ai/Step-Audio-2-mini"
 DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml")
-TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})]
+TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})]

 pytestmark = [pytest.mark.slow, pytest.mark.tts]

@@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming():
    async def _run():
        async_omni = AsyncOmni(
            model=MODEL,
-            stage_configs_path=STAGE_CONFIG,
+            deploy_config=STAGE_CONFIG,
            stage_init_timeout=300,
            enforce_eager=True,
        )
@@ -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."""
@@ -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)

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

    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-    engine._strip_single_engine_args = lambda kwargs: kwargs

    _, stage_configs = engine._resolve_stage_configs(
        "dummy-model",
        {
-            "stage_configs_path": "dummy.yaml",
+            "deploy_config": "dummy.yaml",
            "additional_config": {"torchair_graph_config": {"enabled": True}},
        },
        trust_remote_code=False,
@@ -516,12 +515,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(
    )

    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-    engine._strip_single_engine_args = lambda kwargs: kwargs

    _, stage_configs = engine._resolve_stage_configs(
        "dummy-model",
        {
-            "stage_configs_path": "dummy.yaml",
+            "deploy_config": "dummy.yaml",
            "quantization_config": {"method": "bitsandbytes"},
        },
        trust_remote_code=False,
@@ -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"
@@ -381,7 +381,6 @@ class TestLoadAndResolveStageConfigs:

        config_path, stage_configs, _ = load_and_resolve_stage_configs(
            model="black-forest-labs/FLUX.2-klein-4B",
-            stage_configs_path=None,
            kwargs=kwargs,
            trust_remote_code=False,
            default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
@@ -390,9 +389,7 @@ class TestLoadAndResolveStageConfigs:
        assert len(stage_configs) == 1
        assert "dtype" in stage_configs[0]["engine_args"]

-    def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
-        self, tmp_path, mocker: MockerFixture
-    ):
+    def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
        deploy_path = tmp_path / "qwen3_multi.yaml"
        deploy_path.write_text(
            'stages:\n  - stage_id: 0\n    devices: "0"\n  - stage_id: 1\n    devices: "1,2,3"\n    num_replicas: 3\n',
@@ -413,18 +410,23 @@ class TestLoadAndResolveStageConfigs:
            "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
            return_value=(returned_stage_configs, None),
        )
+        cli_overrides = {
+            "interleave_mm_strings": True,
+            "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
+            "gpu_memory_utilization": 0.55,
+        }

        config_path, stage_configs, _ = load_and_resolve_stage_configs(
            model="dummy-model",
-            stage_configs_path=str(deploy_path),
-            kwargs={},
+            kwargs=cli_overrides,
            trust_remote_code=True,
+            deploy_config_path=str(deploy_path),
        )

        load_stage_configs.assert_called_once_with(
            "dummy-model",
            trust_remote_code=True,
-            base_engine_args={},
+            base_engine_args=cli_overrides,
            deploy_config_path=str(deploy_path),
            stage_overrides=None,
            strategy_config_path=None,
@@ -2635,7 +2635,7 @@ class OmniRunner:
        # affects the test runner wrapper.
        init_timeout: int = 1800,
        log_stats: bool = False,
-        stage_configs_path: str | None = None,
+        deploy_config: str | None = None,
        **kwargs,
    ) -> None:
        startup_t0 = time.perf_counter()
@@ -2653,7 +2653,7 @@ class OmniRunner:
            stage_init_timeout=stage_init_timeout,
            batch_timeout=batch_timeout,
            init_timeout=init_timeout,
-            stage_configs_path=stage_configs_path,
+            deploy_config=deploy_config,
            **kwargs,
        )
        startup_s = time.perf_counter() - startup_t0
@@ -3179,12 +3179,7 @@ def iter_omni_server(
                print("OmniServer stopping...")
        else:
            if stage_config_path is not None:
-                from vllm_omni.entrypoints.utils import is_new_format_deploy_config
-
-                if is_new_format_deploy_config(stage_config_path):
-                    server_args += ["--deploy-config", stage_config_path]
-                else:
-                    server_args += ["--stage-configs-path", stage_config_path]
+                server_args += ["--deploy-config", stage_config_path]

            with (
                OmniServer(
@@ -3237,7 +3232,7 @@ def iter_omni_runner(
        model = model_prefix + model
        if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
            model = resolve_tiny_model_path(model)
-        with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
+        with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner:
            print("OmniRunner started successfully")
            yield runner
            print("OmniRunner stopping...")
@@ -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):
-        # Image / video / text output types are out of scope for this module.
+    def test_diffusion_path_full(self):
+        stub = _StubModMetrics()
+        stage_metrics = _Bag(
+            stage_gen_time_ms=1000.0,
+            diffusion_metrics={
+                "preprocess_time_s": 0.05,
+                "diffusion_engine_exec_time_s": 1.2,
+                "postprocess_time_s": 0.3,
+                "image_num": 1,
+                "resolution": 1024,
+                "num_inference_steps": 50,
+            },
+        )
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=stage_metrics,
+            engine_outputs=_Bag(),
+        )
+        assert ("observe_diffusion_preprocess", "2", "0", 0.05) in stub.calls
+        assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
+        assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls
+        assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls
+
+    def test_diffusion_per_step_skipped_without_num_inference_steps(self):
+        stub = _StubModMetrics()
+        stage_metrics = _Bag(
+            stage_gen_time_ms=1000.0,
+            diffusion_metrics={
+                "diffusion_engine_exec_time_s": 1.2,
+                "preprocess_time_s": 0.05,
+                "postprocess_time_s": 0.3,
+            },
+        )
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=stage_metrics,
+            engine_outputs=_Bag(),
+        )
+        assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
+        assert not any(c[0] == "observe_diffusion_exec_per_step" for c in stub.calls)
+
+    def test_diffusion_metrics_none_skipped(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=_Bag(stage_gen_time_ms=500.0, diffusion_metrics=None),
+            engine_outputs=_Bag(),
+        )
+        assert not any(c[0].startswith("observe_diffusion") for c in stub.calls)
+
+    def test_diffusion_zero_values_observed(self):
+        stub = _StubModMetrics()
+        stage_metrics = _Bag(
+            stage_gen_time_ms=500.0,
+            diffusion_metrics={
+                "preprocess_time_s": 0.0,
+                "diffusion_engine_exec_time_s": 0.0,
+                "postprocess_time_s": 0.0,
+            },
+        )
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=stage_metrics,
+            engine_outputs=_Bag(),
+        )
+        assert ("observe_diffusion_preprocess", "2", "0", 0.0) in stub.calls
+        assert ("observe_diffusion_exec", "2", "0", 0.0) in stub.calls
+        assert ("observe_diffusion_postprocess", "2", "0", 0.0) in stub.calls
+
+    def test_non_audio_output_type_skips_audio_but_still_checks_diffusion(self):
        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()
@@ -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,
@@ -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}."
            )

@@ -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],
@@ -31,12 +31,12 @@ stages:
    output_connectors:
      to_stage_1: shm
    default_sampling_params:
-      temperature: 1.0
-      top_p: 0.95
-      top_k: 50
+      temperature: 1.7
+      top_p: 0.8
+      top_k: 25
      max_tokens: 4096
      seed: 42
-      repetition_penalty: 1.1
+      repetition_penalty: 1.0

  - stage_id: 1
    max_num_seqs: 64
@@ -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"}, (
-            "vae_parallel_mode must be one of {'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', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
            f"but got {self.vae_parallel_mode!r}."
        )
        if self.allgather_degree > 1:
@@ -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 split_dim is None:
+        if self.distributed_executor.parallel_mode == "tile":
            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(
-                "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
-                "to match the DIT group size; falling back to tiled decode. "
-                "requested=%s dit_group=%s split_dim=%s",
-                requested_size,
-                world_size,
-                split_dim,
-            )
+            if self.distributed_executor.parallel_mode == "auto":
+                logger.debug(
+                    "Wan VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
+                    requested_size,
+                    world_size,
+                )
+            elif self._spatial_shard_decode_split_dim() is not None:
+                logger.warning(
+                    "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
+                    "to match the DIT group size; falling back to tiled decode. "
+                    "requested=%s dit_group=%s split_dim=%s",
+                    requested_size,
+                    world_size,
+                    self._spatial_shard_decode_split_dim(),
+                )
            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,
@@ -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",
-        split_padding: tuple[int, int] | None = None,
+        defer_split_after_padding: bool = False,
    ) -> None:
        super().__init__()
        self.padding = padding
-        self.split_dim = split_dim
-        default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1])
-        self.split_padding = split_padding or default_split_padding
+        self.defer_split_after_padding = defer_split_after_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",
-        split_padding: tuple[int, int] | None = None,
+        deferred_padding: tuple[int, int, 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.split_tensor_dim = _spatial_dim(split_dim)
-        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
-        self.halo_size = (kernel_extent - 1) // 2
-        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
-        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
-        if split_dim == "height":
-            if split_padding is None:
-                split_padding = (pad_h, pad_h)
-            self._non_split_padding = (pad_w, pad_w, 0, 0)
-            self.kernel_extent = self.kernel_size[-2]
-            self.stride_extent = self.stride[-2]
-        else:
-            if split_padding is None:
-                split_padding = (pad_w, pad_w)
-            self._non_split_padding = (0, 0, pad_h, pad_h)
-            self.kernel_extent = self.kernel_size[-1]
-            self.stride_extent = self.stride[-1]
-        self.split_pad_left, self.split_pad_right = split_padding
-        self._halo_recv_top_buf: torch.Tensor | None = None
-        self._halo_recv_bottom_buf: torch.Tensor | None = None
-        self._trim_cache: dict[int, tuple[int, int, int]] = {}
+        self._direct_padding = source.padding
+        self._direct_padding_mode = source.padding_mode
+        self._direct_reversed_padding = source._reversed_padding_repeated_twice
+        self._deferred_padding = deferred_padding
+        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
+        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
+            raise ValueError(f"Wan spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
+        self._spatial_pad_h = int(source_padding[-2])
+        self._spatial_pad_w = int(source_padding[-1])
+        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
+        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
+        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}

    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,
-            split_dim=self.split_dim,
+            halo_size=halo_size,
+            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)
-        return _zero_invalid_extent(out, 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=split_dim)

-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
-        trim_params = self._trim_cache.get(local_extent)
+    def _get_trim_params(
+        self,
+        local_extent: int,
+        *,
+        split_dim: str,
+        halo_size: int,
+        split_pad_left: int,
+        split_pad_right: int,
+        kernel_extent: int,
+        stride_extent: int,
+    ) -> tuple[int, int, int]:
+        cache_key = (split_dim, local_extent)
+        trim_params = self._trim_cache.get(cache_key)
        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,
-                pad_before=self.split_pad_left,
-                pad_after=self.split_pad_right,
-                kernel_extent=self.kernel_extent,
-                stride_extent=self.stride_extent,
+                halo_size=halo_size,
+                pad_before=split_pad_left,
+                pad_after=split_pad_right,
+                kernel_extent=kernel_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.split_pad_left = int(self._source_padding[2])
-            self.split_pad_right = int(self._source_padding[3])
-            self.kernel_extent = self.kernel_size[-2]
-            self.stride_extent = self.stride[-2]
-            padding = (
-                self._source_padding[0],
-                self._source_padding[1],
-                0,
-                0,
-                self._source_padding[4],
-                self._source_padding[5],
-            )
-        else:
-            self.split_pad_left = int(self._source_padding[0])
-            self.split_pad_right = int(self._source_padding[1])
-            self.kernel_extent = self.kernel_size[-1]
-            self.stride_extent = self.stride[-1]
-            padding = (
-                0,
-                0,
-                self._source_padding[2],
-                self._source_padding[3],
-                self._source_padding[4],
-                self._source_padding[5],
-            )
-        self.halo_size = (self.kernel_extent - 1) // 2
-        self._padding = padding if self.halo_size > 0 else self._source_padding
-        self._halo_recv_top_buf: torch.Tensor | None = None
-        self._halo_recv_bottom_buf: torch.Tensor | None = None
-        self._trim_cache: dict[int, tuple[int, int, int]] = {}
+        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
+        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
+        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}

    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,
-            split_dim=self.split_dim,
+            halo_size=halo_size,
+            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)
-        return _zero_invalid_extent(out, 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=split_dim)

-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
-        trim_params = self._trim_cache.get(local_extent)
+    def _get_trim_params(
+        self,
+        local_extent: int,
+        *,
+        split_dim: str,
+        halo_size: int,
+        split_pad_left: int,
+        split_pad_right: int,
+        kernel_extent: int,
+        stride_extent: int,
+    ) -> tuple[int, int, int]:
+        cache_key = (split_dim, local_extent)
+        trim_params = self._trim_cache.get(cache_key)
        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,
-                pad_before=self.split_pad_left,
-                pad_after=self.split_pad_right,
-                kernel_extent=self.kernel_extent,
-                stride_extent=self.stride_extent,
+                halo_size=halo_size,
+                pad_before=split_pad_left,
+                pad_after=split_pad_right,
+                kernel_extent=kernel_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,
-                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
+                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
            ),
        )
        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
-                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
+                deferred_padding = tuple(int(value) for value in prev.padding)
        setattr(
            parent,
            name,
            WanDistConv2d(
                child,
                group,
-                split_dim=split_dim,
-                split_padding=split_padding,
+                deferred_padding=deferred_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
-    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
-    instance and is applied only once (subsequent calls are no-ops). A given
-    instance is bound to a single ``split_dim``; switching between
-    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
-    otherwise.
+    convolutions/padding for context-aware variants and attaching a dedicated
+    spatial forward. Outside that forward the replacements execute their exact
+    local PyTorch behavior, so a later auto-selected tile request remains safe.
+    The same installed decoder can alternate height and width sharding.

    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 %s-sharded decode.", split_dim)
+    logger.info("Installed Wan VAE dynamic-axis spatial-shard decode.")


 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,)
@@ -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
@@ -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
    }
@@ -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
-        # beforehand. The stage CLI exposes the same deploy YAML through
-        # stage_configs_path.
-        deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path")
+        # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand.
+        deploy_config_path = kwargs.get("deploy_config")
        # ``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,
-            base_kwargs,
+            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,
@@ -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
-              configurations. If None, configurations are resolved from model
-              pipeline factory.
+            - deploy_config: Optional path to a deploy YAML. If None,
+              configurations are resolved from the model pipeline factory.
            - log_stats: Whether to enable statistics logging.
            - stage_init_timeout: Timeout for per-stage initialization.
            - init_timeout: Total timeout for orchestrator startup.
@@ -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.
@@ -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)

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

@@ -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).
-            Mutually exclusive with ``stage_configs_path``.
+        deploy_config_path: Optional path to a deploy YAML.
        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:
-        raise ValueError(
-            "--stage-configs-path and --deploy-config are mutually exclusive: "
-            "they use different path resolution rules and loading paths. "
-            "Use --deploy-config for new-format YAMLs (preferred); "
-            "--stage-configs-path is kept only for the legacy `stage_args` format "
-            "and will be removed in a future release."
-        )
-    if stage_configs_path is not None and deploy_config_path is None:
-        if not os.path.exists(stage_configs_path):
-            raise FileNotFoundError(
-                f"--stage-configs-path {stage_configs_path!r} does not exist. "
-                "Legacy `stage_configs/` yamls were replaced by `vllm_omni/deploy/<model>.yaml`; "
-                "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
+    config_path = deploy_config_path if deploy_config_path is not None else 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,
+        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:
-            logger.warning(
-                "--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 = []

    stage_configs = filter_stages(config_path, stage_configs, kwargs)
    logger.debug(f"stage_configs: {stage_configs}")
@@ -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)
 # ============================================================================
@@ -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)
-    # `stage_metrics.audio_generated_frames` is the legacy per-chunk
-    # accumulator field on StageRequestStats. No production path currently
-    # fills it, so the fallback below is the live source — but we leave the
-    # field lookup in place in case the accumulator gets re-wired upstream.
-    n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
-    if n_frames == 0:
-        n_frames = count_audio_frames(mm_out)
-    mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
-    duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
-    if duration_s > 0:
-        mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
-        mod_metrics.observe_audio_rtf(
-            stage_label,
-            replica_label,
-            defs.compute_audio_rtf(gen_time_s, duration_s),
-        )
-    else:
-        # Request completed (finish_reason ∈ {stop, length} — error paths
-        # don't reach finalize) but no audio samples were produced. Covers
-        # silent `return None` skips in the talker→code2wav stage
-        # processors and the `parsed.append((0,0))` malformed-length path
-        # in qwen3-tts code2wav. raise-paths surface via the upstream
-        # vllm:request_success_total{finished_reason="error"} channel and
-        # never reach this branch.
-        mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
-    # audio_underrun / continuity are emitted from the streaming path in
-    # observe_audio_streaming_finalize; finalize is too late for the
-    # per-chunk timeline they need.
+    if output_type == "audio":
+        sample_rate = defs.resolve_audio_sample_rate(mm_out)
+        n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
+        if n_frames == 0:
+            n_frames = count_audio_frames(mm_out)
+        mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
+        duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
+        if duration_s > 0:
+            mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
+            mod_metrics.observe_audio_rtf(
+                stage_label,
+                replica_label,
+                defs.compute_audio_rtf(gen_time_s, duration_s),
+            )
+        else:
+            mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
+
+    dm = getattr(stage_metrics, "diffusion_metrics", None)
+    if dm:
+        _key_map = {
+            "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec,
+            "preprocess_time_s": mod_metrics.observe_diffusion_preprocess,
+            "postprocess_time_s": mod_metrics.observe_diffusion_postprocess,
+        }
+        for key, observe_fn in _key_map.items():
+            val = dm.get(key)
+            if val is not None:
+                observe_fn(stage_label, replica_label, float(val))
+
+        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(
@@ -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,
@@ -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,
        ):
@@ -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(
@@ -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(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_padded,
-                num_tokens_across_dp=num_tokens_across_dp,
-                in_profile_run=is_profile,
-                num_actual_tokens=num_tokens_padded,
-                aclgraph_runtime_mode=cudagraph_runtime_mode,
-                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,
+            with (
+                self.maybe_randomize_inputs(input_ids, inputs_embeds, randomize_inputs=randomize_inputs),
+                set_ascend_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_padded,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    in_profile_run=is_profile,
+                    num_actual_tokens=num_tokens_padded,
+                    aclgraph_runtime_mode=cudagraph_runtime_mode,
+                    batch_descriptor=batch_desc,
+                    model_instance=self.model,
+                    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:
Author	SHA1	Message	Date
Hongsheng Liu	089851e55a	Merge origin/main into codex-vae-spatial-decode-auto Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 23:49:11 +08:00
Yukim1andHongsheng Liu	4b66f62e10	[Refactor] Remove internal stage_configs_path plumbing (#5741 ) Signed-off-by: zwhzzz0821 <2831474076@qq.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 21:29:03 +08:00
Weiming Liao	f41451f64c	[NPU] Upgrade to v0.27.0 (#6096 ) Signed-off-by: Weiming Liao <liaowm5@gmail.com>	2026-08-13 20:01:01 +08:00
Hongsheng Liu	bbe6ccc512	[Doc] Refresh community documentation (#6141 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 17:04:06 +08:00
Canlin Guo	bbb709cfb5	[BugFix] Use MOSS-TTS-Local official sample params (#6156 ) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>	2026-08-13 08:56:26 +00:00
vraitiandClaude Opus 4.6	c98517ce7f	[Metrics] Add Diffusion Metrics (#4755 ) Signed-off-by: vraiti <vraiti@redhat.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>	2026-08-13 06:59:13 +00:00
Hongsheng Liu	3466ff9b98	docs: standardize VAE Parallel naming Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-12 23:06:06 +08:00
Hongsheng Liu	e16473c6b5	docs: align VAE parallel file convention Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-12 22:24:58 +08:00
Hongsheng Liu	14d67bcd60	docs: align VAE parallelism file names Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-12 22:15:51 +08:00
Hongsheng Liu	5234e3599c	docs: generalize VAE parallelism design Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-12 22:07:54 +08:00
Hongsheng Liu	3aefb31727	docs(diffusion): remove VAE policy attribution Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-12 12:01:32 +08:00
Hongsheng Liu	c8210b8b5c	fix(diffusion): address Wan auto decode review feedback Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-12 11:49:47 +08:00
Hongsheng Liu	879d19c4c2	Merge branch 'main' into codex-vae-spatial-decode-auto	2026-08-12 11:38:12 +08:00
Hongsheng Liu	af1825e07d	feat(diffusion): select Wan VAE decode mode automatically Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-11 17:36:35 +08:00