Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32)

Add a speed_mode knob (off by default, so the render path stays bit-identical):
default applies channels_last VAE + regional torch.compile of the denoiser's
repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional
compile is gated off for the GGUF transformer (dequantises per-op) and for families
flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image),
so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed
optims run before placement/offload, per the diffusers composition order. status now
reports speed_mode + the optims actually engaged.

Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last',
'tf32'], compile correctly skipped for GGUF; generation works in every mode.

121 CPU tests pass.

studio/backend/core/inference/diffusion.py

@@ -44,6 +44,7 @@ from .diffusion_memory import (
     plan_diffusion_memory,
     snapshot_device_memory,
 )
+from .diffusion_speed import SPEED_OFF, apply_speed_optims
 
 logger = get_logger(__name__)
 
@@ -64,6 +65,9 @@ class _LoadState:
     offload_policy: str = OFFLOAD_NONE
     vae_tiling: bool = False
     memory_mode: str = "auto"
+    # The opt-in speed profile (Phase 3).
+    speed_mode: str = SPEED_OFF
+    speed_optims: tuple = ()
 
 
 @dataclass
@@ -248,6 +252,7 @@ class DiffusionBackend:
         hf_token: Optional[str] = None,
         cpu_offload: bool = False,
         memory_mode: Optional[str] = None,
+        speed_mode: Optional[str] = None,
     ) -> dict[str, Any]:
         """Validate, then run the (slow) load on a daemon thread. Returns at once."""
         fam = self.validate_load_request(
@@ -277,6 +282,7 @@ class DiffusionBackend:
                 hf_token = hf_token,
                 cpu_offload = cpu_offload,
                 memory_mode = memory_mode,
+                speed_mode = speed_mode,
                 _load_token = token,
             ),
             daemon = True,
@@ -396,6 +402,7 @@ class DiffusionBackend:
         hf_token: Optional[str] = None,
         cpu_offload: bool = False,
         memory_mode: Optional[str] = None,
+        speed_mode: Optional[str] = None,
         _load_token: Optional[int] = None,
     ) -> dict[str, Any]:
         # Validate first (cheap, no torch/diffusers) so a direct call with a bad
@@ -449,6 +456,13 @@ class DiffusionBackend:
                 pipeline_cls = getattr(diffusers, fam.pipeline_class)
                 pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
 
+                # Opt-in speed optims run BEFORE placement (channels_last / compile
+                # must precede CPU offload). Off by default -> bit-identical output.
+                speed_applied = apply_speed_optims(
+                    pipe, target, is_gguf = bool(gguf_filename), family = fam,
+                    speed_mode = speed_mode or SPEED_OFF, logger = logger,
+                )
+
                 # Decide placement from MEASURED free device memory vs the model's
                 # estimated resident size (transformer GGUF dequantised + the
                 # companion text-encoder / VAE already cached for `base`), then
@@ -476,6 +490,8 @@ class DiffusionBackend:
                     offload_policy = effective_policy,
                     vae_tiling = effective_tiling,
                     memory_mode = plan.requested_mode,
+                    speed_mode = (speed_mode or SPEED_OFF),
+                    speed_optims = tuple(k for k, v in speed_applied.items() if v),
                 )
 
         logger.info(
@@ -683,6 +699,8 @@ class DiffusionBackend:
                 "offload_policy": None,
                 "vae_tiling": False,
                 "memory_mode": None,
+                "speed_mode": None,
+                "speed_optims": [],
             }
         return {
             "loaded": True,
@@ -695,6 +713,8 @@ class DiffusionBackend:
             "offload_policy": state.offload_policy,
             "vae_tiling": state.vae_tiling,
             "memory_mode": state.memory_mode,
+            "speed_mode": state.speed_mode,
+            "speed_optims": list(state.speed_optims),
         }
 
 

studio/backend/core/inference/diffusion_families.py

@@ -34,6 +34,10 @@ class DiffusionFamily:
     # (~6.5e4) and produce inf -> NaN latents -> a black image. The backend
     # promotes a resolved float16 to float32 for these at load time.
     fp16_incompatible: bool = False
+    # False for families whose denoiser block doesn't compile cleanly with
+    # regional torch.compile (Z-Image). Only consulted on the non-GGUF path; the
+    # GGUF transformer is never compiled regardless.
+    supports_torch_compile: bool = True
 
 
 # Keyed by architecture, not per model variant: a checkpoint's specific base repo
@@ -76,6 +80,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
         aliases = ("zimage", "z_image"),
         # Z-Image's MLP down-projections peak near 9e5, which overflows float16.
         fp16_incompatible = True,
+        # Z-Image's denoiser block is excluded from regional torch.compile.
+        supports_torch_compile = False,
     ),
 )
 

studio/backend/core/inference/diffusion_speed.py

