kernel.py

import functools
import math
import warnings

import torch
import torch.utils.checkpoint
from torch.nn import functional as F


try:
    import config as _alpha_config
except ImportError:
    _alpha_config = None


def _kernel_config_attr(name: str, default):
    
    if _alpha_config is None:
        return default
    if not hasattr(_alpha_config, name):
        warnings.warn(
            f"[kernel.py] config module has no attribute '{name}' -- falling back to "
            f"default {default!r}. If this is unexpected, check for a casing mismatch "
            f"or rename in your config.py.",
            stacklevel=3,
        )
        return default
    return getattr(_alpha_config, name)



def _use_sdpa_config() -> bool:
    """Live read of config.use_sdpa -- default True."""
    return bool(_kernel_config_attr('use_sdpa', True))


class _UnsupportedByBackend(Exception):
    """Raised internally to signal 'this backend can't safely do this' -> fall back."""


_warned: set[str] = set()


def _warn_once(key: str, msg: str) -> None:
    if key not in _warned:
        _warned.add(key)
        warnings.warn(f"[kernel.py] {msg}", stacklevel=3)



_backend_stats: dict[str, int] = {}


def _record_backend(name: str) -> None:
    
    _record_backend_eager(name)


@torch.compiler.disable()
def _record_backend_eager(name: str) -> None:
    _backend_stats[name] = _backend_stats.get(name, 0) + 1


def get_backend_stats() -> dict[str, int]:
    """Returns a copy of accumulated counts, e.g.
    {"splash_success": 12, "splash_fallback": 3, "flash_success": 3, ...}.
    Counts accumulate across the process lifetime (or since the last
    reset_backend_stats() call) -- reset before the segment you want to
    measure if you're isolating one phase of a script."""
    return dict(_backend_stats)


def reset_backend_stats() -> None:
    _backend_stats.clear()


# ---------------------------------------------------------------------------
# Availability checks -- each cached so we only pay the import/probe cost once
# ---------------------------------------------------------------------------

@torch.compiler.allow_in_graph
@functools.lru_cache(maxsize=1)
def _flash_attn_available() -> bool:
    try:
        import flash_attn  # noqa: F401
        return True
    except ImportError:
        return False


@torch.compiler.allow_in_graph
@functools.lru_cache(maxsize=1)
def _xformers_available() -> bool:
    try:
        import xformers.ops  # noqa: F401
        return True
    except ImportError:
        return False


@torch.compiler.allow_in_graph
@functools.lru_cache(maxsize=1)
def _tpu_kernel_available() -> bool:
    try:
        from torch_xla.experimental import custom_kernel  # noqa: F401
        return True
    except ImportError:
        return False


@torch.compiler.allow_in_graph
@functools.lru_cache(maxsize=None)
def _detect_backend(device_key: str) -> str:
    
    if device_key.startswith("xla"):
        if _tpu_kernel_available():
            return "tpu"
        _warn_once("tpu_missing", "torch_xla not importable on an XLA device -- using reference attention.")
        return "reference"

    if device_key.startswith("cuda"):
        idx = int(device_key.split(":")[1]) if ":" in device_key else torch.cuda.current_device()
        major, _minor = torch.cuda.get_device_capability(idx)
        if major >= 8 and _flash_attn_available():
            return "flash"
        if _xformers_available():
            return "xformers"
        if major >= 8:
            _warn_once("flash_missing", "Ampere+ GPU detected but `flash_attn` isn't installed, and neither is `xformers` -- using reference attention (slow).")
        else:
            _warn_once("xformers_missing", "Pre-Ampere GPU (e.g. T4) detected and `xformers` isn't installed -- using reference attention (slow, high memory).")
        return "reference"

    return "reference"


def _backend_for(device: torch.device) -> str:
    if device.type == "cuda":
        return _detect_backend(f"cuda:{device.index if device.index is not None else torch.cuda.current_device()}")
    return _detect_backend(device.type)


# ---------------------------------------------------------------------------
# Reference implementation -- the ground truth every backend is validated
# against, and the universal fallback.
# ---------------------------------------------------------------------------

