"""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()