[Feature] Spatially-sharded (SP) decode for the Wan VAE (#4620)

Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/.nav.yml

@@ -68,7 +68,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 Patch Parallel: user_guide/diffusion/parallelism/vae_patch_parallel.md
+        - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
       - LoRA: user_guide/diffusion/lora.md
       - Custom Pipeline: features/custom_pipeline.md

docs/design/feature/vae_parallel.md

@@ -448,6 +448,42 @@ Complete examples in the codebase:
 
 ---
 
+## Spatially-Sharded Decode (Wan)
+
+The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).
+
+### How it differs from tile parallel
+
+| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+|--------|--------------------------|-----------------------------------------------|
+| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Scope | Decode + encode | Decode only |
+
+Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.
+
+### Wiring
+
+`vae_parallel_mode` flows through the same path as `vae_patch_parallel_size`:
+
+```text
+serve.py (--vae-parallel-mode) / OmniEngineArgs
+  -> DiffusionParallelConfig.vae_parallel_mode
+  -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...)
+  -> DistributedVaeExecutor.parallel_mode
+  -> 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.
+
+### 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.
+- 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`).
+
+---
+
 ## Summary
 
 Adding VAE Patch Parallel support to diffusion model:

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 Patch Parallelism](vae_patch_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 Patch Parallelism Guide
+# VAE Parallelism Guide
 
 
 ## Table of Content
@@ -15,7 +15,7 @@
 
 ## Overview
 
-VAE Patch Parallelism distributes the VAE (Variational AutoEncoder) decode/encode computation across multiple GPUs by splitting the latent space into spatial tiles or patches. Each GPU processes a subset of tiles in parallel, significantly reducing peak memory consumption during the VAE decode stage while maintaining output quality.
+VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width.
 
 This is particularly useful for:
 - **High-resolution image generation** where VAE decode can become a memory bottleneck
@@ -25,7 +25,7 @@ This is particularly useful for:
 See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
 
 
-VAE Patch Parallelism uses two strategies based on image size:
+VAE patch parallelism uses two strategies based on image size:
 
 | Strategy | Use Case | How It Works | Overlap Handling | Output Quality |
 |----------|----------|--------------|------------------|----------------|
@@ -113,6 +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). |
 
 Additional requirements:
 
@@ -125,6 +126,51 @@ 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"`.
+
+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.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    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"
+    ),
+)
+```
+
+Or from the CLI / serving entrypoint:
+
+```bash
+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
+```
+
+**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).
+
+For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \
+    --height 480 --width 832 --num-frames 17 --max-concurrency 1
+```
+
+---
+
 ## Best Practices
 
 ### When to Use

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_diffusion.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 Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling 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
 
@@ -270,7 +270,7 @@ Measured on NVIDIA H800:
 **Memory Optimization:**
 
 - **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
-- **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images
+- **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
 - **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
 
 **Extensions:**

tests/config/test_omni_config.py

@@ -423,6 +423,7 @@ def test_sub_config_fields_match_rfc_scopes():
         "ulysses_mode",
         "cfg_parallel_size",
         "vae_patch_parallel_size",
+        "vae_parallel_mode",
         "use_hsdp",
         "mask_sp_padding",
         "hsdp_shard_size",

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -0,0 +1,417 @@
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+    DistributedAutoencoderKLWan,
+)
+from vllm_omni.platforms import current_omni_platform
+
+# Module-level markers apply to the CPU unit tests below. The multi-GPU
+# correctness test at the end of the file adds the nightly/distributed markers
+# and a skip guard so it only runs where >= 2 CUDA devices are available.
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_split_for_parallel_decode_pads_uneven_height():
+    x = torch.arange(1 * 1 * 1 * 5 * 2, dtype=torch.float32).reshape(1, 1, 1, 5, 2)
+
+    local, expected_height = wan_spatial_shard.split_for_parallel_decode(
+        x,
+        upsample_count=2,
+        rank=2,
+        world_size=3,
+    )
+
+    assert expected_height == 20
+    assert local.shape == (1, 1, 1, 2, 2)
+    assert torch.equal(local[..., 0, :], x[..., 4, :])
+    assert torch.equal(local[..., 1, :], torch.zeros_like(local[..., 1, :]))
+
+
+def test_split_for_parallel_decode_pads_uneven_width():
+    x = torch.arange(1 * 1 * 1 * 2 * 5, dtype=torch.float32).reshape(1, 1, 1, 2, 5)
+
+    local, expected_width = wan_spatial_shard.split_for_parallel_decode(
+        x,
+        upsample_count=2,
+        split_dim="width",
+        rank=2,
+        world_size=3,
+    )
+
+    assert expected_width == 20
+    assert local.shape == (1, 1, 1, 2, 2)
+    assert torch.equal(local[..., :, 0], x[..., :, 4])
+    assert torch.equal(local[..., :, 1], torch.zeros_like(local[..., :, 1]))
+
+
+def test_split_for_parallel_decode_rejects_invalid_split_dim():
+    x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="split_dim"):
+        wan_spatial_shard.split_for_parallel_decode(
+            x,
+            upsample_count=1,
+            split_dim="depth",
+            rank=0,
+            world_size=2,
+        )
+
+
+def test_split_for_parallel_decode_rejects_zero_world_size():
+    x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="world_size"):
+        wan_spatial_shard.split_for_parallel_decode(
+            x,
+            upsample_count=1,
+            rank=0,
+            world_size=0,
+        )
+
+
+def test_split_for_parallel_decode_rejects_rank_out_of_range():
+    x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="rank"):
+        wan_spatial_shard.split_for_parallel_decode(
+            x,
+            upsample_count=1,
+            rank=3,
+            world_size=3,
+        )
+
+
+def test_gather_and_trim_height(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3))
+
+    def fake_all_gather(gathered, x, group=None):
+        for idx, output in enumerate(gathered):
+            output.copy_(x + idx)
+
+    monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather)
+
+    x = torch.zeros((1, 1, 1, 2, 1), dtype=torch.float32)
+    out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object())
+
+    assert out.shape == (1, 1, 1, 5, 1)
+    assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0]))
+
+
+def test_gather_and_trim_width(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3))
+
+    def fake_all_gather(gathered, x, group=None):
+        for idx, output in enumerate(gathered):
+            output.copy_(x + idx)
+
+    monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather)
+
+    x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
+    out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="width", group=object())
+
+    assert out.shape == (1, 1, 1, 1, 5)
+    assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0]))
+
+
+def test_gather_and_trim_rank0_only_assembles_on_rank0(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3))
+
+    def fake_all_gather(gathered, x, group=None):
+        for idx, output in enumerate(gathered):
+            output.copy_(x + idx)
+
+    monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather)
+
+    x = torch.zeros((1, 1, 1, 2, 1), dtype=torch.float32)
+    out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object(), dst=0)
+
+    assert out.shape == (1, 1, 1, 5, 1)
+    assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0]))
+
+
+def test_gather_and_trim_rank0_only_returns_empty_on_non_zero_rank(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
+
+    gathered_sizes = []
+
+    def fake_all_gather(gathered, x, group=None):
+        # Every rank must still take part in the collective even when it discards the result.
+        gathered_sizes.append(len(gathered))
+        for output in gathered:
+            output.copy_(x)
+
+    monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather)
+
+    x = torch.ones((1, 1, 1, 2, 1), dtype=torch.float32)
+    out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object(), dst=0)
+
+    assert gathered_sizes == [3]
+    assert out.numel() == 0
+
+
+def test_reshard_from_trimmed_height_pads_invalid_rows(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
+
+    x = torch.arange(5, dtype=torch.float32).reshape(1, 1, 1, 5, 1)
+    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
+        wan_spatial_shard.SpatialShardContext(
+            input_extent=5,
+            local_input_extent=2,
+            split_dim="height",
+            rank=2,
+            world_size=3,
+        )
+    )
+    try:
+        out = wan_spatial_shard.reshard_from_trimmed_extent(
+            x,
+            local_extent=2,
+            split_dim="height",
+            group=object(),
+        )
+    finally:
+        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
+
+    assert out.shape == (1, 1, 1, 2, 1)
+    assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
+
+
+def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
+
+    x = torch.arange(5, dtype=torch.float32).reshape(1, 1, 1, 1, 5)
+    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
+        wan_spatial_shard.SpatialShardContext(
+            input_extent=5,
+            local_input_extent=2,
+            split_dim="width",
+            rank=2,
+            world_size=3,
+        )
+    )
+    try:
+        out = wan_spatial_shard.reshard_from_trimmed_extent(
+            x,
+            local_extent=2,
+            split_dim="width",
+            group=object(),
+        )
+    finally:
+        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
+
+    assert out.shape == (1, 1, 1, 1, 2)
+    assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
+
+
+def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
+
+    x = torch.randn((1, 1, 1, 4, 2))
+    out, recv_top, recv_bottom = wan_spatial_shard.halo_exchange(
+        x,
+        group=object(),
+        halo_size=1,
+    )
+
+    assert out is x
+    assert recv_top is None
+    assert recv_bottom is None
+
+
+def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.MonkeyPatch):
+    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((0, 1, 1, 1), group=object())
+    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)
+    assert last.shape == (1, 1, 3, 3)
+
+
+def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.MonkeyPatch):
+    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)
+    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)
+    assert last.shape == (1, 1, 2, 3)
+
+
+def test_spatial_shard_height_gate_falls_back_for_partial_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=2, parallel_mode="spatial_shard_height")
+    vae.is_distributed_enabled = lambda: True
+
+    z = torch.zeros((1, 16, 1, 8, 8))
+
+    assert vae._spatial_shard_decode_split_dim() == "height"
+    assert vae._spatial_shard_decode_enabled(z) is False
+
+
+def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
+        lambda group=None: 2,
+    )
+
+    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
+    vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="spatial_shard_width")
+    vae.is_distributed_enabled = lambda: True
+
+    z = torch.zeros((1, 16, 1, 8, 8))
+
+    assert vae._spatial_shard_decode_split_dim() == "width"
+    assert vae._spatial_shard_decode_enabled(z) is True
+
+
+def test_tile_mode_disables_spatial_shard_decode():
+    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
+    vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="tile")
+    vae.is_distributed_enabled = lambda: True
+
+    z = torch.zeros((1, 16, 1, 8, 8))
+
+    assert vae._spatial_shard_decode_split_dim() is None
+    assert vae._spatial_shard_decode_enabled(z) is False
+
+
+# =============================================================================
+# Multi-GPU numerical-correctness test (nightly / full_model, distributed CUDA)
+#
+# Spawns a small process group and verifies that spatial_shard_height/spatial_shard_width decode match
+# a single-process (non-distributed) reference decode of the same latent within
+# tolerance. Requires >= 2 accelerator devices and downloads a Wan VAE.
+# =============================================================================
+
+_SPATIAL_SHARD_MODEL = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+_SPATIAL_SHARD_SUBFOLDER = "vae"
+_SPATIAL_SHARD_WORLD_SIZE = 2
+_SPATIAL_SHARD_LATENT_FRAMES = 5
+_SPATIAL_SHARD_LATENT_HEIGHT = 60
+_SPATIAL_SHARD_LATENT_WIDTH = 104
+_SPATIAL_SHARD_TOLERANCE = 3e-2
+
+
+def _spatial_shard_decode_worker(rank: int, split_dim: str, 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)
+    dtype = torch.float32
+
+    backend = current_omni_platform.dist_backend
+    init_distributed_environment(world_size=_SPATIAL_SHARD_WORLD_SIZE, rank=rank, local_rank=rank, backend=backend)
+    initialize_model_parallel(
+        sequence_parallel_size=_SPATIAL_SHARD_WORLD_SIZE, ulysses_degree=_SPATIAL_SHARD_WORLD_SIZE, backend=backend
+    )
+
+    try:
+        vae = DistributedAutoencoderKLWan.from_pretrained(
+            _SPATIAL_SHARD_MODEL, subfolder=_SPATIAL_SHARD_SUBFOLDER, torch_dtype=dtype
+        )
+        vae.to(device=device, dtype=dtype)
+        vae.eval()
+
+        generator = torch.Generator(device=device).manual_seed(0)
+        latents = torch.randn(
+            (
+                1,
+                vae.config.z_dim,
+                _SPATIAL_SHARD_LATENT_FRAMES,
+                _SPATIAL_SHARD_LATENT_HEIGHT,
+                _SPATIAL_SHARD_LATENT_WIDTH,
+            ),
+            generator=generator,
+            device=device,
+            dtype=dtype,
+        )
+
+        with torch.inference_mode():
+            # Ground-truth reference: standard non-parallel, untiled decode (computed identically on
+            # every rank). Tiling must be OFF so neither the tile-parallel nor the single-GPU tiled
+            # path is exercised; otherwise we would be comparing SP against tiled decode.
+            vae.use_tiling = False
+            vae.set_parallel_size(1, mode="tile")
+            reference = vae.decode(latents, return_dict=False)[0].float()
+
+            # Spatially-sharded decode across the full group (requires tiling to enter tiled_decode).
+            vae.use_tiling = True
+            vae.set_parallel_size(_SPATIAL_SHARD_WORLD_SIZE, mode=f"spatial_shard_{split_dim}")
+            sharded = vae.decode(latents, return_dict=False)[0].float()
+
+        # Only rank 0 assembles the full decoded sample (matching broadcast_result=False);
+        # non-zero ranks return an empty placeholder, so the comparison runs on rank 0 only.
+        if rank == 0:
+            diff = (sharded - reference).abs()
+            return_dict["max_abs_diff"] = diff.max().item()
+            return_dict["mean_abs_diff"] = diff.mean().item()
+            return_dict["shape"] = tuple(sharded.shape)
+    finally:
+        destroy_model_parallel()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.distributed_cuda
+@pytest.mark.skipif(
+    current_omni_platform.get_device_count() < _SPATIAL_SHARD_WORLD_SIZE,
+    reason="Requires >= 2 accelerator devices",
+)
+@pytest.mark.parametrize("split_dim", ["height", "width"])
+def test_spatial_shard_decode_matches_reference(split_dim: str):
+    manager = mp.get_context("spawn").Manager()
+    return_dict = manager.dict()
+    # Use a per-split-dim port to avoid collisions across parametrized runs.
+    master_port = str(29500 + (1 if split_dim == "width" else 0))
+
+    mp.spawn(
+        _spatial_shard_decode_worker,
+        args=(split_dim, return_dict, master_port),
+        nprocs=_SPATIAL_SHARD_WORLD_SIZE,
+        join=True,
+    )
+
+    assert "max_abs_diff" in return_dict, "rank 0 did not report a result"
+    max_abs_diff = return_dict["max_abs_diff"]
+    mean_abs_diff = return_dict["mean_abs_diff"]
+    print(
+        f"spatial_shard_{split_dim} vs reference: max_abs_diff={max_abs_diff:.6e} "
+        f"mean_abs_diff={mean_abs_diff:.6e} shape={return_dict.get('shape')}"
+    )
+    assert max_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
+        f"spatial_shard_{split_dim} max_abs_diff {max_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
+    )
+    assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
+        f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
+    )

