From 7c8f4919d62b1d17ca26c21f4af1bd0c65e14bd0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 08:45:17 +0000 Subject: [PATCH] diffusion: fix step-cache Off, kernel-install dep clobber, retry-under-lock, effective steps - Step cache 'Off' is preserved: the frontend defaulted to 'off' and mapped it to an omitted transformer_cache, which the backend now reads as 'auto', so leaving the control at Off silently enabled FBCache on 20+ step families. Default the control to Auto, add an explicit Auto option, and send auto -> omitted so Off maps to an explicit cache-off. - Kernel auto-install adds --no-deps: 'pip install --only-binary :all: xformers' resolves xformers' pinned torch and replaces the running torch/triton. --no-deps installs only the best-effort kernel wheel; an ABI mismatch just fails to import and falls back to native, never clobbering core deps. - Do not retry a failed kernel install under the load lock: the pre-install runs outside the locks, then the in-lock resolve re-attempts pip (up to 600s) while holding _generate_lock/_lock and blocking unload/cancel/new loads. Record the attempt in a process-level set so the in-lock call short-circuits to native. - Cache auto-toggle keys on effective denoise steps: an image-conditioned run with strength < 1 (upscale default 0.35) denoises a fraction of the requested steps, so a 28-step request runs ~10 steps. Compute the effective count the way diffusers get_timesteps does and gate FBCache on it, only when strength is actually applied. --- studio/backend/core/inference/diffusion.py | 22 +++++++++- .../core/inference/diffusion_attention.py | 28 +++++++++++- .../backend/core/inference/diffusion_cache.py | 18 ++++++++ .../backend/tests/test_diffusion_attention.py | 43 +++++++++++++++++++ studio/backend/tests/test_diffusion_cache.py | 34 +++++++++++++++ .../src/features/images/images-page.tsx | 7 +-- 6 files changed, 146 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index ef62f4358d..8b423299ff 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -78,6 +78,7 @@ from .diffusion_cache import ( TC_AUTO, TC_FBCACHE, apply_step_cache, + effective_denoise_steps, maybe_toggle_step_cache, normalize_transformer_cache, ) @@ -2273,9 +2274,26 @@ class DiffusionBackend: # 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: + # Key the policy on the EFFECTIVE denoise steps: an img2img/upscale/ + # inpaint request at strength < 1 only denoises a fraction of `steps` + # (e.g. a 28-step upscale at strength 0.35 runs ~10 steps), so passing + # the raw request would wrongly engage FBCache on exactly the short + # trajectory the policy keeps uncached. Only fold in `strength` when it + # is ACTUALLY applied to the pipe (same gate as the kwarg below), so a + # stray strength on a txt2img request never shortens the count. + strength_applied = ( + strength + if ( + strength is not None + and init_pil is not None + and "strength" in call_params + ) + else None + ) + denoise_steps = effective_denoise_steps(steps, strength_applied) toggled = maybe_toggle_step_cache( state.pipe, - steps = steps, + steps = denoise_steps, quant_active = state.cache_quant_active, threshold = state.cache_threshold, logger = logger, @@ -2289,7 +2307,7 @@ class DiffusionBackend: if isinstance(entry, dict): entry["value"] = toggled or "off" entry["reason"] = ( - f"auto: {steps}-step generation " + f"auto: {denoise_steps}-step generation " + ("reaches" if toggled else "is below") + f" {FBCACHE_MIN_STEPS}" ) diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 7377f8d919..2f00a09269 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -178,6 +178,15 @@ _INSTALLABLE_BACKENDS: dict[str, tuple[str, str]] = { # 0 - never install; a missing kernel falls back to native _ATTENTION_INSTALL_ENV = "UNSLOTH_DIFFUSION_ATTENTION_INSTALL" +# Packages a pip install has already been attempted for in THIS process (success or +# failure). The loader pre-installs the kernel OUTSIDE its locks and then re-resolves the +# same backend under _generate_lock, where apply_attention_backend would otherwise call +# pip a SECOND time -- for a package with no matching wheel / an offline host that repeat +# runs the full (up to 600s) install while holding the load lock, blocking unload/cancel/ +# new loads for exactly the failure the pre-install was added to keep off the lock. Record +# each attempt so a retry is a no-op and set_attention_backend falls back to native at once. +_INSTALL_ATTEMPTED: set[str] = set() + def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> None: """Best-effort wheel-only install of the package ``backend`` needs, when allowed. @@ -202,6 +211,14 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non return except Exception: # noqa: BLE001 — a broken install probes as missing; try the install pass + # Only ever attempt each package's install once per process. The loader pre-installs + # this backend outside its locks; if that failed (no wheel / offline) the module is + # still missing here, so without this guard the in-lock apply path would re-run the + # whole install under _generate_lock and block unload/cancel. A recorded attempt makes + # the retry a no-op -> set_attention_backend raises on the missing package -> native. + if package in _INSTALL_ATTEMPTED: + return + _INSTALL_ATTEMPTED.add(package) import subprocess import sys @@ -211,7 +228,16 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non ) try: subprocess.run( - [sys.executable, "-m", "pip", "install", "--only-binary", ":all:", package], + # --no-deps: install ONLY this best-effort kernel wheel, never its declared + # dependencies. xformers/flash-attn pin an exact torch (e.g. torch==2.x), so + # normal resolution would upgrade/replace the running torch/triton and leave + # later loads on a different, possibly CUDA-mismatched dependency stack. Without + # its deps an ABI-incompatible kernel simply fails to import -> native fallback, + # which is the same best-effort outcome as an uninstallable wheel. + [ + sys.executable, "-m", "pip", "install", + "--only-binary", ":all:", "--no-deps", package, + ], capture_output = True, timeout = 600, check = True, diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index c68fc8894b..b5bc2341a1 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -126,6 +126,24 @@ def apply_step_cache( return None +def effective_denoise_steps(steps: int, strength: Optional[float]) -> int: + """The number of steps diffusers ACTUALLY denoises for a request. + + An image-conditioned workflow with ``strength`` < 1 (img2img / upscale / inpaint) runs + only a fraction of ``num_inference_steps``: diffusers' ``get_timesteps`` computes + ``init_timestep = min(num_inference_steps * strength, num_inference_steps)`` and denoises + ``num_inference_steps - int(num_inference_steps - init_timestep)`` steps. The auto + step-cache policy must key on THIS count -- e.g. a 28-step upscale at strength 0.35 runs + ~10 real steps, exactly the short trajectory FBCache should stay off (each skipped step + is a large quality hit). ``strength`` None (txt2img / reference) or >= 1 -> the full count. + """ + s = int(steps) + if strength is None or float(strength) >= 1.0: + return s + init = min(s * float(strength), s) + return max(1, s - int(max(s - init, 0))) + + def maybe_toggle_step_cache( pipe: Any, *, diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 5abaa86fa7..4becbf7a4b 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -241,6 +241,10 @@ def _no_real_installs(monkeypatch): # backends (sage/flash*), so hard-disable the gate; install tests re-enable it # with a stubbed subprocess. monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "0") + # The install once-per-process memo is module state; clear it so each test starts + # with a fresh "not yet attempted" set (otherwise an earlier test's attempt would + # make a later install a no-op). + att._INSTALL_ATTEMPTED.clear() class _Recorder: @@ -290,6 +294,45 @@ def test_install_runs_wheel_only_for_missing_kernel(monkeypatch): assert "--only-binary" in cmd and ":all:" in cmd and "sageattention" in cmd +def test_install_uses_no_deps_to_protect_core_deps(monkeypatch): + # A kernel add-on (xformers/flash-attn) pins an exact torch, so a normal install would + # upgrade/replace the running torch/triton. --no-deps installs only the kernel wheel; + # an ABI-incompatible one fails to import and falls back to native rather than clobbering + # the environment's core deps. + 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("xformers") + assert len(run.calls) == 1 + assert "--no-deps" in run.calls[0] + + +def test_failed_install_not_retried_in_same_process(monkeypatch): + # The loader pre-installs the kernel OUTSIDE its locks and then re-resolves the same + # backend under _generate_lock; if the pre-install failed (no wheel / offline) the + # in-lock apply path must NOT re-run pip (a second up-to-600s install holding the load + # lock blocks unload/cancel). The once-per-process memo makes the retry a no-op. + monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto") + import importlib.util + import subprocess as sp + + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) # stays missing + + calls: list[list[str]] = [] + + def _boom(cmd, **kwargs): + calls.append(list(cmd)) + raise sp.CalledProcessError(returncode = 1, cmd = cmd) + + _stub_subprocess(monkeypatch, _boom) + att._ensure_attention_backend_installed("sage") # pre-install attempt (outside lock) + att._ensure_attention_backend_installed("sage") # in-lock retry -> must be skipped + assert len(calls) == 1 + + def test_install_invalidates_import_caches_on_success(monkeypatch): # A wheel written to site-packages after the finder cached that directory can be # missed by the very next import, so a successful install must invalidate the caches diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 8df218101c..aa0e9d354e 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -179,10 +179,44 @@ def test_diffusers_unavailable_runs_uncached(monkeypatch): from core.inference.diffusion_cache import ( # noqa: E402 FBCACHE_MIN_STEPS, TC_AUTO, + effective_denoise_steps, maybe_toggle_step_cache, ) +# ── effective_denoise_steps (strength-aware step count for the auto policy) ───────── +def test_effective_steps_txt2img_is_full_count(): + # No strength (txt2img / reference) -> the full requested step count. + assert effective_denoise_steps(28, None) == 28 + assert effective_denoise_steps(28, 1.0) == 28 # full redraw denoises every step + + +def test_effective_steps_low_strength_shrinks_below_the_bar(): + # A 28-step upscale at strength 0.35 denoises ~10 steps (diffusers get_timesteps), + # which is below FBCACHE_MIN_STEPS -> the auto policy must NOT engage FBCache there. + eff = effective_denoise_steps(28, 0.35) + assert eff == 10 + assert eff < FBCACHE_MIN_STEPS + + +def test_effective_steps_matches_diffusers_get_timesteps(): + # Mirror diffusers exactly: num_inference_steps - int(num_inference_steps - + # min(num_inference_steps * strength, num_inference_steps)). + for steps, strength in [(28, 0.35), (28, 0.8), (50, 0.5), (20, 0.99), (30, 0.1)]: + init = min(steps * strength, steps) + expected = max(1, steps - int(max(steps - init, 0))) + assert effective_denoise_steps(steps, strength) == expected + + +def test_toggle_stays_off_for_low_strength_workflow(monkeypatch): + # End to end: a 28-step request would engage FBCache, but at strength 0.35 the + # effective ~10 steps keep it uncached. + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + mode = maybe_toggle_step_cache(_pipe(t), steps = effective_denoise_steps(28, 0.35)) + assert mode is None and t.enables == 0 + + class _ToggleTransformer(_MixinTransformer): """CacheMixin-style fake with the disable side too, counting transitions.""" diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 6ab4529652..f6c539a919 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -992,7 +992,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { "auto", ); const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto"); - const [transformerCache, setTransformerCache] = useState<"off" | "fbcache">("off"); + const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache">("auto"); const [cpuOffload, setCpuOffload] = useState(false); // The last load descriptor, so "Reapply" can reload the same model with new advanced // options without the user re-picking it from the dropdown. @@ -1475,7 +1475,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant, attention_backend: attentionBackend === "auto" ? undefined : attentionBackend, memory_mode: memoryMode === "auto" ? undefined : memoryMode, - transformer_cache: transformerCache === "off" ? undefined : transformerCache, + transformer_cache: transformerCache === "auto" ? undefined : transformerCache, }); } catch (err) { dismissLoadToast(); @@ -1925,10 +1925,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) { /> setTransformerCache(v as typeof transformerCache)} options={[ + ["auto", "Auto"], ["off", "Off"], ["fbcache", "First-Block-Cache"], ]}