[Docs] Redirect the former PD guide URL

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

docker/Dockerfile.npu

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

docker/Dockerfile.npu.a3

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

docker/Dockerfile.npu.ci

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

docker/Dockerfile.npu.ci.a3

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

docs/.nav.yml

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

docs/community/contact_us.md

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

docs/community/governance.md

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

docs/community/meetings.md

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

docs/community/meetups.md

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

docs/community/volunteers.md

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

docs/configuration/README.md

@@ -14,10 +14,6 @@ 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

docs/configuration/pd_disaggregation.md

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

docs/contributing/README.md

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

docs/design/architecture_overview.md

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

docs/design/feature/vae_parallel.md

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

docs/design/index.md

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

docs/design/metrics.md

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

docs/features/README.md

@@ -20,13 +20,19 @@ 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:
 
-- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
+- Generic, OmniConnector-based
+  [Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
+  yet a standalone User Guide feature. Its
   [OmniConnector implementations](../design/index.md#communication) describe
-  topology and transport contracts. Practical deployment configuration remains
+  transport contracts, while practical connector configuration remains under
-  under [Pipeline and deploy configurations](../configuration/stage_configs.md).
+  [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

docs/features/pd_disaggregation.md

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

docs/usage/faq.md

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

docs/usage/metrics.md

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

docs/user_guide/diffusion/parallelism/overview.md

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

docs/user_guide/diffusion/parallelism/vae_parallelism.md

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

docs/user_guide/diffusion_features.md

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

examples/offline_inference/covo_audio/end2end.py

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

examples/offline_inference/dynin_omni/end2end.py

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

examples/offline_inference/hunyuan_image3/end2end.py

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

examples/offline_inference/minicpmo/end2end.py

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

examples/offline_inference/qwen3_omni/end2end.py

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

examples/offline_inference/step_audio2/end2end.py

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

examples/offline_inference/text_to_image/text_to_image.py

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

examples/offline_inference/text_to_speech/fish_speech/end2end.py

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

examples/offline_inference/text_to_speech/glm_tts/end2end.py

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

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

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

mkdocs.yml

@@ -73,12 +73,11 @@ 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

tests/config/test_config_factory.py

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

tests/config/test_omni_config.py

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

tests/diffusion/distributed/test_wan_spatial_shard.py

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

tests/diffusion/quantization/test_quantization_fp8.py

@@ -199,7 +199,7 @@ def _generate_bagel_image(
     config_path = get_deploy_config_path("ci/bagel.yaml")
     omni_kwargs: dict[str, Any] = {
         "model": "ByteDance-Seed/BAGEL-7B-MoT",
-        "deploy_config": config_path,
+        "stage_configs_path": config_path,
         "stage_init_timeout": 300,
     }
     if diffusion_quantization_config:

tests/diffusion/test_diffusion_output_formatter.py

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

tests/distributed/omni_connectors/test_bagel_mooncake_connector.py

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

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

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

tests/e2e/offline_inference/test_audiox_model_expansion.py

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

tests/e2e/offline_inference/test_bagel_expansion.py

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

tests/e2e/offline_inference/test_cosyvoice3_expansion.py

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

tests/e2e/offline_inference/test_ernie_image_expansion.py

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

tests/e2e/offline_inference/test_glm_tts_expansion.py

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

tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py

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

tests/e2e/offline_inference/test_longcat_video_avatar.py

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

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

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

tests/e2e/offline_inference/test_minicpmo_4_5.py

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

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

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

tests/e2e/offline_inference/test_omnivoice_expansion.py

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

tests/e2e/offline_inference/test_sensenova_u1_expansion.py

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

tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py

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

tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py

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

tests/e2e/offline_inference/test_soulxsinger_expansion.py

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

tests/e2e/offline_inference/test_step_audio2_expansion.py

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

tests/e2e/offline_inference/test_voxtral_tts_expansion.py

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

tests/engine/test_arg_utils.py

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

tests/engine/test_async_omni_engine_abort.py

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

tests/engine/test_single_stage_mode.py

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

tests/entrypoints/test_async_omni.py

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

tests/entrypoints/test_async_omni_diffusion_config.py

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

tests/entrypoints/test_serve.py

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

tests/entrypoints/test_utils.py

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

tests/helpers/runtime.py

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

tests/metrics/test_modality.py

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

tools/baselines/omni_hash_baseline.py

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

vllm_omni/config/omni_config.py

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

vllm_omni/config/stage_config.py

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

vllm_omni/deploy/moss_tts_local.yaml

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

vllm_omni/diffusion/data.py

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

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

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

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

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

vllm_omni/diffusion/output_formatter.py

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

vllm_omni/engine/arg_utils.py

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

vllm_omni/engine/async_omni_engine.py

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

vllm_omni/entrypoints/async_omni.py

@@ -123,8 +123,9 @@ class AsyncOmni(EngineClient, OmniBase):
     Args:
         model: Model name or path to load.
         **kwargs: Additional keyword arguments.
-            - deploy_config: Optional path to a deploy YAML. If None,
+            - stage_configs_path: Optional path to YAML file containing stage
-              configurations are resolved from the model pipeline factory.
+              configurations. If None, configurations are resolved from 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.

vllm_omni/entrypoints/cli/serve.py

@@ -790,10 +790,9 @@ class OmniServeCommand(CLISubcommand):
             "--vae-parallel-mode",
             type=str,
             default="tile",
-            choices=["tile", "auto", "spatial_shard_height", "spatial_shard_width"],
+            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
-            "'tile' (default) uses patch/tile decode; 'auto' selects the longer spatial-shard axis for "
+            "'tile' (default) uses patch/tile parallel decode; "
-            "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. "
@@ -946,6 +945,8 @@ def run_headless(args: TrackingNamespace) -> None:
 
     config_path, stage_configs, _ = load_and_resolve_stage_configs(
         model,
+        # The serve CLI no longer accepts legacy stage_args YAMLs.
+        None,
         args_dict,
         # store_true cannot express an explicit False: absent maps to None
         # ("not specified") so the deploy yaml's per-stage value applies.

vllm_omni/entrypoints/omni_base.py

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

vllm_omni/entrypoints/openai/api_server.py

@@ -679,7 +679,7 @@ async def build_async_omni_from_stage_config(
         EngineClient instance (AsyncOmni) ready for use
 
     Note:
-        Stage configurations are loaded from ``args.deploy_config`` when provided,
+        Stage configurations are loaded from args.stage_configs_path if 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, "deploy_config", None),
+        config_path=getattr(args, "stage_configs_path", None) or getattr(args, "deploy_config", None),
     ):
         from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler
 

vllm_omni/entrypoints/utils.py

@@ -6,6 +6,7 @@ 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
@@ -29,6 +30,15 @@ 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",
 }
@@ -596,6 +606,7 @@ 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,
@@ -604,15 +615,17 @@ 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 a deploy YAML or model defaults.
+    """Load stage configurations from model or YAML file with fallback to 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 a deploy YAML.
+        deploy_config_path: Optional path to deploy YAML (new format).
+            Mutually exclusive with ``stage_configs_path``.
         stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
             Keys are stage_id strings, values are dicts of overrides.
         strategy_config_path: Optional path to a composable-parallel
@@ -623,21 +636,65 @@ 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.
     """
-    config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model)
+    if stage_configs_path is not None and deploy_config_path is not None:
-    stage_configs, omni_lb_policy = load_stage_configs_from_model(
+        raise ValueError(
-        model,
+            "--stage-configs-path and --deploy-config are mutually exclusive: "
-        trust_remote_code=trust_remote_code,
+            "they use different path resolution rules and loading paths. "
-        base_engine_args=kwargs,
+            "Use --deploy-config for new-format YAMLs (preferred); "
-        deploy_config_path=deploy_config_path,
+            "--stage-configs-path is kept only for the legacy `stage_args` format "
-        stage_overrides=stage_overrides,
+            "and will be removed in a future release."
-        strategy_config_path=strategy_config_path,
+        )
-    )
+    if stage_configs_path is not None and deploy_config_path is None:
-    if not stage_configs:
+        if not os.path.exists(stage_configs_path):
-        if default_stage_cfg_factory is not None:
+            raise FileNotFoundError(
-            default_stage_cfg = default_stage_cfg_factory()
+                f"--stage-configs-path {stage_configs_path!r} does not exist. "
-            stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
+                "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
         else:
-            stage_configs = []
+            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 = filter_stages(config_path, stage_configs, kwargs)
     logger.debug(f"stage_configs: {stage_configs}")

vllm_omni/metrics/definitions.py

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

vllm_omni/metrics/modality.py

@@ -78,35 +78,6 @@ _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
@@ -168,44 +139,6 @@ 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,
@@ -228,49 +161,43 @@ 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)
 
-    if output_type == "audio":
+    sample_rate = defs.resolve_audio_sample_rate(mm_out)
-        sample_rate = defs.resolve_audio_sample_rate(mm_out)
+    # `stage_metrics.audio_generated_frames` is the legacy per-chunk
-        n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
+    # accumulator field on StageRequestStats. No production path currently
-        if n_frames == 0:
+    # fills it, so the fallback below is the live source — but we leave the
-            n_frames = count_audio_frames(mm_out)
+    # field lookup in place in case the accumulator gets re-wired upstream.
-        mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
+    n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
-        duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
+    if n_frames == 0:
-        if duration_s > 0:
+        n_frames = count_audio_frames(mm_out)
-            mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
+    mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
-            mod_metrics.observe_audio_rtf(
+    duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
-                stage_label,
+    if duration_s > 0:
-                replica_label,
+        mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
-                defs.compute_audio_rtf(gen_time_s, duration_s),
+        mod_metrics.observe_audio_rtf(
-            )
+            stage_label,
-        else:
+            replica_label,
-            mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
+            defs.compute_audio_rtf(gen_time_s, duration_s),
-
+        )
-    dm = getattr(stage_metrics, "diffusion_metrics", None)
+    else:
-    if dm:
+        # Request completed (finish_reason ∈ {stop, length} — error paths
-        _key_map = {
+        # don't reach finalize) but no audio samples were produced. Covers
-            "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec,
+        # silent `return None` skips in the talker→code2wav stage
-            "preprocess_time_s": mod_metrics.observe_diffusion_preprocess,
+        # processors and the `parsed.append((0,0))` malformed-length path
-            "postprocess_time_s": mod_metrics.observe_diffusion_postprocess,
+        # in qwen3-tts code2wav. raise-paths surface via the upstream
-        }
+        # vllm:request_success_total{finished_reason="error"} channel and
-        for key, observe_fn in _key_map.items():
+        # never reach this branch.
-            val = dm.get(key)
+        mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
-            if val is not None:
+    # audio_underrun / continuity are emitted from the streaming path in
-                observe_fn(stage_label, replica_label, float(val))
+    # observe_audio_streaming_finalize; finalize is too late for the
-
+    # per-chunk timeline they need.
-        exec_time = dm.get("diffusion_engine_exec_time_s")
-        num_steps = dm.get("num_inference_steps")
-        if exec_time is not None and num_steps and num_steps > 0:
-            mod_metrics.observe_diffusion_exec_per_step(
-                stage_label,
-                replica_label,
-                float(exec_time) / int(num_steps),
-            )
 
 
 def observe_audio_first_packet(

vllm_omni/metrics/stats.py

@@ -44,7 +44,7 @@ class StageRequestStats:
     final_output_type: str | None = None
     request_id: str | None = None
     postprocess_time_ms: float = 0.0
-    diffusion_metrics: dict[str, float] = None
+    diffusion_metrics: dict[str, int] = 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: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
+            {k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
             if req_id in self.diffusion_metrics
             else None
         )
@@ -644,19 +644,8 @@ 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():
-                if value is None:
+                self.diffusion_metrics[req_id][key] += value
-                    continue
-                if key in _MS_TO_S:
-                    self.diffusion_metrics[req_id][_MS_TO_S[key]] = float(value) / 1000.0
-                else:
-                    self.diffusion_metrics[req_id][key] = value
 
     def on_forward(
         self,

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -730,9 +730,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
         has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
 
         # Run forward pass
-        defer_kv_connector_finalize = self.speculative_config is not None and (
+        clear_kv_metadata = self.speculative_config is None
-            get_pp_group().is_last_rank or self.broadcast_pp_output
-        )
         with (
             record_function_or_nullcontext("forward"),
             set_ascend_forward_context(
@@ -746,12 +744,13 @@ 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": defer_kv_connector_finalize}
+                    {"defer_finalize": not clear_kv_metadata}
                 ),
             ) as kv_connector_output,
         ):

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -381,9 +381,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
         has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
 
         # Run forward pass
-        defer_kv_connector_finalize = self.speculative_config is not None and (
+        clear_kv_metadata = self.speculative_config is None
-            get_pp_group().is_last_rank or self.broadcast_pp_output
-        )
         with (
             record_function_or_nullcontext("forward"),
             set_ascend_forward_context(
@@ -397,11 +395,12 @@ 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": defer_kv_connector_finalize}),
+                **({"defer_finalize": not clear_kv_metadata}),
             ) as kv_connector_output,
         ):
             if self.cache_config.mamba_cache_mode == "align":
@@ -655,10 +654,9 @@ 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.is_valid_runtime_mode()
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
         # 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.
@@ -919,6 +917,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
                 batch_descriptor=batch_desc,
                 model_instance=self.model,
                 has_sinks=self._has_sinks,
+                input_ids=input_ids,
                 eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
             ):
                 outputs = self.model(

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -125,10 +125,9 @@ 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.is_valid_runtime_mode()
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
         # 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.
@@ -368,21 +367,19 @@ 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 (
+            with set_ascend_forward_context(
-                self.maybe_randomize_inputs(input_ids, inputs_embeds, randomize_inputs=randomize_inputs),
+                attn_metadata,
-                set_ascend_forward_context(
+                self.vllm_config,
-                    attn_metadata,
+                num_tokens=num_tokens_padded,
-                    self.vllm_config,
+                num_tokens_across_dp=num_tokens_across_dp,
-                    num_tokens=num_tokens_padded,
+                in_profile_run=is_profile,
-                    num_tokens_across_dp=num_tokens_across_dp,
+                num_actual_tokens=num_tokens_padded,
-                    in_profile_run=is_profile,
+                aclgraph_runtime_mode=cudagraph_runtime_mode,
-                    num_actual_tokens=num_tokens_padded,
+                batch_descriptor=batch_desc,
-                    aclgraph_runtime_mode=cudagraph_runtime_mode,
+                model_instance=self.model,
-                    batch_descriptor=batch_desc,
+                has_sinks=self._has_sinks,
-                    model_instance=self.model,
+                input_ids=input_ids,
-                    has_sinks=self._has_sinks,
+                eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False,
-                    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: