Warm-save the compile cache by default, compile U-Net denoisers whole-module
diffusion_compile_cache: auto mode now saves the Mega-cache bundle after the first compiled generation (UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE=0 opts out), so users get warm restarts without the distributor env; a bundle hit starts clean (no pointless rewrite of the just-loaded artifacts) and explicit mode 1/on keeps the distributor-style re-save. New register_shape + manifest shape coverage: a STATIC compile produces new artifacts per (width, height, batch), so the generate path registers each generation's shape and an uncovered shape re-dirties the context, growing the bundle to cover every shape the session used. Measured (B200, real backend): Qwen-Image deferred gen-3 hitch 29.1 -> 22.2 s warm with bit-identical output (7.9 MB bundle, ~0.5 s save); SDXL gen-3 115.7 -> 24.7 s and a mid-session 768px recompile 65.8 -> 12.6 s (bundle 63.6 -> 98.7 MB after the 768 re-save). diffusion_speed: U-Net denoisers (UNet2DConditionModel; no _repeated_blocks, so the regional compile never reached them) now get a whole-module STATIC torch.compile on the default tier, plus fused QKV projections and a compiled VAE decode. Measured on SDXL (30 steps / 7.0 / 1024px, 4 prompts, LPIPS vs the bit-exact reference): 6.16 -> 3.14 s end to end (1.96x) at LPIPS 0.035, steady state 0.70-0.88 s/image through the real backend. Rejected on measurement: dynamic=True whole-module (366 s compile for 39.3 ms/step vs static's 73 s for 26.9), regional BasicTransformerBlock only (45.0 ms/step; ResNet convs stay eager), max-autotune + inductor flags (25.9 ms/step for a 445 s warmup), channels-last UNet alone (neutral). DiT tiers unchanged: fused QKV measured exactly neutral under the regional compile (Qwen-Image 6.53 vs 6.52 s), so it stays max-only there, and the DiT VAE decode stays eager (a few % of a DiT generation). compiled_shapes_are_static tells the cache layer which loads are per-shape (max tier, U-Net whole-module). diffusion: register each generation's shape with the compile cache before the save, pass pipe.unet to the cache fingerprint when the pipe has no transformer, and correct the transformer_quant resolved reason on dense loads (it claimed a GGUF transformer was loaded on every non-quantized pipeline load). Tests: 333 passing across the related suites (speed 42, compile_cache 27, cache 40, precision 20, backend, base_precision, transformer_quant, memory); ruff clean. Full measurement record: outputs/image_optim_round2_audit.md.
This commit is contained in:
Daniel Han
parent
6cb44270fc
commit
352fb40089
5 changed files with 372 additions and 17 deletions
|
|
@ -67,6 +67,7 @@ from .diffusion_speed import (
|
|||
SPEED_OFF,
|
||||
apply_speed_optims,
|
||||
compile_eligible,
|
||||
compiled_shapes_are_static,
|
||||
normalize_speed_mode,
|
||||
resolve_speed_mode,
|
||||
restore_backend_flags,
|
||||
|
|
@ -1670,7 +1671,10 @@ class DiffusionBackend:
|
|||
):
|
||||
compile_ctx = compile_cache.begin(
|
||||
family = fam.name,
|
||||
transformer = getattr(pipe, "transformer", None),
|
||||
# U-Net families (SDXL) carry the denoiser as pipe.unet; the
|
||||
# fingerprint needs the module actually compiled.
|
||||
transformer = getattr(pipe, "transformer", None)
|
||||
or getattr(pipe, "unet", None),
|
||||
dtype = getattr(target, "dtype", None),
|
||||
quant = transformer_quant_engaged,
|
||||
attention_backend = attention_engaged,
|
||||
|
|
@ -1757,7 +1761,14 @@ class DiffusionBackend:
|
|||
"transformer_quant": (
|
||||
transformer_quant,
|
||||
transformer_quant_engaged or "off",
|
||||
"not engaged (GGUF transformer loaded)"
|
||||
# The None reason must match the load kind: only a GGUF
|
||||
# load has a GGUF transformer; a dense pipeline /
|
||||
# single-file load simply keeps its dense weights.
|
||||
(
|
||||
"not engaged (GGUF transformer loaded)"
|
||||
if kind == "gguf"
|
||||
else "dense transformer kept unquantized"
|
||||
)
|
||||
if transformer_quant_engaged is None
|
||||
else "re-planned resident for the quantised artifact"
|
||||
if quant_plan is not None
|
||||
|
|
@ -2384,7 +2395,9 @@ class DiffusionBackend:
|
|||
if compile_eligible(target, is_gguf = gguf_transformer, family = state.family):
|
||||
compile_ctx = compile_cache.begin(
|
||||
family = state.family.name,
|
||||
transformer = getattr(state.pipe, "transformer", None),
|
||||
# U-Net families (SDXL) carry the denoiser as pipe.unet.
|
||||
transformer = getattr(state.pipe, "transformer", None)
|
||||
or getattr(state.pipe, "unet", None),
|
||||
dtype = getattr(target, "dtype", None),
|
||||
quant = state.transformer_quant,
|
||||
attention_backend = attention_engaged,
|
||||
|
|
@ -2871,9 +2884,20 @@ class DiffusionBackend:
|
|||
if cancel.is_set():
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# The first compiled generation just paid the compile cost; persist the
|
||||
# warm torch.compile cache bundle when saving is enabled (distributor /
|
||||
# first-run warm). Idempotent + best-effort -- never fails a generation.
|
||||
# warm torch.compile cache bundle when saving is enabled (the auto
|
||||
# default / distributor mode). A STATIC compile (max tier, U-Net
|
||||
# whole-module) produces new artifacts per (width, height, batch), so
|
||||
# register this generation's shape first: a shape the bundle does not
|
||||
# cover re-dirties the context and the save below rewrites the bundle
|
||||
# with the enriched set. Idempotent + best-effort -- never fails a
|
||||
# generation.
|
||||
try:
|
||||
compile_cache.register_shape(
|
||||
state.compile_cache_ctx,
|
||||
(int(width), int(height), int(batch_size)),
|
||||
static = "compiled" in (state.speed_optims or ())
|
||||
and compiled_shapes_are_static(state.pipe, state.speed_mode),
|
||||
)
|
||||
compile_cache.save(state.compile_cache_ctx, logger = logger)
|
||||
except Exception: # noqa: BLE001 — cache persistence is best-effort
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -24,7 +24,8 @@ Lifecycle (driven by the caller, around ``_compile_repeated_blocks``):
|
|||
exists. Must run BEFORE the first compiled forward.
|
||||
2. (compile + one warmup forward happen as usual; on a hit they reuse the cache.)
|
||||
3. ``save(...)`` -> ``save_cache_artifacts`` to the bundle + manifest, AFTER the
|
||||
warmup forward, when in distributor/save mode.
|
||||
warmup forward. On by default (first-run warm; a bundle hit
|
||||
skips the rewrite), disable with the SAVE env knob.
|
||||
4. ``restore(...)`` -> put ``TORCHINDUCTOR_CACHE_DIR`` back on unload.
|
||||
|
||||
Everything is env-gated and best-effort; torch is imported lazily.
|
||||
|
|
@ -42,11 +43,18 @@ from typing import Any, Optional
|
|||
|
||||
# ----------------------------------------------------------------------------- env knobs
|
||||
# UNSLOTH_DIFFUSION_COMPILE_CACHE: auto (default) | 0 | 1
|
||||
# auto -> load a matching bundle if present (no automatic save).
|
||||
# 1 -> load AND save (distributor / first-run warm).
|
||||
# auto -> load a matching bundle if present AND save one after the first compiled
|
||||
# generation (first-run warm: every later session/restart skips the inductor
|
||||
# codegen + Triton compile + autotune part of the warmup). Measured on
|
||||
# Qwen-Image (B200, deferred 3rd-generation engage, FBCache armed): the
|
||||
# compile hitch drops 29.1 -> 22.2 s with bit-identical output; the bundle is
|
||||
# 7.9 MB and the one-time save costs ~0.5 s. The residual warmup is dynamo
|
||||
# tracing + guards, which the Mega-cache deliberately does not capture.
|
||||
# 1 -> same as auto, and also re-saves on a bundle hit (distributor refresh).
|
||||
# 0 -> disabled (plain local compile, no cache dir override).
|
||||
# UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR: root dir for bundles (default under the workspace).
|
||||
# UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE: 1 -> force-enable save even in "auto".
|
||||
# UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE: 0 -> disable the auto save (load-only);
|
||||
# 1 -> keep saving (explicit; same as the auto default).
|
||||
_ENV_MODE = "UNSLOTH_DIFFUSION_COMPILE_CACHE"
|
||||
_ENV_DIR = "UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR"
|
||||
_ENV_SAVE = "UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE"
|
||||
|
|
@ -73,8 +81,10 @@ def _save_enabled(mode: str) -> bool:
|
|||
return False
|
||||
if mode == "on":
|
||||
return True
|
||||
# auto: save only if explicitly opted in.
|
||||
return (os.environ.get(_ENV_SAVE) or "").strip().lower() in ("1", "on", "true", "yes")
|
||||
# auto: save by default (first-run warm -- without a saved bundle no user ever gets
|
||||
# a warm restart, since only distributors ran with mode "on"). The SAVE env stays as
|
||||
# an explicit override: "0" turns the auto save off (load-only), "1" keeps it on.
|
||||
return (os.environ.get(_ENV_SAVE) or "").strip().lower() not in ("0", "off", "false", "no")
|
||||
|
||||
|
||||
def cache_root() -> Path:
|
||||
|
|
@ -165,7 +175,14 @@ def cache_key(env_fp: dict[str, Any], model_fp: dict[str, Any]) -> str:
|
|||
# ----------------------------------------------------------------------------- lifecycle
|
||||
@dataclasses.dataclass
|
||||
class CacheContext:
|
||||
"""Carries the per-load cache state between ``begin`` and ``save``/``restore``."""
|
||||
"""Carries the per-load cache state between ``begin`` and ``save``/``restore``.
|
||||
|
||||
``shapes`` tracks the (width, height, batch) tuples whose STATIC-compile artifacts
|
||||
the bundle covers (persisted in the manifest). A dynamic-shape compile never needs
|
||||
it; a static compile (the max tier's regional compile, the U-Net whole-module
|
||||
compile) produces NEW artifacts per shape, so the caller registers each generation's
|
||||
shape and clears ``saved`` when it sees a new one -- the next ``save`` then rewrites
|
||||
the bundle with the enriched artifact set."""
|
||||
|
||||
key: str
|
||||
dir: Path
|
||||
|
|
@ -178,6 +195,7 @@ class CacheContext:
|
|||
saved: bool = False
|
||||
prev_inductor_dir: Optional[str] = None
|
||||
prev_inductor_dir_set: bool = False
|
||||
shapes: set = dataclasses.field(default_factory = set)
|
||||
|
||||
|
||||
def begin(
|
||||
|
|
@ -245,11 +263,36 @@ def begin(
|
|||
# Try an exact-match load. A miss/mismatch is normal and non-fatal.
|
||||
if ctx.bundle.exists() and ctx.manifest_path.exists():
|
||||
ctx.hit = _try_load(ctx, logger)
|
||||
if ctx.hit and mode != "on":
|
||||
# The artifacts on disk are exactly the ones just loaded, so there is
|
||||
# nothing to save (the write costs ~0.5 s on the first generation for no
|
||||
# change). A NEW static-compile shape later clears ``saved`` via
|
||||
# register_shape; explicit mode "on" (distributor refresh) keeps saving.
|
||||
ctx.saved = True
|
||||
else:
|
||||
_info(logger, f"compile-cache: no bundle for key {key} (will compile locally)")
|
||||
return ctx
|
||||
|
||||
|
||||
def register_shape(ctx: Optional[CacheContext], shape: Any, *, static: bool) -> None:
|
||||
"""Record a generation's (width, height, batch) against the bundle coverage.
|
||||
|
||||
Only meaningful for a STATIC compile (``static=True``): each new shape triggers its
|
||||
own compile, so the bundle written earlier this session (or loaded from disk) lacks
|
||||
those artifacts -- clear ``saved`` so the caller's next ``save`` rewrites the bundle
|
||||
with the enriched artifact set. Dynamic compiles reuse one artifact across shapes,
|
||||
so they never dirty the context. Never raises."""
|
||||
if ctx is None or not static:
|
||||
return
|
||||
try:
|
||||
key = tuple(shape)
|
||||
if key not in ctx.shapes:
|
||||
ctx.shapes.add(key)
|
||||
ctx.saved = False
|
||||
except Exception: # noqa: BLE001 — bookkeeping only
|
||||
pass
|
||||
|
||||
|
||||
def _try_load(ctx: CacheContext, logger: Any) -> bool:
|
||||
try:
|
||||
manifest = json.loads(ctx.manifest_path.read_text())
|
||||
|
|
@ -281,6 +324,12 @@ def _try_load(ctx: CacheContext, logger: Any) -> bool:
|
|||
if info is None:
|
||||
_warn(logger, "compile-cache: load_cache_artifacts returned None (no hit)")
|
||||
return False
|
||||
# The static-compile shapes this bundle covers (see register_shape); a
|
||||
# generation at a shape already here does not dirty the context.
|
||||
try:
|
||||
ctx.shapes = {tuple(s) for s in manifest.get("shapes", [])}
|
||||
except Exception: # noqa: BLE001 — coverage bookkeeping only
|
||||
ctx.shapes = set()
|
||||
_info(logger, f"compile-cache: loaded bundle for key {ctx.key}")
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001
|
||||
|
|
@ -291,7 +340,11 @@ def _try_load(ctx: CacheContext, logger: Any) -> bool:
|
|||
def save(ctx: Optional[CacheContext], *, logger: Any = None) -> bool:
|
||||
"""Persist the compiled artifacts to the bundle + manifest, AFTER a warmup forward.
|
||||
|
||||
No-op unless save is enabled (mode ``on`` or ``UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE``).
|
||||
No-op unless save is enabled (the ``auto``/``on`` default; disable with
|
||||
``UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE=0``) and the context is dirty: a load that
|
||||
HIT a bundle starts clean (``begin`` marks it saved -- rewriting the just-loaded
|
||||
artifacts costs ~0.5 s for no change), and a NEW static-compile shape re-dirties it
|
||||
via ``register_shape`` so the bundle grows to cover every shape the session used.
|
||||
Returns True if a bundle was written.
|
||||
"""
|
||||
if ctx is None or not _save_enabled(ctx.mode) or ctx.saved:
|
||||
|
|
@ -318,6 +371,9 @@ def save(ctx: Optional[CacheContext], *, logger: Any = None) -> bool:
|
|||
"sha256": hashlib.sha256(data).hexdigest(),
|
||||
"env": ctx.env_fp,
|
||||
"model": ctx.model_fp,
|
||||
# Static-compile shape coverage (register_shape); informational for a
|
||||
# dynamic compile (empty or the shapes generated, either way unused).
|
||||
"shapes": sorted(list(s) for s in ctx.shapes),
|
||||
}
|
||||
ctx.manifest_path.write_text(json.dumps(manifest, indent = 2, sort_keys = True, default = str))
|
||||
ctx.saved = True
|
||||
|
|
|
|||
|
|
@ -21,7 +21,9 @@ near-lossless speedups in the order the diffusers guides recommend
|
|||
one-time compile (~7.5-10.4s) and ZERO extra VRAM, resolution-invariant
|
||||
(the dequant inputs are fixed-shape weights). For a dense (non-GGUF) model
|
||||
there is no dequant, so ``default`` falls back to regional torch.compile of
|
||||
the denoiser's repeated block (the only compile lever a dense model has).
|
||||
the denoiser's repeated block; a U-Net denoiser (SDXL) has no repeated-block
|
||||
list, so it gets a whole-module STATIC compile instead (1.61x at LPIPS 0.034
|
||||
on SDXL, see ``_UNET_WHOLE_COMPILE``).
|
||||
max - the FULL torch.compile: regional max-autotune compile of the denoiser's
|
||||
repeated block (which fuses the GGUF dequant AND the matmul/norm/elementwise
|
||||
in one graph -- ~3.2x on the GGUF Z-Image transformer, PSNR ~36 dB vs eager,
|
||||
|
|
@ -209,6 +211,7 @@ def apply_speed_optims(
|
|||
"fused_qkv": False,
|
||||
"compiled": False,
|
||||
"compiled_dequant": False,
|
||||
"compiled_vae_decode": False,
|
||||
}
|
||||
mode = normalize_speed_mode(speed_mode)
|
||||
# TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max,
|
||||
|
|
@ -258,6 +261,12 @@ def apply_speed_optims(
|
|||
if is_gguf and on_cuda and family_allows_compile:
|
||||
applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger)
|
||||
elif compile_eligible(target, is_gguf = is_gguf, family = family):
|
||||
# A U-Net denoiser (SDXL) fuses QKV BEFORE its whole-module compile:
|
||||
# measured 36.3 vs 39.3 ms/step on SDXL under the compile (LPIPS unchanged
|
||||
# at 0.033). DiTs measured exactly neutral under the regional compile
|
||||
# (Qwen-Image 6.53 vs 6.52 s), so they keep the fuse on the max tier only.
|
||||
if _denoiser_unet(pipe) is not None:
|
||||
applied["fused_qkv"] = _fuse_qkv(pipe, logger)
|
||||
applied["compiled"] = _compile_repeated_blocks(
|
||||
pipe,
|
||||
logger,
|
||||
|
|
@ -274,6 +283,13 @@ def apply_speed_optims(
|
|||
offload_active = offload_active,
|
||||
)
|
||||
|
||||
# A compiled U-Net family also compiles the VAE decode: at SDXL's fast step rate the
|
||||
# decode is a real share of each image (measured 4.98 -> 4.25 s over 4 images, LPIPS
|
||||
# unchanged). DiT families skip it (the decode is a few % of their generation).
|
||||
# dynamic=True keeps it resolution-robust; fullgraph=False tolerates offload hooks.
|
||||
if applied["compiled"] and _denoiser_unet(pipe) is not None:
|
||||
applied["compiled_vae_decode"] = _compile_vae_decode(pipe, logger)
|
||||
|
||||
if mode == SPEED_MAX:
|
||||
# Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed.
|
||||
if on_cuda:
|
||||
|
|
@ -296,6 +312,40 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
# U-Net denoisers ship no ``_repeated_blocks`` (their block mix is heterogeneous), so the
|
||||
# regional compile below cannot reach them; these classes instead get a WHOLE-module
|
||||
# ``torch.compile`` with STATIC shapes, the one flavor measured worth its warmup. On SDXL
|
||||
# (B200, 30 steps / 7.0 / 1024px, 4 prompts; scripts/image_speedmem_bench.py levers +
|
||||
# probe): static whole-UNet compile runs 26.9 ms/step vs the 45.9 ms/step bit-exact
|
||||
# reference -- **1.61x end to end (6.16 -> 3.83 s) at LPIPS 0.034** -- while dynamic=True
|
||||
# compiles 5x slower (366 s vs 73 s cold) for less win (39.3 ms/step), and a regional
|
||||
# BasicTransformerBlock compile only reaches 45.0 ms/step (the ResNet convs stay eager).
|
||||
# Static shapes mean a recompile per new (height, width, batch); the Mega-cache bundle
|
||||
# (diffusion_compile_cache) carries each compiled shape across restarts.
|
||||
_UNET_WHOLE_COMPILE: frozenset[str] = frozenset({"UNet2DConditionModel"})
|
||||
|
||||
|
||||
def _denoiser_unet(pipe: Any) -> Any:
|
||||
"""The pipe's U-Net denoiser when its class is on the whole-compile list, else None."""
|
||||
unet = getattr(pipe, "unet", None)
|
||||
if unet is not None and type(unet).__name__ in _UNET_WHOLE_COMPILE:
|
||||
return unet
|
||||
return None
|
||||
|
||||
|
||||
def compiled_shapes_are_static(pipe: Any, speed_mode: Optional[str]) -> bool:
|
||||
"""Whether this load's compiled denoiser artifacts are per-(width, height, batch).
|
||||
|
||||
The ``max`` tier compiles the regional blocks with dynamic=False, and the U-Net
|
||||
whole-module compile is always static; the ``default`` DiT tier compiles
|
||||
dynamic=True (one artifact across shapes). The compile-cache layer keys on this to
|
||||
re-save its bundle when a session generates at a shape it has not covered yet."""
|
||||
mode = normalize_speed_mode(speed_mode)
|
||||
if mode == SPEED_MAX:
|
||||
return True
|
||||
return mode == SPEED_DEFAULT and _denoiser_unet(pipe) is not None
|
||||
|
||||
|
||||
def _denoiser_dits(pipe: Any) -> list:
|
||||
"""Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second
|
||||
expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch
|
||||
|
|
@ -321,7 +371,8 @@ def _compile_repeated_blocks(
|
|||
dits = [
|
||||
t for t in _denoiser_dits(pipe) if callable(getattr(t, "compile_repeated_blocks", None))
|
||||
]
|
||||
if not dits:
|
||||
unet = _denoiser_unet(pipe) if not dits else None
|
||||
if not dits and unet is None:
|
||||
return False
|
||||
# default: mode="default" + dynamic=True -- fast cold start, robust to resolution
|
||||
# changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False --
|
||||
|
|
@ -374,6 +425,23 @@ def _compile_repeated_blocks(
|
|||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "compile_repeated_blocks", exc)
|
||||
return False
|
||||
if unet is not None:
|
||||
# Whole-module static compile for the U-Net classes above. fullgraph mirrors the
|
||||
# regional decision (an active cache or offload hook graph-breaks, though U-Net
|
||||
# pipelines have no CacheMixin so in practice only offload lowers it); dynamic is
|
||||
# ALWAYS False -- the measured recipe -- so each new (height, width, batch) pays
|
||||
# its own compile, carried across restarts by the Mega-cache bundle.
|
||||
# ``Module.compile`` keeps the module identity (in-place ``_compiled_call_impl``),
|
||||
# so unload/status/LoRA gating see the same object the eager path had.
|
||||
unet_kwargs: dict[str, Any] = {"fullgraph": kwargs["fullgraph"], "dynamic": False}
|
||||
if max_autotune:
|
||||
unet_kwargs["mode"] = "max-autotune-no-cudagraphs"
|
||||
try:
|
||||
unet.compile(**unet_kwargs)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "unet whole-module compile", exc)
|
||||
return False
|
||||
# Compile every denoiser DiT (a dual-DiT family such as Ideogram runs both each step); a
|
||||
# per-DiT failure degrades that one to eager without dropping the others.
|
||||
engaged = False
|
||||
|
|
@ -400,6 +468,23 @@ def _compile_repeated_blocks(
|
|||
return engaged
|
||||
|
||||
|
||||
def _compile_vae_decode(pipe: Any, logger: Any) -> bool:
|
||||
"""torch.compile the VAE ``decode`` bound method in place (U-Net families only; the
|
||||
caller gates). Instance-level assignment: the pipe owns it, unload drops it with the
|
||||
pipe, and the module object itself is untouched."""
|
||||
vae = getattr(pipe, "vae", None)
|
||||
decode = getattr(vae, "decode", None) if vae is not None else None
|
||||
if not callable(decode):
|
||||
return False
|
||||
try:
|
||||
import torch
|
||||
vae.decode = torch.compile(decode, fullgraph = False, dynamic = True)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "vae decode compile", exc)
|
||||
return False
|
||||
|
||||
|
||||
def _enable_cudnn_benchmark(logger: Any) -> bool:
|
||||
try:
|
||||
import torch
|
||||
|
|
|
|||
|
|
@ -190,15 +190,83 @@ def test_save_then_load_roundtrip(monkeypatch, tmp_path, fake_megacache):
|
|||
assert ctx2.key == ctx.key
|
||||
|
||||
|
||||
def test_no_save_in_auto_mode(monkeypatch, tmp_path, fake_megacache):
|
||||
def test_auto_mode_saves_by_default(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "auto")
|
||||
monkeypatch.delenv(cc._ENV_SAVE, raising = False)
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert cc.save(ctx) is False # auto without SAVE opt-in does not write
|
||||
assert cc.save(ctx) is True # first-run warm: auto saves the bundle
|
||||
assert ctx.bundle.exists() and ctx.manifest_path.exists()
|
||||
|
||||
# The next load with the same fingerprint hits the just-saved bundle...
|
||||
ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx2.hit is True
|
||||
# ...and does NOT rewrite it under auto (the artifacts on disk are the ones loaded).
|
||||
before = ctx2.bundle.stat().st_mtime_ns
|
||||
assert cc.save(ctx2) is False
|
||||
assert ctx2.bundle.stat().st_mtime_ns == before
|
||||
|
||||
|
||||
def test_save_env_zero_disables_auto_save(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "auto")
|
||||
monkeypatch.setenv(cc._ENV_SAVE, "0")
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert cc.save(ctx) is False # explicit load-only override
|
||||
assert not ctx.bundle.exists()
|
||||
|
||||
|
||||
def test_on_mode_resaves_after_hit(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "on")
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert cc.save(ctx) is True
|
||||
ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx2.hit is True
|
||||
# Distributor mode refreshes the bundle even on a hit (new variants get captured).
|
||||
assert cc.save(ctx2) is True
|
||||
|
||||
|
||||
def test_new_static_shape_redirties_a_hit(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "auto")
|
||||
monkeypatch.delenv(cc._ENV_SAVE, raising = False)
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
|
||||
# Cold session at 1024: the save records the shape coverage in the manifest.
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
cc.register_shape(ctx, (1024, 1024, 1), static = True)
|
||||
assert cc.save(ctx) is True
|
||||
manifest = json.loads(ctx.manifest_path.read_text())
|
||||
assert manifest["shapes"] == [[1024, 1024, 1]]
|
||||
|
||||
# Warm session: the covered shape does not dirty the context...
|
||||
ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx2.hit is True and ctx2.saved is True
|
||||
assert ctx2.shapes == {(1024, 1024, 1)}
|
||||
cc.register_shape(ctx2, (1024, 1024, 1), static = True)
|
||||
assert cc.save(ctx2) is False
|
||||
# ...but a NEW static shape (its compile just produced new artifacts) does, and the
|
||||
# rewritten manifest covers both.
|
||||
cc.register_shape(ctx2, (768, 768, 1), static = True)
|
||||
assert ctx2.saved is False
|
||||
assert cc.save(ctx2) is True
|
||||
manifest = json.loads(ctx2.manifest_path.read_text())
|
||||
assert manifest["shapes"] == [[768, 768, 1], [1024, 1024, 1]]
|
||||
|
||||
|
||||
def test_dynamic_compile_never_dirties(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "auto")
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
cc.save(ctx)
|
||||
ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW)
|
||||
assert ctx2.hit is True
|
||||
# A dynamic-shape compile reuses one artifact across shapes: no re-save.
|
||||
cc.register_shape(ctx2, (768, 768, 1), static = False)
|
||||
assert cc.save(ctx2) is False
|
||||
cc.register_shape(None, (768, 768, 1), static = True) # no context: no-op
|
||||
|
||||
|
||||
def test_fingerprint_mismatch_falls_back(monkeypatch, tmp_path, fake_megacache):
|
||||
monkeypatch.setenv(cc._ENV_MODE, "on")
|
||||
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
|
||||
|
|
|
|||
|
|
@ -68,6 +68,8 @@ def _stub_torch(monkeypatch):
|
|||
cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(allow_tf32 = False)),
|
||||
cudnn = types.SimpleNamespace(allow_tf32 = False, benchmark = False),
|
||||
)
|
||||
# The VAE-decode compile wraps a bound method; identity wrap is enough for tests.
|
||||
torch.compile = lambda fn, **kwargs: fn
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
return torch
|
||||
|
||||
|
|
@ -225,6 +227,7 @@ def test_speed_off_applies_nothing(monkeypatch):
|
|||
"fused_qkv": False,
|
||||
"compiled": False,
|
||||
"compiled_dequant": False,
|
||||
"compiled_vae_decode": False,
|
||||
"fp16_accum": False,
|
||||
}
|
||||
assert pipe.vae.mem_format is None and pipe.compiled is False
|
||||
|
|
@ -359,6 +362,125 @@ def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch):
|
|||
assert pipe.compile_kwargs["dynamic"] is False
|
||||
|
||||
|
||||
# ── U-Net whole-module compile fallback (SDXL) ─────────────────────────────────
|
||||
|
||||
|
||||
class UNet2DConditionModel:
|
||||
"""Fake with the diffusers class NAME the fallback keys on: no
|
||||
compile_repeated_blocks (U-Nets ship no _repeated_blocks), but Module.compile."""
|
||||
|
||||
def __init__(self):
|
||||
self.compile_kwargs = None
|
||||
|
||||
def compile(self, **kwargs):
|
||||
self.compile_kwargs = kwargs
|
||||
|
||||
|
||||
class _SomeOtherUNet(UNet2DConditionModel):
|
||||
pass
|
||||
|
||||
|
||||
class _UNetPipe:
|
||||
def __init__(self, unet = None):
|
||||
self.mem_format = None
|
||||
self.fused = False
|
||||
self.vae = types.SimpleNamespace(to = self._vae_to, decode = lambda z: z)
|
||||
self.unet = UNet2DConditionModel() if unet is None else unet
|
||||
|
||||
def _vae_to(self, *, memory_format):
|
||||
self.mem_format = memory_format
|
||||
|
||||
def fuse_qkv_projections(self):
|
||||
self.fused = True
|
||||
|
||||
|
||||
def test_unet_whole_compile_default_tier(monkeypatch):
|
||||
# SDXL's UNet has no _repeated_blocks, so `default` falls back to a whole-module
|
||||
# STATIC compile (measured 1.61x at LPIPS 0.034 on SDXL): fullgraph on, dynamic OFF.
|
||||
# The U-Net recipe also fuses QKV (36.3 vs 39.3 ms/step) and compiles the VAE decode
|
||||
# (4.98 -> 4.25 s over 4 images) on the same tier.
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _UNetPipe()
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert pipe.unet.compile_kwargs == {"fullgraph": True, "dynamic": False}
|
||||
assert applied["fused_qkv"] is True and pipe.fused is True
|
||||
assert applied["compiled_vae_decode"] is True
|
||||
|
||||
|
||||
def test_dit_default_tier_keeps_fuse_and_vae_decode_off(monkeypatch):
|
||||
# The DiT default tier is unchanged: fused QKV measured exactly neutral there
|
||||
# (Qwen-Image 6.53 vs 6.52 s) so it stays max-only, and the VAE decode is a few
|
||||
# percent of a DiT generation so it stays eager.
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _Pipe(with_compile = True, with_fuse = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert applied["fused_qkv"] is False and pipe.fused is False
|
||||
assert applied["compiled_vae_decode"] is False
|
||||
|
||||
|
||||
def test_unet_whole_compile_offload_drops_fullgraph(monkeypatch):
|
||||
# Offload hooks graph-break exactly as on the regional path.
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _UNetPipe()
|
||||
applied = apply_speed_optims(
|
||||
pipe,
|
||||
_target(),
|
||||
is_gguf = False,
|
||||
family = _family(),
|
||||
speed_mode = SPEED_DEFAULT,
|
||||
offload_active = True,
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert pipe.unet.compile_kwargs == {"fullgraph": False, "dynamic": False}
|
||||
|
||||
|
||||
def test_unet_whole_compile_max_tier_mode(monkeypatch):
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _UNetPipe()
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert pipe.unet.compile_kwargs == {
|
||||
"fullgraph": True,
|
||||
"dynamic": False,
|
||||
"mode": "max-autotune-no-cudagraphs",
|
||||
}
|
||||
|
||||
|
||||
def test_unet_whole_compile_gated_by_class_name(monkeypatch):
|
||||
# An unlisted U-Net class (unmeasured architecture) stays eager rather than paying
|
||||
# an unvalidated whole-module compile.
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _UNetPipe(unet = _SomeOtherUNet())
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is False
|
||||
assert pipe.unet.compile_kwargs is None
|
||||
|
||||
|
||||
def test_unet_whole_compile_failure_degrades_to_eager(monkeypatch):
|
||||
_stub_torch(monkeypatch)
|
||||
|
||||
class _Boom(UNet2DConditionModel):
|
||||
def compile(self, **kwargs):
|
||||
raise RuntimeError("no dynamo on this build")
|
||||
|
||||
_Boom.__name__ = "UNet2DConditionModel"
|
||||
pipe = _UNetPipe(unet = _Boom())
|
||||
applied = apply_speed_optims(
|
||||
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
|
||||
)
|
||||
assert applied["compiled"] is False # best-effort: load proceeds eager
|
||||
|
||||
|
||||
def test_speed_max_tf32_only_on_cuda(monkeypatch):
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _Pipe()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue