Studio diffusion: eager patches + torch.compile cache speed phase

Adds the opt-in speed path for the GGUF diffusion transformer behind a
selectable speed mode (default off, so output is unchanged until a profile
is chosen):

- diffusion_eager_patches.py: shared eager fast-paths (channels_last,
  attention/backend selection, fused norms and QKV) installed at load and
  rolled back on unload or failed load.
- diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent
  torch.compile cache and the GGUF-transformer compile wiring.
- diffusion_arch_patches.py: architecture-specific patches.
- diffusion_patch_backend.py: shared install/restore plumbing.
- diffusion_speed.py: speed-profile planning.

Tests for each module plus the benchmarking and probe scripts used to
measure speed, memory, and accuracy of the path.

studio/backend/tests/test_diffusion_compile_cache.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Unit tests for the pre-warmed torch.compile cache (``diffusion_compile_cache.py``).
+
+The Mega-cache API (``torch.compiler.save_cache_artifacts`` / ``load_cache_artifacts``)
+is monkeypatched with deterministic in-memory fakes so the fingerprint / exact-match /
+integrity / fallback / lifecycle logic is exercised without a real compile. The
+fingerprint helpers run against the real torch on this box.
+"""
+
+from __future__ import annotations
+
+import json
+import types
+
+import pytest
+
+from core.inference import diffusion_compile_cache as cc
+
+
+def _transformer(blocks=("FluxTransformerBlock", "FluxSingleTransformerBlock")):
+    return types.SimpleNamespace(_repeated_blocks=list(blocks))
+
+
+_BEGIN_KW = dict(
+    family="flux.1",
+    dtype="torch.bfloat16",
+    quant=None,
+    attention_backend="_native_cudnn",
+    compile_kwargs={"fullgraph": True, "dynamic": True},
+    shape_bucket="1024x1024",
+)
+
+
+# --------------------------------------------------------------------------- fingerprint
+def test_environment_fingerprint_has_hard_dimensions():
+    fp = cc.environment_fingerprint()
+    for k in ("torch", "torch_cuda", "triton", "diffusers", "gpu_name", "gpu_capability"):
+        assert k in fp
+
+
+def test_cache_key_stable_across_kwarg_order():
+    efp = cc.environment_fingerprint()
+    t = _transformer()
+    a = cc.model_fingerprint(family="flux.1", transformer=t, dtype="bf16", quant=None,
+                             attention_backend="x", compile_kwargs={"fullgraph": True, "dynamic": True})
+    b = cc.model_fingerprint(family="flux.1", transformer=t, dtype="bf16", quant=None,
+                             attention_backend="x", compile_kwargs={"dynamic": True, "fullgraph": True})
+    assert cc.cache_key(efp, a) == cc.cache_key(efp, b)
+
+
+@pytest.mark.parametrize("field,value", [
+    ("family", "qwen-image"),
+    ("dtype", "torch.float16"),
+    ("quant", "int8"),
+    ("attention_backend", "native"),
+    ("shape_bucket", "512x512"),
+])
+def test_cache_key_sensitive_to_model_dims(field, value):
+    efp = cc.environment_fingerprint()
+    t = _transformer()
+    base = dict(family="flux.1", transformer=t, dtype="bf16", quant=None,
+                attention_backend="x", compile_kwargs={"fullgraph": True}, shape_bucket="1024x1024")
+    k0 = cc.cache_key(efp, cc.model_fingerprint(**base))
+    base[field] = value
+    assert cc.cache_key(efp, cc.model_fingerprint(**base)) != k0
+
+
+def test_repeated_blocks_change_key():
+    efp = cc.environment_fingerprint()
+    k1 = cc.cache_key(efp, cc.model_fingerprint(family="f", transformer=_transformer(("A",)),
+                      dtype="bf16", quant=None, attention_backend="x", compile_kwargs={}))
+    k2 = cc.cache_key(efp, cc.model_fingerprint(family="f", transformer=_transformer(("B",)),
+                      dtype="bf16", quant=None, attention_backend="x", compile_kwargs={}))
+    assert k1 != k2
+
+
+# ----------------------------------------------------------------------------- env knobs
+@pytest.mark.parametrize("raw,expected", [
+    ("0", "off"), ("off", "off"), ("1", "on"), ("on", "on"),
+    ("auto", "auto"), ("", "auto"), ("garbage", "auto"),
+])
+def test_cache_mode(monkeypatch, raw, expected):
+    monkeypatch.setenv(cc._ENV_MODE, raw)
+    assert cc.cache_mode() == expected
+
+
+def test_cache_mode_default_auto(monkeypatch):
+    monkeypatch.delenv(cc._ENV_MODE, raising=False)
+    assert cc.cache_mode() == "auto"
+
+
+# ------------------------------------------------------------------------------ disabled
+def test_begin_returns_none_when_disabled(monkeypatch):
+    monkeypatch.setenv(cc._ENV_MODE, "0")
+    assert cc.begin(transformer=_transformer(), **_BEGIN_KW) is None
+
+
+def test_begin_returns_none_without_megacache_api(monkeypatch):
+    monkeypatch.setenv(cc._ENV_MODE, "auto")
+    fake_torch = types.ModuleType("torch")
+    fake_torch.compiler = types.SimpleNamespace()  # no save/load attrs
+    monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch)
+    assert cc.begin(transformer=_transformer(), **_BEGIN_KW) is None
+
+
+# ----------------------------------------------------------------- megacache fake + flow
+@pytest.fixture
+def fake_megacache(monkeypatch):
+    """Patch torch.compiler save/load with deterministic in-memory behaviour."""
+    import torch
+
+    state = {"saved": None, "loaded_with": None}
+
+    def fake_save():
+        return (b"ARTIFACT-BYTES", None)
+
+    def fake_load(data: bytes):
+        state["loaded_with"] = data
+        return object() if data == b"ARTIFACT-BYTES" else None
+
+    monkeypatch.setattr(torch.compiler, "save_cache_artifacts", fake_save, raising=False)
+    monkeypatch.setattr(torch.compiler, "load_cache_artifacts", fake_load, raising=False)
+    return state
+
+
+def test_save_then_load_roundtrip(monkeypatch, tmp_path, fake_megacache):
+    monkeypatch.setenv(cc._ENV_MODE, "on")          # load + save
+    monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
+
+    # First load: cold (no bundle yet).
+    ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    assert ctx is not None and ctx.hit is False
+    assert cc.save(ctx) is True
+    assert ctx.bundle.exists() and ctx.manifest_path.exists()
+
+    # Second load with the SAME fingerprint: warm hit.
+    ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    assert ctx2 is not None and ctx2.hit is True
+    assert fake_megacache["loaded_with"] == b"ARTIFACT-BYTES"
+    assert ctx2.key == ctx.key
+
+
+def test_no_save_in_auto_mode(monkeypatch, tmp_path, fake_megacache):
+    monkeypatch.setenv(cc._ENV_MODE, "auto")
+    monkeypatch.delenv(cc._ENV_SAVE, raising=False)
+    monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
+    ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    assert cc.save(ctx) is False           # auto without SAVE opt-in does not write
+    assert not ctx.bundle.exists()
+
+
+def test_fingerprint_mismatch_falls_back(monkeypatch, tmp_path, fake_megacache):
+    monkeypatch.setenv(cc._ENV_MODE, "on")
+    monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
+    ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    cc.save(ctx)
+
+    # Tamper the manifest's env fingerprint -> exact-match guard must reject the bundle.
+    manifest = json.loads(ctx.manifest_path.read_text())
+    manifest["env"]["torch"] = "0.0.0-other"
+    ctx.manifest_path.write_text(json.dumps(manifest))
+
+    ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    assert ctx2.hit is False               # mismatch -> local compile, non-fatal
+
+
+def test_corrupt_bundle_rejected(monkeypatch, tmp_path, fake_megacache):
+    monkeypatch.setenv(cc._ENV_MODE, "on")
+    monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
+    ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    cc.save(ctx)
+    ctx.bundle.write_bytes(b"CORRUPTED")   # manifest sha256 no longer matches
+
+    ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    assert ctx2.hit is False
+
+
+# ------------------------------------------------------------------------------- restore
+def test_restore_inductor_dir(monkeypatch, tmp_path, fake_megacache):
+    import os
+    monkeypatch.setenv(cc._ENV_MODE, "auto")
+    monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
+    monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", "/tmp/prior-inductor")
+    ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW)
+    assert os.environ["TORCHINDUCTOR_CACHE_DIR"] != "/tmp/prior-inductor"   # redirected
+    cc.restore(ctx)
+    assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == "/tmp/prior-inductor"   # restored