perf(image): compile numeric parity, cache-hook compile arming, FBCache toggle crash fix, TE fp8 zero-row guard

Applies the video round-2 accuracy findings to the image diffusion stack and fixes
two real image-path bugs found while measuring. All numbers B200, production
settings (family default steps/guidance, 1024px, seed 42, 4 fixed prompts), LPIPS
(AlexNet) via the new scripts/image_speedmem_bench.py, which drives the production
lever functions in the loader's own order.

- inductor precision parity: emulate_precision_casts=True on the regional-compile
  path (fused pointwise kernels keep fp32 intermediates where eager rounds to bf16
  between ops). Pairwise LPIPS of the compiled tier vs the same-stack eager tier:
  Qwen-Image 0.019 to 0.006 at identical speed (72.4 vs 72.5 ms/step), FLUX.1-dev
  0.046 to 0.029 at +2% step time (69.8 vs 68.3, reproduced), FLUX.2-klein-4B
  0.018 to 0.017 at identical speed. Snapshot/restored with the other process-wide
  backend flags so an off load never inherits it.
- cache x compile composition: re-point each cache hook's fn_ref.original_forward
  at a torch.compile'd wrapper of the same bound method (armed only where the
  speed layer compiled the block; restored before every disable_cache and before
  the partial-hook cleanup). Qwen-Image FBCache computed steps 91.8 to 71.2 ms
  (back at the uncached compiled rate), 1.21x end to end (7.36 to 6.06 s per 4
  images); FLUX.1-dev already traced through its FBCache hook and is measured
  neutral (same-process armed vs unarmed latents bit-identical). Skip counts
  within noise (13 vs 11 of 76; pairwise LPIPS 0.005).
- FBCache mid-session toggle crash: diffusers 0.39 caches the HookRegistry child
  list on first cache_context use, so an uncached generation followed by a
  20+-step generation (the auto toggle path) enabled hooks the context never
  reached and crashed with "No context is set" (reproduced live on FLUX.1-dev).
  Invalidate the stale child cache after every enable_cache.
- TE fp8_dynamic zero-row guard: torchao per-row fp8 derives a per-output-channel
  scale from the row amax, so an all-zero weight row is 0/0 = NaN. SDXL's
  text_encoder_2 (OpenCLIP bigG) ships exactly such a row, and every explicit
  fp8_dynamic SDXL render came out black; keep zero-row Linears dense (LPIPS
  0.976 black to 0.096 working). Other families' encoders have no such rows and
  are byte-identical.
- No AUTO TE quant exists on the image branch (text_encoder_quant defaults dense,
  explicit-only), so the video round's auto-dense retune has no image analogue;
  the explicit lever's cost is now measured (TE fp8_dynamic alone, LPIPS vs
  bit-exact: Qwen-Image 0.038, FLUX.1-dev 0.084, SDXL 0.096; no speed win, VRAM
  -6.5 GB on Qwen-Image) for the docs.

Tests: 96 passing across the cache/speed/precision suites (11 new arming, 2
child-registry, 2 zero-row, 4 inductor-flag); ruff clean.

studio/backend/tests/test_diffusion_cache.py

@@ -356,3 +356,206 @@ def test_toggle_noop_without_cache_support(monkeypatch):
 
 def test_toggle_noop_without_transformer():
     assert maybe_toggle_step_cache(types.SimpleNamespace(), steps = 28) is None
