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

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

docs/.nav.yml

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

docs/design/architecture_overview.md

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

docs/design/feature/vae_parallel.md

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

docs/design/index.md

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

docs/user_guide/diffusion/parallelism/overview.md

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

docs/user_guide/diffusion/parallelism/vae_parallel.md

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

docs/user_guide/diffusion_features.md

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

tests/config/test_omni_config.py

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

tests/diffusion/distributed/test_wan_spatial_shard.py

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

tests/entrypoints/test_serve.py

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

vllm_omni/config/omni_config.py

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

vllm_omni/diffusion/data.py

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

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

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

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

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

vllm_omni/entrypoints/cli/serve.py

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