@@ -0,0 +1,158 @@
+# 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 speed optimisations for the local diffusion backend.
+
+Off by default, so the default render path stays bit-identical to a plain run (the
+property the regression harness checks). When the operator opts in, this applies the
+lossless-to-near-lossless speedups in the order the diffusers guides recommend
+(channels_last -> regional compile, with TF32 / fused-QKV under "max"):
+
+  off     - nothing (default).
+  default - lossless: channels_last VAE memory format + regional torch.compile of
+            the denoiser's repeated block WHERE eligible (non-GGUF, bf16, CUDA, and
+            a compile-friendly family).
+  max     - default plus near-lossless TF32 matmul and fused QKV projections.
+
+Regional compile is gated off for the GGUF transformer (it dequantises per-op and
+doesn't compile cleanly) and for families flagged not compile-friendly (Z-Image), so
+on today's GGUF path only channels_last / TF32 engage; the compile path activates
+automatically once a non-GGUF bf16 transformer is loaded. torch is imported lazily.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+SPEED_OFF = "off"
+SPEED_DEFAULT = "default"
+SPEED_MAX = "max"
+SPEED_MODES = (SPEED_OFF, SPEED_DEFAULT, SPEED_MAX)
+
+
+def normalize_speed_mode(value: Optional[str]) -> str:
+    """Lower/strip a requested speed mode (dashes ok); None / "" -> off."""
+    if value is None:
+        return SPEED_OFF
+    normalized = str(value).strip().lower().replace("-", "_")
+    if not normalized:
+        return SPEED_OFF
+    if normalized not in SPEED_MODES:
+        raise ValueError(
+            f"Unsupported diffusion speed_mode '{value}'. Use one of: {', '.join(SPEED_MODES)}."
+        )
+    return normalized
+
+
+def compile_eligible(target: Any, *, is_gguf: bool, family: Any) -> bool:
+    """Whether the denoiser's repeated block should be regionally compiled.
+
+    Only on CUDA (incl. ROCm via supports_default_torch_compile), for a non-GGUF
+    bf16 transformer, on a compile-friendly family. The GGUF transformer is never
+    compiled (it dequantises per-op)."""
+    if is_gguf:
+        return False
+    if not bool(getattr(target, "supports_default_torch_compile", False)):
+        return False
+    if not bool(getattr(family, "supports_torch_compile", True)):
+        return False
+    return _is_bfloat16(getattr(target, "dtype", None))
+
+
+def _is_bfloat16(dtype: Any) -> bool:
+    try:
+        import torch
+
+        return dtype is torch.bfloat16
+    except Exception:
+        return str(dtype).endswith("bfloat16")
+
+
+def apply_speed_optims(
+    pipe: Any,
+    target: Any,
+    *,
+    is_gguf: bool,
+    family: Any,
+    speed_mode: str = SPEED_OFF,
+    logger: Any = None,
+) -> dict[str, bool]:
+    """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline,
+    BEFORE placement / offload. Returns which optimisations actually engaged. Every
+    step is best-effort: a pipeline that doesn't support one is simply skipped."""
+    applied = {"channels_last": False, "tf32": False, "fused_qkv": False, "compiled": False}
+    mode = normalize_speed_mode(speed_mode)
+    if mode == SPEED_OFF:
+        return applied
+
+    # Lossless: a channels-last VAE speeds up its convolutions with no numeric change.
+    applied["channels_last"] = _vae_channels_last(pipe, logger)
+
+    # Lossless-ish: regional compile of the repeated denoiser block, where eligible.
+    if compile_eligible(target, is_gguf = is_gguf, family = family):
+        applied["compiled"] = _compile_repeated_blocks(pipe, logger)
+
+    if mode == SPEED_MAX:
+        # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed.
+        if getattr(target, "device", None) == "cuda":
+            applied["tf32"] = _enable_tf32(logger)
+        applied["fused_qkv"] = _fuse_qkv(pipe, logger)
+
+    return applied
+
+
+def _vae_channels_last(pipe: Any, logger: Any) -> bool:
+    vae = getattr(pipe, "vae", None)
+    if vae is None or not hasattr(vae, "to"):
+        return False
+    try:
+        import torch
+
+        vae.to(memory_format = torch.channels_last)
+        return True
+    except Exception as exc:  # noqa: BLE001 — optimisation only
+        _warn(logger, "channels_last", exc)
+        return False
+
+
+def _compile_repeated_blocks(pipe: Any, logger: Any) -> bool:
+    transformer = getattr(pipe, "transformer", None)
+    fn = getattr(transformer, "compile_repeated_blocks", None)
+    if not callable(fn):
+        return False
+    try:
+        fn(fullgraph = True, dynamic = True)
+        return True
+    except Exception as exc:  # noqa: BLE001 — optimisation only
+        _warn(logger, "compile_repeated_blocks", exc)
+        return False
+
+
+def _enable_tf32(logger: Any) -> bool:
+    try:
+        import torch
+
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        return True
+    except Exception as exc:  # noqa: BLE001 — optimisation only
+        _warn(logger, "tf32", exc)
+        return False
+
+
+def _fuse_qkv(pipe: Any, logger: Any) -> bool:
+    for owner in (pipe, getattr(pipe, "transformer", None)):
+        fn = getattr(owner, "fuse_qkv_projections", None)
+        if callable(fn):
+            try:
+                fn()
+                return True
+            except Exception as exc:  # noqa: BLE001 — optimisation only
+                _warn(logger, "fuse_qkv_projections", exc)
+                return False
+    return False
+
+
+def _warn(logger: Any, what: str, exc: Exception) -> None:
+    if logger is not None:
+        logger.warning("diffusion.speed: %s failed: %s", what, exc)

studio/backend/models/inference.py

@@ -1708,6 +1708,12 @@ class DiffusionLoadRequest(BaseModel):
                       "cut), low_vram (offload every component, lowest VRAM, slower). "
                       "Overrides cpu_offload when set.",
     )
+    speed_mode: Optional[Literal["off", "default", "max"]] = Field(
+        None,
+        description = "Opt-in speed optims (default off -> bit-identical output): "
+                      "default (channels_last + regional torch.compile where eligible), "
+                      "max (also TF32 + fused QKV).",
+    )
 
 
 class DiffusionGenerateRequest(BaseModel):
@@ -1804,7 +1810,11 @@ class DiffusionStatusResponse(BaseModel):
     dtype: Optional[str] = Field(None, description = "Compute dtype")
     cpu_offload: bool = Field(False, description = "Whether CPU offload is engaged")
     offload_policy: Optional[str] = Field(
-        None, description = "Resolved offload policy: none | model | sequential"
+        None, description = "Resolved offload policy: none | group | model | sequential"
     )
     vae_tiling: bool = Field(False, description = "Whether VAE tiling/slicing is enabled")
     memory_mode: Optional[str] = Field(None, description = "Requested memory mode")
+    speed_mode: Optional[str] = Field(None, description = "Requested speed mode")
+    speed_optims: list[str] = Field(
+        default_factory = list, description = "Speed optimisations actually engaged"
+    )

studio/backend/routes/inference.py

@@ -10083,6 +10083,7 @@ async def load_diffusion_model(
             hf_token = request.hf_token,
             cpu_offload = request.cpu_offload,
             memory_mode = request.memory_mode,
+            speed_mode = request.speed_mode,
         )
         return DiffusionStatusResponse(**status_dict)
     except (ValueError, FileNotFoundError) as exc:

studio/backend/tests/test_diffusion_backend.py

