Update kernel.py

kernel.py

@@ -1,147 +1,163 @@
-"""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 torch
-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)
-
-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)
-    # 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)
-
-
-# ---------------------------------------------------------------------------
-# 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()
+"""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()
+