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:
- Fused norm and activation.
WanRMS_normplus SiLU is about seven kernels in eager; one Triton kernel does it in a single pass. - Channels-last-3d throughout, which turns the reference's per-upsample layout conversions into free views.
- Deferred residual adds and conv biases, absorbed into the next fused norm.
- 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
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
e1d0d1b73397782f0e59ae24c1835441f1b904c5c71103cb7753c93336167b17
|
|
| MD5 |
e2aa4de7681a44985b33f2dc5d04651e
|
|
| BLAKE2b-256 |
37d82ea6cd5b719e0968e9c3aef6c6422e6fc9f3cf2b52921aa11f78dcdfd3ce
|
Provenance
The following attestation bundles were made for fastvae-0.1.0.tar.gz:
Publisher:
release.yml on Occipital-Labs/fastvae
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
fastvae-0.1.0.tar.gz -
Subject digest:
e1d0d1b73397782f0e59ae24c1835441f1b904c5c71103cb7753c93336167b17 - Sigstore transparency entry: 2531587140
- Sigstore integration time:
-
Permalink:
Occipital-Labs/fastvae@0e2dbe898cec23e65681ae3f012f9acec46e9974 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/Occipital-Labs
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
release.yml@0e2dbe898cec23e65681ae3f012f9acec46e9974 -
Trigger Event:
push
-
Statement type:
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
a0df1361b72a7050eaa35f1f1b08958c74e02da0790ea60431c69aa25b881e35
|
|
| MD5 |
ed961652f6176bf6cef8e1afc8932f11
|
|
| BLAKE2b-256 |
cca36d65d1790d028d85fbdb43fb550a1b90c9def56faa7aab26aeade16b88ee
|
Provenance
The following attestation bundles were made for fastvae-0.1.0-py3-none-any.whl:
Publisher:
release.yml on Occipital-Labs/fastvae
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
fastvae-0.1.0-py3-none-any.whl -
Subject digest:
a0df1361b72a7050eaa35f1f1b08958c74e02da0790ea60431c69aa25b881e35 - Sigstore transparency entry: 2531587802
- Sigstore integration time:
-
Permalink:
Occipital-Labs/fastvae@0e2dbe898cec23e65681ae3f012f9acec46e9974 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/Occipital-Labs
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
release.yml@0e2dbe898cec23e65681ae3f012f9acec46e9974 -
Trigger Event:
push
-
Statement type: