[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/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