# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Select the diffusion transformer's attention backend.

diffusers exposes a unified ``transformer.set_attention_backend(name)`` dispatcher that
swaps the scaled-dot-product-attention kernel, validating hardware/package requirements at
set time and otherwise leaving the default (``native`` = ``F.scaled_dot_product_attention``).
Attention is memory-bandwidth bound, so a better kernel is a real end-to-end win that is
orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never
touches) and composes with torch.compile.

  auto  - the best *exact* (non-quantized) backend for the device. On NVIDIA CUDA that is
          cuDNN's fused attention (``_native_cudnn``), measured ~1.18x end-to-end on a B200
          with LPIPS ~0.004 vs the default (below the compile/quant noise floor). On
          AMD/Intel/Apple/CPU it stays ``native`` (the dispatcher already routes those).
          ``auto`` only upgrades when a speed profile is active, so ``speed_mode=off`` stays
          bit-identical.
  native - force the default SDPA (bit-identical reference).
  cudnn  - cuDNN fused attention (exact; NVIDIA).
  flash / flash3 / flash4 - FlashAttention 2 / 3 (Hopper) / 4 (SM100); exact, kernel-gated.
  sage   - SageAttention (INT8 QK); quantized, a small quality cost, consumer-friendly.
  xformers / aiter - memory-efficient (NVIDIA) / AITER (AMD ROCm).

