unsloth/studio/backend/core/inference/diffusion_speed.py

# 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
near-lossless speedups in the order the diffusers guides recommend
(channels_last + cudnn.benchmark -> compile, with TF32 / fused-QKV under "max"):

  off     - nothing (default; bit-identical reference).
  eager   - everything lossless EXCEPT torch.compile: channels_last VAE +
            cudnn.benchmark + the attention backend + the shared eager monkey-patches
            (fused RMSNorm / AdaLayerNorm + per-arch addcmul fusions, see
            diffusion_eager_patches.py / diffusion_arch_patches.py). The fast first-image
            / casual-use path -- no compile tax to amortise.
  default - LIGHT compile. For a GGUF model: channels_last + cudnn.benchmark +
            torch.compile of ONLY the dequant op chain
            (``torch.compile(dequantize_gguf_tensor, dynamic=True)``) -- the dequant is
            ~70-80% of eager GGUF time, so fusing it gives ~1.24-1.64x for a small
            one-time compile (~7.5-10.4s) and ZERO extra VRAM, resolution-invariant
            (the dequant inputs are fixed-shape weights). For a dense (non-GGUF) model
            there is no dequant, so ``default`` falls back to regional torch.compile of
            the denoiser's repeated block (the only compile lever a dense model has).
  max     - the FULL torch.compile: regional max-autotune compile of the denoiser's
            repeated block (which fuses the GGUF dequant AND the matmul/norm/elementwise
            in one graph -- ~3.2x on the GGUF Z-Image transformer, PSNR ~36 dB vs eager,
            well above the Q4 noise floor) plus TF32 matmul and fused QKV projections.

Tier rationale: ``default`` is the cheap, always-amortising compile (compile just the
hot GGUF dequant; the block stays eager) so the first image is fast and VRAM is
untouched; ``max`` pays the larger regional-compile tax for the bigger warm speedup.
The compiled dequant is deliberately skipped under ``max`` -- the regional block compile
subsumes the dequant fusion (a separately-compiled dequant would be traced into that
graph and break it), so ``max`` runs the stock dequant and lets the block compile it. The
per-family ``supports_torch_compile`` flag and the bf16/CUDA checks gate regional compile.

The backend flags this layer flips (TF32, cudnn.benchmark) are PROCESS-WIDE, so
``snapshot_backend_flags`` / ``restore_backend_flags`` let the caller capture the
prior values at load and restore them at unload, keeping a later ``off`` load
bit-identical instead of inheriting a previous ``max`` run's globals. torch is
imported lazily.
"""

from __future__ import annotations

from typing import Any, Optional

from . import diffusion_gguf_compile as gguf_compile

SPEED_OFF = "off"
SPEED_EAGER = "eager"
SPEED_DEFAULT = "default"
SPEED_MAX = "max"
SPEED_MODES = (SPEED_OFF, SPEED_EAGER, SPEED_DEFAULT, SPEED_MAX)


def snapshot_backend_flags() -> Optional[dict]:
    """Capture the process-wide torch backend flags this layer may mutate, so the
    caller can restore them on unload. None if torch is unavailable. Each flag is read
    defensively so a build/platform missing one (e.g. no cuda.matmul on CPU/MPS) still
    captures the rest -- otherwise a single missing attribute would skip the whole
    snapshot and a real mutated flag would leak."""
    try:
        import torch
    except Exception:  # noqa: BLE001 — no torch -> nothing to snapshot/restore
        return None
    state: dict[str, bool] = {}
    matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
    if matmul is not None and hasattr(matmul, "allow_tf32"):
        state["matmul_tf32"] = bool(matmul.allow_tf32)
    if matmul is not None and hasattr(matmul, "allow_fp16_accumulation"):
        state["matmul_fp16_accum"] = bool(matmul.allow_fp16_accumulation)
    cudnn = getattr(torch.backends, "cudnn", None)
    if cudnn is not None:
        if hasattr(cudnn, "allow_tf32"):
            state["cudnn_tf32"] = bool(cudnn.allow_tf32)
        if hasattr(cudnn, "benchmark"):
            state["cudnn_benchmark"] = bool(cudnn.benchmark)
    inductor_cfg = _inductor_config()
    if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
        state["inductor_emulate_precision_casts"] = bool(inductor_cfg.emulate_precision_casts)
    return state


def restore_backend_flags(state: Optional[dict]) -> None:
    """Restore the flags captured by ``snapshot_backend_flags``. No-op on None. Each
    flag is restored independently so one failure can't leave the others leaked."""
    if not state:
        return
    try:
        import torch
    except Exception:  # noqa: BLE001 — no torch -> nothing to restore
        return

    def _set(obj: Any, attr: str, key: str) -> None:
        if obj is not None and key in state and hasattr(obj, attr):
            try:
                setattr(obj, attr, state[key])
            except Exception:  # noqa: BLE001 — best-effort per-flag restore
                pass

    matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
    _set(matmul, "allow_tf32", "matmul_tf32")
    _set(matmul, "allow_fp16_accumulation", "matmul_fp16_accum")
    cudnn = getattr(torch.backends, "cudnn", None)
    _set(cudnn, "allow_tf32", "cudnn_tf32")
    _set(cudnn, "benchmark", "cudnn_benchmark")
    _set(_inductor_config(), "emulate_precision_casts", "inductor_emulate_precision_casts")


def _inductor_config() -> Any:
    """``torch._inductor.config`` or None. Resolved as attributes off the imported torch
    module (real torch exposes ``_inductor`` directly after ``import torch``) rather
    than a submodule import, so a stubbed/partial torch (tests, exotic builds) cleanly
    reports None instead of picking a stale real module out of ``sys.modules``."""
    try:
        import torch
        return getattr(getattr(torch, "_inductor", None), "config", None)
    except Exception:  # noqa: BLE001 — no inductor -> nothing to snapshot/set
        return None


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 resolve_speed_mode(
    value: Optional[str],
    *,
    is_gguf: bool,
    dense_default: str = SPEED_OFF,
) -> str:
    """The effective speed mode when the caller leaves it UNSET (``None``).

    A GGUF model defaults to ``default``: it compiles only the hot dequant op chain
    (~70-80% of eager GGUF time) for ~1.24-1.64x at a small one-time compile and zero
    extra VRAM -- a cheap, always-amortising win whose numeric perturbation sits well
    below the quantisation noise floor (the dequant graph is unchanged, just
    Inductor-fused). A dense (non-GGUF) model resolves to ``dense_default``: the image
    backend keeps ``off`` at load (bit-identical first generations, with its own
    deferred engagement after repeated use), while the video backend passes ``default``
    -- a clip denoise runs long enough that the one-time compile always amortises
    within a single generation. An explicit value -- including ``"off"`` -- is always
    honored verbatim."""
    if value is None:
        return SPEED_DEFAULT if is_gguf else dense_default
    return normalize_speed_mode(value)


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 bf16
    transformer, on a compile-friendly family. ``is_gguf`` no longer disqualifies:
    ``compile_repeated_blocks`` runs fine on the GGUF transformer (the per-op
    dequant stays eager, the rest of the block compiles) and is ~2.3x faster, so it
    is kept only for signature/logging compatibility."""
    del is_gguf  # GGUF is compile-eligible now; param kept for call-site compat.
    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,
    cache_active: bool = False,
    offload_active: bool = False,
    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.

    ``offload_active`` is the planned offload policy != none: group/model/sequential
    offloading installs ``@torch.compiler.disable``d onload hooks, so the compile must
    drop ``fullgraph`` (same reason as an active step cache) or it crashes at the first
    denoise step."""
    applied = {
        "channels_last": False,
        "cudnn_benchmark": False,
        "tf32": False,
        "fp16_accum": False,
        "fused_qkv": False,
        "compiled": False,
        "compiled_dequant": False,
    }
    mode = normalize_speed_mode(speed_mode)
    # TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max,
    # cudnn.benchmark on any non-off CUDA load). The caller snapshots them before this
    # call and restores on unload / failed load via snapshot_backend_flags /
    # restore_backend_flags, so a later `off` load -- or chat inference in the same
    # process -- never inherits them. We keep no separate bookkeeping here.
    if mode == SPEED_OFF:
        return applied

    on_cuda = getattr(target, "device", None) == "cuda"
    family_allows_compile = bool(getattr(family, "supports_torch_compile", True))

    # Lossless: a channels-last VAE speeds up its convolutions with no numeric change.
    applied["channels_last"] = _vae_channels_last(pipe, logger)

    # Near-lossless: let cuDNN autotune the fixed-shape VAE convs (CUDA only). It may
    # pick a different conv algorithm, so it is a "default"-tier (not bit-identical) win.
    if on_cuda:
        applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger)

    # Consumer-only: fp16 GEMMs accumulate in fp16 (~2x on GeForce-class parts, whose
    # fp32-accumulate rate is halved; datacenter HBM parts gain nothing and keep the
    # safer fp32 accumulate). Only affects fp16 matmuls -- bf16 loads were measured
    # bit-identical with the flag on across every family (36/36 same-seed A/B cases),
    # so on the quality-neutral tiers the flag engages only when the compute dtype is
    # NOT fp16. On an fp16 pipeline (the pre-Ampere fallback dtype) the same harness
    # measured real same-seed drift (mean 2-5% on SDXL / FLUX), so fp16 compute gets
    # the 2x accumulate only under ``max``, the tier that already trades exactness for
    # speed. Guarded by a per-family deny-list fed by the overflow validation harness
    # and the UNSLOTH_DISABLE_FP16_ACCUM kill switch.
    if on_cuda:
        applied["fp16_accum"] = _enable_fp16_accumulation(
            family, logger, dtype = getattr(target, "dtype", None), speed_mode = mode
        )

    # --- the compile lever, remapped per tier ----------------------------------------
    # default = LIGHT compile: for a GGUF model, compile ONLY the dequant op chain
    #   (~70-80% of eager GGUF time) -- cheap, VRAM-free, resolution-invariant; the
    #   transformer block stays eager. A dense model has no dequant, so default falls
    #   back to the regional block compile (its only compile lever).
    # max = FULL compile: regional max-autotune compile of the repeated denoiser block
    #   (fuses dequant + matmul + norm + elementwise in one graph). It subsumes the
    #   dequant fusion, so we do NOT also install the standalone compiled dequant here.
    # eager = no compile at all.
    if mode == SPEED_DEFAULT:
        if is_gguf and on_cuda and family_allows_compile:
            applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger)
        elif compile_eligible(target, is_gguf = is_gguf, family = family):
            applied["compiled"] = _compile_repeated_blocks(
                pipe,
                logger,
                max_autotune = False,
                cache_active = cache_active,
                offload_active = offload_active,
            )
    elif mode == SPEED_MAX and compile_eligible(target, is_gguf = is_gguf, family = family):
        applied["compiled"] = _compile_repeated_blocks(
            pipe,
            logger,
            max_autotune = True,
            cache_active = cache_active,
            offload_active = offload_active,
        )

    if mode == SPEED_MAX:
        # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed.
        if on_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 _denoiser_dits(pipe: Any) -> list:
    """Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second
    expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch
    CFG, an MoE ``transformer_2``). Speed / attention optims must reach ALL of them -- mirroring
    the offload path (diffusion_memory streams the same set) -- else the second DiT runs
    eager / native for every generation while status over-reports the optimisation as engaged."""
    dits: list = []
    for attr in ("transformer", "transformer_2", "unconditional_transformer"):
        m = getattr(pipe, attr, None)
        if m is not None and m not in dits:
            dits.append(m)
    return dits