+
+
+# ── compiled cache-hook inners (regional compile x step cache composition) ──────────
+import functools  # noqa: E402
+
+from core.inference.diffusion_cache import (  # noqa: E402
+    _compile_hooked_block_inners,
+    _invalidate_child_registry_cache,
+    _restore_hooked_block_inners,
+)
+
+
+class _BoundInner:
+    """Provides a plain bound method for fn_ref.original_forward (__self__ present)."""
+
+    def forward(self, *args, **kwargs):
+        return "eager"
+
+
+def _hooked_block(
+    *,
+    compiled = True,
+    hook_name = "fbc_block_hook",
+    bound = True,
+):
+    inner = _BoundInner()
+    orig = inner.forward if bound else functools.partial(_BoundInner.forward, inner)
+    hook = types.SimpleNamespace(fn_ref = types.SimpleNamespace(original_forward = orig))
+    block = types.SimpleNamespace(
+        _diffusers_hook = types.SimpleNamespace(hooks = {hook_name: hook}),
+        _compiled_call_impl = object() if compiled else None,
+    )
+    return block, hook, orig
+
+
+def _fake_dit(blocks):
+    return types.SimpleNamespace(modules = lambda: [types.SimpleNamespace()] + blocks)
+
+
+def _stub_torch_compile(monkeypatch):
+    compiled_calls = []
+
+    def _compile(fn, **kwargs):
+        compiled_calls.append((fn, kwargs))
+        wrapper = lambda *a, **k: fn(*a, **k)  # noqa: E731
+        wrapper._unsloth_test_compiled_of = fn
+        return wrapper
+
+    torch = types.ModuleType("torch")
+    torch.compile = _compile
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    return compiled_calls
+
+
+def test_arming_swaps_inner_for_compiled_wrapper(monkeypatch):
+    calls = _stub_torch_compile(monkeypatch)
+    block, hook, orig = _hooked_block()
+    assert _compile_hooked_block_inners(_fake_dit([block])) == 1
+    assert hook.fn_ref.original_forward is not orig
+    assert hook.fn_ref.original_forward._unsloth_test_compiled_of is orig
+    assert hook._unsloth_orig_inner is orig
+    # The inner compile must match the cache-active tier: graph-breakable + dynamic.
+    assert calls[0][1] == {"fullgraph": False, "dynamic": True}
+
+
+def test_arming_is_idempotent(monkeypatch):
+    _stub_torch_compile(monkeypatch)
+    block, hook, _ = _hooked_block()
+    dit = _fake_dit([block])
+    assert _compile_hooked_block_inners(dit) == 1
+    once = hook.fn_ref.original_forward
+    assert _compile_hooked_block_inners(dit) == 0  # marker short-circuits
+    assert hook.fn_ref.original_forward is once
+
+
+def test_arming_skips_uncompiled_blocks(monkeypatch):
+    # An eager-tier load has no _compiled_call_impl: the hook must stay untouched
+    # (compiling the inner would ADD compile where the user chose eager).
+    _stub_torch_compile(monkeypatch)
+    block, hook, orig = _hooked_block(compiled = False)
+    assert _compile_hooked_block_inners(_fake_dit([block])) == 0
+    assert hook.fn_ref.original_forward is orig
+
+
+def test_arming_skips_partial_captured_inner(monkeypatch):
+    # A stacked hook chain (e.g. group offload) captures a functools.partial, not the
+    # plain bound method; arming would compile the wrong layer of the chain.
+    _stub_torch_compile(monkeypatch)
+    block, hook, orig = _hooked_block(bound = False)
+    assert _compile_hooked_block_inners(_fake_dit([block])) == 0
+    assert hook.fn_ref.original_forward is orig
+
+
+def test_arming_covers_every_cache_hook_family(monkeypatch):
+    # FBCache is the image cache today, but the hook-name table already covers the
+    # MagCache layout too (same fn_ref shape), so a future mode arms for free.
+    _stub_torch_compile(monkeypatch)
+    names = (
+        "mag_cache_leader_block_hook",
+        "mag_cache_block_hook",
+        "fbc_leader_block_hook",
+        "fbc_block_hook",
+    )
+    blocks = [_hooked_block(hook_name = n)[0] for n in names]
+    assert _compile_hooked_block_inners(_fake_dit(blocks)) == len(names)
+
+
+def test_restore_puts_the_exact_original_back(monkeypatch):
+    _stub_torch_compile(monkeypatch)
+    block, hook, orig = _hooked_block()
+    dit = _fake_dit([block])
+    _compile_hooked_block_inners(dit)
+    _restore_hooked_block_inners(dit)
+    assert hook.fn_ref.original_forward is orig
+    assert hook._unsloth_orig_inner is None
+
+
+def test_restore_tolerates_fakes_without_modules():
+    _restore_hooked_block_inners(_MixinTransformer())  # no .modules(): no-op
+
+
+def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch):
+    # The generation-time toggle engages the cache AFTER the load already compiled the
+    # blocks; apply_step_cache must arm the fresh hooks itself.
+    _stub_diffusers(monkeypatch)
+    _stub_torch_compile(monkeypatch)
+    block, hook, orig = _hooked_block()
+
+    class _T(_MixinTransformer):
+        def modules(self):
+            return [block]
+
+    t = _T()
+    engaged = apply_step_cache(_pipe(t), mode = "fbcache")
+    assert engaged == TC_FBCACHE
+    assert hook.fn_ref.original_forward is not orig
+    assert hook._unsloth_orig_inner is orig
+
+
+def test_toggle_disable_restores_inners_before_disable(monkeypatch):
+    # remove_hook splices fn_ref.original_forward back into module.forward, so the
+    # compiled wrapper must be swapped out BEFORE disable_cache runs.
+    _stub_diffusers(monkeypatch)
+    order = []
+
+    class _T(_ToggleTransformer):
+        def disable_cache(self):
+            super().disable_cache()
+            order.append("disable")
+
+        def modules(self):
+            order.append("restore-walk")
+            return []
+
+    t = _T()
+    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 order[-2:] == ["restore-walk", "disable"]
+
+
+def test_enable_failure_restores_inners_before_partial_disable(monkeypatch):
+    # enable_cache can fail after hooking (and arming) some blocks; the partial-hook
+    # cleanup must un-arm them before disable_cache splices original_forward back.
+    _stub_diffusers(monkeypatch)
+    order = []
+
+    class _T(_ToggleTransformer):
+        def enable_cache(self, config):
+            raise RuntimeError("block signature not recognised")
+
+        def disable_cache(self):
+            super().disable_cache()
+            order.append("disable")
+
+        def modules(self):
+            order.append("restore-walk")
+            return []
+
+    t = _T()
+    assert apply_step_cache(_pipe(t), mode = "fbcache") is None
+    assert order == ["restore-walk", "disable"]
+
+
+# ── stale child-registry cache invalidation (mid-session enable) ────────────────────
+
+
+def test_enable_invalidates_stale_child_registry_cache(monkeypatch):
+    # diffusers 0.39 caches the child-registry list on first cache_context use; an
+    # UNCACHED generation already populates it (empty), so a later toggle-time
+    # enable_cache would install hooks the context never reaches ("No context is set").
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer()
+    t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = ["stale"])
+    assert apply_step_cache(_pipe(t), mode = "fbcache") == TC_FBCACHE
+    assert t._diffusers_hook._child_registries_cache is None
+
+
+def test_invalidate_child_registry_cache_tolerates_absence():
+    _invalidate_child_registry_cache(types.SimpleNamespace())  # no registry: no-op
+    reg = types.SimpleNamespace(_child_registries_cache = None)
+    _invalidate_child_registry_cache(types.SimpleNamespace(_diffusers_hook = reg))
+    assert reg._child_registries_cache is None