The fp8_dynamic conv smoke probe only exercised Conv2d, so a torchao build
whose Conv3d kernel path is missing or broken would pass the probe for a
video VAE (HunyuanVideo-1.5), report it quantized, and crash at the first
decode. The probe now runs per (device, conv ndim) and an explicit request
must pass it for every conv dimensionality the target VAE contains; the
auto ladder gate inspects the VAE the same way when one is provided.
The video benchmark moved the fully dense pipeline to CUDA before applying
the configured quant/optimisation levers, the reverse of the production
loader (video.py quantizes before apply_memory_plan). Dense-oversized
configs could OOM where the shipped quantized path loads fine, and
load_peak_gb recorded the dense placement. The bench now builds on CPU,
applies the levers, then places on CUDA and captures the load peak.
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.
