A B200 speed/memory sweep (new scripts/quant_speedmem_bench.py) shows the VAE quant
win is a video story. Image AutoencoderKLs are ~0.15-0.26 GB, so fp8 saves ~0.1 GB
and only slows their tiny decode (+6-16%); the video Conv3d VAEs are ~2.5 GB and
halve to ~1.2 GB at ~2% decode cost. So VAE auto now only engages above a ~1 GB size
floor: small image VAEs stay dense (faster decode, no quant quality risk), video VAEs
still quantize. An explicit fp8 / fp8_dynamic request skips the gate (opted in).
The same sweep confirmed the text-encoder default is already right: fp8_dynamic is
E2E-neutral (denoise per-step unchanged; +2% one-time encode) and, by hidden-state
cosine vs bf16, marginally more accurate than layerwise fp8 -- so that default is left
as is. Tests cover the gate (small skipped, large quantized, explicit bypasses).
A B200 decoded-image LPIPS/SSIM sweep vs the dense bf16 VAE (new
scripts/quant_accuracy_sweep.py) settles the two VAE schemes:
- Layerwise fp8 (storage-only) holds across families (SSIM >= 0.977 on all but
SDXL), so auto now engages layerwise fp8 ONLY. For a VAE decode (a few percent
of end to end) fp8_dynamic's fp8-matmul speedup over storage fp8 is negligible,
so auto never takes the accuracy risk.
- fp8_dynamic (torchao PerTensor conv compute) is in-bar on only FLUX.2 and
Hunyuan and out-of-bar or catastrophic elsewhere (Qwen-Image SSIM 0.46), so it
is now an explicit opt-in, re-gated by a per-family deny list derived from the
sweep. SDXL denies both schemes (its small VAE stays dense).
Also fixes a real decode-time crash: torchao 0.17's fp8 conv kernel rejects
pointwise (1x1 / 1x1x1) convs ("Activation and filter channels must match"), so
an explicit fp8_dynamic request cast fine then threw at the first decode on most
families. The conv filter now keeps 1x1 convs dense, the smoke probe uses a
spatial 3x3 conv (so it exercises the path that actually runs), and an explicit
fp8_dynamic request runs that probe before casting.
Tests updated for the fp8-only auto ladder, the 1x1 exclusion, the explicit
probe gate, and the shipped deny list.
The transformer and text encoder auto-quantize; the VAE stayed dense. VAEs are
convolutional, so torchao int8 (Linear/2D-only) does not apply, but
Float8DynamicActivationFloat8WeightConfig quantizes Conv2d/Conv3d weights with
PerTensor granularity (auto-skipping convs whose channels are not a multiple of 16,
so the 3-channel RGB head stays dense). New diffusion_vae_quant.py offers two
schemes: fp8_dynamic (torchao conv compute fp8, cc>=8.9, resident) and fp8
(diffusers layerwise storage cast, any conv, survives offload); no int8 (no Conv3d
int8 kernel). select_vae_quant_scheme walks (fp8_dynamic, fp8) with a live conv
smoke probe, an offload gate, a per-family deny list, and a force_fp32 gate; the
image + video loaders map unset vae_quant to auto, skip the vae_force_fp32 Wan
families, and record the engaged scheme. Guards _align_vae_dtype to skip the
img2img/inpaint re-cast when the VAE is quantized (its fp8 tensor subclasses reject
.to(dtype=)). Verified on a B200: %16 Conv2d/Conv3d/Linear -> Float8Tensor, conv_out
dense, forward runs.
