Update kernel.py

kernel.py

@@ -218,6 +218,133 @@ def _turing_flash_attention(q, k, v, causal, softmax_scale):
     return out.transpose(1, 2)  # back to (B,H,T,D) for this codebase's convention
 
 
+# ---------------------------------------------------------------------------
+# FlexAttention (torch.nn.attention.flex_attention). Config-gated, not
+# hardware-gated -- set use_flexattention=True in config.py to opt in.
+# Unlike every other backend in this file, FlexAttention is a torch-level
+# composition of ops, not a hand-fused CUDA kernel -- its speed comes
+# entirely from torch.compile specializing the mask/score-mod into one
+# fused kernel at trace time. Eager flex_attention runs at roughly
+# reference-attention speed, so (unlike the @torch.compiler.disable
+# backends elsewhere in this file) this one is deliberately compiled, once,
+# and cached.
+# ---------------------------------------------------------------------------
+
+@functools.lru_cache(maxsize=1)
+def _flex_attention_available() -> bool:
+    try:
+        from torch.nn.attention.flex_attention import flex_attention  # noqa: F401
+        return True
+    except ImportError:
+        return False
+
+
+_compiled_flex_attention = None
+
+
+def _get_compiled_flex_attention():
+    global _compiled_flex_attention
+    if _compiled_flex_attention is None:
+        from torch.nn.attention.flex_attention import flex_attention
+        _compiled_flex_attention = torch.compile(flex_attention, dynamic=False)
+    return _compiled_flex_attention
+
+@torch.compiler.disable 
+@functools.lru_cache(maxsize=32)
+ 
+def _flex_block_mask(q_len: int, kv_len: int, causal: bool, window_size: "int | None", device_key: str):
+    """Build (and cache) a FlexAttention BlockMask for one (shape, mask-kind,
+    device) combo. Constructing a block mask isn't free, and with a fixed
+    training CONTEXT this pays the cost once instead of every forward call."""
+    from torch.nn.attention.flex_attention import create_block_mask
+
+    if window_size is not None:
+        def mask_mod(b, h, q_idx, kv_idx):
+            return (kv_idx <= q_idx) & (q_idx - kv_idx < window_size)
+    elif causal:
+        def mask_mod(b, h, q_idx, kv_idx):
+            return kv_idx <= q_idx
+    else:
+        return None  # full bidirectional attention -- no block mask needed
+
+    return create_block_mask(mask_mod, B=None, H=None, Q_LEN=q_len, KV_LEN=kv_len, device=device_key)
+
+
+def _flex_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
+    
+    if dropout_p > 0 and training:
+        raise _UnsupportedByBackend(
+            "FlexAttention path here doesn't wire up an attention-weight "
+            "dropout term -- falling back."
+        )
+
+    b, hq, t, d = q.shape
+    _, hkv, tk, _ = k.shape
+
+    block_mask = _flex_block_mask(t, tk, causal, window_size, str(q.device))
+    scale = softmax_scale if softmax_scale is not None else (1.0 / math.sqrt(d))
+    flex_fn = _get_compiled_flex_attention()
+
+    # FlexAttention takes (B, H, T, D) natively -- no transpose needed, unlike
+    # the flash_attn/xformers paths above which want (B, T, H, D).
+    return flex_fn(q, k, v, block_mask=block_mask, scale=scale, enable_gqa=(hkv != hq))
+
+
+# ---------------------------------------------------------------------------
+# Blackwell: FlashAttention-4 (CuTeDSL) with FlashAttention-2 fallback.
+#
+# FA4 (`pip install flash-attn-4`, imported as `flash_attn.cute`) targets
+# Hopper and data-center Blackwell (sm_90 / sm_100 / sm_103 -- H100/B200/B300)
+# via warp-specialized MMA instructions. It does NOT run on desktop Blackwell
+# (sm_120, e.g. RTX 50-series): sm_120 uses the same register-to-register
+# HMMA path NVIDIA GPUs have used since Volta, not the warp-specialized MMA
+# FA4's kernel design requires, and this is a physical silicon difference
+# (confirmed by multiple independent sm_120 FA4 build attempts as of early
+# 2026), not a version-gating issue that a newer release fixes. So this
+# block tries FA4 only on the capabilities it can actually run on, and
+# every other Blackwell-family GPU (sm_120 included) falls through to the
+# existing FlashAttention-2 path below.
+# ---------------------------------------------------------------------------
+
+_BLACKWELL_CAPABILITIES = {(9, 0), (10, 0), (10, 3), (12, 0)}  # Hopper + all Blackwell variants
+_FA4_CAPABILITIES = {(9, 0), (10, 0), (10, 3)}  # NOT (12, 0) -- see note above
+
+try:
+    from flash_attn.cute import flash_attn_func as _fa4_attn_func
+except ImportError:
+    _fa4_attn_func = None
+
+
+def _fa4_eligible(q: torch.Tensor, causal: bool, window_size: "int | None", dropout_p: float) -> bool:
+    if _fa4_attn_func is None:
+        return False
+    if q.device.type != "cuda":
+        return False
+    if torch.cuda.get_device_capability(q.device) not in _FA4_CAPABILITIES:
+        return False
+    if window_size is not None:  # not confirmed supported by FA4's current (beta) public API
+        return False
+    if dropout_p > 0.0:
+        return False
+    return True
+
+
+def _fa4_attention(q, k, v, causal, softmax_scale):
+    # FA4's public surface is still small/beta -- only `causal` is confirmed
+    # from the library's own usage example (flash_attn.cute docs show only
+    # `flash_attn_func(q, k, v, causal=True)`). Anything beyond that
+    # (window_size, dropout) is gated out by _fa4_eligible above rather than
+    # guessed at here. Assumes the same (B, T, H, D) layout as FA2/FA3; if
+    # that assumption is wrong for a given release the call raises and the
+    # dispatcher below falls back to FlashAttention-2 automatically.
+    qf = q.transpose(1, 2).contiguous()
+    kf = k.transpose(1, 2).contiguous()
+    vf = v.transpose(1, 2).contiguous()
+    kwargs = {"softmax_scale": softmax_scale} if softmax_scale is not None else {}
+    out = _fa4_attn_func(qf, kf, vf, causal=causal, **kwargs)
+    return out.transpose(1, 2)
+
+
 # ---------------------------------------------------------------------------
 # FlashAttention-2 (NVIDIA Ampere+)
 # ---------------------------------------------------------------------------
