Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF

- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
  CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
  cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
  transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
  transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
  now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
  quantized FBCache threshold when the caller leaves it unset, instead of the dense default
  that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
  fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
  relative to the script, not a hardcoded path.

scripts/fbcache_flux_probe.py

@@ -18,7 +18,7 @@ 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")
+OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "fbcache_flux_images"
 
 
 _LP = {"fn": None}
@@ -101,8 +101,13 @@ def run(
             from diffusers.hooks import apply_first_block_cache
             apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = threshold))
     if compile_:
+        # FBCache's per-step decision is a graph break, so a cached run must compile with
+        # fullgraph=False (mirroring the production path); fullgraph=True would fail the
+        # warmup compile and the row would silently fall back to an eager cached run,
+        # producing misleading speedup numbers.
+        fullgraph = threshold is None
         try:
-            pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)
+            pipe.transformer.compile_repeated_blocks(fullgraph = fullgraph, dynamic = True)
         except Exception as exc:  # noqa: BLE001
             print(f"    [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush = True)
     try:

studio/backend/core/inference/diffusion.py

@@ -588,7 +588,10 @@ class DiffusionBackend:
                     pipe,
                     mode = transformer_cache,
                     threshold = transformer_cache_threshold,
-                    quant_active = transformer_quant_engaged is not None,
+                    # 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),
                     logger = logger,
                 )
                 speed_applied = apply_speed_optims(

studio/backend/core/inference/diffusion_cache.py

@@ -77,16 +77,22 @@ def apply_step_cache(
         if threshold is not None
         else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD)
     )
+    # Only engage via the transformer's native enable_cache (the diffusers CacheMixin path).
+    # That mixin is present exactly when the pipeline wraps the transformer call in a
+    # cache_context, which the First-Block-Cache hook requires at run time. The lower-level
+    # apply_first_block_cache hook would install on a non-CacheMixin transformer too (e.g.
+    # Z-Image), but its pipeline opens no cache_context, so the first generation would crash
+    # inside the hook -- so a model without enable_cache runs uncached per the best-effort
+    # contract instead of being reported as cached and then failing.
+    enable_cache = getattr(transformer, "enable_cache", None)
+    if not callable(enable_cache):
+        _warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
+        return None
     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)
+        enable_cache(config)
         try:
             transformer._unsloth_step_cache = f"{mode}@{thr}"
         except Exception:  # noqa: BLE001 — marker is best-effort

studio/backend/tests/test_diffusion_cache.py

@@ -4,10 +4,9 @@
 """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.
+``FirstBlockCacheConfig`` lazily), and the pipeline is a fake that records the engaged config.
+So normalisation, the CacheMixin (``enable_cache``) gating, threshold selection, and the
+best-effort failure handling are all exercised without torch or a real diffusers model.
 """
 
 from __future__ import annotations
@@ -62,8 +61,11 @@ class _MixinTransformer:
         self.enabled_with = config
 
 
-class _HookTransformer:
-    """A transformer with no ``enable_cache`` -> the standalone hook is used."""
+class _NonCacheMixinTransformer:
+    """A transformer with no ``enable_cache`` (not a CacheMixin) -> must run uncached.
+
+    Its pipeline opens no ``cache_context``, so installing FBCache would crash at generation;
+    the load runs uncached instead (e.g. Z-Image)."""
 
 
 def _pipe(transformer):
@@ -117,14 +119,14 @@ def test_explicit_threshold_overrides_quant(monkeypatch):
     assert t.enabled_with.threshold == 0.2
 
 
-def test_fallback_to_standalone_hook(monkeypatch):
+def test_non_cachemixin_runs_uncached(monkeypatch):
+    # A transformer without enable_cache (e.g. Z-Image) must NOT install the standalone hook
+    # -- its pipeline opens no cache_context, so it runs uncached instead of crashing at gen.
     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
+    t = _NonCacheMixinTransformer()
+    assert apply_step_cache(_pipe(t), mode = "fbcache") is None
+    assert rec == {}  # the standalone hook was never called
 
 
 def test_incompatible_model_runs_uncached(monkeypatch):