Description
Speeds up VAE decode/encode by applying fused Triton norm+SiLU kernels and channels_last convolution layout. Supported VAEs are auto-detected:
- Wan 2.1/2.2 video VAEs (including Qwen-Image variant): RMSNorm, ~1.4×/1.15× speedup.
- KL image VAEs (Flux/Flux2, SDXL, SD1.5): GroupNorm, ~1.6–1.8× at 2048px.
- LTXV/LTX2 video VAEs: PixelNorm; timestep-conditioned decoder blocks get norm-only fusion.
Other architectures are not supported and will be passed through unchanged. The node applies patches on a cloned patcher, so the optimizations only exist while this VAE is loaded.
Inputs
| Name | Type | Default | Tooltip |
|---|---|---|---|
vae | VAE | – | The VAE model to patch. |
fuse_norm_silu | BOOLEAN | true | Replace norm+SiLU chains (RMSNorm for Wan, GroupNorm for KL VAEs) with fused Triton kernels (single pass, fp32 accumulation). Requires Triton. |
channels_last | BOOLEAN | true | Convert convolution weights to channels_last memory format, removing cuDNN layout transposes around every conv. Required for the fused GroupNorm kernel to engage on KL VAEs. |
int8_conv | BOOLEAN | false | EXPERIMENTAL: run the VAE decoder’s 3×3 convolutions on int8 tensor cores (4× the bf16 rate on Ada+). Weights quantized per-out-channel, activations dynamically per-tensor; ~45–48 dB vs. bf16 decode, minor quality loss possible. Supported on Wan 2.1/2.2 and KL image VAEs (Flux2/SDXL/SD1.5); ignored for others. |
autotune | BOOLEAN | false | Benchmark several kernel block-size configs on first use of each tensor shape and cache the fastest. Brief stutter per new resolution, usually a few percent faster after warmup. |
Outputs
| Type | Description |
|---|---|
VAE | The patched VAE object (a cloned patcher). Use this output in place of the original VAE connection. |
Usage Notes
- All optimizations require Triton to be installed in your ComfyUI environment (typically via
pip install triton). Without Triton,fuse_norm_siluandint8_convwill be silently skipped. - The node auto-detects supported VAE architectures; unsupported VAEs are forwarded unmodified.
- Enable
autotuneafter you have settled on your typical resolutions for best long-run speed, but be aware of the initial stutter on new shapes. - The
int8_convflag is experimental – verify your output quality before relying on it in production. - Because the VAE is patched via a cloned patcher, the original VAE node in your workflow is unchanged. You must route the output of this node to the decode/encode nodes.
- This node is found in the KJNodes/experimental category of the node menu.
Comments
Sign in with GitHub to join the discussion.