From 352fb40089a5722025aa0d435d69db90adce4836 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 06:18:29 +0000 Subject: [PATCH] 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. --- studio/backend/core/inference/diffusion.py | 34 ++++- .../core/inference/diffusion_compile_cache.py | 72 +++++++++-- .../backend/core/inference/diffusion_speed.py | 89 ++++++++++++- .../tests/test_diffusion_compile_cache.py | 72 ++++++++++- studio/backend/tests/test_diffusion_speed.py | 122 ++++++++++++++++++ 5 files changed, 372 insertions(+), 17 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 59767274af..01bf37a418 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_compile_cache.py b/studio/backend/core/inference/diffusion_compile_cache.py index 97812ba9f9..0460e8e819 100644 --- a/studio/backend/core/inference/diffusion_compile_cache.py +++ b/studio/backend/core/inference/diffusion_compile_cache.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 9e2007719c..213d0a077f 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_compile_cache.py b/studio/backend/tests/test_diffusion_compile_cache.py index 7193e9c0f4..1c4ac25795 100644 --- a/studio/backend/tests/test_diffusion_compile_cache.py +++ b/studio/backend/tests/test_diffusion_compile_cache.py @@ -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)) diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 4d7bedacf5..746d05ac1c 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -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()