vllm_omni/config/omni_config.py

@@ -353,6 +353,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     ulysses_mode: str = "strict"
     cfg_parallel_size: int = Field(default=1, ge=1, le=3)
     vae_patch_parallel_size: int = Field(default=1, ge=1)
+    vae_parallel_mode: str = "tile"
     use_hsdp: bool = False
     mask_sp_padding: bool = False
     hsdp_shard_size: int = -1
@@ -362,6 +363,11 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
         self.sequence_parallel_size = self.ulysses_degree * self.ring_degree
         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"}:
+            raise ValueError(
+                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                f"but got {self.vae_parallel_mode!r}."
+            )
 
         other_parallel_world_size = (
             self.pipeline_parallel_size

vllm_omni/config/stage_config.py

@@ -338,6 +338,7 @@ class StageDeployConfig:
     sequence_parallel_size: int | None = None
     cfg_parallel_size: int | None = None
     vae_patch_parallel_size: int | None = None
+    vae_parallel_mode: str | None = None
     use_hsdp: bool | None = None
     hsdp_shard_size: int | None = None
     hsdp_replicate_size: int | None = None

vllm_omni/diffusion/data.py

@@ -183,6 +183,21 @@ class DiffusionParallelConfig:
     vae_patch_parallel_size: int = 1
     """Number of ranks used for VAE patch/tile parallelism (decode/encode)."""
 
+    vae_parallel_mode: str = "tile"
+    """VAE parallel decode strategy.
+
+    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
+      of spatial tiles and the results are stitched on rank 0.
+    - "spatial_shard_height": Spatially-sharded decode that splits decoder
+      feature maps along height and exchanges halo rows around spatial
+      convolutions.
+    - "spatial_shard_width": Same as "spatial_shard_height" but sharded along width.
+
+    The "spatial_shard_*" modes are decode-only and currently require
+    ``vae_patch_parallel_size`` to match the DiT group size; otherwise the VAE
+    falls back to tile-parallel decode at runtime.
+    """
+
     use_hsdp: bool = False
     """Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding."""
 
@@ -214,6 +229,10 @@ class DiffusionParallelConfig:
             f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}"
         )
         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'}, "