def _compile_repeated_blocks(
    pipe: Any,
    logger: Any,
    *,
    max_autotune: bool = False,
    cache_active: bool = False,
    offload_active: bool = False,
) -> bool:
    dits = [
        t for t in _denoiser_dits(pipe) if callable(getattr(t, "compile_repeated_blocks", None))
    ]
    if not dits:
        return False
    # default: mode="default" + dynamic=True -- fast cold start, robust to resolution
    # changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False --
    # Triton autotuning for a few % more on GEMM/conv-heavy models, at a much longer
    # compile and a recompile per new resolution. The CUDA-graph modes (reduce-overhead
    # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled
    # block because its static output buffer is overwritten across denoise steps.
    #
    # fullgraph drops to False when a step cache OR CPU offloading is engaged: both insert
    # an ``@torch.compiler.disable``d function into the forward -- FBCache's per-step
    # decision, and group/model/sequential offload's ``ModuleGroup.onload_`` streaming hook
    # -- i.e. a graph break, which fullgraph=True rejects ("Skip inlining
    # torch.compiler.disable()d function"). The break is cheap and the rest of the block
    # still compiles.
    kwargs: dict[str, Any] = {
        "fullgraph": not (cache_active or offload_active),
        "dynamic": not max_autotune,
    }
    if max_autotune:
        kwargs["mode"] = "max-autotune-no-cudagraphs"
    try:
        import torch

        # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block
        # shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default
        # recompile_limit of 8. Once the limit is hit a resident load hard-errors under
        # fullgraph (and an offload/cache load silently drops the overflow blocks to eager),
        # so raise it well past that (64) for headroom on larger heterogeneous DiTs. This is
        # diffusers' own documented fix for regional-compile recompilation (their guide bumps
        # cache_size_limit). Deliberately NOT force_parameter_static_shapes=False: it doesn't
        # cut the variant count here and makes each compile ~6x slower (24s -> 143s cold).
        dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None)
        if dynamo_cfg is not None:
            for _limit_attr in ("recompile_limit", "cache_size_limit"):  # name varies by torch ver
                if hasattr(dynamo_cfg, _limit_attr):
                    setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
        # Match eager's intermediate rounding inside inductor's fused pointwise kernels:
        # by default they keep chains in fp32 where eager materialises bf16 between ops,
        # a per-forward rounding delta that a multi-step denoise amplifies chaotically.
        # Measured (B200, scripts/image_speedmem_bench.py, pairwise LPIPS of the
        # compiled tier vs the same-stack eager tier): Qwen-Image 0.019 -> 0.006 at
        # identical speed, FLUX.1-dev 0.046 -> 0.029 at +2% step time, FLUX.2-klein
        # 0.018 -> 0.017 at identical speed; on the video DiT (HunyuanVideo-1.5-720p)
        # full-clip LPIPS vs bit-exact drops 0.221 -> 0.052 at zero cost. Process-
        # global, so snapshot_backend_flags carries it and unload restores the prior
        # value.
        inductor_cfg = _inductor_config()
        if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
            inductor_cfg.emulate_precision_casts = True
    except Exception as exc:  # noqa: BLE001 — optimisation only
        _warn(logger, "compile_repeated_blocks", exc)
        return False
    # Compile every denoiser DiT (a dual-DiT family such as Ideogram runs both each step); a
    # per-DiT failure degrades that one to eager without dropping the others.
    engaged = False
    for transformer in dits:
        try:
            transformer.compile_repeated_blocks(**kwargs)
            engaged = True
        except Exception as exc:  # noqa: BLE001 — optimisation only
            _warn(logger, "compile_repeated_blocks", exc)
            continue
        # A step cache engaged BEFORE this compile (the production load order) has
        # already wrapped each block's forward in a @torch.compiler.disable'd hook, so
        # the compute branch would run eager on every non-skipped step and forfeit the
        # regional compile entirely. Re-point the hooks' inner forward at compiled
        # wrappers; no-op when no cache hooks are installed. The toggle path (cache
        # engaged after load) is armed by apply_step_cache instead. Lazy import:
        # diffusion_cache imports nothing from this module, but keep the dependency
        # one-directional at import time.
        try:
            from .diffusion_cache import _compile_hooked_block_inners
            _compile_hooked_block_inners(transformer, logger)
        except Exception as exc:  # noqa: BLE001 — optimisation only
            _warn(logger, "cache-hook inner compile", exc)
    return engaged