Best-effort: an unavailable backend (missing kernel / wrong arch) is caught and the load
falls back to the diffusers default rather than failing. torch/diffusers imported lazily.
"""

from __future__ import annotations

from typing import Any, Optional

ATTN_AUTO = "auto"
ATTN_NATIVE = "native"

# User-facing alias -> the diffusers dispatcher backend name.
_ALIASES: dict[str, str] = {
    "native": "native",
    "sdpa": "native",
    "cudnn": "_native_cudnn",
    "flash": "flash",
    "flash2": "flash",
    "flash3": "_flash_3_hub",
    "flash4": "flash_4_hub",
    "sage": "sage",
    "xformers": "xformers",
    "aiter": "aiter",
}
ATTN_ALIASES = (ATTN_AUTO,) + tuple(dict.fromkeys(_ALIASES))


def normalize_attention_backend(value: Optional[str]) -> Optional[str]:
    """Lower/strip a requested attention backend; None / "" / "auto" -> "auto".

    Raises ValueError for an unsupported alias so a bad request is rejected cheaply."""
    if value is None:
        return ATTN_AUTO
    normalized = str(value).strip().lower()
    if not normalized:
        return ATTN_AUTO
    if normalized not in ATTN_ALIASES:
        raise ValueError(
            f"Unsupported attention_backend '{value}'. Use one of: {', '.join(ATTN_ALIASES)}."
        )
    return normalized


# Backends diffusers validates only by *package* at set time (``_check_attention_backend_
# requirements`` checks the ``kernels`` install, not the GPU), but whose kernels need a
# specific CUDA arch at run time -- so an explicit request on the wrong card loads/sets fine
# and then crashes mid-generation. Gate them up front by a (min, max-exclusive) compute
# capability range. FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel, so it
# needs an upper bound: an explicit flash3 on a B200 (SM100) must drop to native instead of
# setting fine then crashing at generation. FlashAttention 4 is Blackwell+ (no upper bound).
_ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = {
    "_flash_3_hub": ((9, 0), (10, 0)),  # FlashAttention 3 -> Hopper (SM90) only
    "flash_4_hub": ((10, 0), None),  # FlashAttention 4 -> Blackwell (SM100)+
}


def _cuda_capability() -> Optional[tuple[int, int]]:
    """(major, minor) compute capability of the active CUDA device, or None if unknown."""
    try:
        import torch
        if not torch.cuda.is_available():
            return None
        return tuple(torch.cuda.get_device_capability())  # type: ignore[return-value]
    except Exception:  # noqa: BLE001
        return None


def _backend_arch_supported(backend: str) -> bool:
    """False only when ``backend`` needs a CUDA arch outside this device's supported range.

    Unknown capability (no CUDA / detection failure) returns True so we never block on a
    guess -- diffusers' own set-time check still guards the package, and a genuine run-time
    failure falls back to native."""
    bounds = _ARCH_CAPABILITY.get(backend)
    if bounds is None:
        return True
    have = _cuda_capability()
    if have is None:
        return True
    low, high = bounds
    return have >= low and (high is None or have < high)


def _is_cuda_nvidia(target: Any) -> bool:
    """CUDA device on an NVIDIA (non-ROCm) build -- where cuDNN attention applies."""
    if getattr(target, "device", None) != "cuda":
        return False
    try:
        import torch
        return getattr(torch.version, "hip", None) is None
    except Exception:  # noqa: BLE001
        return False


def select_attention_backend(
    target: Any, requested: Optional[str], *, speed_active: bool
) -> Optional[str]:
    """The dispatcher backend name to apply, or None to leave the diffusers default.

    An explicit alias is honored verbatim (apply falls back if its kernel is unavailable).
    ``auto`` upgrades to cuDNN on NVIDIA CUDA only when a speed profile is active (so
    ``off`` stays bit-identical); everywhere else it returns None (native default)."""
    alias = normalize_attention_backend(requested)
    if alias != ATTN_AUTO:
        backend = _ALIASES[alias]
        if backend == "native":
            return None
        # Every explicit kernel here (cuDNN / flash* / sage) is CUDA+NVIDIA-only; on
        # ROCm / MPS / CPU diffusers accepts the name at set time and the first
        # generation crashes, so drop to the native default up front.
        if not _is_cuda_nvidia(target):
            return None
        # An arch-gated kernel (flash3/flash4) on a card that can't run it would set fine
        # then crash mid-generation, so drop it to the native default up front.
        if not _backend_arch_supported(backend):
            return None
        # cuDNN fused SDPA needs Ampere+ (SM80); diffusers accepts it on pre-SM80 cards
        # (T4/V100) then fails at the first generation, so apply the same gate to an
        # explicit cuDNN request as the auto path already does.
        if backend == "_native_cudnn" and not _cudnn_attention_supported():
            return None
        return backend
    # auto
    if speed_active and _is_cuda_nvidia(target) and _cudnn_attention_supported():
        return "_native_cudnn"
    return None


def _cudnn_attention_supported() -> bool:
    """cuDNN fused SDPA needs Ampere+ (SM80). On pre-SM80 NVIDIA cards (T4 SM75 /
    V100 SM70) diffusers accepts ``_native_cudnn`` at set time but the kernel fails at
    the first generation, so gate the auto-cuDNN upgrade on capability. Unknown
    capability allows it (diffusers' set-time check + the run-time fallback still guard)."""
    have = _cuda_capability()
    return have is None or have >= (8, 0)


def apply_attention_backend(
    pipe: Any,
    backend: Optional[str],
    *,
    logger: Any = None,
) -> Optional[str]:
    """Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher.

    Returns the backend actually engaged, or None when left at the native default (either
    because ``backend`` was None or because the requested kernel was unavailable -> graceful
    fallback, never a load failure).

    diffusers keeps a *process-wide* active attention backend that ``set_attention_backend``
    also updates, and a fresh transformer's processors follow it (their ``_attention_backend``
    defaults to None). So a load that wants native must restore it explicitly: otherwise it
    silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile),
    breaking the bit-identical/``off`` guarantee. Best-effort throughout."""
    transformer = getattr(pipe, "transformer", None)
    fn = getattr(transformer, "set_attention_backend", None)
    if not callable(fn):
        return None
    if backend is not None:
        try:
            fn(backend)
            # set_attention_backend also pins the backend in diffusers' process-wide
            # registry. This transformer's own processors keep it locally (their
            # _attention_backend is now explicit), so reset the global default back to
            # native -- otherwise a later component whose processors are unconfigured
            # (backend None) silently inherits this kernel.
            _reset_global_backend_to_native(logger)
            if logger is not None:
                logger.info("diffusion.attention: backend=%s", backend)
            return backend
        except Exception as exc:  # noqa: BLE001 — unavailable kernel -> restore native below
            _warn(logger, backend, exc)
    # No backend requested, or the requested one failed: pin the native default so a stale
    # process-wide backend from a previous load can't leak into this one.
    _restore_native_backend(fn, logger)
    return None


def _active_attention_backend() -> Optional[str]:
    """The diffusers process-wide active attention backend name, or None if undeterminable."""
    try:
        from diffusers.models.attention_dispatch import _AttentionBackendRegistry

        # get_active_backend() returns a (AttentionBackendName, fn) tuple (or None), so
        # take element 0 and read its .value (e.g. "native"); reading .value off the
        # tuple itself would yield a junk string that never compares equal to a name.
        active = _AttentionBackendRegistry.get_active_backend()
        if active is None:
            return None
        name = active[0] if isinstance(active, tuple) else active
        return getattr(name, "value", str(name))
    except Exception:  # noqa: BLE001
        return None


def _reset_global_backend_to_native(logger: Any) -> None:
    """Reset diffusers' process-wide active attention backend to native after a
    successful per-transformer set, so a later component whose processors are
    unconfigured (backend None) does not inherit this transformer's kernel. The
    transformer's own processors keep the backend just set. Best-effort and silent:
    if the diffusers internals move, the prior (leaking) behavior is unchanged."""
    if _active_attention_backend() == ATTN_NATIVE:
        return
    try:
        from diffusers.models.attention_dispatch import (
            AttentionBackendName,
            _AttentionBackendRegistry,
        )
        _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE)
    except Exception:  # noqa: BLE001 — best-effort; leave the global as-is on any change
        pass


def _restore_native_backend(set_backend_fn: Any, logger: Any) -> None:
    """Force the native default when the global active backend isn't already native."""
    if _active_attention_backend() == ATTN_NATIVE:
        return  # already native -> avoid redundant work and an extra dispatcher warning
    try:
        set_backend_fn(ATTN_NATIVE)
    except Exception as exc:  # noqa: BLE001 — best-effort restore
        _warn(logger, ATTN_NATIVE, exc)


def _warn(logger: Any, what: str, exc: Exception) -> None:
    if logger is not None:
        logger.warning("diffusion.attention: %s unavailable (%s); using default", what, exc)