diff --git a/scripts/image_speedmem_bench.py b/scripts/image_speedmem_bench.py new file mode 100644 index 0000000000..ee2269d103 --- /dev/null +++ b/scripts/image_speedmem_bench.py @@ -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() diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index ad5eb50127..b79e148859 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -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: diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 24df29d26e..1f7b3d6ebb 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -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) diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 9a1ec7ecc5..9e2007719c 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 1794482239..eb8db319d4 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 8871f05ef0..b835434b37 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 594de55983..4d7bedacf5 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -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]