From ca8f415a7e11668a53153cba3575720fbc08250b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 10:04:46 +0000 Subject: [PATCH] Auto-quantize the VAE (image/video decoder) to fp8 (default on, gated) 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. --- studio/backend/core/inference/diffusion.py | 55 ++- .../core/inference/diffusion_vae_quant.py | 303 ++++++++++++++ studio/backend/core/inference/video.py | 42 ++ studio/backend/models/inference.py | 26 ++ studio/backend/routes/inference.py | 1 + studio/backend/routes/video.py | 2 + .../backend/tests/test_diffusion_vae_quant.py | 379 ++++++++++++++++++ 7 files changed, 805 insertions(+), 3 deletions(-) create mode 100644 studio/backend/core/inference/diffusion_vae_quant.py create mode 100644 studio/backend/tests/test_diffusion_vae_quant.py diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index f7db4a745a..da04a4ceee 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -91,6 +91,7 @@ from .diffusion_cache import ( normalize_transformer_cache, ) from .diffusion_precision import TE_QUANT_AUTO, normalize_te_quant, quantize_text_encoders +from .diffusion_vae_quant import VAE_QUANT_AUTO, normalize_vae_quant, quantize_vae from .diffusion_prequant import ( load_prequantized_transformer, resolve_prequant_source, @@ -388,6 +389,10 @@ class _LoadState: backend_flags_before: Optional[dict] = None # Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C). text_encoder_quant: Optional[str] = None + # VAE quantisation actually engaged: "fp8" (layerwise storage) | "fp8_dynamic" (torchao + # conv compute) | None. When set, the resident VAE holds fp8 tensor subclasses, so the + # img2img/inpaint _align_vae_dtype re-cast must be skipped (it would corrupt them). + vae_quant: Optional[str] = None # Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8" # | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded. transformer_quant: Optional[str] = None @@ -801,6 +806,7 @@ class DiffusionBackend: memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, + vae_quant: Optional[str] = None, transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, @@ -849,6 +855,7 @@ class DiffusionBackend: memory_mode = memory_mode, speed_mode = speed_mode, text_encoder_quant = text_encoder_quant, + vae_quant = vae_quant, transformer_quant = transformer_quant, transformer_quant_fast_accum = transformer_quant_fast_accum, transformer_prequant_path = transformer_prequant_path, @@ -1167,6 +1174,7 @@ class DiffusionBackend: memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, + vae_quant: Optional[str] = None, transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, @@ -1208,12 +1216,18 @@ class DiffusionBackend: normalize_attention_backend(attention_backend) normalize_transformer_cache(transformer_cache) normalize_te_quant(text_encoder_quant) + normalize_vae_quant(vae_quant) # text_encoder_quant tri-state, mirroring transformer_quant: UNSET (None / "") -> auto, # which picks the best accurate TE scheme for this GPU + family (fp8_dynamic / int8 / # layerwise fp8) or stays dense when none qualifies. An explicit "none"/"off" pins the # encoder dense; an explicit scheme forces it. So the shipped default is auto. if text_encoder_quant is None or str(text_encoder_quant).strip() == "": text_encoder_quant = TE_QUANT_AUTO + # vae_quant tri-state, same contract: UNSET -> auto (fp8_dynamic conv compute on resident + # fp8-GEMM silicon that passes the conv probe, else layerwise fp8, else dense); none/off -> + # dense; an explicit scheme forces it. + if vae_quant is None or str(vae_quant).strip() == "": + vae_quant = VAE_QUANT_AUTO # For a full pipeline the repo itself supplies every component, so it is its # own base; the single-file kinds resolve the companion base diffusers repo. base = ( @@ -1731,6 +1745,18 @@ class DiffusionBackend: offload_active = plan.offload_policy != OFFLOAD_NONE, logger = logger, ) + # Quantise the dense VAE (opt-in fp8 layerwise / fp8_dynamic torchao conv), + # also before placement so the offload hooks move the smaller weights. The + # image families do not force-fp32 their VAE; auto skips torchao under offload. + vae_quant_engaged = quantize_vae( + pipe, + target, + mode = vae_quant, + family = fam.name, + offload_active = plan.offload_policy != OFFLOAD_NONE, + force_fp32 = False, + logger = logger, + ) # Apply the placement planned above (from MEASURED free device memory vs # the model's estimated resident size). apply_memory_plan returns the @@ -1778,6 +1804,15 @@ class DiffusionBackend: if text_encoder_quant == TE_QUANT_AUTO else "requested", ), + "vae_quant": ( + vae_quant, + vae_quant_engaged or "off", + "dense (no accurate scheme for this GPU / disabled)" + if vae_quant_engaged is None + else "auto-selected for this GPU + family" + if vae_quant == VAE_QUANT_AUTO + else "requested", + ), "attention_backend": ( attention_backend, attention_engaged or "native", @@ -1823,6 +1858,7 @@ class DiffusionBackend: speed_optims = tuple(k for k, v in speed_applied.items() if v), backend_flags_before = backend_flags_before, text_encoder_quant = te_quant, + vae_quant = vae_quant_engaged, transformer_quant = transformer_quant_engaged, attention_backend = attention_engaged, attention_request = attention_backend, @@ -2231,7 +2267,9 @@ class DiffusionBackend: return pipe @staticmethod - def _align_vae_dtype(pipe: Any, denoiser_attr: str = "transformer") -> None: + def _align_vae_dtype( + pipe: Any, denoiser_attr: str = "transformer", vae_quant: Optional[str] = None + ) -> None: """Cast the VAE to the denoiser's compute dtype before an image-conditioned call. The img2img/inpaint pipelines VAE-encode the input image at the text- encoder dtype (bf16), but a prior txt2img DECODE may have left the shared VAE @@ -2239,7 +2277,14 @@ class DiffusionBackend: (bf16 image vs fp32 VAE). Re-aligning here is safe: our families run bf16 or fp32 only (the fp16 guard promotes fp16), and a later txt2img decode re-upcasts as needed. ``denoiser_attr`` is ``pipe.transformer`` for DiT families and - ``pipe.unet`` for SDXL. Best-effort; a no-op when already aligned.""" + ``pipe.unet`` for SDXL. Best-effort; a no-op when already aligned. + + Skipped when ``vae_quant`` engaged: a quantised VAE holds fp8 tensor subclasses + that mishandle ``.to(dtype=...)`` (torchao rejects it), so the re-cast would + corrupt the weights. The VAE already runs at the compute dtype under fp8, so the + alignment is unnecessary there anyway.""" + if vae_quant is not None: + return denoiser = getattr(pipe, denoiser_attr, None) vae = getattr(pipe, "vae", None) if denoiser is None or vae is None: @@ -2724,7 +2769,9 @@ class DiffusionBackend: if init_pil is not None: # Keep the VAE encode dtype consistent with the input image. # state.family is always a DiffusionFamily, which defines denoiser_attr. - self._align_vae_dtype(pipe, state.family.denoiser_attr) + # A quantised VAE (fp8 tensor subclasses) must NOT be re-cast, so pass the + # engaged scheme through to skip the re-align in that case. + self._align_vae_dtype(pipe, state.family.denoiser_attr, state.vae_quant) # Pipelines vary in which kwargs they accept (img2img derives size from the # input image and may reject width/height; a distilled pipe may take no @@ -3007,6 +3054,7 @@ class DiffusionBackend: "speed_mode": None, "speed_optims": [], "text_encoder_quant": None, + "vae_quant": None, "transformer_quant": None, "attention_backend": None, "transformer_cache": None, @@ -3032,6 +3080,7 @@ class DiffusionBackend: "speed_mode": state.speed_mode, "speed_optims": list(state.speed_optims), "text_encoder_quant": state.text_encoder_quant, + "vae_quant": state.vae_quant, "transformer_quant": state.transformer_quant, "attention_backend": state.attention_backend, "transformer_cache": state.transformer_cache, diff --git a/studio/backend/core/inference/diffusion_vae_quant.py b/studio/backend/core/inference/diffusion_vae_quant.py new file mode 100644 index 0000000000..92ccd0e985 --- /dev/null +++ b/studio/backend/core/inference/diffusion_vae_quant.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto low-precision casting of the diffusion pipeline's VAE (image / video decoder). + +The transformer and the companion text encoder both quantise, but the VAE loads dense +and is a real resident chunk (the FLUX.2 / Qwen conv VAEs run to a few GB, the video +Conv3d VAEs more). Unlike the encoders it is CONVOLUTIONAL, so the Linear-only torchao +int8 path does not apply -- there is no int8 Conv3d kernel. Exactly two schemes fit: + + fp8_dynamic - torchao dynamic fp8 COMPUTE with PER-TENSOR granularity. Its + Float8DynamicActivationFloat8WeightConfig quantises Conv2d (4D) and + Conv3d (5D) weights, not just Linear, and torchao auto-skips any conv + whose C_out / C_in is not a multiple of 16 (so the 3-channel RGB + ``conv_out`` head stays dense). Needs fp8-GEMM silicon (cc >= 8.9) and a + resident (non-offloaded) VAE -- the fp8 tensor subclasses reject the + Module.to() an offload hook uses. + fp8 - diffusers layerwise casting: 8-bit (e4m3) STORAGE, upcast per layer to the + compute dtype. Storage-only, so it runs on ANY conv (2D / 3D) on any + fp8-capable card (cc >= 8.9) and survives group offload. + +There is no int8 (no Conv3d int8 kernel) and no nvfp4 for the VAE. ``auto`` (the loader +default) walks (fp8_dynamic, fp8): fp8_dynamic on resident data-center / Ada+ silicon that +passes a live conv smoke probe, else layerwise fp8, else dense. ``none``/``off`` keeps the +VAE dense bf16; an explicit scheme forces it (re-gated). A quantised VAE MUST NOT be +``.to(dtype=...)``'d afterwards (the fp8 tensor subclasses mishandle it), so the loader +skips the img2img/inpaint VAE re-align when a scheme engaged. torch / diffusers / torchao +are imported lazily so the module stays importable in a no-torch runtime. +""" + +from __future__ import annotations + +from typing import Any, Optional + +VAE_QUANT_FP8 = "fp8" +VAE_QUANT_FP8_DYNAMIC = "fp8_dynamic" +VAE_QUANT_AUTO = "auto" +# Concrete schemes (excludes "auto"): "auto" resolves to one of these via +# select_vae_quant_scheme, and these are the values the casters dispatch on. +VAE_QUANT_MODES = (VAE_QUANT_FP8, VAE_QUANT_FP8_DYNAMIC) + +# The RGB decoder head + output normalisations stay dense. torchao's fp8_dynamic path +# already skips the 3-channel ``conv_out`` via its C_out/C_in %16 rule, but the layerwise +# fp8 path has no such rule, so name them explicitly for both. Substring-matched against the +# module fqn: "conv_out"/"proj_out" (the pixel head), "conv_norm_out"/"norm_out" (the head's +# group-norm), whose low-magnitude outputs the coarse fp8 grid would band. +_VAE_KEEP_DENSE_TOKENS = ("conv_out", "proj_out", "conv_norm_out", "norm_out") + +# Best-first ``auto`` order: fp8_dynamic (compute fp8 on the conv tensor cores) leads; layerwise +# ``fp8`` (storage-only) is the universal fallback and the sole scheme that survives group offload. +_VAE_AUTO_LADDER = (VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8) + +# VAEs whose activation ranges break a scheme at the MODEL level (measured decoded-image +# LPIPS / SSIM vs the dense bf16 VAE). Populated from the accuracy sweep; a denied scheme is +# skipped by ``auto`` and refused when requested explicitly. Empty by default -- the +# vae_force_fp32 video families are gated separately at the loader (they never quantise). +_VAE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {} + +# Cache of device -> bool for the fp8_dynamic conv smoke probe (run once per device). +_VAE_DYNAMIC_PROBE_CACHE: dict[str, bool] = {} + + +def normalize_vae_quant(value: Optional[str]) -> Optional[str]: + """Lower/strip a requested VAE quant; None / "" / "none" / "off" -> None, + "auto" -> "auto" (resolved later by select_vae_quant_scheme). + + Raises ValueError for an unsupported value so a bad request is rejected cheaply.""" + if value is None: + return None + normalized = str(value).strip().lower().replace("-", "_") + if not normalized or normalized in ("none", "off"): + return None + if normalized == VAE_QUANT_AUTO: + return VAE_QUANT_AUTO + if normalized not in VAE_QUANT_MODES: + raise ValueError( + f"Unsupported vae_quant '{value}'. Use one of: " + f"{', '.join((VAE_QUANT_AUTO,) + VAE_QUANT_MODES)}, none/off." + ) + return normalized + + +def vae_quant_supported(target: Any, mode: str) -> bool: + """Whether ``mode`` is usable for ``target``: a CUDA device with a bf16 compute dtype, plus + the fp8 dtype (fp8 layerwise storage) and, for fp8_dynamic, fp8-GEMM silicon (Ada sm_89+ / + Hopper / Blackwell). There is no int8 / nvfp4 VAE scheme, so both modes need fp8.""" + if getattr(target, "device", None) != "cuda": + return False + try: + import torch + + if getattr(target, "dtype", None) is not torch.bfloat16: + return False + if mode == VAE_QUANT_FP8: + return hasattr(torch, "float8_e4m3fn") + if mode == VAE_QUANT_FP8_DYNAMIC: + # Compute fp8 (torch._scaled_mm on the conv weights) needs fp8-GEMM silicon. + return hasattr(torch, "float8_e4m3fn") and torch.cuda.get_device_capability() >= (8, 9) + except Exception: + return False + return False + + +def _vae_family_denied(family: Optional[str], scheme: str) -> bool: + return scheme in _VAE_FAMILY_SCHEME_DENY.get((family or "").strip().lower(), frozenset()) + + +def _vae_fp8_dynamic_probe(device: str) -> bool: + """True iff torchao's fp8_dynamic CONV path runs on this build: quantise a tiny + Conv2d(16, 16, 1) (channels a multiple of 16 so torchao does not skip it) with the + PerTensor fp8 config and run one forward. Cached per device. This makes ``auto`` robust + to a torchao build whose Float8 config lacks the conv path (the Linear-only fp8 the + transformer probes does not prove the conv path works) -- it fails here and the ladder + falls to layerwise fp8 rather than crashing at the first decode.""" + if device in _VAE_DYNAMIC_PROBE_CACHE: + return _VAE_DYNAMIC_PROBE_CACHE[device] + ok = False + try: + import torch + from torch import nn + from torchao.quantization import ( + Float8DynamicActivationFloat8WeightConfig, + PerTensor, + quantize_, + ) + + conv = nn.Conv2d(16, 16, 1).to(device = device, dtype = torch.bfloat16) + quantize_( + conv, + Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()), + filter_fn = lambda m, fqn = "": isinstance(m, nn.Conv2d), + ) + x = torch.randn(1, 16, 4, 4, device = device, dtype = torch.bfloat16) + with torch.no_grad(): + conv(x) + torch.cuda.synchronize() + ok = True + except Exception: + ok = False + _VAE_DYNAMIC_PROBE_CACHE[device] = ok + return ok + + +def select_vae_quant_scheme( + target: Any, + requested: Optional[str], + *, + family: Optional[str] = None, + offload_active: bool = False, + force_fp32: bool = False, +) -> Optional[str]: + """Resolve the concrete VAE scheme to apply, or None to stay dense bf16. + + An explicit scheme is returned as-is unless family-denied (``quantize_vae`` re-gates it). + ``force_fp32`` (a family that runs its VAE in fp32) always stays dense. ``auto`` walks + (fp8_dynamic, fp8) and returns the first that: survives the active offload policy (torchao + fp8_dynamic tensors reject the Module.to() an offload hook uses -> only layerwise fp8 under + offload), is not family-denied, is hardware-supported, and (fp8_dynamic only) passes a real + conv smoke probe. Returns None when nothing qualifies (e.g. no CUDA / pre-Ada).""" + requested = normalize_vae_quant(requested) + if requested is None or force_fp32: + return None + if requested != VAE_QUANT_AUTO: + return None if _vae_family_denied(family, requested) else requested + from .diffusion_transformer_quant import _capability + + if _capability() is None: + return None + device = str(getattr(target, "device", "cuda")) + for scheme in _VAE_AUTO_LADDER: + # torchao fp8_dynamic produces tensor subclasses that reject Module.to(); an offload hook + # moves the VAE that way and hard-crashes, so under offload only layerwise fp8 engages. + if offload_active and scheme == VAE_QUANT_FP8_DYNAMIC: + continue + if _vae_family_denied(family, scheme): + continue + if not vae_quant_supported(target, scheme): + continue + # fp8_dynamic additionally needs the torchao CONV fp8 kernel to actually run on this build. + if scheme == VAE_QUANT_FP8_DYNAMIC and not _vae_fp8_dynamic_probe(device): + continue + return scheme + return None + + +def quantize_vae( + pipe: Any, + target: Any, + *, + mode: Optional[str], + family: Optional[str] = None, + offload_active: bool = False, + force_fp32: bool = False, + logger: Any = None, +) -> Optional[str]: + """Quantise ``pipe.vae`` in place with ``mode`` (auto / fp8 / fp8_dynamic). Returns the scheme + actually engaged, or None when disabled, unsupported, force-fp32, family-denied, or no VAE was + cast. ``auto`` resolves to the best scheme for the GPU + family via ``select_vae_quant_scheme``; + an explicit scheme is re-gated the same way. Under offload the torchao fp8_dynamic mode is + skipped (its tensor subclasses reject the ``Module.to()`` an offload hook uses); layerwise fp8 + still engages. Best-effort: any failure leaves the VAE dense.""" + mode = normalize_vae_quant(mode) + if mode is None: + return None + if mode == VAE_QUANT_AUTO: + mode = select_vae_quant_scheme( + target, + VAE_QUANT_AUTO, + family = family, + offload_active = offload_active, + force_fp32 = force_fp32, + ) + if mode is None: + return None + else: + # An explicit scheme re-runs the same gates ``auto`` applies in select_vae_quant_scheme. + if force_fp32: + _note(logger, f"vae '{mode}' skipped: family runs its VAE in fp32") + return None + if _vae_family_denied(family, mode): + _note(logger, f"vae '{mode}' denied for family '{family}' (out-of-bar; staying dense)") + return None + # Layerwise fp8 is storage-only and streams fine under offload; the torchao fp8_dynamic + # tensors reject Module.to(), so skip only that one when an offload policy is active. + if offload_active and mode == VAE_QUANT_FP8_DYNAMIC: + _note( + logger, + f"vae '{mode}' skipped under offload (torchao tensors reject Module.to()); " + "pin a resident memory mode or use fp8", + ) + return None + if not vae_quant_supported(target, mode): + return None + vae = getattr(pipe, "vae", None) + if vae is None: + return None + try: + if mode == VAE_QUANT_FP8_DYNAMIC: + _cast_vae_fp8_dynamic(vae, target) + else: + _cast_vae_fp8(vae, target) + return mode + except Exception as exc: # noqa: BLE001 — leave the VAE dense + _warn(logger, mode, exc) + return None + + +def _cast_vae_fp8_dynamic(vae: Any, target: Any) -> None: + # torchao dynamic fp8 COMPUTE with PER-TENSOR granularity (NOT the DiT's per-row config): + # Float8DynamicActivationFloat8WeightConfig quantises Conv2d (4D) / Conv3d (5D) weights, not + # just Linear. Both channel dims must be a multiple of 16 (torchao skips the conv otherwise, + # which already leaves the 3-channel RGB head dense), and the decoder head / output norms are + # excluded by name for good measure. A weight.dim() < 2 (e.g. a bias-only leaf) never matches. + from torch import nn + from torchao.quantization import ( + Float8DynamicActivationFloat8WeightConfig, + PerTensor, + quantize_, + ) + + def filter_fn(module: Any, fqn: str = "") -> bool: + if not isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv3d)): + return False + weight = getattr(module, "weight", None) + if weight is None or weight.dim() < 2: + return False + if weight.shape[0] % 16 != 0 or weight.shape[1] % 16 != 0: + return False + name = fqn.lower() if fqn else "" + return not any(tok in name for tok in _VAE_KEEP_DENSE_TOKENS) + + quantize_(vae, Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()), filter_fn = filter_fn) + + +def _cast_vae_fp8(vae: Any, target: Any) -> None: + import re + + import torch + from diffusers.hooks import apply_layerwise_casting + from diffusers.hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN + + # diffusers' layerwise casting stores each supported leaf module's weights in fp8 and upcasts + # them per forward. Storage-only, so it works on any conv (2D / 3D) and survives offload. Keep + # the RGB decoder head / output norms dense (their low-magnitude pixels band on the coarse fp8 + # grid); the diffusers default pattern only covers pos/patch embeds. Names are literal + # substrings (re.escape) matched against the module fqn. + skip = tuple(DEFAULT_SKIP_MODULES_PATTERN) + tuple(re.escape(t) for t in _VAE_KEEP_DENSE_TOKENS) + apply_layerwise_casting( + vae, + storage_dtype = torch.float8_e4m3fn, + compute_dtype = target.dtype, + skip_modules_pattern = skip, + ) + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.vae_quant: (%s) failed: %s", what, exc) + + +def _note(logger: Any, msg: str) -> None: + if logger is not None: + logger.info("diffusion.vae_quant: %s", msg) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index cac8e13511..84d31decbd 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -79,6 +79,7 @@ from .diffusion_transformer_quant import ( select_transformer_quant_scheme, ) from .diffusion_precision import TE_QUANT_AUTO, normalize_te_quant, quantize_text_encoders +from .diffusion_vae_quant import VAE_QUANT_AUTO, normalize_vae_quant, quantize_vae from .video_families import ( VIDEO_CANCELLED_MSG, VIDEO_NOT_LOADED_MSG, @@ -289,6 +290,10 @@ class _VideoLoadState: # The companion text encoder (UMT5 / Gemma3 / Qwen2.5-VL) loads dense bf16 and is often the # largest resident component; this shrinks it in place, mirroring the image backend. text_encoder_quant: Optional[str] = None + # VAE quant actually engaged ("fp8" layerwise | "fp8_dynamic" torchao conv) or None. The + # convolutional decoder (Conv2d/Conv3d) shrinks in place; the vae_force_fp32 families + # (Wan) never quantise (force_fp32 -> dense). Mirrors the image backend's _LoadState. + vae_quant: Optional[str] = None resolved: Optional[dict] = None @@ -394,6 +399,7 @@ class VideoBackend: model_kind: Optional[str] = None, transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None, + vae_quant: Optional[str] = None, ) -> VideoFamily: """Cheap, network-free validation shared by the route and the load path.""" kind = resolve_video_model_kind(gguf_filename, model_kind) @@ -517,6 +523,8 @@ class VideoBackend: # Reject a malformed text_encoder_quant the same way (applies to any load kind: the dense # text encoder is resident for pipeline / gguf / single_file alike). normalize_te_quant(text_encoder_quant) + # Same for vae_quant (the dense VAE is resident for every load kind). + normalize_vae_quant(vae_quant) _ensure_mp4_encoder_available() return fam @@ -537,6 +545,7 @@ class VideoBackend: transformer_cache_threshold: Optional[float] = None, transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None, + vae_quant: Optional[str] = None, model_kind: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" @@ -549,6 +558,7 @@ class VideoBackend: model_kind = model_kind, transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant, + vae_quant = vae_quant, ) with self._lock: if self._loading is not None and self._loading.error is None: @@ -573,6 +583,7 @@ class VideoBackend: transformer_cache_threshold = transformer_cache_threshold, transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant, + vae_quant = vae_quant, model_kind = model_kind, _load_token = token, ), @@ -886,6 +897,7 @@ class VideoBackend: transformer_cache_threshold: Optional[float] = None, transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None, + vae_quant: Optional[str] = None, model_kind: Optional[str] = None, _load_token: Optional[int] = None, _base_local_dir: Optional[str] = None, @@ -901,6 +913,7 @@ class VideoBackend: model_kind = model_kind, transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant, + vae_quant = vae_quant, ) kind = resolve_video_model_kind(gguf_filename, model_kind) # text_encoder_quant tri-state (mirrors the image backend + transformer_quant): UNSET @@ -908,6 +921,10 @@ class VideoBackend: # "none"/"off" pins the encoder dense; a scheme forces it. So the shipped default is auto. if text_encoder_quant is None or str(text_encoder_quant).strip() == "": text_encoder_quant = TE_QUANT_AUTO + # vae_quant tri-state, same contract. The vae_force_fp32 families (Wan) keep the VAE dense + # regardless (quantize_vae's force_fp32 gate), so auto is safe as the shipped default. + if vae_quant is None or str(vae_quant).strip() == "": + vae_quant = VAE_QUANT_AUTO base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo) with self._lock: @@ -1206,6 +1223,19 @@ class VideoBackend: offload_active = plan.offload_policy != "none", logger = logger, ) + # Quantise the dense convolutional VAE (opt-in fp8 layerwise / fp8_dynamic torchao conv). + # The vae_force_fp32 families (Wan) run the VAE in fp32 for numerical stability, so + # force_fp32 pins them dense (quantising the fp32 VAE bands the decode); auto skips the + # torchao mode under offload. Best-effort: a failure leaves the VAE dense. + vae_quant_engaged = quantize_vae( + pipe, + target, + mode = vae_quant, + family = fam.name, + offload_active = plan.offload_policy != "none", + force_fp32 = getattr(fam, "vae_force_fp32", False), + logger = logger, + ) # ── optimisation layers, in the image backend's order: step cache FIRST # (compile keys its fullgraph decision off an active cache: FBCache hooks @@ -1383,6 +1413,15 @@ class VideoBackend: if text_encoder_quant_engaged is not None else "not engaged (dense bf16 text encoder loaded)", ), + "vae_quant": ( + vae_quant, + vae_quant_engaged or "off", + "dense VAE quantised in place" + if vae_quant_engaged is not None + else "not engaged (fp32 VAE family / offload / disabled -> dense)" + if getattr(fam, "vae_force_fp32", False) + else "not engaged (dense VAE loaded)", + ), } ) @@ -1416,6 +1455,7 @@ class VideoBackend: cache_threshold = transformer_cache_threshold, transformer_quant = transformer_quant_engaged, text_encoder_quant = text_encoder_quant_engaged, + vae_quant = vae_quant_engaged, resolved = resolved, ) # Ownership of the globals transferred to _state / _teardown_state. @@ -1770,6 +1810,7 @@ class VideoBackend: "transformer_cache": None, "transformer_quant": None, "text_encoder_quant": None, + "vae_quant": None, "has_audio": False, "defaults": None, "resolved": None, @@ -1798,6 +1839,7 @@ class VideoBackend: "transformer_cache": state.transformer_cache, "transformer_quant": state.transformer_quant, "text_encoder_quant": state.text_encoder_quant, + "vae_quant": state.vae_quant, "has_audio": fam.has_audio, "defaults": { "steps": default_steps, diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0be18d2e93..35aada65a7 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1758,6 +1758,15 @@ class DiffusionLoadRequest(BaseModel): "cc>=8.0), or nvfp4 (~4x smaller, Blackwell sm_100+). none/off keeps it dense bf16. A " "memory-vs-quality tradeoff (shifts fine detail); pairs well with balanced mode.", ) + vae_quant: Optional[Literal["auto", "none", "off", "fp8", "fp8_dynamic"]] = Field( + None, + description = "Quantise the VAE (image decoder). auto (the default when unset) picks the " + "best accurate scheme for this GPU + family: fp8_dynamic (torchao PER-TENSOR fp8 COMPUTE " + "on the conv tensor cores, cc >= 8.9, resident only) where a conv probe passes, else " + "layerwise fp8 (diffusers 8-bit storage, cc >= 8.9, survives offload), else dense. The " + "3-channel RGB head + output norms always stay dense. There is no int8 (no Conv3d int8 " + "kernel). none/off keeps the VAE dense bf16; an explicit scheme forces it.", + ) transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = ( Field( None, @@ -2117,6 +2126,9 @@ class DiffusionStatusResponse(BaseModel): text_encoder_quant: Optional[str] = Field( None, description = "Text-encoder quantisation engaged: fp8 | nvfp4 | null" ) + vae_quant: Optional[str] = Field( + None, description = "VAE quantisation engaged: fp8 | fp8_dynamic | null" + ) transformer_quant: Optional[str] = Field( None, description = "Transformer quant engaged on the dense fast path: int8 | fp8 | " @@ -2364,6 +2376,15 @@ class VideoLoadRequest(BaseModel): "for a family without a measured schedule); nvfp4 = torchao 4-bit weight-only (Blackwell " "sm_100+). none/off keeps the encoder dense. Mirrors the image backend's field.", ) + vae_quant: Optional[Literal["auto", "none", "off", "fp8", "fp8_dynamic"]] = Field( + None, + description = "Quantise the VAE (video decoder). auto (the default when unset) picks the " + "best accurate scheme for this GPU + family: fp8_dynamic (torchao PER-TENSOR fp8 COMPUTE " + "on the Conv2d/Conv3d tensor cores, cc >= 8.9, resident only), else layerwise fp8 " + "(diffusers 8-bit storage, cc >= 8.9, survives offload), else dense. The fp32-VAE families " + "(Wan) always stay dense. No int8 (no Conv3d int8 kernel). none/off keeps it dense; an " + "explicit scheme forces it. Mirrors the image backend's field.", + ) @field_validator("attention_backend", mode = "before") @classmethod @@ -2535,6 +2556,11 @@ class VideoStatusResponse(BaseModel): "(null = the dense bf16 encoder is loaded). An int8 request without a per-family " "keep-bf16 schedule is reported as the fp8 it fell back to.", ) + vae_quant: Optional[str] = Field( + None, + description = "VAE quant engaged: fp8 | fp8_dynamic | null (null = the dense VAE is " + "loaded; the fp32-VAE families always report null).", + ) has_audio: bool = Field( False, description = "Whether the loaded family produces a synchronized audio track" ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 75c913e01c..6e89b5dc08 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -12515,6 +12515,7 @@ async def load_diffusion_model( memory_mode = request.memory_mode, speed_mode = request.speed_mode, text_encoder_quant = request.text_encoder_quant, + vae_quant = request.vae_quant, transformer_quant = request.transformer_quant, transformer_quant_fast_accum = request.transformer_quant_fast_accum, transformer_prequant_path = request.transformer_prequant_path, diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index a053c9fac7..cf421a6923 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -96,6 +96,7 @@ async def load_video_model( model_kind = request.model_kind, transformer_quant = request.transformer_quant, text_encoder_quant = request.text_encoder_quant, + vae_quant = request.vae_quant, ) # Refuse while training is running: a multi-GB video pipeline would compete # with the training subprocess for VRAM. Mirrors the image-load guard. @@ -125,6 +126,7 @@ async def load_video_model( transformer_cache_threshold = request.transformer_cache_threshold, transformer_quant = request.transformer_quant, text_encoder_quant = request.text_encoder_quant, + vae_quant = request.vae_quant, model_kind = request.model_kind, ) return VideoStatusResponse(**status_dict) diff --git a/studio/backend/tests/test_diffusion_vae_quant.py b/studio/backend/tests/test_diffusion_vae_quant.py new file mode 100644 index 0000000000..ba621d6526 --- /dev/null +++ b/studio/backend/tests/test_diffusion_vae_quant.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for VAE quantisation (``diffusion_vae_quant.py``). + +Hermetic: torch + the diffusers / torchao casters are stubbed via ``sys.modules`` so +gating, the conv filter, and the apply path run without a GPU, real diffusers, or real +torchao. Mirrors tests/test_diffusion_precision.py's stubbing style. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +import core.inference.diffusion_vae_quant as vq +from core.inference.diffusion_vae_quant import ( + VAE_QUANT_AUTO, + VAE_QUANT_FP8, + VAE_QUANT_FP8_DYNAMIC, + _cast_vae_fp8, + _cast_vae_fp8_dynamic, + normalize_vae_quant, + quantize_vae, + select_vae_quant_scheme, + vae_quant_supported, +) + + +def _target(*, device = "cuda", dtype = "bfloat16", cc = (10, 0)): + return types.SimpleNamespace(device = device, dtype = dtype, _cc = cc) + + +class _Weight: + """A stand-in for a conv / linear weight tensor: exposes ``.shape`` and ``.dim()``.""" + + def __init__(self, shape): + self.shape = shape + + def dim(self): + return len(self.shape) + + +def _stub_torch(monkeypatch, *, with_fp8 = True, cc = (10, 0)): + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.float16 = "float16" + if with_fp8: + torch.float8_e4m3fn = "float8_e4m3fn" + # The conv filter isinstance-checks nn.Linear / nn.Conv2d / nn.Conv3d, and the layerwise + # caster reads torch.float8_e4m3fn, so the stub torch must expose both. + torch.nn = types.SimpleNamespace( + Linear = type("Linear", (), {}), + Conv2d = type("Conv2d", (), {}), + Conv3d = type("Conv3d", (), {}), + ) + torch.cuda = types.SimpleNamespace( + get_device_capability = lambda *a: cc, + synchronize = lambda *a, **k: None, + is_available = lambda: True, + ) + monkeypatch.setitem(sys.modules, "torch", torch) + return torch + + +def _stub_torchao(monkeypatch, captured): + # torchao's fp8_dynamic conv config + quantize_. Records the (config, filter_fn) so the + # PerTensor granularity and the conv filter closure can be asserted. + tq = types.ModuleType("torchao.quantization") + tq.quantize_ = lambda module, config, filter_fn = None: captured.update( + module = module, config = config, filter_fn = filter_fn + ) + tq.Float8DynamicActivationFloat8WeightConfig = lambda granularity = None: ("fp8dyn", granularity) + tq.PerTensor = lambda: "pertensor" + monkeypatch.setitem(sys.modules, "torchao.quantization", tq) + return tq + + +def _stub_diffusers(monkeypatch, recorder): + hooks = types.ModuleType("diffusers.hooks") + casting = types.ModuleType("diffusers.hooks.layerwise_casting") + casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",) + hooks.apply_layerwise_casting = lambda module, **kw: recorder.append(("fp8", module, kw)) + monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks) + monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting) + + +def _stub_capability(monkeypatch, cc): + """Stub the transformer module's ``_capability`` that select_vae_quant_scheme imports.""" + dtq = types.ModuleType("core.inference.diffusion_transformer_quant") + dtq._capability = lambda: cc + monkeypatch.setitem(sys.modules, "core.inference.diffusion_transformer_quant", dtq) + return dtq + + +def _allow_vae(monkeypatch, allowed): + """Force vae_quant_supported to accept only ``allowed`` (simulates the hardware gate).""" + monkeypatch.setattr(vq, "vae_quant_supported", lambda target, mode: mode in allowed) + + +# ── normalisation ───────────────────────────────────────────────────────────── + + +def test_normalize_vae_quant(): + assert normalize_vae_quant(None) is None + assert normalize_vae_quant("") is None + assert normalize_vae_quant("none") is None + # "off" disables (like the TE / transformer normalisers) -> dense. + assert normalize_vae_quant("off") is None + # "auto" passes through for select_vae_quant_scheme to resolve. + assert normalize_vae_quant("AUTO") == VAE_QUANT_AUTO + assert normalize_vae_quant("FP8") == VAE_QUANT_FP8 + # Hyphens fold to underscores so "fp8-dynamic" is accepted. + assert normalize_vae_quant("FP8-Dynamic") == VAE_QUANT_FP8_DYNAMIC + # int8 / nvfp4 have no VAE scheme -> rejected. + with pytest.raises(ValueError): + normalize_vae_quant("int8") + with pytest.raises(ValueError): + normalize_vae_quant("nvfp4") + + +# ── gating ──────────────────────────────────────────────────────────────────── + + +def test_vae_quant_supported_fp8_requires_cuda_bf16_and_fp8(monkeypatch): + _stub_torch(monkeypatch, with_fp8 = True, cc = (8, 9)) + assert vae_quant_supported(_target(), VAE_QUANT_FP8) is True + assert vae_quant_supported(_target(device = "cpu"), VAE_QUANT_FP8) is False + assert vae_quant_supported(_target(dtype = "float16"), VAE_QUANT_FP8) is False + # No fp8 dtype at all -> unsupported. + _stub_torch(monkeypatch, with_fp8 = False, cc = (8, 9)) + assert vae_quant_supported(_target(), VAE_QUANT_FP8) is False + + +def test_vae_quant_supported_fp8_dynamic_requires_sm89(monkeypatch): + # Compute fp8 conv (torch._scaled_mm) needs fp8-GEMM silicon: Ada sm_89+ / Hopper / Blackwell. + _stub_torch(monkeypatch, cc = (8, 9)) + assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is True + _stub_torch(monkeypatch, cc = (9, 0)) + assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is True + # Ampere (8.0) has no fp8 GEMM. + _stub_torch(monkeypatch, cc = (8, 0)) + assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is False + + +# ── auto ladder (select_vae_quant_scheme) ─────────────────────────────────────── + + +def test_select_datacenter_prefers_fp8_dynamic(monkeypatch): + # Data-center fp8-GEMM silicon with a passing conv probe: fp8_dynamic leads the ladder. + _stub_capability(monkeypatch, (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True) + assert select_vae_quant_scheme(_target(), "auto", family = "flux.1") == VAE_QUANT_FP8_DYNAMIC + + +def test_select_offload_uses_layerwise_fp8(monkeypatch): + # Under offload the torchao fp8_dynamic mode (rejects Module.to()) is skipped BEFORE the + # probe -> layerwise fp8. The probe must not even run. + _stub_capability(monkeypatch, (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + monkeypatch.setattr( + vq, "_vae_fp8_dynamic_probe", lambda device: pytest.fail("probe must not run under offload") + ) + assert ( + select_vae_quant_scheme(_target(), "auto", offload_active = True) == VAE_QUANT_FP8 + ) + + +def test_select_force_fp32_stays_dense(monkeypatch): + # A force-fp32 (Wan) family never quantises, for auto or an explicit request. + _stub_capability(monkeypatch, (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + assert select_vae_quant_scheme(_target(), "auto", force_fp32 = True) is None + assert select_vae_quant_scheme(_target(), "fp8", force_fp32 = True) is None + + +def test_select_family_deny_skips_scheme(monkeypatch): + _stub_capability(monkeypatch, (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True) + monkeypatch.setattr( + vq, "_VAE_FAMILY_SCHEME_DENY", {"badfam": frozenset({VAE_QUANT_FP8_DYNAMIC})} + ) + # fp8_dynamic denied for this family -> falls to layerwise fp8. + assert select_vae_quant_scheme(_target(), "auto", family = "badfam") == VAE_QUANT_FP8 + + +def test_select_no_capability_is_none(monkeypatch): + _stub_capability(monkeypatch, None) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + assert select_vae_quant_scheme(_target(), "auto") is None + + +def test_select_support_gate_falls_to_fp8(monkeypatch): + # fp8_dynamic hardware-unsupported (only fp8 allowed) -> layerwise fp8. + _stub_capability(monkeypatch, (8, 9)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8}) + monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True) + assert select_vae_quant_scheme(_target(), "auto") == VAE_QUANT_FP8 + + +def test_select_smoke_failure_falls_to_fp8(monkeypatch): + # fp8_dynamic is hardware-supported but its conv probe fails -> layerwise fp8. + _stub_capability(monkeypatch, (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: False) + assert select_vae_quant_scheme(_target(), "auto") == VAE_QUANT_FP8 + + +def test_select_explicit_passthrough(monkeypatch): + # An explicit request is returned as-is (quantize_vae re-gates it); no ladder walk, + # so no transformer-module stub is needed. + assert select_vae_quant_scheme(_target(), "fp8") == VAE_QUANT_FP8 + assert select_vae_quant_scheme(_target(), "fp8_dynamic") == VAE_QUANT_FP8_DYNAMIC + assert select_vae_quant_scheme(_target(), None) is None + assert select_vae_quant_scheme(_target(), "none") is None + + +def test_select_explicit_family_deny_returns_none(monkeypatch): + monkeypatch.setattr(vq, "_VAE_FAMILY_SCHEME_DENY", {"badfam": frozenset({VAE_QUANT_FP8})}) + assert select_vae_quant_scheme(_target(), "fp8", family = "badfam") is None + + +# ── conv-aware filter (_cast_vae_fp8_dynamic) ─────────────────────────────────── + + +def test_fp8_dynamic_conv_filter(monkeypatch): + # The real filter closure the PerTensor fp8 conv config receives: %16-channel Conv2d/Conv3d + # and Linear quantise; off-%16 channels and the conv_out / norm_out head stay dense. + torch = _stub_torch(monkeypatch) + captured: dict = {} + _stub_torchao(monkeypatch, captured) + _cast_vae_fp8_dynamic(object(), _target()) + # PerTensor granularity (NOT the DiT's per-row) is what the config was built with. + assert captured["config"] == ("fp8dyn", "pertensor") + ff = captured["filter_fn"] + nn = torch.nn + + def _mod(cls, shape): + m = cls() + m.weight = _Weight(shape) + return m + + # Conv2d (4D) / Conv3d (5D) with both channel dims a multiple of 16 quantise. + assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.up.0.resnets.0.conv1") is True + assert ff(_mod(nn.Conv3d, (64, 64, 3, 3, 3)), "decoder.mid_block.conv3d") is True + # nn.Linear (mid-block attention projections) also quantise. + assert ff(_mod(nn.Linear, (512, 512)), "decoder.mid_block.attentions.0.to_q") is True + # Channels not a multiple of 16 excluded (torchao would skip them regardless): the RGB + # in/out head (C=3) and any off-16 dim. + assert ff(_mod(nn.Conv2d, (128, 3, 3, 3)), "encoder.conv_in") is False + assert ff(_mod(nn.Conv2d, (24, 128, 3, 3)), "decoder.up.1.upsamplers.0.conv") is False + # conv_out / proj_out / norm_out excluded by NAME even with %16 channels. + assert ff(_mod(nn.Conv2d, (16, 128, 3, 3)), "decoder.conv_out") is False + assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.conv_norm_out") is False + assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.norm_out.conv") is False + assert ff(_mod(nn.Linear, (512, 512)), "decoder.proj_out") is False + # A non-conv/linear module (e.g. a GroupNorm) is excluded outright. + class _GroupNorm: + pass + + gn = _GroupNorm() + gn.weight = _Weight((128,)) + assert ff(gn, "decoder.mid_block.resnets.0.norm1") is False + # A weight with dim() < 2 (a 1D param) is excluded. + assert ff(_mod(nn.Conv2d, (128,)), "decoder.some.bias_only") is False + + +def test_cast_vae_fp8_layerwise_skips_head_and_norms(monkeypatch): + # The layerwise storage cast passes the decoder head + norm tokens through to diffusers' + # skip_modules_pattern (on top of the diffusers default), so they stay dense. + _stub_torch(monkeypatch) + recorder: list = [] + _stub_diffusers(monkeypatch, recorder) + vae = object() + _cast_vae_fp8(vae, _target()) + assert len(recorder) == 1 + _, mod, kw = recorder[0] + assert mod is vae + assert kw["storage_dtype"] == "float8_e4m3fn" + assert kw["compute_dtype"] == "bfloat16" + skip = kw["skip_modules_pattern"] + # The diffusers default is preserved and the keep-dense tokens are appended. + assert "norm" in skip + for tok in ("conv_out", "proj_out", "conv_norm_out", "norm_out"): + assert tok in skip + + +# ── apply (quantize_vae) ──────────────────────────────────────────────────────── + + +def test_quantize_vae_disabled_returns_none(monkeypatch): + pipe = types.SimpleNamespace(vae = object()) + assert quantize_vae(pipe, _target(), mode = None) is None + assert quantize_vae(pipe, _target(), mode = "none") is None + + +def test_quantize_vae_force_fp32_stays_dense(monkeypatch): + # A force-fp32 (Wan) family never casts, for an explicit scheme or auto. + _stub_torch(monkeypatch, cc = (10, 0)) + monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("must not cast")) + monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("must not cast")) + pipe = types.SimpleNamespace(vae = object()) + assert quantize_vae(pipe, _target(), mode = "fp8", force_fp32 = True) is None + assert quantize_vae(pipe, _target(), mode = "auto", force_fp32 = True) is None + + +def test_quantize_vae_offload_skips_fp8_dynamic(monkeypatch): + # Explicit fp8_dynamic under offload is skipped (torchao tensors reject Module.to()); + # layerwise fp8 still engages. + _stub_torch(monkeypatch, cc = (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + monkeypatch.setattr( + vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("torchao must not run under offload") + ) + pipe = types.SimpleNamespace(vae = object()) + assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", offload_active = True) is None + fp8_calls: list = [] + monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: fp8_calls.append(v)) + assert quantize_vae(pipe, _target(), mode = "fp8", offload_active = True) == VAE_QUANT_FP8 + assert len(fp8_calls) == 1 + + +def test_quantize_vae_unsupported_hw_is_noop(monkeypatch): + # An explicit scheme on hardware that does not support it applies nothing. + _stub_torch(monkeypatch, cc = (8, 0)) + _allow_vae(monkeypatch, set()) + monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("must not cast")) + pipe = types.SimpleNamespace(vae = object()) + assert quantize_vae(pipe, _target(), mode = "fp8") is None + + +def test_quantize_vae_none_vae_is_noop(monkeypatch): + # A pipeline with no VAE attribute is a best-effort no-op even when the mode is supported. + _stub_torch(monkeypatch, cc = (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8}) + pipe = types.SimpleNamespace() # no .vae + assert quantize_vae(pipe, _target(), mode = "fp8") is None + + +def test_quantize_vae_explicit_fp8_applies(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8}) + calls: list = [] + monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v)) + vae = object() + pipe = types.SimpleNamespace(vae = vae) + assert quantize_vae(pipe, _target(), mode = "fp8") == VAE_QUANT_FP8 + assert calls == [vae] + + +def test_quantize_vae_tolerates_caster_failure(monkeypatch): + # The caster raising leaves the VAE dense (best-effort) -> None. + _stub_torch(monkeypatch, cc = (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8}) + + def _boom(v, t): + raise RuntimeError("fp8 unsupported for this layer") + + monkeypatch.setattr(vq, "_cast_vae_fp8", _boom) + pipe = types.SimpleNamespace(vae = object()) + assert quantize_vae(pipe, _target(), mode = "fp8") is None + + +def test_quantize_vae_auto_resolves_and_applies(monkeypatch): + # End-to-end: mode="auto" resolves via the ladder then applies the resolved caster. + _stub_torch(monkeypatch, cc = (10, 0)) + _stub_capability(monkeypatch, (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) + monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True) + calls: list = [] + monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: calls.append(v)) + vae = object() + pipe = types.SimpleNamespace(vae = vae) + assert quantize_vae(pipe, _target(), mode = "auto", family = "flux.1") == VAE_QUANT_FP8_DYNAMIC + assert calls == [vae]