+            f"but got {self.vae_parallel_mode!r}."
+        )
         assert self.sequence_parallel_size == self.ulysses_degree * self.ring_degree, (
             "Sequence parallel size must be equal to the product of ulysses degree and ring degree,"
             f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}"

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -5,11 +5,13 @@ from contextlib import nullcontext
 from typing import Any
 
 import torch
+import torch.distributed as dist
 from diffusers.models.autoencoders import AutoencoderKLWan
 from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedOperator,
     DistributedVaeMixin,
@@ -268,10 +270,54 @@ 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:
+        mode = self.distributed_executor.parallel_mode
+        if mode == "spatial_shard_width":
+            return "width"
+        if mode == "spatial_shard_height":
+            return "height"
+        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:
+            return False
+        if z.ndim != 5:
+            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
+
+        group = self.distributed_executor.group
+        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,
+            )
+            return False
+        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):
+            split_dim = self._spatial_shard_decode_split_dim()
+            logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim)
+            return wan_spatial_shard.spatial_shard_decode(
+                self,
+                z,
+                group=self.distributed_executor.group,
+                return_dict=return_dict,
+                split_dim=split_dim,
+            )
+
         logger.debug("Decode running with distributed executor")
         result = self.distributed_executor.execute(
             z,

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -49,9 +49,11 @@ class DistributedVaeExecutor:
         self.world_size = dist.get_world_size(self.group)
         self.rank = dist.get_rank(self.group)
         self.parallel_size = 1
+        self.parallel_mode = "tile"
 
-    def set_parallel_size(self, parallel_size: int):
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile"):
         self.parallel_size = parallel_size
+        self.parallel_mode = mode
 
     def gather_tensors(self, tensor: torch.Tensor):
         gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
@@ -170,8 +172,8 @@ class DistributedVaeMixin:
     def init_distributed(self):
         self.distributed_executor = DistributedVaeExecutor()
 
-    def set_parallel_size(self, parallel_size: int) -> None:
-        self.distributed_executor.set_parallel_size(parallel_size)
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
+        self.distributed_executor.set_parallel_size(parallel_size, mode=mode)
 
     def is_distributed_enabled(self) -> bool:
         if (

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -0,0 +1,816 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# The halo-exchange spatial-parallel decode here is adapted from SGLang's
+# spatial-parallel VAE decode
+#   https://github.com/sgl-project/sglang
+#   (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py)
+# which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo).
+# This version generalizes the height-only sharding to shard along height or
+# width and adds the Wan causal-conv ``feat_cache`` handling.
+"""Spatially-sharded Wan VAE decode.
+
+The existing distributed Wan VAE path shards *tiles*.  This module adds an
+opt-in decode backend that shards decoder feature maps along height or width and
+exchanges boundary rows/columns before spatial convolutions.  It is
+intentionally decode-only and keeps checkpoint loading unchanged by patching the
+already-loaded decoder.
+"""
+
+from __future__ import annotations
+
+import math
+from contextlib import nullcontext
+from contextvars import ContextVar
+from dataclasses import dataclass
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class SpatialShardContext:
+    input_extent: int
+    local_input_extent: int
+    split_dim: str
+    rank: int
+    world_size: int
+
+
+_SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar(
+    "wan_vae_spatial_shard_context",
+    default=None,
+)
+
+
+def _spatial_dim(split_dim: str) -> int:
+    if split_dim == "height":
+        return -2
+    if split_dim == "width":
+        return -1
+    raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+
+
+def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor:
+    if dim < 0:
+        dim += x.dim()
+    return x.narrow(dim, start, length)
+
+
+def _global_rank(group: dist.ProcessGroup, group_rank: int) -> int:
+    try:
+        return dist.get_global_rank(group, group_rank)
+    except Exception:
+        return group_rank
+
+
+def _rank_world(group: dist.ProcessGroup) -> tuple[int, int]:
+    return dist.get_rank(group), dist.get_world_size(group)
+
+
+def _pad_along_dim(x: torch.Tensor, pad: int, dim: int, value: float = 0.0) -> torch.Tensor:
+    if pad <= 0:
+        return x
+    shape = list(x.shape)
+    shape[dim] = pad
+    padding = torch.full(shape, value, dtype=x.dtype, device=x.device)
+    return torch.cat([x, padding], dim=dim)
+
+
+def _maybe_contiguous_for_shard_gather(x: torch.Tensor) -> torch.Tensor:
+    if (
+        x.dim() == 5
+        and hasattr(torch, "channels_last_3d")
+        and x.is_contiguous(memory_format=torch.channels_last_3d)
+        and not x.is_contiguous()
+    ):
+        return x.contiguous()
+    return x
+
+
+def _halo_memory_format(reference: torch.Tensor) -> torch.memory_format:
+    if reference.dim() > 1 and reference.stride(1) == 1:
+        if reference.dim() == 5 and hasattr(torch, "channels_last_3d"):
+            return torch.channels_last_3d
+        if reference.dim() == 4:
+            return torch.channels_last
+    return torch.contiguous_format
+
+
+def _current_full_extent(local_extent: int) -> int | None:
+    ctx = _SPATIAL_SHARD_CONTEXT.get()
+    if ctx is None:
+        return None
+    if ctx.local_input_extent <= 0:
+        return None
+    scale = local_extent / ctx.local_input_extent
+    rounded_scale = round(scale)
+    if not math.isclose(scale, rounded_scale, rel_tol=0.0, abs_tol=1e-6):
+        return None
+    return ctx.input_extent * rounded_scale
+
+
+def _local_valid_extent(local_extent: int) -> int:
+    ctx = _SPATIAL_SHARD_CONTEXT.get()
+    full_extent = _current_full_extent(local_extent)
+    if ctx is None or full_extent is None:
+        return local_extent
+    start = ctx.rank * local_extent
+    return max(0, min(local_extent, full_extent - start))
+
+
+def _zero_invalid_extent(x: torch.Tensor, *, split_dim: str) -> torch.Tensor:
+    dim = _spatial_dim(split_dim)
+    dim_size = x.shape[dim]
+    valid_extent = _local_valid_extent(dim_size)
+    if valid_extent >= dim_size:
+        return x
+    x = x.clone()
+    invalid = _narrow_along_dim(x, dim, valid_extent, dim_size - valid_extent)
+    invalid.zero_()
+    return x
+
+
+def split_for_parallel_decode(
+    x: torch.Tensor,
+    *,
+    upsample_count: int,
+    split_dim: str = "height",
+    group: dist.ProcessGroup | None = None,
+    rank: int | None = None,
+    world_size: int | None = None,
+) -> tuple[torch.Tensor, int]:
+    """Shard latent/feature spatial extent and return expected full output extent."""
+    if group is not None:
+        rank, world_size = _rank_world(group)
+    rank = 0 if rank is None else int(rank)
+    world_size = 1 if world_size is None else int(world_size)
+    if world_size < 1:
+        raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.")
+    if not 0 <= rank < world_size:
+        raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+
+    dim = _spatial_dim(split_dim)
+    expected_extent = x.shape[dim] * (2**upsample_count)
+    if world_size <= 1:
+        return x, expected_extent
+
+    pad = (world_size - (x.shape[dim] % world_size)) % world_size
+    if pad:
+        x = _pad_along_dim(x, pad, dim=dim)
+    chunk_size = x.shape[dim] // world_size
+    return _narrow_along_dim(x, dim, rank * chunk_size, chunk_size).contiguous(), expected_extent
+
+
+def all_gather_along_dim(
+    x: torch.Tensor,
+    *,
+    group: dist.ProcessGroup,
+    dim: int,
+    dst: int | None = None,
+) -> torch.Tensor:
+    rank, world_size = _rank_world(group)
+    if world_size <= 1:
+        return x
+    x = _maybe_contiguous_for_shard_gather(x)
+    gathered = [torch.empty_like(x) for _ in range(world_size)]
+    # NCCL has no rank-local gather, so every rank joins the collective; only ``dst``
+    # keeps the assembled tensor while the rest drop their copies.
+    dist.all_gather(gathered, x.contiguous(), group=group)
+    if dst is not None and rank != dst:
+        return x.new_zeros(0)
+    return torch.cat(gathered, dim=dim)
+
+
+def reshard_from_trimmed_extent(
+    x: torch.Tensor,
+    *,
+    local_extent: int,
+    split_dim: str,
+    group: dist.ProcessGroup,
+) -> torch.Tensor:
+    rank, world_size = _rank_world(group)
+    if world_size <= 1:
+        return x
+
+    dim = _spatial_dim(split_dim)
+    valid_extent = _local_valid_extent(local_extent)
+    start = rank * local_extent
+    local = _narrow_along_dim(x, dim, start, valid_extent).contiguous()
+    if valid_extent < local_extent:
+        local = _pad_along_dim(local, local_extent - valid_extent, dim=dim)
+    return local
+
+
+def gather_and_trim_extent(
+    x: torch.Tensor,
+    *,
+    expected_extent: int | None,
+    split_dim: str,
+    group: dist.ProcessGroup,
+    dst: int | None = None,
+) -> torch.Tensor:
+    dim = _spatial_dim(split_dim)
+    rank, _ = _rank_world(group)
+    out = all_gather_along_dim(x, group=group, dim=dim, dst=dst)
+    if dst is not None and rank != dst:
+        return out
+    if expected_extent is not None and out.shape[dim] != expected_extent:
+        out = _narrow_along_dim(out, dim, 0, expected_extent).contiguous()
+    return out
+
+
+def _ensure_recv_buf(recv_buf: torch.Tensor | None, reference: torch.Tensor) -> torch.Tensor:
+    memory_format = _halo_memory_format(reference)
+    if (
+        recv_buf is None
+        or recv_buf.shape != reference.shape
+        or recv_buf.dtype != reference.dtype
+        or recv_buf.device != reference.device
+        or not recv_buf.is_contiguous(memory_format=memory_format)
+    ):
+        return torch.empty(
+            reference.shape,
+            dtype=reference.dtype,
+            device=reference.device,
+            memory_format=memory_format,
+        )
+    return recv_buf
+
+
+def _halo_exchange_p2p(
+    *,
+    rank: int,
+    world_size: int,
+    group: dist.ProcessGroup,
+    top_row_ref: torch.Tensor,
+    bottom_row_ref: torch.Tensor,
+    recv_top_buf: torch.Tensor,
+    recv_bottom_buf: torch.Tensor,
+) -> None:
+    p2p_ops = []
+    if rank > 0:
+        prev_rank = _global_rank(group, rank - 1)
+        top_row = top_row_ref.contiguous(memory_format=_halo_memory_format(top_row_ref))
+        p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group))
+        p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group))
+    else:
+        recv_top_buf.zero_()
+
+    if rank < world_size - 1:
+        next_rank = _global_rank(group, rank + 1)
+        bottom_row = bottom_row_ref.contiguous(memory_format=_halo_memory_format(bottom_row_ref))
+        p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group))
+        p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group))
+    else:
+        recv_bottom_buf.zero_()
+
+    if p2p_ops:
+        reqs = dist.batch_isend_irecv(p2p_ops)
+        for req in reqs:
+            req.wait()
+
+
+def halo_exchange(
+    x: torch.Tensor,
+    *,
+    group: dist.ProcessGroup,
+    halo_size: int,
+    split_dim: str = "height",
+    recv_top_buf: torch.Tensor | None = None,
+    recv_bottom_buf: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+    if halo_size <= 0:
+        return x, recv_top_buf, recv_bottom_buf
+
+    rank, world_size = _rank_world(group)
+    if world_size <= 1:
+        return x, recv_top_buf, recv_bottom_buf
+
+    dim = _spatial_dim(split_dim)
+    top_row_ref = _narrow_along_dim(x, dim, 0, halo_size)
+    bottom_row_ref = _narrow_along_dim(x, dim, x.shape[dim] - halo_size, halo_size)
+    recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row_ref)
+    recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row_ref)
+
+    _halo_exchange_p2p(
+        rank=rank,
+        world_size=world_size,
+        group=group,
+        top_row_ref=top_row_ref,
+        bottom_row_ref=bottom_row_ref,
+        recv_top_buf=recv_top_buf,
+        recv_bottom_buf=recv_bottom_buf,
+    )
+
+    return torch.cat([recv_top_buf, x, recv_bottom_buf], dim=dim), recv_top_buf, recv_bottom_buf
+
+
+class WanDistZeroPad2d(nn.Module):
+    """Apply ZeroPad2d only at global split-dimension boundaries."""
+
+    def __init__(
+        self,
+        padding: tuple[int, int, int, int],
+        group: dist.ProcessGroup,
+        *,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
+    ) -> 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.group = group
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rank, world_size = _rank_world(self.group)
+        left, right, top, bottom = self.padding
+        if world_size > 1:
+            if self.split_dim == "height":
+                top = top if rank == 0 else 0
+                bottom = bottom if rank == world_size - 1 else 0
+            else:
+                left = left if rank == 0 else 0
+                right = right if rank == world_size - 1 else 0
+        return F.pad(x, (left, right, top, bottom))
+
+
+class WanDistConv2d(nn.Conv2d):
+    def __init__(
+        self,
+        source: nn.Conv2d,
+        group: dist.ProcessGroup,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
+    ):
+        super().__init__(
+            source.in_channels,
+            source.out_channels,
+            source.kernel_size,
+            stride=source.stride,
+            padding=0,
+            dilation=source.dilation,
+            groups=source.groups,
+            bias=source.bias is not None,
+            padding_mode=source.padding_mode,
+            device=source.weight.device,
+            dtype=source.weight.dtype,
+        )
+        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]] = {}
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.pad(x, self._non_split_padding)
+        x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
+            x,
+            group=self.group,
+            halo_size=self.halo_size,
+            split_dim=self.split_dim,
+            recv_top_buf=self._halo_recv_top_buf,
+            recv_bottom_buf=self._halo_recv_bottom_buf,
+        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
+        if shift:
+            x_padded = _narrow_along_dim(
+                x_padded,
+                self.split_tensor_dim,
+                shift,
+                x_padded.shape[self.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)
+
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
+        if trim_params is None:
+            rank, world_size = _rank_world(self.group)
+            trim_params = _compute_conv_trim_params(
+                local_extent=local_extent,
+                rank=rank,
+                world_size=world_size,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
+            )
+            self._trim_cache[local_extent] = trim_params
+        return trim_params
+
+
+class WanDistCausalConv3d(nn.Conv3d):
+    def __init__(
+        self,
+        source: nn.Conv3d,
+        group: dist.ProcessGroup,
+        split_dim: str = "height",
+    ):
+        super().__init__(
+            source.in_channels,
+            source.out_channels,
+            source.kernel_size,
+            stride=source.stride,
+            padding=0,
+            dilation=source.dilation,
+            groups=source.groups,
+            bias=source.bias is not None,
+            padding_mode=source.padding_mode,
+            device=source.weight.device,
+            dtype=source.weight.dtype,
+        )
+        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]] = {}
+
+    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+        padding = list(self._padding)
+        if cache_x is not None and padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+
+        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,
+            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])
+        if shift:
+            x_padded = _narrow_along_dim(
+                x_padded,
+                self.split_tensor_dim,
+                shift,
+                x_padded.shape[self.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)
+
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
+        if trim_params is None:
+            rank, world_size = _rank_world(self.group)
+            trim_params = _compute_conv_trim_params(
+                local_extent=local_extent,
+                rank=rank,
+                world_size=world_size,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
+            )
+            self._trim_cache[local_extent] = trim_params
+        return trim_params
+
+
+def _compute_conv_trim_params(
+    *,
+    local_extent: int,
+    rank: int,
+    world_size: int,
+    halo_size: int,
+    pad_before: int,
+    pad_after: int,
+    kernel_extent: int,
+    stride_extent: int,
+) -> tuple[int, int, int]:
+    global_start = rank * local_extent
+    shift = 0
+    if halo_size > 0 and stride_extent > 1:
+        shift = (global_start - halo_size + pad_before) % stride_extent
+        if shift:
+            global_start += shift
+
+    global_extent = local_extent * world_size
+    min_i = math.ceil(((-pad_before) - (global_start - halo_size)) / stride_extent)
+    max_i = math.floor(
+        ((global_extent - 1 + pad_after) - (kernel_extent - 1) - (global_start - halo_size)) / stride_extent
+    )
+    return shift, max(min_i, 0), max_i + 1
+
+
+def _trim_local_conv_output(
+    out: torch.Tensor,
+    halo_size: int,
+    start: int,
+    upper_bound: int,
+    *,
+    split_dim: str,
+) -> torch.Tensor:
+    if halo_size <= 0:
+        return out
+    dim = _spatial_dim(split_dim)
+    end = min(upper_bound, out.shape[dim])
+    if start != 0 or end != out.shape[dim]:
+        out = _narrow_along_dim(out, dim, start, end - start)
+    return out
+
+
+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:
+        _, world_size = _rank_world(group)
+        if world_size <= 1:
+            return orig_forward(x, *args, **kwargs)
+        dim = _spatial_dim(split_dim)
+        local_extent = x.shape[dim]
+        gathered = all_gather_along_dim(x, group=group, dim=dim).contiguous()
+        full_extent = _current_full_extent(local_extent)
+        if full_extent is not None:
+            gathered = _narrow_along_dim(gathered, dim, 0, full_extent).contiguous()
+        out = orig_forward(gathered, *args, **kwargs)
+        return reshard_from_trimmed_extent(out, local_extent=local_extent, split_dim=split_dim, group=group)
+
+    module.forward = MethodType(_forward, module)
+    module._vllm_omni_spatial_shard_attention = True  # type: ignore[attr-defined]
+
+
+def _replace_child(
+    parent: nn.Module,
+    name: str,
+    child: nn.Module,
+    group: dist.ProcessGroup,
+    split_dim: str,
+) -> None:
+    if child.__class__.__name__ == "WanCausalConv3d":
+        setattr(
+            parent,
+            name,
+            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,
+                group,
+                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):
+        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):
+                split_padding = prev.split_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]))
+        setattr(
+            parent,
+            name,
+            WanDistConv2d(
+                child,
+                group,
+                split_dim=split_dim,
+                split_padding=split_padding,
+            ),
+        )
+
+
+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)
+        inside_attention = True
+
+    for name, child in list(module.named_children()):
+        if child.__class__.__name__ == "WanCausalConv3d":
+            _replace_child(module, name, child, group, split_dim)
+            continue
+        if isinstance(child, nn.Conv2d) and not inside_attention:
+            _replace_child(module, name, child, group, split_dim)
+            continue
+        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
+
+
+def _decoder_upsample_count(decoder: nn.Module) -> int:
+    count = 0
+    for block in getattr(decoder, "up_blocks", []):
+        if getattr(block, "upsampler", None) is not None or getattr(block, "upsamplers", None) is not None:
+            count += 1
+    return count
+
+
+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.
+
+    Only group-relative rank 0 assembles the final decoded frame, mirroring the
+    distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
+    take part in the collectives but return an empty placeholder.
+    """
+    _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)
+    upsample_count = _decoder_upsample_count(decoder)
+    orig_forward = decoder.forward
+
+    def _forward(
+        self: nn.Module,
+        x: torch.Tensor,
+        feat_cache: list[torch.Tensor] | None = None,
+        feat_idx: list[int] | None = None,
+        first_chunk: bool = False,
+    ) -> torch.Tensor:
+        if feat_idx is None:
+            feat_idx = [0]
+        tensor_dim = _spatial_dim(split_dim)
+        input_extent = x.shape[tensor_dim]
+        x, expected_extent = split_for_parallel_decode(
+            x,
+            upsample_count=upsample_count,
+            split_dim=split_dim,
+            group=group,
+        )
+        rank, world_size = _rank_world(group)
+        token = _SPATIAL_SHARD_CONTEXT.set(
+            SpatialShardContext(
+                input_extent=input_extent,
+                local_input_extent=x.shape[tensor_dim],
+                split_dim=split_dim,
+                rank=rank,
+                world_size=world_size,
+            )
+        )
+        try:
+            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
+        finally:
+            _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)
+    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)
+
+
+def spatial_shard_decode(
+    vae: Any,
+    z: torch.Tensor,
+    *,
+    group: dist.ProcessGroup,
+    return_dict: bool = True,
+    split_dim: str = "height",
+) -> DecoderOutput | tuple[torch.Tensor]:
+    install_wan_spatial_shard_decode(vae, group, split_dim=split_dim)
+
+    if z.shape[2] == 0:
+        raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.")
+
+    # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with
+    # the halo/all-gather collectives; they just skip keeping/assembling the output.
+    rank, world_size = _rank_world(group)
+    produce_output = world_size <= 1 or rank == 0
+
+    vae.clear_cache()
+    try:
+        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
+        with context:
+            x = vae.post_quant_conv(z)
+            decoded_chunks = []
+            for i in range(z.shape[2]):
+                vae._conv_idx = [0]
+                chunk = vae.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=vae._feat_map,
+                    feat_idx=vae._conv_idx,
+                    first_chunk=(i == 0),
+                )
+                if produce_output:
+                    decoded_chunks.append(chunk)
+
+            if produce_output:
+                out = torch.cat(decoded_chunks, dim=2)
+                if vae.config.patch_size is not None:
+                    out = unpatchify(out, patch_size=vae.config.patch_size)
+                out = torch.clamp(out, min=-1.0, max=1.0)
+            else:
+                out = z.new_zeros(0)
+    finally:
+        vae.clear_cache()
+
+    if not return_dict:
+        return (out,)
+    return DecoderOutput(sample=out)

vllm_omni/diffusion/registry.py

@@ -390,7 +390,7 @@ def initialize_model(
             model.vae.use_tiling = od_config.vae_use_tiling
 
         if is_distributed_vae:
-            model.vae.set_parallel_size(vae_pp_size)
+            model.vae.set_parallel_size(vae_pp_size, mode=od_config.parallel_config.vae_parallel_mode)
 
         # Apply sequence parallelism if enabled
         # This follows diffusers' pattern where enable_parallelism() is called

vllm_omni/engine/arg_utils.py

@@ -477,6 +477,7 @@ class OrchestratorArgs:
     diffusion_kv_cache_skip_layers: str | None = None
     cfg_parallel_size: int = 1
     vae_patch_parallel_size: int = 1
+    vae_parallel_mode: str = "tile"
     default_sampling_params: str | None = None
     max_generated_image_size: int | None = None
     tts_max_instructions_length: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -953,6 +953,7 @@ class AsyncOmniEngine:
             cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
             vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
+            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
             enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
             use_hsdp = normalized_kwargs.get("use_hsdp", False)
             hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1)
@@ -971,6 +972,7 @@ class AsyncOmniEngine:
                 ulysses_mode=ulysses_mode,
                 cfg_parallel_size=cfg_parallel_size,
                 vae_patch_parallel_size=vae_patch_parallel_size,
+                vae_parallel_mode=vae_parallel_mode,
                 use_hsdp=use_hsdp,
                 hsdp_shard_size=hsdp_shard_size,
                 hsdp_replicate_size=hsdp_replicate_size,

vllm_omni/entrypoints/cli/serve.py

@@ -644,6 +644,18 @@ class OmniServeCommand(CLISubcommand):
             "Distributes VAE decode workload across multiple ranks by splitting the latent spatially. "
             "Equivalent to setting DiffusionParallelConfig.vae_patch_parallel_size.",
         )
+        omni_config_group.add_argument(
+            "--vae-parallel-mode",
+            type=str,
+            default="tile",
+            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
+            help="VAE parallel decode strategy for diffusion models. "
+            "'tile' (default) uses patch/tile parallel 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. "
+            "Equivalent to setting DiffusionParallelConfig.vae_parallel_mode.",
+        )
 
         # Default sampling parameters
         omni_config_group.add_argument(