fastvae

A fast video VAE decoder for consumer Blackwell (sm_120). Decodes at full resolution without tiling, 2.4x faster than the reference Wan VAE decoder in eager PyTorch.

Install

pip install fastvae

Pure Python. The kernels are Triton and compile on first call.

Use

Patch an existing diffusers pipeline:

from fastvae import accelerate_vae

accelerate_vae(pipe)   # pipe.vae.decode now runs through fastvae

Or call the decoder directly:

from fastvae import FastVAEDecoder

dec = FastVAEDecoder.from_vae(vae)
frames = dec.decode(latents)          # (B, 3, T, H, W) in [-1, 1]

For ComfyUI, install fastvae into ComfyUI's Python environment, then register the node:

mkdir -p ComfyUI/custom_nodes/fastvae
echo 'from fastvae.integrations.comfy_node import *' > ComfyUI/custom_nodes/fastvae/__init__.py

This adds FastVAEDecode under latent/video, a drop-in replacement for VAEDecode.

Benchmarks

RTX PRO 6000 Blackwell (sm_120), bf16, Wan 2.1 VAE. Median of 3 runs.

Every peak fits a 32 GB RTX 5090 without tiling. torch.compile reaches 1.74x on the shortest case, tiled decode is slower than eager, and the Wan 2.2 VAE speedup is 1.7x.

Reproduce with python -m fastvae.bench.harness --vae <pipeline-dir>. Raw output is in benchmarks/.

How it works

Four changes to the reference decoder:

1. Fused norm and activation. WanRMS_norm plus SiLU is about seven kernels in eager; one Triton kernel does it in a single pass.
2. Channels-last-3d throughout, which turns the reference's per-upsample layout conversions into free views.
3. Deferred residual adds and conv biases, absorbed into the next fused norm.
4. Temporal chunking instead of spatial tiling. Causal conv caching is exact for any chunk length, so chunk size trades peak VRAM for launch count without changing the output.

Ops that only move data are bitwise identical to the reference, asserted in tests. The channel-wise L2 norm is not, since its fp32 reduction order cannot be matched, which diffuses to a max absolute error of 4e-2 on a [-1, 1] output, against 5.2e-2 for torch.compile and 2.0e-1 for tiled decode.

Supported models

AutoencoderKLWan (Wan 2.1 and Wan 2.2). Support is an adapter that rebuilds a reference module tree from fastvae.ops primitives. See fastvae/models/wan.py and register_adapter.

Development

pytest tests/                                    # exactness suite
FASTVAE_TEST_MODEL=<pipeline-dir> pytest tests/  # plus real-checkpoint tests
python -m fastvae.bench.harness --vae <pipeline-dir>
python -m fastvae.bench.m0_profile --model <pipeline-dir>