@@ -512,6 +639,46 @@ def _build_attention_fn(device: torch.device):
     Called once; result is cached in _ATTENTION_FN."""
     backend = _backend_for(device)
 
+    if device.type == "cuda" and _kernel_config_attr("use_flexattention", False) and _flex_attention_available():
+        def _flex_dispatch(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
+            try:
+                result = _flex_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
+                _record_backend("flex_success")
+                return result
+            except _UnsupportedByBackend as e:
+                _record_backend(f"flex_fallback:{type(e).__name__}")
+                _warn_once(f"flex_unsupported_{type(e).__name__}",
+                           f"FlexAttention can't handle this call ({e!r}) -- falling back to xformers.")
+            except Exception as e:  # noqa: BLE001
+                _record_backend(f"flex_fallback:{type(e).__name__}")
+                _warn_once(f"flex_fail_{type(e).__name__}",
+                           f"FlexAttention failed ({e!r}) -- falling back to xformers.")
+            return _xformers_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
+        return "flex", _flex_dispatch
+
+    if device.type == "cuda" and torch.cuda.get_device_capability(device) in _BLACKWELL_CAPABILITIES:
+        def _blackwell_dispatch(q, k, v, causal, window_size, dropout_p, training, softmax_scale):
+            if _fa4_eligible(q, causal, window_size, dropout_p if training else 0.0):
+                try:
+                    result = _fa4_attention(q, k, v, causal, softmax_scale)
+                    _record_backend("fa4_success")
+                    return result
+                except Exception as e:  # noqa: BLE001
+                    _record_backend(f"fa4_fallback:{type(e).__name__}")
+                    _warn_once(f"fa4_fail_{type(e).__name__}",
+                               f"FlashAttention-4 failed ({e!r}) -- falling back to FlashAttention-2.")
+            if _flash_attn_available():
+                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"FlashAttention-2 failed ({e!r}) -- falling back to xformers.")
+            return _xformers_attention(q, k, v, causal, window_size, dropout_p, training, softmax_scale)
+        return "blackwell", _blackwell_dispatch
+
     if (
         device.type == "cuda"
         and _turing_flash_attn_func is not None
@@ -633,6 +800,26 @@ def print_attention_backend(device: "torch.device | str | None" = None) -> None:
             )
             lines.append(f"           turing-flash    not eligible ({reason})")
 
+        # FlexAttention: config-gated (use_flexattention in config.py), not hardware-gated
+        flex_cfg_on = _kernel_config_attr("use_flexattention", False)
+        flex_active = flex_cfg_on and _flex_attention_available()
+        if flex_active:
+            lines.append("           flex-attention  ACTIVE (use_flexattention=True in config.py)")
+        elif flex_cfg_on:
+            lines.append("           flex-attention  requested but not importable (torch < 2.5?)")
+        else:
+            lines.append("           flex-attention  off (set use_flexattention=True in config.py to enable)")
+
+        # Blackwell family (Hopper sm_90 + all Blackwell variants): FA4 with FA2 fallback
+        blackwell_eligible = (major, minor) in _BLACKWELL_CAPABILITIES
+        if blackwell_eligible:
+            if (major, minor) in _FA4_CAPABILITIES and _fa4_attn_func is not None:
+                lines.append(f"           fa4             AVAILABLE (sm_{major}{minor})")
+            elif (major, minor) in _FA4_CAPABILITIES:
+                lines.append(f"           fa4             not installed (pip install flash-attn-4)")
+            else:
+                lines.append(f"           fa4             not eligible (sm_120 lacks FA4's required warp-specialized MMA -- uses FlashAttention-2 instead)")
+
         # flash_attn 2.x
         if _flash_attn_available():
             lines.append(f"           flash_attn      available (Ampere+ path, backend={backend!r})")
@@ -646,7 +833,13 @@ def print_attention_backend(device: "torch.device | str | None" = None) -> None:
             lines.append("           xformers        not installed")
 
         # Summarise what will actually fire first.
-        if turing_eligible:
+        if flex_active:
+            active = "flex-attention  (config-enabled, torch.compile'd)"
+        elif blackwell_eligible and (major, minor) in _FA4_CAPABILITIES and _fa4_attn_func is not None:
+            active = f"fa4  (sm_{major}{minor})"
+        elif blackwell_eligible and _flash_attn_available():
+            active = f"flash_attn 2.x  (sm_{major}{minor}, FA4 not eligible/installed)"
+        elif turing_eligible:
             active = "turing-flash  (sm_75 + flash_attention_interface)"
         elif backend == "flash":
             active = "flash_attn 2.x"