def _enable_cudnn_benchmark(logger: Any) -> bool:
    try:
        import torch
        torch.backends.cudnn.benchmark = True
        return True
    except Exception as exc:  # noqa: BLE001 — optimisation only
        _warn(logger, "cudnn_benchmark", 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


# Families the overflow validation harness (scripts/fp16_accum_validate.py) found to
# produce non-finite activations or NEW black frames under fp16 accumulation. Empty by
# measurement: across all six families the harness found no overflow anywhere -- bf16
# loads are bit-identical with the flag on, and fp16 loads stay finite (their same-seed
# drift is why fp16 compute is additionally gated to the ``max`` tier below).
_FP16_ACCUM_DENY: frozenset[str] = frozenset()


def _enable_fp16_accumulation(
    family: Any,
    logger: Any,
    *,
    dtype: Any = None,
    speed_mode: Optional[str] = None,
) -> bool:
    """Turn on fp16-accumulated fp16 GEMMs for consumer GPUs, where they run ~2x the
    fp32-accumulate rate (datacenter HBM parts are not throughput-nerfed, so they keep
    the safer default). Gated on: the torch build exposing the flag (2.10+), a
    consumer-class device, the family not being deny-listed by the overflow harness,
    the UNSLOTH_DISABLE_FP16_ACCUM kill switch being unset, and -- when the pipeline
    compute dtype IS fp16, the only case where the accumulator width changes results --
    the ``max`` tier (measured same-seed drift: mean 2-5%; bf16 loads are bit-identical
    so they engage on any tier). The caller's snapshot/restore pair returns the
    process-wide flag to its prior value on unload."""
    import os

    if os.environ.get("UNSLOTH_DISABLE_FP16_ACCUM", "").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    ):
        return False
    name = str(getattr(family, "name", family or "")).lower()
    if name in _FP16_ACCUM_DENY:
        return False
    if str(dtype).replace("torch.", "") == "float16" and speed_mode != SPEED_MAX:
        return False
    try:
        import torch

        matmul = torch.backends.cuda.matmul
        if not hasattr(matmul, "allow_fp16_accumulation"):
            return False
        from .diffusion_transformer_quant import _is_consumer_gpu

        if not _is_consumer_gpu():
            return False
        matmul.allow_fp16_accumulation = True
        return True
    except Exception as exc:  # noqa: BLE001 — optimisation only
        _warn(logger, "fp16_accum", exc)
        return False


def _fuse_qkv(pipe: Any, logger: Any) -> bool:
    # Prefer the pipe-level fuse (it covers every component the pipe knows about); else fuse each
    # denoiser DiT directly so a dual-DiT family (Ideogram) fuses BOTH experts, not just the first.
    fn = getattr(pipe, "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
    engaged = False
    for transformer in _denoiser_dits(pipe):
        tfn = getattr(transformer, "fuse_qkv_projections", None)
        if callable(tfn):
            try:
                tfn()
                engaged = True
            except Exception as exc:  # noqa: BLE001 — optimisation only
                _warn(logger, "fuse_qkv_projections", exc)
    return engaged


def _warn(logger: Any, what: str, exc: Exception) -> None:
    if logger is not None:
        logger.warning("diffusion.speed: %s failed: %s", what, exc)