William LinandClaude Opus 5.5
bec159cac0
- Attention: on sm_100+ with flash_attn.cute installed and FASTVIDEO_FA4 unset, the PAN2 pipeline turns FA4 on before its modules load. One B300, 241 frames: 288 s instead of 670 s with the SDPA fallback. FASTVIDEO_FA4=0 keeps the previous backend. - First frame: an image_path is read as its stored pixels, without the EXIF orientation rotation, as the reference sampler reads it (load_image gains exif_transpose=True). All 11 sign-off inputs, one EXIF-rotated, now reach the VAE bit-identical to the reference. - Wan2.2 VAE decode kernel: each process builds in a temporary directory of its own instead of a shared rank<LOCAL_RANK> directory, and finished builds go through the loader's keyed cache, added by atomic rename. Two processes building at once both succeed; a later one fetches the build in about 2 s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
100 lines
6.5 KiB
Markdown
100 lines
6.5 KiB
Markdown
# 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
|
|
|
|
```python
|
|
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 calls `decode` with the full latents; rank 0 gets the video and the
|
|
other ranks `None`. 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`). Each process builds in a directory of its own, so concurrent processes do
|
|
not collide, and puts the finished build in the `cache/` beside it, from which later processes copy it.
|
|
- `WanVAEDecoder(..., frames_per_pass=4)` sets the latent frames each GPU carries through one pass (see below).
|
|
- `WanVAEDecoder` prepares the weights once (layout, phase and Winograd transforms of the weights). The raw entry
|
|
points are `prepare(weights, rank, world, pg)`, `run(z, weights, rank, world, pg, out)` (4 frames per pass) and
|
|
`run_with(..., out, frames_per_pass)` in `solution/binding.cpp`, destination-passing, on the current stream;
|
|
`weights` are the 110 tensors of `weight_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).
|
|
|
|
```bash
|
|
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): `tcgen05` with 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_kernel` routes `AutoencoderKLWan.decode` here 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.json` holds 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_weights` and `weight_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.
|