Skip to main content

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.

case eager fastvae speedup peak VRAM
480p / 81f 2817 ms 1171 ms 2.41x 6.82 GiB
480p / 161f 5613 ms 2340 ms 2.40x 7.18 GiB
720p / 81f 6570 ms 2782 ms 2.36x 15.21 GiB
720p / 161f 13069 ms 5419 ms 2.41x 16.06 GiB

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>

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

fastvae-0.1.0.tar.gz (34.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

fastvae-0.1.0-py3-none-any.whl (34.5 kB view details)

Uploaded Python 3

File details

Details for the file fastvae-0.1.0.tar.gz.

File metadata

  • Download URL: fastvae-0.1.0.tar.gz
  • Upload date:
  • Size: 34.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for fastvae-0.1.0.tar.gz
Algorithm Hash digest
SHA256 e1d0d1b73397782f0e59ae24c1835441f1b904c5c71103cb7753c93336167b17
MD5 e2aa4de7681a44985b33f2dc5d04651e
BLAKE2b-256 37d82ea6cd5b719e0968e9c3aef6c6422e6fc9f3cf2b52921aa11f78dcdfd3ce

See more details on using hashes here.

Provenance

The following attestation bundles were made for fastvae-0.1.0.tar.gz:

Publisher: release.yml on Occipital-Labs/fastvae

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file fastvae-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: fastvae-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 34.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for fastvae-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 a0df1361b72a7050eaa35f1f1b08958c74e02da0790ea60431c69aa25b881e35
MD5 ed961652f6176bf6cef8e1afc8932f11
BLAKE2b-256 cca36d65d1790d028d85fbdb43fb550a1b90c9def56faa7aab26aeade16b88ee

See more details on using hashes here.

Provenance

The following attestation bundles were made for fastvae-0.1.0-py3-none-any.whl:

Publisher: release.yml on Occipital-Labs/fastvae

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page