From 7f44a98ad8d6b3ba7165f83065ca7ef9a681c127 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 07:33:07 +0000 Subject: [PATCH] Auto-install optional attention kernels and toggle the step cache per generation Attention: apply_attention_backend now best-effort installs the package an explicitly requested optional backend needs (sage -> sageattention, flash -> flash-attn, flash3/flash4 -> kernels, xformers), wheel-only via pip --only-binary=:all: so a host without a CUDA toolchain never starts a source build. Gated by UNSLOTH_DIFFUSION_ATTENTION_INSTALL (auto|0), mirroring the sd.cpp prebuilt installer gate, and only reached after the arch gating in select_attention_backend, so no install is attempted for a kernel this card cannot run. Any failure keeps today's native fallback. Step cache: transformer_cache gains a real auto state (unset or "auto"). At load the policy engages FBCache when the model's default schedule reaches FBCACHE_MIN_STEPS = 20 (dev-style 28-step models win ~1.4x; 4-9-step distilled models never engage, a skipped step costs too much there). generate() then re-checks the ACTUAL step count and toggles the cache idempotently across the bar, so one resident load serves both a 28-step and a 4-step request with the right cache state, and status/resolved provenance follow the toggle. An explicit off or fbcache request is pinned and never toggled. Compile drops fullgraph when an auto cache could still engage on a cache-capable transformer, since enabling FBCache under a fullgraph-compiled transformer would crash. Verified on GPU: flux.1-schnell load starts uncached (4-step default), engages fbcache at 24 steps, disengages at 4, re-engages at 28, with images at each step and the provenance record tracking each transition. --- studio/backend/core/inference/diffusion.py | 107 ++++++++++++++++-- .../core/inference/diffusion_attention.py | 67 +++++++++++ .../backend/core/inference/diffusion_cache.py | 63 ++++++++++- .../backend/tests/test_diffusion_attention.py | 85 ++++++++++++++ studio/backend/tests/test_diffusion_cache.py | 94 +++++++++++++++ 5 files changed, 402 insertions(+), 14 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 37630f6670..8037dde174 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -32,6 +32,7 @@ from .diffusion_families import ( DIFFUSION_CANCELLED_MSG, DIFFUSION_NOT_LOADED_MSG, DiffusionFamily, + default_generation_params, detect_family_for_pick, resolve_base_repo, resolve_local_gguf_child, @@ -70,7 +71,14 @@ from .diffusion_attention import ( ) from . import diffusion_compile_cache as compile_cache from . import diffusion_gguf_compile as gguf_compile -from .diffusion_cache import apply_step_cache +from .diffusion_cache import ( + FBCACHE_MIN_STEPS, + TC_AUTO, + TC_FBCACHE, + apply_step_cache, + maybe_toggle_step_cache, + normalize_transformer_cache, +) from .diffusion_precision import quantize_text_encoders from .diffusion_prequant import ( load_prequantized_transformer, @@ -277,6 +285,13 @@ class _LoadState: attention_backend: Optional[str] = None # Step cache engaged ("fbcache") or None. Opt-in, for many-step models. transformer_cache: Optional[str] = None + # True when the cache decision was AUTO on a cache-capable transformer: generate() + # then re-checks the actual step count and toggles FBCache across FBCACHE_MIN_STEPS. + # An explicit request (off / fbcache) is never toggled. + cache_auto: bool = False + # Inputs the generation-time toggle re-applies (quantised threshold + override). + cache_quant_active: bool = False + cache_threshold: Optional[float] = None # Shared eager monkey-patches (diffusion_eager_patches) installed for this load (any # non-off speed tier). Uninstalled on unload so a later `off` load is bit-identical. eager_patched: bool = False @@ -1142,20 +1157,58 @@ class DiffusionBackend: ), logger = logger, ) - # Opt-in step caching (First-Block-Cache), also before compile. OFF by - # default; for many-step models it reuses the transformer tail across steps - # (~1.4x on Flux at LPIPS ~0.08). When engaged, compile must drop fullgraph - # (the cache's per-step decision is a graph break), so pass it through. + # Step caching (First-Block-Cache), also before compile. For many-step + # models it reuses the transformer tail across steps (~1.4x on Flux at + # LPIPS ~0.08). When engaged, compile must drop fullgraph (the cache's + # per-step decision is a graph break), so pass it through. + # Tri-state request: unset / "auto" lets the step-count policy decide + # (engage when this model's DEFAULT schedule reaches FBCACHE_MIN_STEPS, + # then re-check against the actual step count on every generation); + # explicit "off" / "fbcache" are pinned and never toggled. + cache_request = normalize_transformer_cache(transformer_cache) + cache_auto = transformer_cache is None or cache_request == TC_AUTO + cache_quant_active = ( + transformer_quant_engaged is not None or bool(gguf_filename) + ) + default_steps: Optional[int] = None + if cache_auto: + default_steps, _ = default_generation_params( + gguf_filename, repo_id, base, fam.name + ) + cache_request = ( + TC_FBCACHE if default_steps >= FBCACHE_MIN_STEPS else None + ) cache_engaged = apply_step_cache( pipe, - mode = transformer_cache, + mode = cache_request, threshold = transformer_cache_threshold, # GGUF transformers are quantized too (the default Studio path), so the # cache needs the higher quantized threshold to still trigger -- not just # the dense-quant fast path. - quant_active = transformer_quant_engaged is not None or bool(gguf_filename), + quant_active = cache_quant_active, logger = logger, ) + # An auto decision can flip at generation time, but only on a transformer + # that supports caching at all; a non-CacheMixin transformer (e.g. + # Z-Image) can never engage, so compile keeps fullgraph there. + cache_may_toggle = cache_auto and callable( + getattr(getattr(pipe, "transformer", None), "enable_cache", None) + ) + if cache_auto: + if cache_engaged: + cache_reason = ( + f"auto: {default_steps}-step default schedule reaches " + f"{FBCACHE_MIN_STEPS}; re-checked per generation" + ) + elif cache_request is not None: + cache_reason = "auto: model does not support step caching" + else: + cache_reason = ( + f"auto: {default_steps}-step default schedule is below " + f"{FBCACHE_MIN_STEPS}; re-checked per generation" + ) + else: + cache_reason = "requested" # Install the shared compile-safe eager patches (fused RMSNorm / # AdaLayerNorm) for any active speed tier. They are class-level, idempotent # and math-equivalent (FMA / fused -> neutral under compile, equal-or-more @@ -1221,8 +1274,12 @@ class DiffusionBackend: compile_kwargs = { # Mirrors apply_speed_optims' fullgraph decision: an active # step cache OR a planned offload graph-breaks, so the cached - # bundle must be keyed on the same fullgraph setting. + # bundle must be keyed on the same fullgraph setting. An auto + # cache that could still engage mid-session also drops + # fullgraph: enabling FBCache under a fullgraph-compiled + # transformer would crash the first cached generation. "fullgraph": cache_engaged is None + and not cache_may_toggle and plan.offload_policy == OFFLOAD_NONE, "dynamic": effective_speed != SPEED_MAX, "mode": "max-autotune-no-cudagraphs" @@ -1238,7 +1295,7 @@ class DiffusionBackend: is_gguf = gguf_transformer, family = fam, speed_mode = effective_speed, - cache_active = cache_engaged is not None, + cache_active = cache_engaged is not None or cache_may_toggle, # The planned offload policy: group/model/sequential offload installs # compiler-disabled onload hooks, so compile must drop fullgraph. offload_active = plan.offload_policy != OFFLOAD_NONE, @@ -1312,9 +1369,9 @@ class DiffusionBackend: "planned from measured free VRAM vs estimated footprint", ), "transformer_cache": ( - transformer_cache, + None if cache_auto else transformer_cache, cache_engaged or "off", - "off by default" if transformer_cache is None else "requested", + cache_reason, ), "cpu_offload": ( True if cpu_offload else None, @@ -1343,6 +1400,9 @@ class DiffusionBackend: transformer_quant = transformer_quant_engaged, attention_backend = attention_engaged, transformer_cache = cache_engaged, + cache_auto = cache_may_toggle, + cache_quant_active = cache_quant_active, + cache_threshold = transformer_cache_threshold, eager_patched = eager_patched, compile_cache_ctx = compile_ctx, hf_token = hf_token, @@ -2102,6 +2162,31 @@ class DiffusionBackend: if "callback_on_step_end" in call_params: kwargs["callback_on_step_end"] = _on_step + # An AUTO cache decision is re-checked against the ACTUAL step count: + # a 28-step dev-style request gains FBCache even when the load's default + # schedule kept it off, and a few-step turbo request drops it (skipping + # a step there is a large quality hit). Explicit choices never toggle. + if state.cache_auto: + toggled = maybe_toggle_step_cache( + state.pipe, + steps = steps, + quant_active = state.cache_quant_active, + threshold = state.cache_threshold, + logger = logger, + ) + if toggled != state.transformer_cache: + # _LoadState is frozen (loads swap it as one unit); this is the + # one deliberate in-place update, tracking the pipe-level toggle + # that already happened so status() reports the true cache state. + object.__setattr__(state, "transformer_cache", toggled) + entry = (state.resolved or {}).get("transformer_cache") + if isinstance(entry, dict): + entry["value"] = toggled or "off" + entry["reason"] = ( + f"auto: {steps}-step generation " + + ("reaches" if toggled else "is below") + + f" {FBCACHE_MIN_STEPS}" + ) # Start each generation from a clean step cache: FBCache residuals from # a prior request on this resident pipe would otherwise be compared # against this generation's first step (shape mismatch on a resolution/ diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index e652b0068c..912b20e73e 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -154,6 +154,72 @@ def _cudnn_attention_supported() -> bool: return have is None or have >= (8, 0) +# Optional-kernel backends the loader may install on demand: dispatcher name -> +# (probe module, pip package). Only wheels are ever installed (--only-binary=:all:): +# a source build of flash-attn or sageattention takes tens of minutes and needs a +# CUDA toolchain, which a Studio host cannot be assumed to have -- no wheel for this +# python/torch/cuda combo means the request falls back to the native default exactly +# as an uninstallable kernel does today. cuDNN/native need nothing (ship with torch). +_INSTALLABLE_BACKENDS: dict[str, tuple[str, str]] = { + "sage": ("sageattention", "sageattention"), + "flash": ("flash_attn", "flash-attn"), + "_flash_3_hub": ("kernels", "kernels"), # FA3/FA4 stream from the HF kernels hub + "flash_4_hub": ("kernels", "kernels"), + "xformers": ("xformers", "xformers"), +} + +# Gate for the on-demand install, mirroring UNSLOTH_DIFFUSION_SD_CPP_INSTALL: +# auto (default) / 1 - install the missing package when a gated backend is requested +# 0 - never install; a missing kernel falls back to native +_ATTENTION_INSTALL_ENV = "UNSLOTH_DIFFUSION_ATTENTION_INSTALL" + + +def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> None: + """Best-effort wheel-only install of the package ``backend`` needs, when allowed. + + Called after arch gating (select_attention_backend already dropped kernels this + card cannot run), so an install attempt is only made for a backend that could + actually work here. Failure is logged and swallowed: the subsequent + set_attention_backend raises on the still-missing package and the load falls + back to the native default, same as before this hook existed.""" + import importlib.util + import os + + spec = _INSTALLABLE_BACKENDS.get(backend) + if spec is None: + return + module, package = spec + gate = os.environ.get(_ATTENTION_INSTALL_ENV, "auto").strip().lower() + if gate in ("0", "false", "no", "off"): + return + try: + if importlib.util.find_spec(module) is not None: + return + except Exception: # noqa: BLE001 — a broken install probes as missing; try the install + pass + import subprocess + import sys + + if logger is not None: + logger.info( + "diffusion.attention: installing %s for backend=%s (wheel-only)", package, backend + ) + try: + subprocess.run( + [sys.executable, "-m", "pip", "install", "--only-binary", ":all:", package], + capture_output = True, + timeout = 600, + check = True, + ) + except Exception as exc: # noqa: BLE001 — no wheel / no network -> native fallback + if logger is not None: + logger.warning( + "diffusion.attention: could not install %s (%s); falling back to default", + package, + exc, + ) + + def apply_attention_backend( pipe: Any, backend: Optional[str], @@ -176,6 +242,7 @@ def apply_attention_backend( if not callable(fn): return None if backend is not None: + _ensure_attention_backend_installed(backend, logger) try: fn(backend) # set_attention_backend also pins the backend in diffusers' process-wide diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index b7a2b7ac45..a6691836d0 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -28,6 +28,7 @@ from __future__ import annotations from typing import Any, Optional TC_OFF = "off" +TC_AUTO = "auto" TC_FBCACHE = "fbcache" TC_MODES = (TC_FBCACHE,) @@ -37,9 +38,15 @@ TC_MODES = (TC_FBCACHE,) DEFAULT_FBCACHE_THRESHOLD = 0.08 QUANT_FBCACHE_THRESHOLD = 0.12 +# The auto policy's step-count bar: FBCache's win scales with step count (each skipped +# step is a larger quality hit on a short trajectory), so auto engages it only at 20+ +# steps -- full "dev"-style schedules (28+) qualify, distilled turbo models (4-9) never do. +FBCACHE_MIN_STEPS = 20 + def normalize_transformer_cache(value: Optional[str]) -> Optional[str]: - """Lower/strip a requested cache mode; None / "" / "none" / "off" -> None (disabled). + """Lower/strip a requested cache mode; None / "" / "none" / "off" -> None (disabled), + "auto" -> TC_AUTO (the loader decides from the step count). Raises ValueError for an unsupported value so a bad request is rejected cheaply.""" if value is None: @@ -47,9 +54,12 @@ def normalize_transformer_cache(value: Optional[str]) -> Optional[str]: normalized = str(value).strip().lower().replace("-", "_") if not normalized or normalized in ("none", "off"): return None + if normalized == TC_AUTO: + return TC_AUTO if normalized not in TC_MODES: raise ValueError( - f"Unsupported transformer_cache '{value}'. Use one of: off, {', '.join(TC_MODES)}." + f"Unsupported transformer_cache '{value}'. Use one of: off, auto, " + f"{', '.join(TC_MODES)}." ) return normalized @@ -67,7 +77,9 @@ def apply_step_cache( the default; ``quant_active`` raises the default so the cache still triggers on a quantised transformer. Best-effort: never raises for an incompatible model.""" mode = normalize_transformer_cache(mode) - if mode is None: + if mode is None or mode == TC_AUTO: + # AUTO must be resolved by the loader (step-count policy) before reaching the + # engage call; treat a stray auto as off rather than crashing the load. return None transformer = getattr(pipe, "transformer", None) if transformer is None: @@ -105,6 +117,51 @@ def apply_step_cache( return None +def maybe_toggle_step_cache( + pipe: Any, + *, + steps: int, + quant_active: bool = False, + threshold: Optional[float] = None, + logger: Any = None, +) -> Optional[str]: + """Generation-time enable/disable for an AUTO cache decision, keyed on the actual + step count: engage FBCache at ``FBCACHE_MIN_STEPS`` or more, run uncached below it. + Idempotent (the ``_unsloth_step_cache`` marker tracks the engaged state), so calling + it on every generation is cheap. Only the loader's auto path calls this; an explicit + user choice is never toggled. Returns the mode now active (or None when uncached).""" + transformer = getattr(pipe, "transformer", None) + if transformer is None: + return None + engaged = getattr(transformer, "_unsloth_step_cache", None) + want = int(steps) >= FBCACHE_MIN_STEPS + if want and not engaged: + return apply_step_cache( + pipe, + mode = TC_FBCACHE, + threshold = threshold, + quant_active = quant_active, + logger = logger, + ) + if not want and engaged: + disable_cache = getattr(transformer, "disable_cache", None) + if callable(disable_cache): + try: + disable_cache() + transformer._unsloth_step_cache = None + if logger is not None: + logger.info( + "diffusion.cache: fbcache disengaged (auto: %s steps < %s)", + steps, + FBCACHE_MIN_STEPS, + ) + return None + except Exception as exc: # noqa: BLE001 — keep the cache rather than crash + _warn(logger, "fbcache disable", exc) + return TC_FBCACHE + return TC_FBCACHE if engaged else None + + def _warn(logger: Any, what: str, exc: Exception) -> None: if logger is not None: logger.warning("diffusion.cache: %s unavailable (%s); running uncached", what, exc) diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 3e43aab8a9..48647bb0b4 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -223,3 +223,88 @@ def test_active_attention_backend_reads_tuple_return(): _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE) assert att._active_attention_backend() == "native" + + +# ── on-demand wheel-only install of optional kernels ───────────────────────────── +@pytest.fixture(autouse = True) +def _no_real_installs(monkeypatch): + # Unit tests must never shell out to pip: the apply path probes installable + # backends (sage/flash*), so hard-disable the gate; install tests re-enable it + # with a stubbed subprocess. + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "0") + + +class _Recorder: + def __init__(self): + self.calls = [] + + def __call__(self, cmd, **kwargs): + self.calls.append(list(cmd)) + return types.SimpleNamespace(returncode = 0) + + +def _stub_subprocess(monkeypatch, run): + import subprocess + + monkeypatch.setattr(subprocess, "run", run) + + +def test_install_skipped_when_gate_disabled(monkeypatch): + run = _Recorder() + _stub_subprocess(monkeypatch, run) + att._ensure_attention_backend_installed("sage") + assert run.calls == [] + + +def test_install_skipped_when_module_present(monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + import importlib.util + + monkeypatch.setattr( + importlib.util, "find_spec", lambda name: object() if name == "sageattention" else None + ) + run = _Recorder() + _stub_subprocess(monkeypatch, run) + att._ensure_attention_backend_installed("sage") + assert run.calls == [] + + +def test_install_runs_wheel_only_for_missing_kernel(monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + import importlib.util + + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + run = _Recorder() + _stub_subprocess(monkeypatch, run) + att._ensure_attention_backend_installed("sage") + assert len(run.calls) == 1 + cmd = run.calls[0] + assert "--only-binary" in cmd and ":all:" in cmd and "sageattention" in cmd + + +def test_install_never_attempted_for_builtin_backends(monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + run = _Recorder() + _stub_subprocess(monkeypatch, run) + att._ensure_attention_backend_installed("_native_cudnn") + att._ensure_attention_backend_installed("native") + assert run.calls == [] + + +def test_install_failure_falls_back_to_native(monkeypatch): + # pip failing (no wheel for this platform) must not break the load: the apply + # path proceeds, set_attention_backend raises on the missing package, and the + # dispatcher is restored to native -- same contract as before the hook. + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + import importlib.util + import subprocess as sp + + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + + def _boom(cmd, **kwargs): + raise sp.CalledProcessError(returncode = 1, cmd = cmd) + + _stub_subprocess(monkeypatch, _boom) + monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") + t = _FakeTransformer(fail = True) + assert apply_attention_backend(_pipe(t), "sage") is None diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 62071d9aa6..40537529ef 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -147,3 +147,97 @@ def test_diffusers_unavailable_runs_uncached(monkeypatch): monkeypatch.setitem(sys.modules, "diffusers", None) t = _MixinTransformer() assert apply_step_cache(_pipe(t), mode = "fbcache") is None + + +# ── the auto policy: normalize("auto") + generation-time toggling ────────────────── +from core.inference.diffusion_cache import ( # noqa: E402 + FBCACHE_MIN_STEPS, + TC_AUTO, + maybe_toggle_step_cache, +) + + +class _ToggleTransformer(_MixinTransformer): + """CacheMixin-style fake with the disable side too, counting transitions.""" + + def __init__(self): + super().__init__() + self.enables = 0 + self.disables = 0 + + def enable_cache(self, config): + super().enable_cache(config) + self.enables += 1 + + def disable_cache(self): + self.disables += 1 + + +def test_normalize_auto_is_a_distinct_state(): + assert normalize_transformer_cache("auto") == TC_AUTO + assert normalize_transformer_cache(" AUTO ") == TC_AUTO + + +def test_apply_treats_stray_auto_as_off(monkeypatch): + # AUTO must be resolved by the loader; if it ever reaches the engage call the + # load runs uncached instead of crashing. + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + assert apply_step_cache(_pipe(t), mode = "auto") is None + assert t.enabled_with is None + + +def test_toggle_engages_at_the_step_bar(monkeypatch): + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + mode = maybe_toggle_step_cache(_pipe(t), steps = FBCACHE_MIN_STEPS) + assert mode == TC_FBCACHE and t.enables == 1 + assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD + assert t._unsloth_step_cache + + +def test_toggle_uses_quant_threshold(monkeypatch): + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + maybe_toggle_step_cache(_pipe(t), steps = 28, quant_active = True) + assert t.enabled_with.threshold == QUANT_FBCACHE_THRESHOLD + + +def test_toggle_is_idempotent_when_engaged(monkeypatch): + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + maybe_toggle_step_cache(_pipe(t), steps = 28) + mode = maybe_toggle_step_cache(_pipe(t), steps = 28) + assert mode == TC_FBCACHE and t.enables == 1 and t.disables == 0 + + +def test_toggle_disengages_below_the_bar(monkeypatch): + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + 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 not t._unsloth_step_cache + # and it stays off on repeat calls (no flapping disable calls). + assert maybe_toggle_step_cache(_pipe(t), steps = 8) is None + assert t.disables == 1 + + +def test_toggle_reengages_after_a_disable(monkeypatch): + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + maybe_toggle_step_cache(_pipe(t), steps = 28) + maybe_toggle_step_cache(_pipe(t), steps = 8) + mode = maybe_toggle_step_cache(_pipe(t), steps = 24) + assert mode == TC_FBCACHE and t.enables == 2 + + +def test_toggle_noop_without_cache_support(monkeypatch): + _stub_diffusers(monkeypatch) + t = _NonCacheMixinTransformer() + assert maybe_toggle_step_cache(_pipe(t), steps = 28) is None + assert maybe_toggle_step_cache(_pipe(t), steps = 8) is None + + +def test_toggle_noop_without_transformer(): + assert maybe_toggle_step_cache(types.SimpleNamespace(), steps = 28) is None