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

studio/backend/tests/test_diffusion_precision.py

@@ -393,3 +393,81 @@ def test_nvfp4_filter_keeps_vision_tower_dense(monkeypatch):
     assert ff(object(), "lm_head") is False
     assert ff(object(), "model.decoder.wo") is False
     assert ff(object(), "model.layers.5.self_attn.q_proj") is True
+
+
+# ── zero-output-row guard (per-row fp8 NaN protection) ───────────────────────────
+
+
+class _FakeAmaxVec:
+    def __init__(self, vals):
+        self._vals = vals
+
+    def __eq__(self, other):  # noqa: PLW0642 -- tensor-style elementwise compare
+        return _FakeAmaxVec([v == other for v in self._vals])
+
+    def any(self):
+        return _FakeScalar(any(self._vals))
+
+
+class _FakeScalar:
+    def __init__(self, v):
+        self._v = v
+
+    def item(self):
+        return self._v
+
+
+class _FakeWeight:
+    """Tensor-shaped stand-in supporting the exact chain the guard runs:
+    ``weight.abs().amax(dim = -1) == 0 -> .any().item()``."""
+
+    ndim = 2
+
+    def __init__(self, rows):
+        self._rows = rows
+
+    def abs(self):
+        return _FakeWeight([[abs(v) for v in r] for r in self._rows])
+
+    def amax(self, dim = -1):
+        return _FakeAmaxVec([max(r) for r in self._rows])
+
+
+def test_weight_zero_output_row_detection():
+    # A dead output row NaNs torchao's per-row fp8 (scale 0 -> 0/0); SDXL's
+    # text_encoder_2 (OpenCLIP bigG) really ships one in layers.2.self_attn.out_proj --
+    # measured: every fp8_dynamic SDXL render was black until the row is kept dense.
+    zero_row = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.0, 0.0]]))
+    dense = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.3, 0.0]]))
+    assert dp._weight_has_zero_output_row(zero_row) is True
+    assert dp._weight_has_zero_output_row(dense) is False
+    # Non-2D / absent weights are not the per-row scheme's input: never flagged.
+    w3 = _FakeWeight([[1.0]])
+    w3.ndim = 3
+    assert dp._weight_has_zero_output_row(types.SimpleNamespace(weight = w3)) is False
+    assert dp._weight_has_zero_output_row(types.SimpleNamespace()) is False
+
+    # An unreadable weight falls through to quantize_'s own handling.
+    class _Boom:
+        @property
+        def weight(self):
+            raise RuntimeError("meta tensor")
+
+    assert dp._weight_has_zero_output_row(_Boom()) is False
+
+
+def test_fp8_dynamic_filter_skips_zero_row_linear(monkeypatch):
+    # The fp8_dynamic caster must leave a zero-output-row Linear dense while the rest
+    # of the encoder still quantises (a family-wide deny would forfeit the whole win).
+    _stub_torch(monkeypatch)
+    captured: dict = {}
+    _stub_transformer_quant(monkeypatch, captured)
+    enc = types.SimpleNamespace(_keep_in_fp32_modules = [])
+
+    dp._cast_fp8_dynamic(enc, _target())
+
+    ff = captured["filter_fn"]
+    dead = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.0, 0.0]]))
+    live = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.5, 0.5]]))
+    assert ff(dead, "text_model.encoder.layers.2.self_attn.out_proj") is False
+    assert ff(live, "text_model.encoder.layers.2.mlp.fc1") is True

studio/backend/tests/test_diffusion_speed.py

@@ -531,3 +531,85 @@ def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch):
     )
     assert applied["fp16_accum"] is True
     assert torch.backends.cuda.matmul.allow_fp16_accumulation is True
+
+
+# ── inductor precision-cast emulation (compile-vs-eager numeric parity) ─────────
+
+
+def _stub_inductor_config(
+    monkeypatch,
+    torch,
+    *,
+    emulate = False,
+):
+    """Attach a fake ``_inductor.config`` to the stubbed torch module (diffusion_speed
+    resolves it as attributes off the imported torch, never via sys.modules -- so the
+    real torch._inductor lingering in sys.modules cannot leak into stubbed tests)."""
+    cfg = types.SimpleNamespace(emulate_precision_casts = emulate)
+    torch._inductor = types.SimpleNamespace(config = cfg)
+    return cfg
+
+
+def test_regional_compile_enables_emulate_precision_casts(monkeypatch):
+    # Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds
+    # to bf16 between ops; over a multi-step denoise that compounds to a visible drift.
+    # emulate_precision_casts restores eager's rounding at zero measured speed cost, so
+    # the regional compile path must switch it on.
+    torch = _stub_torch(monkeypatch)
+    _stub_gguf_accel(monkeypatch)
+    cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
+    pipe = _Pipe(with_compile = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
+    )
+    assert applied["compiled"] is True
+    assert cfg.emulate_precision_casts is True
+
+
+def test_snapshot_restores_emulate_precision_casts(monkeypatch):
+    # The flag is process-global, so the unload path must restore the pre-load value
+    # exactly like the TF32 / cudnn.benchmark globals.
+    torch = _stub_torch(monkeypatch)
+    cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
+    snap = snapshot_backend_flags()
+    assert snap["inductor_emulate_precision_casts"] is False
+    cfg.emulate_precision_casts = True
+    restore_backend_flags(snap)
+    assert cfg.emulate_precision_casts is False
+
+
+def test_missing_inductor_config_is_tolerated(monkeypatch):
+    # A build without torch._inductor (or with the flag renamed) must neither break the
+    # snapshot nor the compile path.
+    _stub_torch(monkeypatch)  # the stub torch has no _inductor attribute
+    _stub_gguf_accel(monkeypatch)
+    snap = snapshot_backend_flags()
+    assert "inductor_emulate_precision_casts" not in snap
+    pipe = _Pipe(with_compile = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
+    )
+    assert applied["compiled"] is True
+
+
+def test_regional_compile_arms_cache_hook_inners(monkeypatch):
+    # The production load order engages the step cache BEFORE compile, so the regional
+    # compile pass must re-arm the already-installed cache hooks with compiled inner
+    # forwards (otherwise every computed step runs eager under the hook's
+    # torch.compiler.disable and forfeits the regional compile).
+    _stub_torch(monkeypatch)
+    _stub_gguf_accel(monkeypatch)
+    from core.inference import diffusion_cache as dc_mod
+
+    armed = []
+    monkeypatch.setattr(
+        dc_mod,
+        "_compile_hooked_block_inners",
+        lambda transformer, logger = None: armed.append(transformer) or 1,
+    )
+    pipe = _Pipe(with_compile = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
+    )
+    assert applied["compiled"] is True
+    assert armed == [pipe.transformer]