ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

429
scripts/image_speedmem_bench.py Normal file
View file

@ -0,0 +1,429 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Speed + accuracy lever benchmark for the IMAGE diffusion backend (per-lever LPIPS).
Drives the SAME production lever functions the image loader calls -- ``apply_step_cache``,
``apply_attention_backend``, ``apply_speed_optims``, ``quantize_text_encoders``, the
compile-safe eager patches -- with the loader's own default arguments and order, so each
measured configuration reflects a real load. For each config it loads the pipeline fresh
(quant/compile mutate irreversibly), warms up (to pay the one-time compile), renders a
fixed prompt set at a fixed seed, and reports total latency, median per-step ms, peak
resident GB, and mean LPIPS(AlexNet) vs the bit-exact reference config (speed off,
native attention, uncached, dense) rendered at the same seed/settings.
Lever isolation knobs (for before/after measurement of shipped fixes):
--no-epc force torch._inductor.config.emulate_precision_casts back off after
the speed layer enables it (the pre-fix compile numerics).
--unarm-cache restore the cache hooks' eager inner forwards after the speed layer
arms them (the pre-fix cache x compile composition).
Example:
CUDA_VISIBLE_DEVICES=3 python scripts/image_speedmem_bench.py --family flux.1-dev \\
--config compile --out outputs/image_speedmem
"""
from __future__ import annotations
import argparse
import gc
import json
import os
import sys
import time
import types
from pathlib import Path
from typing import Any, Optional
os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
_REPO_ROOT = Path(__file__).resolve().parent.parent
_BACKEND_ROOT = _REPO_ROOT / "studio" / "backend"
for _p in (str(_BACKEND_ROOT), str(_REPO_ROOT / "scripts")):
if _p not in sys.path:
sys.path.insert(0, _p)
# Fixed prompt set (the diffusion_quality.py defaults + one photographic subject) so the
# LPIPS mean is not hostage to a single composition.
PROMPTS = [
"A cozy reading nook by a rain-streaked window, warm lamplight, a cat asleep on a stack of books",
"A lone lighthouse on a rocky cliff at sunset, dramatic clouds, crashing waves, highly detailed",
"A bustling night market street in the rain, neon signs reflected in puddles, cinematic",
"A photograph of an astronaut riding a horse on the surface of the moon, detailed, 8k",
]
# Production defaults per family (diffusion_families.default_generation_params).
_FAMILIES: dict[str, dict[str, Any]] = {
"qwen-image": {"repo": "Qwen/Qwen-Image", "family": "qwen-image"},
"flux.1-dev": {"repo": "black-forest-labs/FLUX.1-dev", "family": "flux.1"},
"flux.2-klein-4b": {"repo": "black-forest-labs/FLUX.2-klein-4B", "family": "flux.2-klein"},
"sdxl": {"repo": "stabilityai/stable-diffusion-xl-base-1.0", "family": "sdxl"},
}
# te speed attn cache
_CONFIGS: dict[str, dict[str, Any]] = {
# bit-exact reference: everything off / native / dense.
"reference": dict(te = "none", speed = "off", attn = "native", cache = "off"),
# the non-compile floor: eager patches + attention auto-upgrade, no compile.
"eager": dict(te = "none", speed = "eager", attn = "auto", cache = "off"),
# the default dense tier (regional compile), uncached.
"compile": dict(te = "none", speed = "default", attn = "auto", cache = "off"),
# max tier (max-autotune regional compile + TF32 + fused QKV), uncached.
"speedmax": dict(te = "none", speed = "max", attn = "auto", cache = "off"),
# the default tier + FBCache (the auto path for 20+ step schedules).
"fbcache": dict(te = "none", speed = "default", attn = "auto", cache = "fbcache"),
# FBCache without compile (isolates the cache's own drift from the compile floor).
"fbcache_eager": dict(te = "none", speed = "eager", attn = "auto", cache = "fbcache"),
# TE quant isolation on the bit-exact stack: the conditioning perturbation ALONE.
"te_fp8dyn": dict(te = "fp8_dynamic", speed = "off", attn = "native", cache = "off"),
"te_fp8": dict(te = "fp8", speed = "off", attn = "native", cache = "off"),
}
def _sync() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.synchronize()
def _reset_peak() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
def _alloc_gb() -> float:
import torch
return torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
def _peak_gb() -> float:
import torch
return torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
def _empty() -> None:
import torch
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
_LP: dict = {}
def _lpips_alex(ref_arr, arr) -> Optional[float]:
"""LPIPS(AlexNet) between two HxWx3 uint8 images (net on CPU). None if lpips missing."""
try:
import lpips
import torch
fn = _LP.get("fn")
if fn is None:
fn = lpips.LPIPS(net = "alex", verbose = False).eval()
_LP["fn"] = fn
def _t(a):
import torch as _torch
return _torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
with torch.no_grad():
return float(fn(_t(ref_arr), _t(arr)).item())
except Exception:
return None
def _import_diffusers():
import torch # noqa: F401
import torchao # noqa: F401
import diffusers.utils.import_utils as iu
iu._bitsandbytes_available = False
import diffusers
return diffusers
def _target():
"""Stand-in for DiffusionDeviceTarget: what the real lever functions read."""
import torch
return types.SimpleNamespace(
device = "cuda",
dtype = torch.bfloat16,
supports_default_torch_compile = True,
)
def _find_family(name: str):
from core.inference.diffusion_families import _FAMILIES as ALL
for fam in ALL:
if fam.name == name:
return fam
raise SystemExit(f"unknown family '{name}'")
def _apply_levers(
pipe,
cfg: dict,
*,
fam_obj,
no_epc: bool = False,
unarm_cache: bool = False,
logger = None,
) -> dict:
"""Apply the configured levers with the loader's own argument values, in the loader's
order (diffusion.py): TE quant -> attention -> step cache -> eager patches -> speed."""
from core.inference.diffusion_precision import quantize_text_encoders
from core.inference.diffusion_attention import (
apply_attention_backend,
select_attention_backend,
)
from core.inference.diffusion_cache import apply_step_cache, _restore_hooked_block_inners
from core.inference.diffusion_eager_patches import (
install_compile_safe_patches,
uninstall_patches,
)
from core.inference.diffusion_arch_patches import (
install_arch_patches,
uninstall_arch_patches,
)
from core.inference.diffusion_speed import apply_speed_optims
tgt = _target()
engaged: dict[str, Any] = {"te": None, "attn": None, "cache": None, "speed_optims": {}}
if cfg["te"] != "none":
engaged["te"] = quantize_text_encoders(
pipe, tgt, mode = cfg["te"], family = fam_obj.name, logger = logger
)
speed_mode = cfg["speed"]
engaged["attn"] = apply_attention_backend(
pipe,
select_attention_backend(
tgt, None if cfg["attn"] == "auto" else cfg["attn"], speed_active = speed_mode != "off"
),
logger = logger,
)
if cfg["cache"] != "off":
engaged["cache"] = apply_step_cache(
pipe, mode = cfg["cache"], quant_active = False, logger = logger
)
if speed_mode != "off":
install_compile_safe_patches()
install_arch_patches()
else:
uninstall_patches()
uninstall_arch_patches()
engaged["speed_optims"] = apply_speed_optims(
pipe,
tgt,
is_gguf = False,
family = fam_obj,
speed_mode = speed_mode,
cache_active = engaged["cache"] is not None,
offload_active = False,
)
if no_epc:
import torch
cfg_ind = getattr(getattr(torch, "_inductor", None), "config", None)
if cfg_ind is not None and hasattr(cfg_ind, "emulate_precision_casts"):
cfg_ind.emulate_precision_casts = False
engaged["epc_forced_off"] = True
if unarm_cache:
transformer = getattr(pipe, "transformer", None)
if transformer is not None:
_restore_hooked_block_inners(transformer)
engaged["cache_unarmed"] = True
return engaged
def _generate(
pipe,
fam_obj,
*,
steps: int,
guidance: float,
size: int,
seed: int,
limit: Optional[int] = None,
) -> tuple:
"""Render every prompt at a fixed per-prompt seed; returns (arrays, total_s, step_ms)."""
import numpy as np
import torch
call_params = {}
try:
import inspect
call_params = inspect.signature(pipe.__call__).parameters
except (TypeError, ValueError):
pass
step_times: list[float] = []
last: dict[str, float] = {}
def _cb(p, i, t, kw):
now = time.perf_counter()
if "t" in last:
step_times.append(now - last["t"])
last["t"] = now
return kw
arrs = []
total = 0.0
for idx, prompt in enumerate(PROMPTS[: limit or len(PROMPTS)]):
kwargs: dict[str, Any] = {
"prompt": prompt,
"num_inference_steps": steps,
"width": size,
"height": size,
"generator": torch.Generator("cuda").manual_seed(seed + idx),
}
if fam_obj.cfg_kwarg in call_params:
kwargs[fam_obj.cfg_kwarg] = guidance
if "callback_on_step_end" in call_params:
kwargs["callback_on_step_end"] = _cb
last.clear()
_sync()
t0 = time.perf_counter()
with torch.inference_mode():
image = pipe(**kwargs).images[0]
_sync()
total += time.perf_counter() - t0
arrs.append(np.array(image.convert("RGB")))
med_step = sorted(step_times)[len(step_times) // 2] * 1000.0 if step_times else None
return arrs, total, med_step, step_times
def main() -> None:
ap = argparse.ArgumentParser(description = __doc__.splitlines()[0])
ap.add_argument("--family", required = True, choices = sorted(_FAMILIES))
ap.add_argument("--config", required = True, choices = sorted(_CONFIGS))
ap.add_argument("--steps", type = int, default = None, help = "override the family default")
ap.add_argument("--size", type = int, default = 1024)
ap.add_argument("--seed", type = int, default = 42)
ap.add_argument("--out", default = "outputs/image_speedmem")
ap.add_argument("--no-epc", action = "store_true")
ap.add_argument("--unarm-cache", action = "store_true")
ap.add_argument("--tag", default = None, help = "output row name (default: config name)")
args = ap.parse_args()
import logging
logging.basicConfig(level = logging.INFO, format = "%(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("image_speedmem")
fam_spec = _FAMILIES[args.family]
cfg = _CONFIGS[args.config]
tag = args.tag or args.config
import numpy as np
import torch
diffusers = _import_diffusers()
from core.inference.diffusion_families import default_generation_params
fam_obj = _find_family(fam_spec["family"])
steps, guidance = default_generation_params(fam_spec["repo"])
if args.steps is not None:
steps = args.steps
out_dir = Path(args.out) / args.family
out_dir.mkdir(parents = True, exist_ok = True)
ref_npz = out_dir / f"ref_seed{args.seed}_st{steps}_{args.size}.npz"
logger.info(
"family=%s config=%s steps=%d guidance=%s size=%d seed=%d",
args.family,
args.config,
steps,
guidance,
args.size,
args.seed,
)
_reset_peak()
t0 = time.perf_counter()
pipe = diffusers.DiffusionPipeline.from_pretrained(fam_spec["repo"], torch_dtype = torch.bfloat16)
load_s = time.perf_counter() - t0
engaged = _apply_levers(
pipe,
cfg,
fam_obj = fam_obj,
no_epc = args.no_epc,
unarm_cache = args.unarm_cache,
logger = logger,
)
pipe.to("cuda")
weights_gb = _alloc_gb()
# Warmup: pays the one-time compile (and the cuDNN autotune) outside the timed runs.
wt0 = time.perf_counter()
_generate(
pipe,
fam_obj,
steps = steps,
guidance = guidance,
size = args.size,
seed = args.seed + 1000,
limit = 1,
)
warmup_s = time.perf_counter() - wt0
_reset_peak()
arrs, total_s, med_step_ms, step_times = _generate(
pipe, fam_obj, steps = steps, guidance = guidance, size = args.size, seed = args.seed
)
gen_peak = _peak_gb()
# Persist / score against the reference.
lpips_vals: list[float] = []
if args.config == "reference" and not (args.no_epc or args.unarm_cache):
np.savez_compressed(ref_npz, *arrs)
if ref_npz.exists():
ref = np.load(ref_npz)
refs = [ref[k] for k in ref.files]
for r, a in zip(refs, arrs):
v = _lpips_alex(r, a)
if v is not None:
lpips_vals.append(v)
from PIL import Image
for i, a in enumerate(arrs):
Image.fromarray(a).save(out_dir / f"{tag}_p{i}.png")
row = {
"family": args.family,
"config": args.config,
"tag": tag,
"steps": steps,
"guidance": guidance,
"size": args.size,
"seed": args.seed,
"engaged": {k: v for k, v in engaged.items()},
"load_s": round(load_s, 2),
"warmup_s": round(warmup_s, 2),
"total_gen_s": round(total_s, 2),
"per_image_s": round(total_s / len(PROMPTS), 3),
"median_step_ms": round(med_step_ms, 1) if med_step_ms else None,
"step_times_s": [round(t, 4) for t in step_times],
"weights_gb": round(weights_gb, 2),
"gen_peak_gb": round(gen_peak, 2),
"lpips_vs_ref_mean": round(sum(lpips_vals) / len(lpips_vals), 4) if lpips_vals else None,
"lpips_vs_ref_per_prompt": [round(v, 4) for v in lpips_vals] or None,
}
(out_dir / f"{tag}.json").write_text(json.dumps(row, indent = 2, default = str))
print(json.dumps(row, indent = 2, default = str))
del pipe
_empty()
if __name__ == "__main__":
main()