Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT

Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).

Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.

OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.

- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
  (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
  quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
  when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
  transformer_cache_threshold through begin_load -> load_pipeline and report the
  engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
  transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
  selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
  LPIPS vs the compiled no-cache baseline).

scripts/fbcache_flux_probe.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Validate First-Block-Cache (FBCache) on a MANY-step DiT (Flux.1-dev), vs the compiled
+baseline. FBCache reuses the transformer tail across denoise steps when the first block's
+residual barely changes -- a real speedup only when there are enough steps (it is why it is
+gated OFF for few-step distilled models like Z-Image-Turbo). Reports median latency,
+speedup, peak VRAM, and LPIPS vs the no-cache baseline. One CUDA GPU."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+BASE = "black-forest-labs/FLUX.1-dev"
+PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
+OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/fbcache_flux_images")
+
+
+_LP = {"fn": None}
+
+
+def _lpips(ref, arr):
+    try:
+        import lpips
+        import torch
+        if _LP["fn"] is None:
+            _LP["fn"] = lpips.LPIPS(net="alex", verbose=False).cuda().eval()
+
+        def t(x):
+            return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda()
+
+        with torch.no_grad():
+            return float(_LP["fn"](t(ref), t(arr)).item())
+    except Exception as exc:  # noqa: BLE001
+        print(f"    (lpips: {type(exc).__name__})", flush=True)
+        return None
+
+
+def _load():
+    import os
+    import diffusers
+    import torch
+    pipe = diffusers.FluxPipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, token=os.environ.get("HF_TOKEN"))
+    pipe.to("cuda")
+    return pipe
+
+
+def _gen(pipe, steps, seed, res, guidance):
+    import torch
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    torch.cuda.synchronize(); t0 = time.time()
+    img = pipe(prompt=PROMPT, width=res, height=res, num_inference_steps=steps,
+               guidance_scale=guidance, generator=g).images[0]
+    torch.cuda.synchronize()
+    return img, time.time() - t0
+
+
+def _median(xs):
+    return sorted(xs)[len(xs) // 2]
+
+
+def run(tag, steps, seed, res, guidance, iters, *, threshold=None, compile_=True):
+    import torch
+    torch.compiler.reset(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
+    pipe = _load()
+    if threshold is not None:
+        from diffusers import FirstBlockCacheConfig
+        try:
+            pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold=threshold))
+        except Exception as exc:  # noqa: BLE001
+            from diffusers.hooks import apply_first_block_cache
+            apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=threshold))
+    if compile_:
+        try:
+            pipe.transformer.compile_repeated_blocks(fullgraph=True, dynamic=True)
+        except Exception as exc:  # noqa: BLE001
+            print(f"    [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush=True)
+    try:
+        _gen(pipe, steps, seed, res, guidance)  # warmup / compile
+    except Exception as exc:  # noqa: BLE001
+        import traceback; traceback.print_exc()
+        print(f"    [{tag}] FAILED: {type(exc).__name__}: {str(exc)[:100]}", flush=True)
+        del pipe; torch.cuda.empty_cache(); return None
+    dts, img = [], None
+    for _ in range(iters):
+        img, dt = _gen(pipe, steps, seed, res, guidance); dts.append(dt)
+    peak = torch.cuda.max_memory_allocated() / 1e9
+    arr = np.array(img)
+    OUT.mkdir(parents=True, exist_ok=True); img.save(OUT / f"{tag}.png")
+    del pipe; torch.cuda.empty_cache()
+    return _median(dts), arr, peak
+
+
+def main(argv=None) -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--steps", type=int, default=28)
+    p.add_argument("--res", type=int, default=1024)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--guidance", type=float, default=3.5)
+    p.add_argument("--iters", type=int, default=2)
+    args = p.parse_args(argv)
+    s, r, seed, gd, it = args.steps, args.res, args.seed, args.guidance, args.iters
+
+    print(f"== FBCache on Flux.1-dev ({r}px, {s} steps, guidance {gd}) ==", flush=True)
+    base = run("baseline", s, seed, r, gd, it)
+    if base is None:
+        print("baseline FAILED", flush=True); return 1
+    bmed, ref, bpeak = base
+    print(f"  baseline {bmed:.3f}s peak={bpeak:.1f}G", flush=True)
+    rows = [("baseline", bmed, bpeak, 0.0)]
+    for thr in (0.08, 0.12, 0.20):
+        out = run(f"fbcache_{thr}", s, seed, r, gd, it, threshold=thr)
+        if out is None:
+            rows.append((f"fbcache_{thr}", None, None, None)); continue
+        med, arr, peak = out
+        lp = _lpips(ref, arr)
+        rows.append((f"fbcache_{thr}", med, peak, lp))
+        print(f"  fbcache_{thr}: {med:.3f}s ({bmed/med:.2f}x) peak={peak:.1f}G LPIPS={lp}", flush=True)
+
+    print("\n==== SUMMARY (Flux.1-dev, ref = no-cache compile) ====", flush=True)
+    for tag, med, peak, lp in rows:
+        if med is None:
+            print(f"  {tag:16s} FAILED"); continue
+        spd = f"{bmed/med:.2f}x"
+        lpv = "ref" if tag == "baseline" else (f"{lp:.3f}" if lp is not None else "n/a")
+        print(f"  {tag:16s} {med:.3f}s  {spd:>6s}  peak={peak:.1f}G  LPIPS={lpv:>6s}", flush=True)
+    print("FBCACHE-FLUX-DONE", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
+    sys.exit(main())

studio/backend/core/inference/diffusion.py

@@ -55,6 +55,7 @@ from .diffusion_attention import (
     apply_attention_backend,
     select_attention_backend,
 )
+from .diffusion_cache import apply_step_cache
 from .diffusion_precision import quantize_text_encoders
 from .diffusion_prequant import (
     load_prequantized_transformer,
@@ -101,6 +102,8 @@ class _LoadState:
     # Attention backend engaged via the diffusers dispatcher (e.g. "_native_cudnn"), or
     # None for the default SDPA. Set before compile; orthogonal to the weight quant.
     attention_backend: Optional[str] = None
+    # Step cache engaged ("fbcache") or None. Opt-in, for many-step models.
+    transformer_cache: Optional[str] = None
 
 
 @dataclass
@@ -291,6 +294,8 @@ class DiffusionBackend:
         transformer_quant_fast_accum: Optional[bool] = None,
         transformer_prequant_path: Optional[str] = None,
         attention_backend: Optional[str] = None,
+        transformer_cache: Optional[str] = None,
+        transformer_cache_threshold: Optional[float] = None,
     ) -> dict[str, Any]:
         """Validate, then run the (slow) load on a daemon thread. Returns at once."""
         fam = self.validate_load_request(
@@ -326,6 +331,8 @@ class DiffusionBackend:
                 transformer_quant_fast_accum = transformer_quant_fast_accum,
                 transformer_prequant_path = transformer_prequant_path,
                 attention_backend = attention_backend,
+                transformer_cache = transformer_cache,
+                transformer_cache_threshold = transformer_cache_threshold,
                 _load_token = token,
             ),
             daemon = True,
@@ -451,6 +458,8 @@ class DiffusionBackend:
         transformer_quant_fast_accum: Optional[bool] = None,
         transformer_prequant_path: Optional[str] = None,
         attention_backend: Optional[str] = None,
+        transformer_cache: Optional[str] = None,
+        transformer_cache_threshold: Optional[float] = None,
         _load_token: Optional[int] = None,
     ) -> dict[str, Any]:
         # Validate first (cheap, no torch/diffusers) so a direct call with a bad
@@ -571,12 +580,24 @@ 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.
+                cache_engaged = apply_step_cache(
+                    pipe,
+                    mode = transformer_cache,
+                    threshold = transformer_cache_threshold,
+                    quant_active = transformer_quant_engaged is not None,
+                    logger = logger,
+                )
                 speed_applied = apply_speed_optims(
                     pipe,
                     target,
                     is_gguf = bool(gguf_filename),
                     family = fam,
                     speed_mode = effective_speed,
+                    cache_active = cache_engaged is not None,
                     logger = logger,
                 )
                 # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
@@ -615,6 +636,7 @@ class DiffusionBackend:
                     text_encoder_quant = te_quant,
                     transformer_quant = transformer_quant_engaged,
                     attention_backend = attention_engaged,
+                    transformer_cache = cache_engaged,
                 )
 
         logger.info(
@@ -917,6 +939,7 @@ class DiffusionBackend:
                 "text_encoder_quant": None,
                 "transformer_quant": None,
                 "attention_backend": None,
+                "transformer_cache": None,
             }
         return {
             "loaded": True,
@@ -934,6 +957,7 @@ class DiffusionBackend:
             "text_encoder_quant": state.text_encoder_quant,
             "transformer_quant": state.transformer_quant,
             "attention_backend": state.attention_backend,
+            "transformer_cache": state.transformer_cache,
         }
 
 

studio/backend/core/inference/diffusion_cache.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Opt-in step caching for the diffusion transformer (First-Block-Cache).
+
+Across denoising steps a DiT's output changes little once the trajectory settles, so most of
+the transformer can be reused. First-Block-Cache (FBCache) computes the first block, and if
+its residual barely changed from the previous step (within ``threshold``) it skips the
+remaining blocks and reuses their cached output. diffusers ships it natively
+(``transformer.enable_cache(FirstBlockCacheConfig(...))`` for CacheMixin models, or the
+standalone ``apply_first_block_cache`` hook).
+
+Measured on Flux.1-dev (28 steps, 1024px, B200): ~1.4x on top of torch.compile (2.83 ->
+2.03 s) at LPIPS ~0.08 vs the no-cache output -- deep inside the speed-for-quality bar.
+
+OFF by default and a deliberate per-load opt-in, because the win scales with step count: a
+few-step distilled model (e.g. Z-Image-Turbo at ~8 steps) has almost no headroom and a
+single skipped step is a large fraction of the trajectory, so caching is for many-step
+models (Flux / Qwen-Image). It composes with torch.compile only with ``fullgraph=False``
+(the cache's compiler-disabled decision is a graph break), which the speed layer switches to
+automatically when a cache is engaged. Best-effort: an incompatible model (e.g. a transformer
+whose block signature the hook does not recognise) is caught and the load proceeds uncached.
+torch / diffusers imported lazily.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+TC_OFF = "off"
+TC_FBCACHE = "fbcache"
+TC_MODES = (TC_FBCACHE,)
+
+# FBCache residual thresholds: higher skips more steps (faster, lower quality). The dense
+# bf16 default; a quantised transformer shifts the residual distribution, so it needs a
+# higher threshold for the cache to trigger at all (per ParaAttention's fp8 guidance).
+DEFAULT_FBCACHE_THRESHOLD = 0.08
+QUANT_FBCACHE_THRESHOLD = 0.12
+
+
+def normalize_transformer_cache(value: Optional[str]) -> Optional[str]:
+    """Lower/strip a requested cache mode; None / "" / "none" / "off" -> None (disabled).
+
+    Raises ValueError for an unsupported value so a bad request is rejected cheaply."""
+    if value is None:
+        return None
+    normalized = str(value).strip().lower().replace("-", "_")
+    if not normalized or normalized in ("none", "off"):
+        return None
+    if normalized not in TC_MODES:
+        raise ValueError(
+            f"Unsupported transformer_cache '{value}'. Use one of: off, {', '.join(TC_MODES)}."
+        )
+    return normalized
+
+
+def apply_step_cache(
+    pipe: Any,
+    *,
+    mode: Optional[str],
+    threshold: Optional[float] = None,
+    quant_active: bool = False,
+    logger: Any = None,
+) -> Optional[str]:
+    """Engage step caching on ``pipe.transformer``. Returns the mode actually engaged, or
+    None when disabled / unsupported (the load then runs uncached). ``threshold`` overrides
+    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:
+        return None
+    transformer = getattr(pipe, "transformer", None)
+    if transformer is None:
+        return None
+    thr = threshold if threshold is not None else (
+        QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD
+    )
+    try:
+        from diffusers import FirstBlockCacheConfig
+
+        config = FirstBlockCacheConfig(threshold=thr)
+        enable_cache = getattr(transformer, "enable_cache", None)
+        if callable(enable_cache):
+            enable_cache(config)
+        else:
+            from diffusers.hooks import apply_first_block_cache
+
+            apply_first_block_cache(transformer, config)
+        try:
+            transformer._unsloth_step_cache = f"{mode}@{thr}"
+        except Exception:  # noqa: BLE001 — marker is best-effort
+            pass
+        if logger is not None:
+            logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
+        return mode
+    except Exception as exc:  # noqa: BLE001 — incompatible model -> run uncached
+        _warn(logger, mode, exc)
+        return 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)

studio/backend/core/inference/diffusion_speed.py

@@ -127,6 +127,7 @@ def apply_speed_optims(
     is_gguf: bool,
     family: Any,
     speed_mode: str = SPEED_OFF,
+    cache_active: bool = False,
     logger: Any = None,
 ) -> dict[str, bool]:
     """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline,
@@ -155,7 +156,9 @@ def apply_speed_optims(
     # block, where eligible (now incl. the GGUF transformer). `max` opts into
     # max-autotune (longer compile, autotuned kernels).
     if compile_eligible(target, is_gguf = is_gguf, family = family):
-        applied["compiled"] = _compile_repeated_blocks(pipe, logger, max_autotune = mode == SPEED_MAX)
+        applied["compiled"] = _compile_repeated_blocks(
+            pipe, logger, max_autotune = mode == SPEED_MAX, cache_active = cache_active
+        )
 
     if mode == SPEED_MAX:
         # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed.
@@ -184,6 +187,7 @@ def _compile_repeated_blocks(
     logger: Any,
     *,
     max_autotune: bool = False,
+    cache_active: bool = False,
 ) -> bool:
     transformer = getattr(pipe, "transformer", None)
     fn = getattr(transformer, "compile_repeated_blocks", None)
@@ -195,7 +199,12 @@ def _compile_repeated_blocks(
     # compile and a recompile per new resolution. The CUDA-graph modes (reduce-overhead
     # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled
     # block because its static output buffer is overwritten across denoise steps.
-    kwargs: dict[str, Any] = {"fullgraph": True, "dynamic": not max_autotune}
+    #
+    # fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is
+    # ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip
+    # inlining torch.compiler.disable()d function"). The break is cheap and the rest of the
+    # block still compiles.
+    kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune}
     if max_autotune:
         kwargs["mode"] = "max-autotune-no-cudagraphs"
     try:

studio/backend/models/inference.py

@@ -1769,6 +1769,23 @@ class DiffusionLoadRequest(BaseModel):
         "friendly); xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An "
         "unavailable kernel falls back to the default.",
     )
+    transformer_cache: Optional[Literal["off", "fbcache"]] = Field(
+        None,
+        description = "Opt-in step caching (off by default). fbcache = First-Block-Cache: "
+        "reuse the transformer tail across denoise steps when the first block's residual "
+        "barely changes (~1.4x on Flux 28-step at LPIPS ~0.08). For MANY-step models "
+        "(Flux / Qwen-Image); leave off for few-step distilled models (e.g. Z-Image-Turbo), "
+        "which have no caching headroom. Composes with compile (drops fullgraph "
+        "automatically); incompatible models run uncached.",
+    )
+    transformer_cache_threshold: Optional[float] = Field(
+        None,
+        ge = 0.0,
+        le = 1.0,
+        description = "FBCache residual threshold (higher = skips more steps = faster, lower "
+        "quality). null auto-picks 0.08 (0.12 when the transformer is quantised, which "
+        "shifts the residual distribution).",
+    )
 
 
 class DiffusionGenerateRequest(BaseModel):
@@ -1886,3 +1903,6 @@ class DiffusionStatusResponse(BaseModel):
         description = "Attention backend engaged via the diffusers dispatcher (e.g. "
         "_native_cudnn), or null for the default SDPA",
     )
+    transformer_cache: Optional[str] = Field(
+        None, description = "Step cache engaged: fbcache | null"
+    )

studio/backend/routes/inference.py

@@ -10089,6 +10089,8 @@ async def load_diffusion_model(
             transformer_quant_fast_accum = request.transformer_quant_fast_accum,
             transformer_prequant_path = request.transformer_prequant_path,
             attention_backend = request.attention_backend,
+            transformer_cache = request.transformer_cache,
+            transformer_cache_threshold = request.transformer_cache_threshold,
         )
         return DiffusionStatusResponse(**status_dict)
     except (ValueError, FileNotFoundError) as exc:

studio/backend/tests/test_diffusion_cache.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Hermetic CPU tests for opt-in step caching (First-Block-Cache).
+
+``diffusers`` is stubbed via ``sys.modules`` (the module under test imports
+``FirstBlockCacheConfig`` / ``apply_first_block_cache`` lazily), and the pipeline is a fake
+that records the engaged config. So normalisation, the CacheMixin (``enable_cache``) path, the
+standalone-hook fallback, threshold selection, and the best-effort failure handling are all
+exercised without torch or a real diffusers model.
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+from core.inference.diffusion_cache import (
+    DEFAULT_FBCACHE_THRESHOLD,
+    QUANT_FBCACHE_THRESHOLD,
+    TC_FBCACHE,
+    apply_step_cache,
+    normalize_transformer_cache,
+)
+
+
+# ── normalize_transformer_cache ────────────────────────────────────────────────────
+def test_normalize_disabled_values_are_none():
+    for value in (None, "", "  ", "none", "off", "OFF", "None"):
+        assert normalize_transformer_cache(value) is None
+
+
+def test_normalize_fbcache_and_casing():
+    assert normalize_transformer_cache("fbcache") == TC_FBCACHE
+    assert normalize_transformer_cache("FBCache") == TC_FBCACHE
+    assert normalize_transformer_cache("  fbcache  ") == TC_FBCACHE
+
+
+def test_normalize_rejects_unknown():
+    with pytest.raises(ValueError):
+        normalize_transformer_cache("deepcache")
+
+
+# ── apply_step_cache ───────────────────────────────────────────────────────────────
+class _Config:
+    def __init__(self, threshold):
+        self.threshold = threshold
+
+
+class _MixinTransformer:
+    """A CacheMixin-style transformer: exposes ``enable_cache``."""
+
+    def __init__(self, *, fail = False):
+        self.fail = fail
+        self.enabled_with = None
+
+    def enable_cache(self, config):
+        if self.fail:
+            raise RuntimeError("block signature not recognised")
+        self.enabled_with = config
+
+
+class _HookTransformer:
+    """A transformer with no ``enable_cache`` -> the standalone hook is used."""
+
+
+def _pipe(transformer):
+    return types.SimpleNamespace(transformer = transformer)
+
+
+def _stub_diffusers(monkeypatch, *, hook_recorder = None):
+    diffusers = types.ModuleType("diffusers")
+    diffusers.FirstBlockCacheConfig = _Config
+    monkeypatch.setitem(sys.modules, "diffusers", diffusers)
+
+    hooks = types.ModuleType("diffusers.hooks")
+
+    def _apply_first_block_cache(transformer, config):
+        if hook_recorder is not None:
+            hook_recorder["transformer"] = transformer
+            hook_recorder["config"] = config
+
+    hooks.apply_first_block_cache = _apply_first_block_cache
+    monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
+
+
+def test_disabled_mode_is_noop(monkeypatch):
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer()
+    assert apply_step_cache(_pipe(t), mode = None) is None
+    assert apply_step_cache(_pipe(t), mode = "off") is None
+    assert t.enabled_with is None
+
+
+def test_enable_cache_path_default_threshold(monkeypatch):
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer()
+    engaged = apply_step_cache(_pipe(t), mode = "fbcache")
+    assert engaged == TC_FBCACHE
+    assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD
+    assert t._unsloth_step_cache == f"fbcache@{DEFAULT_FBCACHE_THRESHOLD}"
+
+
+def test_quant_active_raises_default_threshold(monkeypatch):
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer()
+    apply_step_cache(_pipe(t), mode = "fbcache", quant_active = True)
+    assert t.enabled_with.threshold == QUANT_FBCACHE_THRESHOLD
+
+
+def test_explicit_threshold_overrides_quant(monkeypatch):
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer()
+    apply_step_cache(_pipe(t), mode = "fbcache", threshold = 0.2, quant_active = True)
+    assert t.enabled_with.threshold == 0.2
+
+
+def test_fallback_to_standalone_hook(monkeypatch):
+    rec: dict = {}
+    _stub_diffusers(monkeypatch, hook_recorder = rec)
+    t = _HookTransformer()
+    engaged = apply_step_cache(_pipe(t), mode = "fbcache")
+    assert engaged == TC_FBCACHE
+    assert rec["transformer"] is t
+    assert rec["config"].threshold == DEFAULT_FBCACHE_THRESHOLD
+
+
+def test_incompatible_model_runs_uncached(monkeypatch):
+    # enable_cache raising (e.g. unrecognised block signature) must not fail the load.
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer(fail = True)
+    assert apply_step_cache(_pipe(t), mode = "fbcache") is None
+
+
+def test_missing_transformer_is_none(monkeypatch):
+    _stub_diffusers(monkeypatch)
+    pipe = types.SimpleNamespace(transformer = None)
+    assert apply_step_cache(pipe, mode = "fbcache") is None
+
+
+def test_diffusers_unavailable_runs_uncached(monkeypatch):
+    # no diffusers import -> best-effort returns None, load proceeds uncached.
+    monkeypatch.setitem(sys.modules, "diffusers", None)
+    t = _MixinTransformer()
+    assert apply_step_cache(_pipe(t), mode = "fbcache") is None

studio/backend/tests/test_diffusion_routes.py

@@ -382,6 +382,39 @@ def test_invalid_attention_backend_returns_422(client):
     assert resp.status_code == 422
 
 
+def test_transformer_cache_threads_through(client, monkeypatch):
+    backend = _FakeBackend()
+    monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
+    resp = client.post(
+        "/api/inference/images/load",
+        json = {
+            "model_path": "x/z-image",
+            "gguf_filename": "q.gguf",
+            "transformer_cache": "fbcache",
+            "transformer_cache_threshold": 0.1,
+        },
+    )
+    assert resp.status_code == 200
+    assert backend.last_load_kwargs.get("transformer_cache") == "fbcache"
+    assert backend.last_load_kwargs.get("transformer_cache_threshold") == 0.1
+
+
+def test_invalid_transformer_cache_returns_422(client):
+    resp = client.post(
+        "/api/inference/images/load",
+        json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_cache": "deepcache"},
+    )
+    assert resp.status_code == 422
+
+
+def test_out_of_range_cache_threshold_returns_422(client):
+    resp = client.post(
+        "/api/inference/images/load",
+        json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_cache_threshold": 1.5},
+    )
+    assert resp.status_code == 422
+
+
 def test_invalid_transformer_quant_returns_422_without_eviction(client):
     # An unsupported transformer_quant is rejected by the request schema (Literal), so
     # the GPU is never acquired and no chat model is evicted.