William Lin
cf3ea82e74
- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output. - WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32 accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it); decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the video, with the ranks agreeing on whether every one of them could build the kernel. - PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions, set_fused_kernels). - PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels). - VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame. - The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23.
Wan2.2 VAE decoder (sm_100a, sm_103a)
The decoder of the Wan2.2 VAE (the VAE of Wan2.2-TI2V-5B: 48 latent channels, 16 x 16 spatial and 4x temporal
compression) at 704 x 1248, on one, two or four B200 / GB200 GPUs of a node. It decodes latents [1, 48, T, 44, 78]
into the video [1, 3, 4T - 3, 704, 1248] for any T >= 1, as diffusers' AutoencoderKLWan.decode(z).sample does;
reference.py defines the result exactly, a plain-PyTorch decoder bit-exact with diffusers in fp32 and in bf16.
Use
import sys
sys.path.insert(0, "fastvideo-kernel/wan22_vae_decode")
from wan22_vae_decode import WanVAEDecoder
decoder = WanVAEDecoder(vae.state_dict(), device="cuda") # diffusers' AutoencoderKLWan, or its state dict
video = decoder.decode(latents) # [1, 48, T, 44, 78], as the VAE's decode takes them -> [1, 3, 4T-3, 704, 1248] bf16
- Over several GPUs, every rank builds a decoder with the same process group (
WanVAEDecoder(sd, device, pg), default: torch.distributed's default group) and callsdecodewith the full latents; rank 0 gets the video and the other ranksNone. On one GPU without torch.distributed, the decoder initializes a single-rank NCCL group. - The first call builds the extension from
solution/(about a minute) into$WAN22_VAE_DECODE_BUILD_DIR(default~/.cache/fastvideo/wan22_vae_decode, one directory perLOCAL_RANK), and later processes reuse the build. WanVAEDecoder(..., frames_per_pass=4)sets the latent frames each GPU carries through one pass (see below).WanVAEDecoderprepares the weights once (layout, phase and Winograd transforms of the weights). The raw entry points areprepare(weights, rank, world, pg),run(z, weights, rank, world, pg, out)(4 frames per pass) andrun_with(..., out, frames_per_pass)insolution/binding.cpp, destination-passing, on the current stream;weightsare the 110 tensors ofweight_names()in bf16.
Performance and memory
At PAN2's default 704 x 1248, 24 fps, on GB200 GPUs, against PAN2's decode today (diffusers' AutoencoderKLWan in bf16
under autocast, one GPU, one latent frame at a time). Memory is the peak allocated and, after it, the peak reserved by
PyTorch's caching allocator (what nvidia-smi shows), per GPU; the default is 4 latent frames per GPU per pass:
| Decoder | 8 s (193 frames) | 12 s (289 frames) | 20 s (481 frames) |
|---|---|---|---|
| Release, 1 GPU | 5.88 s, 12.9 / 17.1 GiB | 8.84 s, 13.9 / 19.4 GiB | 14.76 s, 15.8 / 91.5 GiB |
| This decoder, 1 GPU | 1.17 s (x5.0), 20.7 / 36.3 GiB | 1.74 s (x5.1), 21.2 / 36.8 GiB | 2.89 s (x5.1), 22.2 / 37.8 GiB |
| This decoder, 2 GPUs | 619 ms, 18.7 / 33.0 GiB | 929 ms, 19.2 / 33.5 GiB | 1.54 s, 20.2 / 34.5 GiB |
| This decoder, 4 GPUs | 328 ms, 17.8 / 41.4 GiB | 475 ms, 18.3 / 34.1 GiB | 787 ms (x18.8), 19.2 / 42.9 GiB |
frames_per_pass trades memory for speed and leaves the video unchanged. On one GPU:
| Latent frames per pass | 8 s | 12 s | 20 s |
|---|---|---|---|
| 2 | 1.25 s, 14.4 / 24.4 GiB | 1.88 s, 14.9 / 24.9 GiB | 3.13 s, 15.9 / 25.8 GiB |
| 4 (default) | 1.17 s, 20.7 / 36.3 GiB | 1.74 s, 21.2 / 36.8 GiB | 2.89 s, 22.2 / 37.8 GiB |
| 8 | 1.10 s, 33.7 / 62.1 GiB | 1.64 s, 34.2 / 62.6 GiB | 2.72 s, 35.2 / 63.6 GiB |
| 16 | 1.06 s, 59.8 / 121.0 GiB | 1.58 s, 60.2 / 121.7 GiB | 2.67 s, 61.2 / 122.5 GiB |
- With
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, the default reserves 30.6-32.0 GiB on one GPU and 2 frames per pass 19.6-21.1 GiB, at the same speed. - On the optimization run's nine workloads (
tools/bench.py: 121, 241 and 481 frames at 1, 2 and 4 GPUs, SM clocks 1.6-1.95 GHz), the default's geomean is 758 ms, against 1043 ms for the run's first solution (x1.38) and 698 ms at 16 frames per pass. - The CUTLASS Winograd component products take 60 % of the GPU time, the transforms and epilogues 30 %
(
notes/optimization.md).
Correctness
Against the fp32 decode of real latents, with d = video - reference per clip: every element finite and in [-1, 1],
max |d| at most three times what the release's bf16 decode errs by, and the RMS of d likewise. On this tree, every
clip at worlds 1, 2 and 4 and the short clips f1, f5, f17 at worlds 1 and 4 pass, using 0.30-0.66 of the max
allowance and 0.35-0.42 of the RMS one (PSNR 56-60 dB against fp32; the release's bf16 decode has 58-61 dB).
cd fastvideo-kernel/wan22_vae_decode
DATA=/path/to/data python tools/bench.py --solution solution # every workload this job has the GPUs for
DATA=/path/to/data python tools/bench.py --workload f5x1 f17x1 --no-time # correctness of the short clips
cd .. && WAN22_VAE_WEIGHTS=/path/to/vae python -m pytest tests/test_wan22_vae_decode.py # the wrapper, one GPU
tools/bench.py reads $DATA: the latents (latents_<source>.pt), the weights in weight_names() order
(decoder_weights_bf16.pt, weight_names.json), the fp32 decodes (ref_<clip>.pt) and the release decode's errors
(calibration.json); tools/spec.py documents each. The tests need only the VAE's weights.
Limits
- sm_100a and sm_103a (B200 / GB200, B300):
tcgen05with tensor memory, TMA and 2-CTA clusters; the loader builds for the GPU it runs on. CUDA 13, torch 2.12 tested; the tests pass on B300 too, built for sm_103a. - 704 x 1248 video (latents 44 x 78) and batch 1; any T >= 1.
- 1, 2 or 4 GPUs of one node were checked; the code accepts up to 8 ranks.
- bf16 weights and activations, fp32 accumulation and fp32 RMS-norm statistics.
- In FastVideo,
WanVAEConfig.decode_kernelroutesAutoencoderKLWan.decodehere where it applies (on one GPU per rank,WanVAEDecoder(..., local=True)), and PAN2 sampling turns it on (PAN2PipelineConfig.vae_decode_kernel). On a B300 with PAN2 latents it decodes 7x faster than the fp32 decode, at PSNR 61-63 dB and SSIM 0.998 against it.
Layout
solution/: the kernels and their binding;solution.jsonholds the build flags. The CUTLASS and CuTe headers come from fastvideo-kernel's CUTLASS submodule (include/cutlass, or$CUTLASS_DIR); the decoder's structure follows diffusers' (WAN_LICENSE.txt).wan22_vae_decode.py:WanVAEDecoder,decoder_weightsandweight_names.reference.py: the definition of the result.tools/:bench.py(check, audit and time),loader.py(the build),spec.py(workloads, weights, data),audit.py,bounds.py.notes/optimization.md: the design, what each merge bought, what did not pay, and where to go next.