Patch Qwen VAE Triton W8A8

If you've ever watched a Krea 2 or Qwen-Image run finish sampling and then hang for another four minutes, you've met the Qwen-Image VAE decoder. It's Wan 2.1's autoencoder with the encoder frozen and the decoder retuned for text legibility, and it's the most expensive thing in a short workflow. The author's own benchmark is honest about it: an 8-step 1024×1024 Krea 2 Turbo run takes 46 seconds of sampling and 257 seconds of VAE load and decode. The picture is already done; the codec is the bottleneck.

Patch Qwen VAE Triton W8A8 attacks exactly that. It quantizes the decoder's CausalConv3d layers to INT8 and runs them as Triton implicit-GEMM kernels, which in the tested setup cut the VAE phase from 257.675s to 208.415s and the total workflow from 343s to 291s. On a resolution change (1024×1024 → 840×1256) it took 43 seconds off. A real 15–16% end-to-end, not a rounding error.

What it actually does

Nothing exotic - just a selective lie about precision. The node walks the decoder, picks the causal 3D convolutions it can handle (3×3×3 kernel, stride 1, dilation 1, groups 1, at least 64 in/out channels, in-channels divisible by 32), and swaps them for a Triton implementation: weights symmetric INT8 with a per-output-channel scale, activations dynamically quantized to INT8 per call, INT8×INT8 multiplied with INT32 accumulation, rescaled in FP32, cast back to the incoming dtype. Latent-facing input and RGB-facing output layers stay native on purpose.

Two details worth knowing: it clones the VAE and registers the replacements through ComfyUI's object patcher rather than mutating your VAE in place, so other branches of the graph keep the original - and it keeps the VAE resident while the INT8 buffers are live, so expect a bit more VRAM pinned during decode, not less.

The two controls you actually touch

Preset is the whole user interface, and it's architecture-adaptive rather than a magic number. The node lists the distinct channel tiers it finds in your decoder and picks from the top down:

• Conservative - top channel tier only. On the tested Qwen topology, ~15 of 28 eligible layers.
• Balanced (the default) - top two tiers, ~22 of 28, leaving the lowest 96-channel tier native.
• Aggressive - everything eligible except two layers the author found produce unacceptable noise when quantized, decoder.upsamples.14.residual.2 and .6. ~26 W8A8 layers plus 2 native.

disable_node is a passthrough toggle, and it's the right way to A/B this: enable it, run the same seed twice, compare. There's a second output, report, a plain STRING listing the preset, the channel tiers, every patched layer and every layer left native. When output looks wrong, the report tells you instantly whether the layer you care about got quantized.

Wire it the boring way: VAELoader → Patch Qwen VAE Triton W8A8 → VAEDecode. It outputs a normal VAE, so no custom decode node is needed.

Installing it

ComfyUI Manager, search Qwen VAE Triton W8A8 (registry package qwen-vae-triton, publisher puppet-vision), or from a shell:

cd ComfyUI/custom_nodes
git clone https://github.com/AllenCraigBarnard/ComfyUI-Qwen-VAE-Triton

Then restart, and start ComfyUI with Triton enabled - this is not optional, the module raises on import without it:

python main.py --listen 0.0.0.0 --port 8188 --enable-triton-backend

If your Docker or launcher passes args via an env var, put the flag there instead (CLI_ARGS=--enable-triton-backend). There's no model download and no pip dependency in requirements.txt - by design, it's a comment telling you Triton must come from your existing PyTorch/ROCm/CUDA environment. And don't swap a working ROCm Triton build for an arbitrary PyPI wheel. The pack was developed on ROCm 7 / Strix Halo (gfx1151), so the GPU env-var block in the README is a working example for one machine, not a config to copy - HSA_OVERRIDE_GFX_VERSION=11.5.1 on a different card is a good way to break things.

Where people get burned

Wrong VAE, hard error. With no eligible layers it raises No eligible Qwen/Wan decoder CausalConv3d layers were found - you pointed it at something that isn't the Wan/Qwen causal autoencoder. An SDXL or Flux VAE won't work. Load qwen_image_vae.safetensors (Krea 2 uses the same file) and try again.

You see no speedup, and the log explains why. The kernel silently falls back to native convolution for anything it can't take: non-accelerator tensors, non-5D inputs, batch size above 1, unsupported kernel/stride, or any exception mid-kernel. Batch-decoding several images at once is the easy one to trip - the log prints W8A8 fallback [layer]: batch size 4. Not broken, just unaccelerated.

Triton on Windows. Historically painful, and nothing here changes that. Get triton working in your Comfy environment first.

Judge quality on a T2I image, not on Krea's over-smoothing. The Qwen VAE's texture ceiling is a known, contested complaint, and people swap in the Wan 2.1 decoder to fight it. That's a separate lever: this node makes the decode faster, it does not make the decoder better. If Aggressive shows speckle in flat areas, drop to Balanced - the author's own note is that Balanced matched Aggressive's end-to-end time after a resolution change with more layers left native, which is why Balanced is the default.

One last caveat, stated plainly: this is a brand-new pack from one author, and those benchmarks are his measurements on his hardware. Flip disable_node once and confirm the gain yourself before rebuilding a workflow around it.