def _attention_block(q_block, k_block, v_block, q_start, k_start, causal, window_size, dropout_p, training, scale):
    """One tile's worth of causal/windowed attention. Kept as a standalone
    function (not a closure) so torch.utils.checkpoint can call it directly."""
    qb, kb = q_block.shape[2], k_block.shape[2]
    scores = torch.matmul(q_block, k_block.transpose(-2, -1)) * scale

    if causal or window_size is not None:
        q_idx = torch.arange(q_start, q_start + qb, device=q_block.device).view(qb, 1)
        k_idx = torch.arange(k_start, k_start + kb, device=q_block.device).view(1, kb)
        allowed = (k_idx <= q_idx) if causal else torch.ones(qb, kb, dtype=torch.bool, device=q_block.device)
        if window_size is not None:
            allowed = allowed & (q_idx - k_idx < window_size)
        mask = torch.where(allowed, torch.zeros(1, device=q_block.device), torch.full((1,), -1e4, device=q_block.device))
        scores = scores + mask.to(scores.dtype)

    weights = torch.softmax(scores.float(), dim=-1).to(scores.dtype)
    if training and dropout_p > 0:
        weights = torch.nn.functional.dropout(weights, p=dropout_p)
    return torch.matmul(weights, v_block)


@torch.compiler.disable
def _reference_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale, query_block_size=128):
    
    b, hq, t, d = q.shape
    _, hkv, tk, _ = k.shape
    if hkv != hq:
        reps = hq // hkv
        k = k.repeat_interleave(reps, dim=1)
        v = v.repeat_interleave(reps, dim=1)

    scale = softmax_scale if softmax_scale is not None else (1.0 / math.sqrt(d))

    outputs = []
    for q_start in range(0, t, query_block_size):
        q_end = min(q_start + query_block_size, t)
        q_block = q[:, :, q_start:q_end, :]

        k_lo = max(0, q_start - window_size + 1) if window_size is not None else 0
        k_hi = min(q_end, tk) if causal else tk
        k_block = k[:, :, k_lo:k_hi, :]
        v_block = v[:, :, k_lo:k_hi, :]

        
        if q_block.device.type == "xla":
            out_block = _attention_block(
                q_block, k_block, v_block, q_start, k_lo,
                causal, window_size, dropout_p, training, scale,
            )
        else:
            out_block = torch.utils.checkpoint.checkpoint(
                _attention_block, q_block, k_block, v_block, q_start, k_lo,
                causal, window_size, dropout_p, training, scale,
                use_reentrant=False,
            )
        outputs.append(out_block)

    return torch.cat(outputs, dim=2)


try:
    from flash_attention_interface import flash_attn_func as _turing_flash_attn_func
except ImportError:
    _turing_flash_attn_func = None

_TURING_SUPPORTED_HEAD_DIMS = (64, 128)


def _turing_flash_eligible(q: torch.Tensor, k: torch.Tensor, causal: bool,
                            window_size: int | None, dropout_p: float) -> bool:
    if _turing_flash_attn_func is None:
        return False
    if q.device.type != "cuda":
        return False
    major, minor = torch.cuda.get_device_capability(q.device)
    if (major, minor) != (7, 5):
        return False
    if q.shape[-1] not in _TURING_SUPPORTED_HEAD_DIMS:
        return False
    if window_size is not None:
        return False
    if dropout_p > 0.0:
        return False
    return True


def _turing_flash_attention(q, k, v, causal, softmax_scale):
    
    q_ = q.transpose(1, 2)  # (B,H,T,D) -> (B,T,H,D), this repo's expected layout
    k_ = k.transpose(1, 2)
    v_ = v.transpose(1, 2)
    out = _turing_flash_attn_func(q_, k_, v_, softmax_scale=softmax_scale, causal=causal)
    return out.transpose(1, 2)  # back to (B,H,T,D) for this codebase's convention


# ---------------------------------------------------------------------------
# FlashAttention-2 (NVIDIA Ampere+)
# ---------------------------------------------------------------------------

def _flash_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
    from flash_attn import flash_attn_func

    # flash_attn wants (B, T, H, D); we work in (B, H, T, D) throughout.
    qf = q.transpose(1, 2).contiguous()
    kf = k.transpose(1, 2).contiguous()
    vf = v.transpose(1, 2).contiguous()

    # window_size=(-1,-1) means "unbounded" (plain causal/full) per flash-attn's
    # own convention; (W-1, 0) means "attend to self + W-1 previous tokens".
    ws = (window_size - 1, 0) if window_size is not None else (-1, -1)

    out = flash_attn_func(
        qf, kf, vf,
        dropout_p=dropout_p if training else 0.0,
        softmax_scale=softmax_scale,
        causal=causal,
        window_size=ws,
    )  # GQA/MQA handled natively by flash_attn_func (kf/vf may have fewer heads than qf)
    return out.transpose(1, 2)  # back to (B, H, T, D)


