perf(image): compile numeric parity, cache-hook compile arming, FBCache toggle crash fix, TE fp8 zero-row guard

Applies the video round-2 accuracy findings to the image diffusion stack and fixes
two real image-path bugs found while measuring. All numbers B200, production
settings (family default steps/guidance, 1024px, seed 42, 4 fixed prompts), LPIPS
(AlexNet) via the new scripts/image_speedmem_bench.py, which drives the production
lever functions in the loader's own order.

- inductor precision parity: emulate_precision_casts=True on the regional-compile
  path (fused pointwise kernels keep fp32 intermediates where eager rounds to bf16
  between ops). Pairwise LPIPS of the compiled tier vs the same-stack eager tier:
  Qwen-Image 0.019 to 0.006 at identical speed (72.4 vs 72.5 ms/step), FLUX.1-dev
  0.046 to 0.029 at +2% step time (69.8 vs 68.3, reproduced), FLUX.2-klein-4B
  0.018 to 0.017 at identical speed. Snapshot/restored with the other process-wide
  backend flags so an off load never inherits it.
- cache x compile composition: re-point each cache hook's fn_ref.original_forward
  at a torch.compile'd wrapper of the same bound method (armed only where the
  speed layer compiled the block; restored before every disable_cache and before
  the partial-hook cleanup). Qwen-Image FBCache computed steps 91.8 to 71.2 ms
  (back at the uncached compiled rate), 1.21x end to end (7.36 to 6.06 s per 4
  images); FLUX.1-dev already traced through its FBCache hook and is measured
  neutral (same-process armed vs unarmed latents bit-identical). Skip counts
  within noise (13 vs 11 of 76; pairwise LPIPS 0.005).
- FBCache mid-session toggle crash: diffusers 0.39 caches the HookRegistry child
  list on first cache_context use, so an uncached generation followed by a
  20+-step generation (the auto toggle path) enabled hooks the context never
  reached and crashed with "No context is set" (reproduced live on FLUX.1-dev).
  Invalidate the stale child cache after every enable_cache.
- TE fp8_dynamic zero-row guard: torchao per-row fp8 derives a per-output-channel
  scale from the row amax, so an all-zero weight row is 0/0 = NaN. SDXL's
  text_encoder_2 (OpenCLIP bigG) ships exactly such a row, and every explicit
  fp8_dynamic SDXL render came out black; keep zero-row Linears dense (LPIPS
  0.976 black to 0.096 working). Other families' encoders have no such rows and
  are byte-identical.
- No AUTO TE quant exists on the image branch (text_encoder_quant defaults dense,
  explicit-only), so the video round's auto-dense retune has no image analogue;
  the explicit lever's cost is now measured (TE fp8_dynamic alone, LPIPS vs
  bit-exact: Qwen-Image 0.038, FLUX.1-dev 0.084, SDXL 0.096; no speed win, VRAM
  -6.5 GB on Qwen-Image) for the docs.

Tests: 96 passing across the cache/speed/precision suites (11 new arming, 2
child-registry, 2 zero-row, 4 inductor-flag); ruff clean.
This commit is contained in:
Daniel Han 2026-07-10 16:07:46 +00:00
commit de2f22df2b
7 changed files with 997 additions and 2 deletions

View file

@ -217,6 +217,24 @@ def _cast_int8_selective(encoder: Any, target: Any, skip_first: int, skip_last:
quantize_(encoder, _make_quant_config(TQ_INT8), filter_fn = filter_fn)
def _weight_has_zero_output_row(module: Any) -> bool:
"""True when a Linear's weight contains an all-zero OUTPUT row. torchao's per-row
fp8 scheme derives a per-output-channel scale from that row's amax, so a dead row
yields scale 0 -> 0/0 = NaN through the whole forward. Real checkpoints ship such
rows: SDXL's text_encoder_2 (OpenCLIP ViT-bigG) has one in
``text_model.encoder.layers.2.self_attn.out_proj`` -- measured on B200: every
fp8_dynamic SDXL render came out black (NaN embeddings) until this Linear is left
dense. Cheap (one amax per Linear, once per load); False on any error so the
caster's own failure handling stays in charge."""
try:
weight = getattr(module, "weight", None)
if weight is None or weight.ndim != 2:
return False
return bool((weight.abs().amax(dim = -1) == 0).any().item())
except Exception: # noqa: BLE001 -- unreadable weight: let quantize_ decide
return False
def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
# torchao dynamic fp8 COMPUTE, per-row (per-token activation + per-output-channel weight ->
# torch._scaled_mm on the fp8 tensor cores). Unlike the layerwise `fp8` backend this keeps the
@ -232,9 +250,15 @@ def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
# require_bf16: scaled_mm asserts a bf16 weight, so skip any stray non-bf16 Linear the encoder
# keeps (belt-and-suspenders over the named T5 wo exclusion) rather than aborting the pass.
filter_fn = make_filter_fn(
base = make_filter_fn(
DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
)
# A Linear with an all-zero output row NaNs under per-row scaling (scale 0 -> 0/0);
# keep exactly those Linears dense so one dead row cannot black out every render.
def filter_fn(module: Any, fqn: str = "") -> bool:
return base(module, fqn) and not _weight_has_zero_output_row(module)
quantize_(encoder, _make_quant_config(TQ_FP8), filter_fn = filter_fn)