@@ -840,6 +840,19 @@ def test_load_explicit_cpu_offload_engages_model_offload_on_cuda(fake_runtime, t
     assert status["offload_policy"] == "model" and status["cpu_offload"] is True
 
 
+def test_load_speed_mode_threads_and_defaults_off(fake_runtime, tmp_path):
+    # No speed_mode -> off, no optimisations engaged (the bit-identical default).
+    (tmp_path / "m.gguf").write_bytes(b"x")
+    backend = DiffusionBackend()
+    status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image")
+    assert status["speed_mode"] == "off" and status["speed_optims"] == []
+    # An explicit speed_mode threads through to status (engaged optims are GPU-verified).
+    status2 = backend.load_pipeline(
+        str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "max"
+    )
+    assert status2["speed_mode"] == "max"
+
+
 def test_load_fast_mode_stays_resident_on_cuda(fake_runtime, tmp_path, monkeypatch):
     (tmp_path / "m.gguf").write_bytes(b"x")
     backend = DiffusionBackend()

studio/backend/tests/test_diffusion_speed.py

@@ -0,0 +1,160 @@
+# 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 opt-in diffusion speed layer (``diffusion_speed.py``).
+
+Hermetic: torch is stubbed via ``sys.modules`` only where a path needs it, so the
+gating logic and the best-effort applier run without a GPU or real diffusers.
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+from core.inference.diffusion_speed import (
+    SPEED_DEFAULT,
+    SPEED_MAX,
+    SPEED_OFF,
+    apply_speed_optims,
+    compile_eligible,
+    normalize_speed_mode,
+)
+
+
+def _target(*, device = "cuda", dtype = "bfloat16", compile_ok = True):
+    return types.SimpleNamespace(
+        device = device, dtype = dtype, supports_default_torch_compile = compile_ok,
+    )
+
+
+def _family(*, compile_ok = True):
+    return types.SimpleNamespace(supports_torch_compile = compile_ok)
+
+
+def _stub_torch(monkeypatch):
+    torch = types.ModuleType("torch")
+    torch.bfloat16 = "bfloat16"  # _is_bfloat16 compares by identity then str fallback
+    torch.channels_last = "channels_last"
+    torch.backends = types.SimpleNamespace(
+        cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(allow_tf32 = False)),
+        cudnn = types.SimpleNamespace(allow_tf32 = False),
+    )
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    return torch
+
+
+# ── normalisation ─────────────────────────────────────────────────────────────
+
+
+def test_normalize_speed_mode():
+    assert normalize_speed_mode(None) == SPEED_OFF
+    assert normalize_speed_mode("") == SPEED_OFF
+    assert normalize_speed_mode("MAX") == SPEED_MAX
+    with pytest.raises(ValueError):
+        normalize_speed_mode("ludicrous")
+
+
+# ── compile gating ────────────────────────────────────────────────────────────
+
+
+def test_compile_eligible_requires_non_gguf_bf16_cuda_friendly(monkeypatch):
+    _stub_torch(monkeypatch)
+    # The happy path: non-GGUF, bf16, CUDA, compile-friendly family.
+    assert compile_eligible(_target(), is_gguf = False, family = _family()) is True
+    # GGUF is never compiled.
+    assert compile_eligible(_target(), is_gguf = True, family = _family()) is False
+    # fp16 (non-bf16) is excluded.
+    assert compile_eligible(_target(dtype = "float16"), is_gguf = False, family = _family()) is False
+    # A family flagged not compile-friendly (Z-Image) is excluded.
+    assert compile_eligible(_target(), is_gguf = False, family = _family(compile_ok = False)) is False
+    # No compile support (e.g. ROCm/XPU/MPS) is excluded.
+    assert compile_eligible(_target(compile_ok = False), is_gguf = False, family = _family()) is False
+
+
+# ── applier ───────────────────────────────────────────────────────────────────
+
+
+class _Pipe:
+    def __init__(self, *, with_compile = False, with_fuse = False) -> None:
+        self.vae = types.SimpleNamespace(mem_format = None, to = self._vae_to)
+        self.transformer = types.SimpleNamespace()
+        if with_compile:
+            self.transformer.compile_repeated_blocks = self._compile
+        if with_fuse:
+            self.fuse_qkv_projections = self._fuse
+        self.compiled = False
+        self.fused = False
+
+    def _vae_to(self, *, memory_format):
+        self.vae.mem_format = memory_format
+
+    def _compile(self, *, fullgraph, dynamic):
+        self.compiled = True
+
+    def _fuse(self):
+        self.fused = True
+
+
+def test_speed_off_applies_nothing(monkeypatch):
+    _stub_torch(monkeypatch)
+    pipe = _Pipe(with_compile = True, with_fuse = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_OFF
+    )
+    assert applied == {"channels_last": False, "tf32": False, "fused_qkv": False, "compiled": False}
+    assert pipe.vae.mem_format is None and pipe.compiled is False
+
+
+def test_speed_default_channels_last_and_compile_when_eligible(monkeypatch):
+    torch = _stub_torch(monkeypatch)
+    pipe = _Pipe(with_compile = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
+    )
+    assert applied["channels_last"] is True and pipe.vae.mem_format == torch.channels_last
+    assert applied["compiled"] is True and pipe.compiled is True
+    # default does not flip TF32 or fuse QKV.
+    assert applied["tf32"] is False and applied["fused_qkv"] is False
+
+
+def test_speed_default_skips_compile_for_gguf(monkeypatch):
+    _stub_torch(monkeypatch)
+    pipe = _Pipe(with_compile = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT
+    )
+    assert applied["channels_last"] is True  # lossless layout still applies
+    assert applied["compiled"] is False and pipe.compiled is False  # GGUF never compiles
+
+
+def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch):
+    torch = _stub_torch(monkeypatch)
+    pipe = _Pipe(with_compile = True, with_fuse = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX
+    )
+    assert applied["tf32"] is True and torch.backends.cuda.matmul.allow_tf32 is True
+    assert applied["fused_qkv"] is True and pipe.fused is True
+
+
+def test_speed_max_tf32_only_on_cuda(monkeypatch):
+    _stub_torch(monkeypatch)
+    pipe = _Pipe()
+    applied = apply_speed_optims(
+        pipe, _target(device = "mps", compile_ok = False), is_gguf = True,
+        family = _family(), speed_mode = SPEED_MAX,
+    )
+    assert applied["tf32"] is False  # not CUDA -> no TF32
+
+
+def test_apply_tolerates_missing_optims(monkeypatch):
+    _stub_torch(monkeypatch)
+    # A bare pipe (no vae.to, no compile, no fuse) must not crash.
+    bare = types.SimpleNamespace(vae = None, transformer = types.SimpleNamespace())
+    applied = apply_speed_optims(
+        bare, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX
+    )
+    assert applied["channels_last"] is False and applied["fused_qkv"] is False