# ---------------------------------------------------------------------------
# xFormers (NVIDIA pre-Ampere, e.g. T4)
# ---------------------------------------------------------------------------

@torch.compiler.disable
def _xformers_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
    
    import xformers.ops as xops
    from xformers.ops.fmha.attn_bias import LowerTriangularMask, BlockDiagonalCausalMask

    qf = q.transpose(1, 2).contiguous()  # (B, T, H, D)
    kf = k.transpose(1, 2).contiguous()
    vf = v.transpose(1, 2).contiguous()

    b, t, hq, d = qf.shape
    _, tk, hkv, _ = kf.shape
    if hkv != hq:
        # xFormers' native grouped-query path (5D inputs) is fragile across
        # versions; repeat_interleave is a small, version-stable cost and
        # correctness matters more here than shaving a few heads of memory.
        reps = hq // hkv
        kf = kf.repeat_interleave(reps, dim=2)
        vf = vf.repeat_interleave(reps, dim=2)

    if window_size is not None:
        if not causal:
            raise _UnsupportedByBackend("non-causal windowed attention not implemented for the xFormers path")
        
        bias = BlockDiagonalCausalMask.from_seqlens(
            q_seqlen=[t] * b, kv_seqlen=[tk] * b,
        ).make_local_attention(window_size)
        qf_packed = qf.reshape(1, b * t, hq, d)
        kf_packed = kf.reshape(1, b * tk, hq, d)
        vf_packed = vf.reshape(1, b * tk, hq, d)
        out = xops.memory_efficient_attention(
            qf_packed, kf_packed, vf_packed, attn_bias=bias,
            p=dropout_p if training else 0.0,
            scale=softmax_scale,
        )
        out = out.reshape(b, t, hq, d)
        return out.transpose(1, 2)

    bias = LowerTriangularMask() if causal else None
    out = xops.memory_efficient_attention(
        qf, kf, vf, attn_bias=bias,
        p=dropout_p if training else 0.0,
        scale=softmax_scale,
    )
    return out.transpose(1, 2)


# ---------------------------------------------------------------------------
# TPU (torch_xla)
# ---------------------------------------------------------------------------

@functools.lru_cache(maxsize=1)
def _splash_attention_available() -> bool:
    try:
        from jax.experimental.pallas.ops.tpu.splash_attention import (  # noqa: F401
            splash_attention_kernel,
            splash_attention_mask,
        )
        import torch_xla.core.xla_builder  # noqa: F401
        return True
    except ImportError:
        return False



_splash_disabled_reason: str | None = None


def reset_splash_circuit_breaker() -> None:
    global _splash_disabled_reason
    _splash_disabled_reason = None


def get_splash_disabled_reason() -> str | None:
    
    return _splash_disabled_reason




class _SplashAttentionFn(torch.autograd.Function):
   
    @staticmethod
    def forward(ctx, q, k, v, num_heads, q_len, k_len, scale):
        import torch_xla
        import torch_xla.core.xla_builder as xb

        orig_dtype = q.dtype
        q_b, k_b, v_b = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)

        def fwd_jax(qj, kj, vj):
            from jax.experimental.pallas.ops.tpu.splash_attention import (
                splash_attention_kernel, splash_attention_mask,
            )
            import jax
            mask = splash_attention_mask.MultiHeadMask(
                masks=[splash_attention_mask.CausalMask(shape=(q_len, k_len)) for _ in range(num_heads)]
            )
            kernel_fn = splash_attention_kernel.make_splash_mha(mask=mask, head_shards=1, q_seq_shards=1)
            return jax.vmap(kernel_fn)(q=qj * scale, k=kj, v=vj)

        out = xb.call_jax(fwd_jax, (q_b, k_b, v_b), {}, "arya_splash_attention_fwd")

        
        torch_xla.sync(reset_scope=False)

        ctx.save_for_backward(q_b, k_b, v_b)
        ctx.num_heads, ctx.q_len, ctx.k_len, ctx.scale, ctx.orig_dtype = num_heads, q_len, k_len, scale, orig_dtype
        return out.to(orig_dtype)

    @staticmethod
    def backward(ctx, grad_output):
        import torch_xla
        import torch_xla.core.xla_builder as xb
        q_b, k_b, v_b = ctx.saved_tensors
        num_heads, q_len, k_len, scale = ctx.num_heads, ctx.q_len, ctx.k_len, ctx.scale
        grad_output_b = grad_output.to(torch.bfloat16)

        def bwd_jax(qj, kj, vj, gj):
            from jax.experimental.pallas.ops.tpu.splash_attention import (
                splash_attention_kernel, splash_attention_mask,
            )
            import jax

            def raw(qj_, kj_, vj_):
                mask = splash_attention_mask.MultiHeadMask(
                    masks=[splash_attention_mask.CausalMask(shape=(q_len, k_len)) for _ in range(num_heads)]
                )
                kernel_fn = splash_attention_kernel.make_splash_mha(mask=mask, head_shards=1, q_seq_shards=1)
                return jax.vmap(kernel_fn)(q=qj_ * scale, k=kj_, v=vj_)

            _, vjp_fn = jax.vjp(raw, qj, kj, vj)
            return vjp_fn(gj)  # (dq, dk, dv) -- call_jax supports PyTree returns

        dq, dk, dv = xb.call_jax(bwd_jax, (q_b, k_b, v_b, grad_output_b), {}, "arya_splash_attention_bwd")
        torch_xla.sync(reset_scope=False)  # same reasoning as forward() -- must fail here, not later, elsewhere
        orig_dtype = ctx.orig_dtype
        return dq.to(orig_dtype), dk.to(orig_dtype), dv.to(orig_dtype), None, None, None, None


