[optimization] Added Wan 2.2 VAE decode optimizations (#7056)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>

docs/user_guide/diffusion/vae_fast_path.md

@@ -0,0 +1,102 @@
+# Wan VAE Decoder Fast Path
+
+Video pipelines built on the Wan 2.1/2.2 causal VAE (Cosmos3, Wan2.2, LingBot,
+Helios, SANA-Video, LongCat-Video, DreamZero) spend a large share of their
+end-to-end time decoding latents to pixels. The decoder is a chunked 3D
+convolutional network whose reference implementation in diffusers moves every
+activation through memory several times per layer (normalization, activation,
+causal padding, feature-cache bookkeeping, shortcut upsampling). vLLM-Omni
+installs a fast path on every loaded Wan VAE that fuses this data movement into
+a handful of Triton kernels while leaving the convolutions themselves untouched.
+
+## Levels
+
+The fast path is controlled by `--vae-fast-path` (engine argument
+`vae_fast_path`, deploy-config key `vae_fast_path`):
+
+| Level | Default | Output vs. diffusers | What it does |
+|---|---|---|---|
+| `lossless` | yes | bit-identical | Fused RMSNorm epilogue, fused causal-conv input assembly and cache refresh, fused shortcut upsampling and residual adds (with the neighbouring convolution biases folded in), single-pass nearest 2x upsampling, preallocated output assembly. |
+| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus decoder convolution weights converted to channels-last memory format so cuDNN picks its channels-last kernels, and a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the bias of the preceding `conv1`. |
+| `off` | no | bit-identical | Reference diffusers decoder. |
+
+```bash
+# Default: bit-exact fast path
+vllm serve nvidia/Cosmos3-Nano --omni
+
+# Fastest: channels-last decoder (not bit-exact)
+vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path channels_last
+
+# Reference implementation, e.g. to bisect a quality issue
+vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path off
+```
+
+The equivalent per-stage deploy configuration is:
+
+```yaml
+stages:
+  - stage_id: 0
+    vae_fast_path: channels_last
+```
+
+## Behavior and limitations
+
+- Only CUDA is supported. On other platforms the reference decoder runs.
+- The fast path is installed once per VAE instance when the pipeline is
+  initialized. It rebinds the forwards of the loaded decoder modules; parameter
+  names, `state_dict` keys and weight loading are unchanged.
+- Installation is skipped, with a logged reason, if a custom forward wrapper
+  or forward hook would be replaced or bypassed by the fused paths. Wrappers
+  and hooks on modules that the fast path still calls normally are preserved.
+- Every fused kernel validates its inputs and falls back to the exact PyTorch
+  expression for anything it does not support (unusual dtypes or layouts, CPU
+  tensors, autograd enabled, `torch.compile` tracing).
+- The SiLU activation is folded into the normalization kernel only after an
+  exhaustive self-test over all bf16/fp16 values proves the fused epilogue
+  bit-identical to `F.silu` on the running toolkit; otherwise SiLU stays a
+  separate operation. The startup log reports `fused_silu=...`.
+- Spatially sharded VAE decode (`--vae-parallel-mode spatial_shard_height` or
+  `spatial_shard_width`) is not combined with the fast path; the installer skips
+  the VAE and logs why. Tiled and tile-parallel decode work unchanged.
+- The `channels_last` level changes the order in which cuDNN accumulates
+  convolutions, so outputs differ from the reference in the last bits. Use
+  `lossless` when bitwise reproducibility against diffusers matters.
+- For VAEs kept in fp32, cuDNN's channels-last convolution algorithms run in
+  TF32 under PyTorch's default `torch.backends.cudnn.allow_tf32 = True`, which
+  dominates the difference to the reference (about 1e-3). Set
+  `torch.backends.cudnn.allow_tf32 = False` if full fp32 convolutions are
+  required.
+
+## Measured
+
+Cosmos3-Nano VAE, 1280x720 x 189 frames, bf16, one GB200 GPU, `bench_wan_vae_decode.py`
+(best of 2 runs after warmup):
+
+| `--vae-fast-path` | Decode time | Speedup | Output vs. `off` |
+|---|-------------|---------|---|
+| `off` | 6.00 s      | 1.00x   | reference |
+| `lossless` | 3.26 s      | 1.84x   | bit-identical |
+| `channels_last` | 2.51 s      | 2.39x   | PSNR 62.6 dB, max abs diff 4.6e-2 |
+
+Peak decode memory was unchanged (about 13 GiB): it is set by the largest
+activations and the cuDNN workspace, not by the output assembly.
+
+## Benchmarking
+
+`benchmarks/diffusion/bench_wan_vae_decode.py` decodes seeded latents with the
+real VAE at each level and reports decode time, peak memory, bitwise equality
+and PSNR against the `off` level. `--profile` prints a per-kernel table with the
+convolution share and any layout-transpose kernels:
+
+```bash
+python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+    --size 1280x720 --frames 189 --fast-path off,lossless,channels_last --profile
+```
+
+The same script benchmarks multi-GPU VAE decode when launched with `torchrun`;
+the decode is timed across all ranks and rank 0 reports:
+
+```bash
+torchrun --nproc-per-node 2 benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+    --vae-patch-parallel-size 2 --vae-parallel-mode tile
+```