diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 45104bd1c3..76e381fb87 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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), } diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 805ea7f04c..ba9125fabf 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/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, ), ) diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py new file mode 100644 index 0000000000..a731c3124f --- /dev/null +++ b/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) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index f30f8a31f6..892f25172f 100644 --- a/studio/backend/models/inference.py +++ b/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" + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 83726516bf..0265ac6a9a 100644 --- a/studio/backend/routes/inference.py +++ b/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: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 57be08fb4d..8c157c8875 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/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() diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py new file mode 100644 index 0000000000..72a4f7899d --- /dev/null +++ b/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