def _splash_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
    
    if not causal:
        raise _UnsupportedByBackend("Splash Attention wiring here only supports causal=True -- falling back.")
    if window_size is not None:
        raise _UnsupportedByBackend("Splash Attention wiring here doesn't support window_size -- falling back.")
    if dropout_p > 0 and training:
        raise _UnsupportedByBackend("Splash Attention wiring here doesn't support dropout -- falling back.")

    q_, k_, v_ = q, k, v
    hq, hkv = q.shape[1], k.shape[1]
    if hkv != hq:
        reps = hq // hkv
        k_ = k.repeat_interleave(reps, dim=1)
        v_ = v.repeat_interleave(reps, dim=1)

    num_heads = q_.shape[1]
    q_len, k_len = q_.shape[2], k_.shape[2]
    scale = softmax_scale if softmax_scale is not None else (1.0 / math.sqrt(q_.shape[-1]))

    return _SplashAttentionFn.apply(q_, k_, v_, num_heads, q_len, k_len, scale)


@torch.compiler.disable
def _tpu_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
    
    from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention

    global _splash_disabled_reason
    if _splash_disabled_reason is not None:
        _record_backend("splash_circuit_broken")
    elif _splash_attention_available():
        try:
            result = _splash_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
            _record_backend("splash_success")
            return result
        except _UnsupportedByBackend as e:
            
            _record_backend(f"splash_fallback:{type(e).__name__}")
            _warn_once(
                f"splash_fail_{type(e).__name__}",
                f"Splash Attention call failed ({e!r}) -- falling back to flash_attention.",
            )
        except Exception as e:  # noqa: BLE001 -- a REAL execution failure (post-sync) -- trip the breaker
            _record_backend(f"splash_fallback:{type(e).__name__}")
            _splash_disabled_reason = repr(e)
            _warn_once(
                f"splash_fail_{type(e).__name__}",
                f"Splash Attention call failed with a real execution error ({e!r}) -- "
                f"this usually means a jax/jaxlib/libtpu version mismatch in this "
                f"environment (e.g. 'Unsupported version: expected <= N but got M' is "
                f"Mosaic IR version skew between jax and libtpu -- fix by aligning "
                f"`pip install -U \"jax[tpu]\" jaxlib` as a matched pair). Disabling "
                f"Splash Attention for the rest of this run and falling back to "
                f"flash_attention -- call kernel.reset_splash_circuit_breaker() to "
                f"retry after fixing the environment.",
            )
    else:
        _record_backend("splash_unavailable")

    if window_size is not None:
        raise _UnsupportedByBackend(
            "torch_xla's built-in flash_attention wrapper doesn't expose a "
            "sliding-window argument -- falling back to reference."
        )
    if dropout_p > 0 and training:
        raise _UnsupportedByBackend(
            "torch_xla's built-in flash_attention wrapper doesn't take a "
            "dropout argument -- falling back to reference."
        )

    q_, k_, v_ = q, k, v
    hq, hkv = q.shape[1], k.shape[1]
    if hkv != hq:
        reps = hq // hkv
        k_ = k.repeat_interleave(reps, dim=1)
        v_ = v.repeat_interleave(reps, dim=1)

    
    orig_dtype = q_.dtype
    q_b = q_.to(torch.bfloat16)
    k_b = k_.to(torch.bfloat16)
    v_b = v_.to(torch.bfloat16)
    result = xla_flash_attention(q_b, k_b, v_b, causal=causal)
    result = result.to(orig_dtype)
    _record_backend("tpu_flash_success")
    return result





@torch.compiler.disable
def fused_block_sparse_attention(
    q_blocks: torch.Tensor,
    k_sel: torch.Tensor,
    v_sel: torch.Tensor,
    bias: torch.Tensor,
    dropout_p: float = 0.0,
    training: bool = True,
    softmax_scale: float | None = None,
) -> torch.Tensor | None:
    
    if q_blocks.device.type == "cuda" and _xformers_available():
        try:
            return _xformers_block_sparse_attention(q_blocks, k_sel, v_sel, bias, dropout_p, training, softmax_scale)
        except Exception as e:  # noqa: BLE001 -- must never crash training, just fall back
            _warn_once(
                f"block_sparse_xformers_fail_{type(e).__name__}",
                f"xFormers block-sparse attention failed ({e!r}) -- falling back to reference.",
            )
    return None





def _sdpa_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
    
    if causal or window_size is not None:
        q_idx = torch.arange(t, device=q.device).view(t, 1)
        k_idx = torch.arange(t, device=q.device).view(1, t)
        
        visible = k_idx <= q_idx
        if window_size is not None:
            visible = visible & (k_idx > q_idx - window_size)
        
        attn_mask = visible.bool().view(1, 1, t, t)  # broadcast over B and H
    else:
        attn_mask = None

    scale = softmax_scale if softmax_scale is not None else (1.0 / math.sqrt(d))

    
    effective_dropout_p = dropout_p * float(training)

   
    try:
        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=effective_dropout_p,
            is_causal=False,  
            scale=scale,
            enable_gqa=(hq != hkv),  
        )
        return out
    except TypeError:
        
        if hq != hkv:
            reps = hq // hkv
            k = k.repeat_interleave(reps, dim=1)
            v = v.repeat_interleave(reps, dim=1)
        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=effective_dropout_p,
            is_causal=False,
            scale=scale,
        )
        return out


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------

def fused_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    causal: bool = True,
    window_size: int | None = None,
    dropout_p: float = 0.0,
    training: bool = True,
    softmax_scale: float | None = None,
) -> torch.Tensor:
    
    backend = _backend_for(q.device)

    
    if _turing_flash_eligible(q, k, causal, window_size, dropout_p if training else 0.0):
        try:
            result = _turing_flash_attention(q, k, v, causal, softmax_scale)
            _record_backend("turing_flash_success")
            return result
        except Exception as e:  # noqa: BLE001
            _record_backend(f"turing_flash_fallback:{type(e).__name__}")
            _warn_once(
                f"turing_flash_fail_{type(e).__name__}",
                f"flash-attention-turing call failed ({e!r}) -- falling back to SDPA/xformers.",
            )

    
    if _use_sdpa_config() and backend in ("flash", "xformers"):
        try:
            result = _sdpa_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
            _record_backend("sdpa_success")
            return result
        except Exception as e:  # noqa: BLE001
            _record_backend(f"sdpa_fallback:{type(e).__name__}")
            _warn_once(f"sdpa_fail_{type(e).__name__}", f"SDPA call failed ({e!r}) -- falling back to flash_attn/xformers.")

    # ── 2. flash_attn 2.x (Ampere+) ───────────────────────────────────────
    if backend == "flash":
        try:
            result = _flash_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
            _record_backend("flash_success")
            return result
        except Exception as e:  # noqa: BLE001
            _record_backend(f"flash_fallback:{type(e).__name__}")
            _warn_once(f"flash_fail_{type(e).__name__}", f"flash_attn call failed ({e!r}) -- falling back.")

    # ── 3. xFormers (T4 fallback if SDPA fails) ───────────────────────────
    if backend in ("flash", "xformers"):
        try:
            result = _xformers_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
            _record_backend("xformers_success")
            return result
        except Exception as e:  # noqa: BLE001
            _record_backend(f"xformers_fallback:{type(e).__name__}")
            _warn_once(f"xformers_fail_{type(e).__name__}", f"xformers call failed ({e!r}) -- falling back.")

    # ── 4. TPU path ───────────────────────────────────────────────────────
    if backend == "tpu":
        try:
            result = _tpu_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
            return result
        except Exception as e:  # noqa: BLE001
            _record_backend(f"tpu_fallback:{type(e).__name__}")
            _warn_once(f"tpu_fail_{type(e).__name__}", f"torch_xla flash_attention call failed ({e!r}) -- falling back.")

    # ── 5. Reference fallback (always correct, always slow) ───────────────
    _record_backend("reference")
    return _reference_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)


def print_attention_backend(device: "torch.device | str | None" = None) -> None:
    """Print the attention backend that will be used for the given device.

    Resolves the full priority chain without running a forward pass:
      turing-flash (sm_75, flash_attention_interface installed, eligible shape)
      -> SDPA (use_sdpa=True, device is cuda/flash/xformers class)
      -> flash_attn 2.x (Ampere+, flash_attn installed)
      -> xformers (any CUDA, xformers installed)
      -> tpu (XLA device, torch_xla installed)
      -> reference (tiled matmul, always-correct fallback)

    Called once at model init so the backend is visible at the top of the log
    before training starts.
    """
    if device is None:
        if torch.cuda.is_available():
            device = torch.device("cuda", torch.cuda.current_device())
        else:
            device = torch.device("cpu")
    elif isinstance(device, str):
        device = torch.device(device)

    dev_str = str(device)
    backend = _backend_for(device)

    lines = [f"[kernel] attention backend  device={dev_str}"]

    if device.type == "cuda":
        major, minor = torch.cuda.get_device_capability(device)
        gpu_name = torch.cuda.get_device_name(device)
        lines.append(f"           gpu             {gpu_name}  (sm_{major}{minor})")

        # Turing flash: only sm_75 + flash_attention_interface installed
        turing_eligible = (
            (major, minor) == (7, 5)
            and _turing_flash_attn_func is not None
        )
        if turing_eligible:
            lines.append("           turing-flash    AVAILABLE (sm_75 + flash_attention_interface)")
        else:
            reason = (
                "sm != 7.5" if (major, minor) != (7, 5)
                else "flash_attention_interface not installed"
            )
            lines.append(f"           turing-flash    not eligible ({reason})")

        # SDPA: gated by use_sdpa config AND backend in (flash, xformers)
        sdpa_config = _use_sdpa_config()
        sdpa_eligible = sdpa_config and backend in ("flash", "xformers")
        if sdpa_eligible:
            lines.append("           sdpa            WILL BE USED  (use_sdpa=True, F.scaled_dot_product_attention)")
        else:
            reason = "use_sdpa=False in config" if not sdpa_config else f"backend={backend!r} not in flash/xformers"
            lines.append(f"           sdpa            skipped ({reason})")

        # flash_attn 2.x
        if _flash_attn_available():
            lines.append(f"           flash_attn      available (Ampere+ path, backend={backend!r})")
        else:
            lines.append("           flash_attn      not installed")

        # xformers
        if _xformers_available():
            lines.append("           xformers        available (fallback if SDPA/flash fail)")
        else:
            lines.append("           xformers        not installed")

        # Summarise what will actually fire first
        if turing_eligible:
            active = "turing-flash  (sm_75 eligible; will bypass SDPA for matching head_dims)"
        elif sdpa_eligible:
            active = "sdpa  (F.scaled_dot_product_attention — dispatches to FlashAttention/cudnn/math kernel)"
        elif backend == "flash":
            active = "flash_attn 2.x"
        elif backend == "xformers":
            active = "xformers memory_efficient_attention"
        else:
            active = "reference  (tiled matmul — slow; install flash_attn or xformers)"

        lines.append(f"           >>> active <<<  {active}")

    elif device.type == "xla":
        splash = _splash_attention_available()
        lines.append(f"           splash-attn     {'available' if splash else 'not available'}")
        lines.append(f"           torch_xla flash {'available' if _tpu_kernel_available() else 'not available'}")
        active = "splash-attention" if splash else ("torch_xla flash_attention" if _tpu_kernel_available() else "reference")
        lines.append(f"           >>> active <<<  {active}")

    else:
        lines.append("           >>> active <<<  reference  (CPU — tiled matmul)")

    print("\n".join(lines))