vllm-omni/docs/user_guide/diffusion/vae_fast_path.md

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. The opt-in channels_last level also specializes first-frame convolutions that have no history.

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 channels-last convolution weights, a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the preceding conv1 bias, and Conv2d for supported single-frame causal convolutions without history.
off	no	bit-identical	Reference diffusers decoder.

# 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:

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 a supported causal convolution with one input frame and no history, the temporal input is [0, 0, x]. The channels_last level uses Conv2d with the last temporal weight slice, preserving the convolution formula for finite inputs and the existing cache update. Calls with history retain Conv3d. The weight slice is taken on each call, so weight updates and device moves do not leave a stale packed copy.
• 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:

python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
    --size 1280x720 --frames 189 --fast-path off,lossless,channels_last --profile

For a single-GPU comparison that isolates first-frame Conv2d within channels_last, use --first-frame-ablation. It alternates the two paths, disables cuDNN TF32 and autotuning, and prints JSON with every timing sample, summary statistics, peak allocated memory and output comparisons. Per-call weight packing is included in the measured decode time.

python benchmarks/diffusion/bench_wan_vae_decode.py \
    --model Wan-AI/Wan2.1-T2V-1.3B-Diffusers --size 832x480 --frames 1 \
    --dtype bf16 --first-frame-ablation --warmup 60 --iters 50 --seed 0

The default input is a seeded latent tensor. --latents /path/to/latents.pt can instead load a saved VAE input after latent mean/std scaling; the benchmark selects the temporal prefix requested by --frames and checks its shape.

The same script benchmarks multi-GPU VAE decode when launched with torchrun; the decode is timed across all ranks and rank 0 reports:

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