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
parent
ec90b8658d
commit
de2f22df2b
7 changed files with 997 additions and 2 deletions
429
scripts/image_speedmem_bench.py
Normal file
429
scripts/image_speedmem_bench.py
Normal 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()
|
||||
|
|
@ -64,6 +64,126 @@ def normalize_transformer_cache(value: Optional[str]) -> Optional[str]:
|
|||
return normalized
|
||||
|
||||
|
||||
def _invalidate_child_registry_cache(transformer: Any) -> None:
|
||||
"""Drop the HookRegistry's cached child-registry list after (un)installing hooks.
|
||||
|
||||
``cache_context`` propagates the state context through ``_get_child_registries``,
|
||||
which diffusers 0.39 caches on first use. An UNCACHED generation already calls
|
||||
``cache_context`` (the pipeline wraps every denoise call), creating the
|
||||
transformer-level registry with an EMPTY cached child list -- so a later
|
||||
``enable_cache`` (the auto step-count toggle engaging FBCache mid-session) installs
|
||||
block hooks that ``_set_context`` never reaches, and the first cached forward dies
|
||||
with "No context is set". Invalidate the stale cache so the next ``cache_context``
|
||||
rebuilds it over the freshly hooked blocks. Best-effort and cheap (one attribute)."""
|
||||
registry = getattr(transformer, "_diffusers_hook", None)
|
||||
if registry is not None and getattr(registry, "_child_registries_cache", None) is not None:
|
||||
try:
|
||||
registry._child_registries_cache = None
|
||||
except Exception: # noqa: BLE001 -- diffusers internals moved; leave as-is
|
||||
pass
|
||||
|
||||
|
||||
# diffusers' cache hook registry names whose compute branch we re-point at a compiled
|
||||
# inner forward (leader = the measuring first block, block = the remaining ones); both
|
||||
# hook families share the fn_ref layout.
|
||||
_CACHE_HOOK_NAMES = (
|
||||
"mag_cache_leader_block_hook",
|
||||
"mag_cache_block_hook",
|
||||
"fbc_leader_block_hook",
|
||||
"fbc_block_hook",
|
||||
)
|
||||
|
||||
|
||||
def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
|
||||
"""Restore the regional compile on cache-hooked blocks' COMPUTED steps.
|
||||
|
||||
``enable_cache`` replaces each block's ``forward`` with the hook's ``new_forward``
|
||||
(stashing the pre-hook bound method in ``fn_ref.original_forward``), whose skip
|
||||
decision is data-dependent Python: MagCache ``@torch.compiler.disable``s the whole
|
||||
``new_forward`` (recursive -- the compute branch runs EAGER), and even FBCache's
|
||||
traceable ``new_forward`` graph-breaks around its disabled threshold decision,
|
||||
which on some archs (measured: Qwen-Image) drops the compute branch's call into
|
||||
``original_forward`` out of the compiled region -- the block's regional compile
|
||||
artifact (``_compiled_call_impl``) is never reached and the cache forfeits the
|
||||
compile win on every non-skipped step. An explicitly ``torch.compile``d callable
|
||||
re-enables
|
||||
dynamo for its own extent even inside a disabled frame, so re-pointing
|
||||
``fn_ref.original_forward`` at a compiled wrapper of the same bound method restores
|
||||
compiled compute steps while the skip decision stays eager exactly as designed.
|
||||
Measured (B200, scripts/image_speedmem_bench.py): Qwen-Image FBCache computed steps
|
||||
91.8 -> 71.2 ms (= the uncached compiled rate), 1.21x end to end; FLUX.1-dev is
|
||||
neutral (its FBCache ``new_forward`` happens to trace, so computed steps were
|
||||
already compiled -- same-process armed vs unarmed latents bit-identical); on the
|
||||
video DiT balanced MagCache went 39.4 -> 26.9 s at 50 steps.
|
||||
|
||||
Only blocks the speed layer actually compiled are armed (``_compiled_call_impl``
|
||||
guard -- eager tiers stay untouched), and only when ``original_forward`` is a plain
|
||||
bound method (a stacked hook chain, e.g. offload, captures a partial and is
|
||||
skipped). Idempotent via the ``_unsloth_orig_inner`` marker; best-effort. Returns
|
||||
the number of hooks armed."""
|
||||
try:
|
||||
import torch
|
||||
except Exception: # noqa: BLE001 -- no torch, nothing to arm
|
||||
return 0
|
||||
armed = 0
|
||||
try:
|
||||
for module in transformer.modules():
|
||||
registry = getattr(module, "_diffusers_hook", None)
|
||||
if registry is None or getattr(module, "_compiled_call_impl", None) is None:
|
||||
continue
|
||||
hooks = getattr(registry, "hooks", None) or {}
|
||||
for name in _CACHE_HOOK_NAMES:
|
||||
hook = hooks.get(name)
|
||||
fn_ref = getattr(hook, "fn_ref", None) if hook is not None else None
|
||||
orig = getattr(fn_ref, "original_forward", None)
|
||||
if orig is None or getattr(hook, "_unsloth_orig_inner", None) is not None:
|
||||
continue
|
||||
if getattr(orig, "__self__", None) is None:
|
||||
continue # not the plain bound method; arming would miss the block
|
||||
# fullgraph=False / dynamic=True: a cache is active by definition (its
|
||||
# decision points graph-break) and this matches the default tier the
|
||||
# regional compile used. Dynamo caches per code object, so re-arming
|
||||
# after a toggle is effectively free (~0.03 s).
|
||||
fn_ref.original_forward = torch.compile(orig, fullgraph = False, dynamic = True)
|
||||
hook._unsloth_orig_inner = orig
|
||||
armed += 1
|
||||
except Exception as exc: # noqa: BLE001 -- best-effort: the cache still works eager
|
||||
_warn(logger, "cache-hook inner compile", exc)
|
||||
return armed
|
||||
if armed and logger is not None:
|
||||
logger.info(
|
||||
"diffusion.cache: %d cache-hooked block(s) armed with compiled inner forwards",
|
||||
armed,
|
||||
)
|
||||
return armed
|
||||
|
||||
|
||||
def _restore_hooked_block_inners(transformer: Any) -> None:
|
||||
"""Undo ``_compile_hooked_block_inners``: put the plain bound methods back and clear
|
||||
the markers. MUST run before ``disable_cache`` -- ``remove_hook`` splices
|
||||
``fn_ref.original_forward`` back into ``module.forward``, and leaving the compiled
|
||||
wrapper there would pin a stale compiled callable onto the uncached path."""
|
||||
try:
|
||||
modules = list(transformer.modules())
|
||||
except Exception: # noqa: BLE001 -- not a torch module (tests/fakes): nothing armed
|
||||
return
|
||||
for module in modules:
|
||||
registry = getattr(module, "_diffusers_hook", None)
|
||||
if registry is None:
|
||||
continue
|
||||
hooks = getattr(registry, "hooks", None) or {}
|
||||
for name in _CACHE_HOOK_NAMES:
|
||||
hook = hooks.get(name)
|
||||
orig = getattr(hook, "_unsloth_orig_inner", None) if hook is not None else None
|
||||
if orig is None:
|
||||
continue
|
||||
try:
|
||||
hook.fn_ref.original_forward = orig
|
||||
hook._unsloth_orig_inner = None
|
||||
except Exception: # noqa: BLE001 -- per-hook best-effort
|
||||
pass
|
||||
|
||||
|
||||
def _pipeline_opens_cache_context(pipe: Any) -> bool:
|
||||
"""Whether the pipeline enters ``transformer.cache_context(...)`` in its denoise loop.
|
||||
The First-Block-Cache hook requires it at run time, and a CacheMixin transformer alone
|
||||
|
|
@ -137,6 +257,15 @@ def apply_step_cache(
|
|||
|
||||
config = FirstBlockCacheConfig(threshold = thr)
|
||||
enable_cache(config)
|
||||
# enable_cache AFTER the pipe has already run leaves a stale cached child-registry
|
||||
# list on the transformer's HookRegistry; the block hooks just installed would then
|
||||
# never receive the cache context. Must follow every enable_cache.
|
||||
_invalidate_child_registry_cache(transformer)
|
||||
# If the blocks are already regionally compiled (the generation-time toggle
|
||||
# path: compile ran at load), re-point the fresh hooks' compute branch at
|
||||
# compiled inners; the load path (cache before compile) is armed by
|
||||
# _compile_repeated_blocks instead. No-op when nothing is compiled.
|
||||
_compile_hooked_block_inners(transformer, logger)
|
||||
try:
|
||||
transformer._unsloth_step_cache = f"{mode}@{thr}"
|
||||
except Exception: # noqa: BLE001 — marker is best-effort
|
||||
|
|
@ -146,7 +275,10 @@ def apply_step_cache(
|
|||
return mode
|
||||
except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached
|
||||
# enable_cache can fail after hooking some blocks; drop any partial hooks so
|
||||
# the reported-uncached model doesn't actually run half-cached.
|
||||
# the reported-uncached model doesn't actually run half-cached. Any armed
|
||||
# compiled inners must be restored FIRST (remove_hook splices original_forward
|
||||
# back into module.forward).
|
||||
_restore_hooked_block_inners(transformer)
|
||||
try:
|
||||
transformer.disable_cache()
|
||||
except Exception: # noqa: BLE001
|
||||
|
|
@ -225,6 +357,10 @@ def maybe_toggle_step_cache(
|
|||
disable_cache = getattr(transformer, "disable_cache", None)
|
||||
if callable(disable_cache):
|
||||
try:
|
||||
# Before remove_hook splices fn_ref.original_forward back into
|
||||
# module.forward: the compiled inner wrappers must not leak onto the
|
||||
# uncached path.
|
||||
_restore_hooked_block_inners(transformer)
|
||||
disable_cache()
|
||||
transformer._unsloth_step_cache = None
|
||||
if logger is not None:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -77,6 +77,9 @@ def snapshot_backend_flags() -> Optional[dict]:
|
|||
state["cudnn_tf32"] = bool(cudnn.allow_tf32)
|
||||
if hasattr(cudnn, "benchmark"):
|
||||
state["cudnn_benchmark"] = bool(cudnn.benchmark)
|
||||
inductor_cfg = _inductor_config()
|
||||
if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
|
||||
state["inductor_emulate_precision_casts"] = bool(inductor_cfg.emulate_precision_casts)
|
||||
return state
|
||||
|
||||
|
||||
|
|
@ -103,6 +106,19 @@ def restore_backend_flags(state: Optional[dict]) -> None:
|
|||
cudnn = getattr(torch.backends, "cudnn", None)
|
||||
_set(cudnn, "allow_tf32", "cudnn_tf32")
|
||||
_set(cudnn, "benchmark", "cudnn_benchmark")
|
||||
_set(_inductor_config(), "emulate_precision_casts", "inductor_emulate_precision_casts")
|
||||
|
||||
|
||||
def _inductor_config() -> Any:
|
||||
"""``torch._inductor.config`` or None. Resolved as attributes off the imported torch
|
||||
module (real torch exposes ``_inductor`` directly after ``import torch``) rather
|
||||
than a submodule import, so a stubbed/partial torch (tests, exotic builds) cleanly
|
||||
reports None instead of picking a stale real module out of ``sys.modules``."""
|
||||
try:
|
||||
import torch
|
||||
return getattr(getattr(torch, "_inductor", None), "config", None)
|
||||
except Exception: # noqa: BLE001 — no inductor -> nothing to snapshot/set
|
||||
return None
|
||||
|
||||
|
||||
def normalize_speed_mode(value: Optional[str]) -> str:
|
||||
|
|
@ -342,6 +358,19 @@ def _compile_repeated_blocks(
|
|||
for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver
|
||||
if hasattr(dynamo_cfg, _limit_attr):
|
||||
setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
|
||||
# Match eager's intermediate rounding inside inductor's fused pointwise kernels:
|
||||
# by default they keep chains in fp32 where eager materialises bf16 between ops,
|
||||
# a per-forward rounding delta that a multi-step denoise amplifies chaotically.
|
||||
# Measured (B200, scripts/image_speedmem_bench.py, pairwise LPIPS of the
|
||||
# compiled tier vs the same-stack eager tier): Qwen-Image 0.019 -> 0.006 at
|
||||
# identical speed, FLUX.1-dev 0.046 -> 0.029 at +2% step time, FLUX.2-klein
|
||||
# 0.018 -> 0.017 at identical speed; on the video DiT (HunyuanVideo-1.5-720p)
|
||||
# full-clip LPIPS vs bit-exact drops 0.221 -> 0.052 at zero cost. Process-
|
||||
# global, so snapshot_backend_flags carries it and unload restores the prior
|
||||
# value.
|
||||
inductor_cfg = _inductor_config()
|
||||
if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
|
||||
inductor_cfg.emulate_precision_casts = True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "compile_repeated_blocks", exc)
|
||||
return False
|
||||
|
|
@ -354,6 +383,20 @@ def _compile_repeated_blocks(
|
|||
engaged = True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "compile_repeated_blocks", exc)
|
||||
continue
|
||||
# A step cache engaged BEFORE this compile (the production load order) has
|
||||
# already wrapped each block's forward in a @torch.compiler.disable'd hook, so
|
||||
# the compute branch would run eager on every non-skipped step and forfeit the
|
||||
# regional compile entirely. Re-point the hooks' inner forward at compiled
|
||||
# wrappers; no-op when no cache hooks are installed. The toggle path (cache
|
||||
# engaged after load) is armed by apply_step_cache instead. Lazy import:
|
||||
# diffusion_cache imports nothing from this module, but keep the dependency
|
||||
# one-directional at import time.
|
||||
try:
|
||||
from .diffusion_cache import _compile_hooked_block_inners
|
||||
_compile_hooked_block_inners(transformer, logger)
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "cache-hook inner compile", exc)
|
||||
return engaged
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -356,3 +356,206 @@ def test_toggle_noop_without_cache_support(monkeypatch):
|
|||
|
||||
def test_toggle_noop_without_transformer():
|
||||
assert maybe_toggle_step_cache(types.SimpleNamespace(), steps = 28) is None
|
||||
|
||||
|
||||
# ── compiled cache-hook inners (regional compile x step cache composition) ──────────
|
||||
import functools # noqa: E402
|
||||
|
||||
from core.inference.diffusion_cache import ( # noqa: E402
|
||||
_compile_hooked_block_inners,
|
||||
_invalidate_child_registry_cache,
|
||||
_restore_hooked_block_inners,
|
||||
)
|
||||
|
||||
|
||||
class _BoundInner:
|
||||
"""Provides a plain bound method for fn_ref.original_forward (__self__ present)."""
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
return "eager"
|
||||
|
||||
|
||||
def _hooked_block(
|
||||
*,
|
||||
compiled = True,
|
||||
hook_name = "fbc_block_hook",
|
||||
bound = True,
|
||||
):
|
||||
inner = _BoundInner()
|
||||
orig = inner.forward if bound else functools.partial(_BoundInner.forward, inner)
|
||||
hook = types.SimpleNamespace(fn_ref = types.SimpleNamespace(original_forward = orig))
|
||||
block = types.SimpleNamespace(
|
||||
_diffusers_hook = types.SimpleNamespace(hooks = {hook_name: hook}),
|
||||
_compiled_call_impl = object() if compiled else None,
|
||||
)
|
||||
return block, hook, orig
|
||||
|
||||
|
||||
def _fake_dit(blocks):
|
||||
return types.SimpleNamespace(modules = lambda: [types.SimpleNamespace()] + blocks)
|
||||
|
||||
|
||||
def _stub_torch_compile(monkeypatch):
|
||||
compiled_calls = []
|
||||
|
||||
def _compile(fn, **kwargs):
|
||||
compiled_calls.append((fn, kwargs))
|
||||
wrapper = lambda *a, **k: fn(*a, **k) # noqa: E731
|
||||
wrapper._unsloth_test_compiled_of = fn
|
||||
return wrapper
|
||||
|
||||
torch = types.ModuleType("torch")
|
||||
torch.compile = _compile
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
return compiled_calls
|
||||
|
||||
|
||||
def test_arming_swaps_inner_for_compiled_wrapper(monkeypatch):
|
||||
calls = _stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block()
|
||||
assert _compile_hooked_block_inners(_fake_dit([block])) == 1
|
||||
assert hook.fn_ref.original_forward is not orig
|
||||
assert hook.fn_ref.original_forward._unsloth_test_compiled_of is orig
|
||||
assert hook._unsloth_orig_inner is orig
|
||||
# The inner compile must match the cache-active tier: graph-breakable + dynamic.
|
||||
assert calls[0][1] == {"fullgraph": False, "dynamic": True}
|
||||
|
||||
|
||||
def test_arming_is_idempotent(monkeypatch):
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, _ = _hooked_block()
|
||||
dit = _fake_dit([block])
|
||||
assert _compile_hooked_block_inners(dit) == 1
|
||||
once = hook.fn_ref.original_forward
|
||||
assert _compile_hooked_block_inners(dit) == 0 # marker short-circuits
|
||||
assert hook.fn_ref.original_forward is once
|
||||
|
||||
|
||||
def test_arming_skips_uncompiled_blocks(monkeypatch):
|
||||
# An eager-tier load has no _compiled_call_impl: the hook must stay untouched
|
||||
# (compiling the inner would ADD compile where the user chose eager).
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block(compiled = False)
|
||||
assert _compile_hooked_block_inners(_fake_dit([block])) == 0
|
||||
assert hook.fn_ref.original_forward is orig
|
||||
|
||||
|
||||
def test_arming_skips_partial_captured_inner(monkeypatch):
|
||||
# A stacked hook chain (e.g. group offload) captures a functools.partial, not the
|
||||
# plain bound method; arming would compile the wrong layer of the chain.
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block(bound = False)
|
||||
assert _compile_hooked_block_inners(_fake_dit([block])) == 0
|
||||
assert hook.fn_ref.original_forward is orig
|
||||
|
||||
|
||||
def test_arming_covers_every_cache_hook_family(monkeypatch):
|
||||
# FBCache is the image cache today, but the hook-name table already covers the
|
||||
# MagCache layout too (same fn_ref shape), so a future mode arms for free.
|
||||
_stub_torch_compile(monkeypatch)
|
||||
names = (
|
||||
"mag_cache_leader_block_hook",
|
||||
"mag_cache_block_hook",
|
||||
"fbc_leader_block_hook",
|
||||
"fbc_block_hook",
|
||||
)
|
||||
blocks = [_hooked_block(hook_name = n)[0] for n in names]
|
||||
assert _compile_hooked_block_inners(_fake_dit(blocks)) == len(names)
|
||||
|
||||
|
||||
def test_restore_puts_the_exact_original_back(monkeypatch):
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block()
|
||||
dit = _fake_dit([block])
|
||||
_compile_hooked_block_inners(dit)
|
||||
_restore_hooked_block_inners(dit)
|
||||
assert hook.fn_ref.original_forward is orig
|
||||
assert hook._unsloth_orig_inner is None
|
||||
|
||||
|
||||
def test_restore_tolerates_fakes_without_modules():
|
||||
_restore_hooked_block_inners(_MixinTransformer()) # no .modules(): no-op
|
||||
|
||||
|
||||
def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch):
|
||||
# The generation-time toggle engages the cache AFTER the load already compiled the
|
||||
# blocks; apply_step_cache must arm the fresh hooks itself.
|
||||
_stub_diffusers(monkeypatch)
|
||||
_stub_torch_compile(monkeypatch)
|
||||
block, hook, orig = _hooked_block()
|
||||
|
||||
class _T(_MixinTransformer):
|
||||
def modules(self):
|
||||
return [block]
|
||||
|
||||
t = _T()
|
||||
engaged = apply_step_cache(_pipe(t), mode = "fbcache")
|
||||
assert engaged == TC_FBCACHE
|
||||
assert hook.fn_ref.original_forward is not orig
|
||||
assert hook._unsloth_orig_inner is orig
|
||||
|
||||
|
||||
def test_toggle_disable_restores_inners_before_disable(monkeypatch):
|
||||
# remove_hook splices fn_ref.original_forward back into module.forward, so the
|
||||
# compiled wrapper must be swapped out BEFORE disable_cache runs.
|
||||
_stub_diffusers(monkeypatch)
|
||||
order = []
|
||||
|
||||
class _T(_ToggleTransformer):
|
||||
def disable_cache(self):
|
||||
super().disable_cache()
|
||||
order.append("disable")
|
||||
|
||||
def modules(self):
|
||||
order.append("restore-walk")
|
||||
return []
|
||||
|
||||
t = _T()
|
||||
maybe_toggle_step_cache(_pipe(t), steps = 28)
|
||||
mode = maybe_toggle_step_cache(_pipe(t), steps = 8)
|
||||
assert mode is None and t.disables == 1
|
||||
assert order[-2:] == ["restore-walk", "disable"]
|
||||
|
||||
|
||||
def test_enable_failure_restores_inners_before_partial_disable(monkeypatch):
|
||||
# enable_cache can fail after hooking (and arming) some blocks; the partial-hook
|
||||
# cleanup must un-arm them before disable_cache splices original_forward back.
|
||||
_stub_diffusers(monkeypatch)
|
||||
order = []
|
||||
|
||||
class _T(_ToggleTransformer):
|
||||
def enable_cache(self, config):
|
||||
raise RuntimeError("block signature not recognised")
|
||||
|
||||
def disable_cache(self):
|
||||
super().disable_cache()
|
||||
order.append("disable")
|
||||
|
||||
def modules(self):
|
||||
order.append("restore-walk")
|
||||
return []
|
||||
|
||||
t = _T()
|
||||
assert apply_step_cache(_pipe(t), mode = "fbcache") is None
|
||||
assert order == ["restore-walk", "disable"]
|
||||
|
||||
|
||||
# ── stale child-registry cache invalidation (mid-session enable) ────────────────────
|
||||
|
||||
|
||||
def test_enable_invalidates_stale_child_registry_cache(monkeypatch):
|
||||
# diffusers 0.39 caches the child-registry list on first cache_context use; an
|
||||
# UNCACHED generation already populates it (empty), so a later toggle-time
|
||||
# enable_cache would install hooks the context never reaches ("No context is set").
|
||||
_stub_diffusers(monkeypatch)
|
||||
t = _MixinTransformer()
|
||||
t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = ["stale"])
|
||||
assert apply_step_cache(_pipe(t), mode = "fbcache") == TC_FBCACHE
|
||||
assert t._diffusers_hook._child_registries_cache is None
|
||||
|
||||
|
||||
def test_invalidate_child_registry_cache_tolerates_absence():
|
||||
_invalidate_child_registry_cache(types.SimpleNamespace()) # no registry: no-op
|
||||
reg = types.SimpleNamespace(_child_registries_cache = None)
|
||||
_invalidate_child_registry_cache(types.SimpleNamespace(_diffusers_hook = reg))
|
||||
assert reg._child_registries_cache is None
|
||||
|
|
|
|||
|
|
@ -393,3 +393,81 @@ def test_nvfp4_filter_keeps_vision_tower_dense(monkeypatch):
|
|||
assert ff(object(), "lm_head") is False
|
||||
assert ff(object(), "model.decoder.wo") is False
|
||||
assert ff(object(), "model.layers.5.self_attn.q_proj") is True
|
||||
|
||||
|
||||
# ── zero-output-row guard (per-row fp8 NaN protection) ───────────────────────────
|
||||
|
||||
|
||||
class _FakeAmaxVec:
|
||||
def __init__(self, vals):
|
||||
self._vals = vals
|
||||
|
||||
def __eq__(self, other): # noqa: PLW0642 -- tensor-style elementwise compare
|
||||
return _FakeAmaxVec([v == other for v in self._vals])
|
||||
|
||||
def any(self):
|
||||
return _FakeScalar(any(self._vals))
|
||||
|
||||
|
||||
class _FakeScalar:
|
||||
def __init__(self, v):
|
||||
self._v = v
|
||||
|
||||
def item(self):
|
||||
return self._v
|
||||
|
||||
|
||||
class _FakeWeight:
|
||||
"""Tensor-shaped stand-in supporting the exact chain the guard runs:
|
||||
``weight.abs().amax(dim = -1) == 0 -> .any().item()``."""
|
||||
|
||||
ndim = 2
|
||||
|
||||
def __init__(self, rows):
|
||||
self._rows = rows
|
||||
|
||||
def abs(self):
|
||||
return _FakeWeight([[abs(v) for v in r] for r in self._rows])
|
||||
|
||||
def amax(self, dim = -1):
|
||||
return _FakeAmaxVec([max(r) for r in self._rows])
|
||||
|
||||
|
||||
def test_weight_zero_output_row_detection():
|
||||
# A dead output row NaNs torchao's per-row fp8 (scale 0 -> 0/0); SDXL's
|
||||
# text_encoder_2 (OpenCLIP bigG) really ships one in layers.2.self_attn.out_proj --
|
||||
# measured: every fp8_dynamic SDXL render was black until the row is kept dense.
|
||||
zero_row = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.0, 0.0]]))
|
||||
dense = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.3, 0.0]]))
|
||||
assert dp._weight_has_zero_output_row(zero_row) is True
|
||||
assert dp._weight_has_zero_output_row(dense) is False
|
||||
# Non-2D / absent weights are not the per-row scheme's input: never flagged.
|
||||
w3 = _FakeWeight([[1.0]])
|
||||
w3.ndim = 3
|
||||
assert dp._weight_has_zero_output_row(types.SimpleNamespace(weight = w3)) is False
|
||||
assert dp._weight_has_zero_output_row(types.SimpleNamespace()) is False
|
||||
|
||||
# An unreadable weight falls through to quantize_'s own handling.
|
||||
class _Boom:
|
||||
@property
|
||||
def weight(self):
|
||||
raise RuntimeError("meta tensor")
|
||||
|
||||
assert dp._weight_has_zero_output_row(_Boom()) is False
|
||||
|
||||
|
||||
def test_fp8_dynamic_filter_skips_zero_row_linear(monkeypatch):
|
||||
# The fp8_dynamic caster must leave a zero-output-row Linear dense while the rest
|
||||
# of the encoder still quantises (a family-wide deny would forfeit the whole win).
|
||||
_stub_torch(monkeypatch)
|
||||
captured: dict = {}
|
||||
_stub_transformer_quant(monkeypatch, captured)
|
||||
enc = types.SimpleNamespace(_keep_in_fp32_modules = [])
|
||||
|
||||
dp._cast_fp8_dynamic(enc, _target())
|
||||
|
||||
ff = captured["filter_fn"]
|
||||
dead = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.0, 0.0]]))
|
||||
live = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.5, 0.5]]))
|
||||
assert ff(dead, "text_model.encoder.layers.2.self_attn.out_proj") is False
|
||||
assert ff(live, "text_model.encoder.layers.2.mlp.fc1") is True
|
||||
|
|
|
|||
|
|
@ -531,3 +531,85 @@ def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch):
|
|||
)
|
||||
assert applied["fp16_accum"] is True
|
||||
assert torch.backends.cuda.matmul.allow_fp16_accumulation is True
|
||||
|
||||
|
||||
# ── inductor precision-cast emulation (compile-vs-eager numeric parity) ─────────
|
||||
|
||||
|
||||
def _stub_inductor_config(
|
||||
monkeypatch,
|
||||
torch,
|
||||
*,
|
||||
emulate = False,
|
||||
):
|
||||
"""Attach a fake ``_inductor.config`` to the stubbed torch module (diffusion_speed
|
||||
resolves it as attributes off the imported torch, never via sys.modules -- so the
|
||||
real torch._inductor lingering in sys.modules cannot leak into stubbed tests)."""
|
||||
cfg = types.SimpleNamespace(emulate_precision_casts = emulate)
|
||||
torch._inductor = types.SimpleNamespace(config = cfg)
|
||||
return cfg
|
||||
|
||||
|
||||
def test_regional_compile_enables_emulate_precision_casts(monkeypatch):
|
||||
# Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds
|
||||
# to bf16 between ops; over a multi-step denoise that compounds to a visible drift.
|
||||
# emulate_precision_casts restores eager's rounding at zero measured speed cost, so
|
||||
# the regional compile path must switch it on.
|
||||
torch = _stub_torch(monkeypatch)
|
||||
_stub_gguf_accel(monkeypatch)
|
||||
cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert cfg.emulate_precision_casts is True
|
||||
|
||||
|
||||
def test_snapshot_restores_emulate_precision_casts(monkeypatch):
|
||||
# The flag is process-global, so the unload path must restore the pre-load value
|
||||
# exactly like the TF32 / cudnn.benchmark globals.
|
||||
torch = _stub_torch(monkeypatch)
|
||||
cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
|
||||
snap = snapshot_backend_flags()
|
||||
assert snap["inductor_emulate_precision_casts"] is False
|
||||
cfg.emulate_precision_casts = True
|
||||
restore_backend_flags(snap)
|
||||
assert cfg.emulate_precision_casts is False
|
||||
|
||||
|
||||
def test_missing_inductor_config_is_tolerated(monkeypatch):
|
||||
# A build without torch._inductor (or with the flag renamed) must neither break the
|
||||
# snapshot nor the compile path.
|
||||
_stub_torch(monkeypatch) # the stub torch has no _inductor attribute
|
||||
_stub_gguf_accel(monkeypatch)
|
||||
snap = snapshot_backend_flags()
|
||||
assert "inductor_emulate_precision_casts" not in snap
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
|
||||
|
||||
def test_regional_compile_arms_cache_hook_inners(monkeypatch):
|
||||
# The production load order engages the step cache BEFORE compile, so the regional
|
||||
# compile pass must re-arm the already-installed cache hooks with compiled inner
|
||||
# forwards (otherwise every computed step runs eager under the hook's
|
||||
# torch.compiler.disable and forfeits the regional compile).
|
||||
_stub_torch(monkeypatch)
|
||||
_stub_gguf_accel(monkeypatch)
|
||||
from core.inference import diffusion_cache as dc_mod
|
||||
|
||||
armed = []
|
||||
monkeypatch.setattr(
|
||||
dc_mod,
|
||||
"_compile_hooked_block_inners",
|
||||
lambda transformer, logger = None: armed.append(transformer) or 1,
|
||||
)
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert armed == [pipe.transformer]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue