kernel.py

"""Attention kernel with two backends, chosen once at import.

* FlexAttention (default): CUDA with compute capability >= 8.0 (sm80+), torch >= 2.5.
  Masks are BlockMasks; the mask logic lives in mask_mod.
* xFormers memory-efficient attention: CUDA GPUs below sm80 (no Triton/Flex support).
  Masks are dense boolean (T, T) or (B, T, T) tensors converted to additive biases.

Every attention layer calls `fused_attention(q, k, v, mask)`; the mask object comes from
`causal_masks` or `doc_masks`, so the backend is invisible to model.py.

Attention-weight dropout is not supported on either backend (DenseMHA refuses dropout > 0).
"""
import functools
import importlib
import warnings

import torch
import torch.nn.functional as F
from torch.nn.attention.flex_attention import create_block_mask, flex_attention

_IS_CUDA = torch.cuda.is_available() and torch.version.cuda is not None
_USE_XFORMERS = _IS_CUDA and torch.cuda.get_device_capability() < (8, 0)

_XFORMERS_OK = _USE_XFORMERS  # cleared on the first failure; then SDPA is used

if _USE_XFORMERS:
    # Loaded dynamically: Transformers scans source files for static imports and would
    # require xformers on every machine, including ones that never use it.
    xops = importlib.import_module("xformers.ops")  # pip install xformers (match your torch)
    _flex = None
else:
    xops = None
    _flex = torch.compile(flex_attention, dynamic=False) if _IS_CUDA else flex_attention


class _XMask:
    """Dense allowed-positions mask for the xFormers path. allowed: bool (T, T) or (B, T, T)."""
    __slots__ = ("allowed",)

    def __init__(self, allowed: torch.Tensor):
        self.allowed = allowed


def fused_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask) -> torch.Tensor:
    """q: (B, Hq, T, D); k, v: (B, Hkv, T, D). GQA is handled natively (Hq a multiple of Hkv)."""
    if isinstance(mask, _XMask):
        return _xformers_attention(q, k, v, mask.allowed)
    return _flex(q, k, v, block_mask=mask, enable_gqa=q.shape[1] != k.shape[1])


def _xformers_attention(q, k, v, allowed):
    B, Hq, T, D = q.shape
    Hkv = k.shape[1]
    if Hq != Hkv:
        # query head h uses kv head h // group, the same grouping FlexAttention's enable_gqa uses
        group = Hq // Hkv
        k = k.repeat_interleave(group, dim=1)
        v = v.repeat_interleave(group, dim=1)
    bias = torch.zeros(allowed.shape, dtype=q.dtype, device=q.device).masked_fill(~allowed, float("-inf"))
    bias = bias[None, None] if bias.dim() == 2 else bias[:, None]
    bias = bias.expand(B, Hq, T, T)
    global _XFORMERS_OK
    if _XFORMERS_OK:
        try:
            # xFormers expects (B, T, H, D)
            out = xops.memory_efficient_attention(
                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_bias=bias)
            _record_backend("xformers")
            return out.transpose(1, 2)
        except Exception as e:
            # Typical cause: xformers built for a different torch/CUDA than the runtime.
            _XFORMERS_OK = False
            warnings.warn(f"xformers memory_efficient_attention unavailable ({type(e).__name__}: {e}); "
                          "using torch scaled_dot_product_attention with the same mask.")
    # Same math with PyTorch's built-in attention (works on any CUDA GPU, no extra package).
    out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias)
    _record_backend("sdpa_masked")
    return out


# ---------------------------------------------------------------------------
# Masks. One rule for every layer: window = window_size or seq_len; since q - kv <= seq_len - 1,
# "no window" is the same expression as "window = seq_len".
# ---------------------------------------------------------------------------

@functools.lru_cache(maxsize=None)
def _causal_mask(seq_len: int, window_size, device_key: str):
    window = window_size or seq_len

    def mask_mod(b, h, q_idx, kv_idx):
        return (kv_idx <= q_idx) & (q_idx - kv_idx < window)

    return create_block_mask(mask_mod, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=device_key)


@functools.lru_cache(maxsize=None)
def _xcausal_mask(seq_len: int, window_size, device_key: str):
    idx = torch.arange(seq_len, device=device_key)
    q_idx = idx[:, None]
    kv_idx = idx[None, :]
    window = window_size or seq_len
    return _XMask((kv_idx <= q_idx) & (q_idx - kv_idx < window))


@torch.compiler.disable
def causal_masks(seq_len: int, windows, device) -> dict:
    """{window_size: mask} for plain causal attention; cached per (seq_len, window, device)."""
    if _USE_XFORMERS:
        return {w: _xcausal_mask(seq_len, w, str(device)) for w in windows}
    return {w: _causal_mask(seq_len, w, str(device)) for w in windows}


@torch.compiler.disable
def doc_masks(doc_ids: torch.Tensor, windows) -> dict:
    """{window_size: mask} for causal + same-document attention. doc_ids: (B, T) int, per batch."""
    b_, t_ = doc_ids.shape

    if _USE_XFORMERS:
        idx = torch.arange(t_, device=doc_ids.device)
        q_idx = idx[:, None]
        kv_idx = idx[None, :]
        same_doc = doc_ids[:, :, None] == doc_ids[:, None, :]          # (B, T, T)

        def build_x(window_size):
            window = window_size or t_
            return _XMask(same_doc & (kv_idx <= q_idx) & (q_idx - kv_idx < window))

        return {w: build_x(w) for w in windows}

    def build(window_size):
        window = window_size or t_

        def mask_mod(b, h, q_idx, kv_idx):
            return (doc_ids[b, q_idx] == doc_ids[b, kv_idx]) & (kv_idx <= q_idx) & (q_idx - kv_idx < window)

        return create_block_mask(mask_mod, B=b_, H=None, Q_LEN=t_, KV_LEN=t_,
                                 device=doc_ids.device, _compile=True)

    return {w: build(w) for w in windows}


# ---------------------------------------------------------------------------
# Telemetry counters (model.py records its Liger/Triton success/fallback events here).
# ---------------------------------------------------------------------------

_backend_stats: dict[str, int] = {}


def _record_backend(name: str) -> None:
    # Inside a compiled region this body runs only at trace time, so counting here under-reports.
    # Record nothing while compiling rather than report wrong numbers.
    if torch.compiler.is_compiling():
        return
    _backend_stats[name] = _backend_stats.get(name, 0) + 1


def get_backend_stats() -> dict[str, int]:
    return dict(_backend_stats)


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