DFlash/Domino: Add support for the FlexAttention FLASH backend

specforge/algorithms/common/dflash_family_model.py

@@ -8,6 +8,7 @@
 import torch.nn.functional as F
 
 from specforge.modeling.draft.dflash import DFlashDraftModel
+from specforge.modeling.draft.flex_attention_backend import flex_attention_backend
 
 try:
     from torch.nn.attention.flex_attention import BlockMask, create_block_mask
@@ -77,6 +78,7 @@ def create_dflash_block_mask(
     S: int,
     block_size: int,
     device: torch.device,
+    flex_block_size=None,
 ):
     """Construct Flex Attention BlockMask for DFlash training.
 
@@ -112,8 +114,17 @@ def dflash_mask_mod(b, h, q_idx, kv_idx):
     Q_LEN = N * block_size
     KV_LEN = S + N * block_size
 
+    kwargs = {}
+    if flex_block_size is not None:
+        kwargs["BLOCK_SIZE"] = flex_block_size
     return create_block_mask(
-        dflash_mask_mod, B=B, H=None, Q_LEN=Q_LEN, KV_LEN=KV_LEN, device=device
+        dflash_mask_mod,
+        B=B,
+        H=None,
+        Q_LEN=Q_LEN,
+        KV_LEN=KV_LEN,
+        device=device,
+        **kwargs,
     )
 
 
@@ -289,6 +300,10 @@ def _forward_draft_blocks(
                 S=seq_len,
                 block_size=self.block_size,
                 device=device,
+                # FLASH requires a minimum of this block size.
+                flex_block_size=(
+                    (256, 128) if flex_attention_backend() == "FLASH" else None
+                ),
             )
         else:
             dflash_attn_mask = create_dflash_sdpa_mask(

specforge/modeling/draft/dflash.py

@@ -4,6 +4,7 @@
 from torch import nn
 from transformers import DynamicCache
 from transformers.cache_utils import Cache
+from transformers.integrations.flex_attention import compile_friendly_flex_attention
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.qwen3.modeling_qwen3 import (
     ALL_ATTENTION_FUNCTIONS,
@@ -18,6 +19,7 @@
 from typing_extensions import Tuple, Unpack
 
 from .dflash_kernels import DEFAULT_DFLASH_KERNELS, DFlashKernels
+from .flex_attention_backend import flex_attention_backend
 from .registry import register_draft
 
 
@@ -60,6 +62,10 @@ def __init__(
         )
         self.scaling = self.head_dim**-0.5
         self.attention_dropout = config.attention_dropout
+        if config._attn_implementation == "flex_attention":
+            assert (
+                config.attention_dropout == 0.0
+            ), "DFlash FlexAttention requires attention_dropout=0.0"
         self.is_causal = False
         self.q_proj = nn.Linear(
             config.hidden_size,
@@ -83,11 +89,6 @@ def __init__(
         )
         self.q_norm = kernels.make_rms_norm(self.head_dim, config.rms_norm_eps)
         self.k_norm = kernels.make_rms_norm(self.head_dim, config.rms_norm_eps)
-        self.sliding_window = (
-            config.sliding_window
-            if config.layer_types[layer_idx] == "sliding_attention"
-            else None
-        )
 
     def forward(
         self,
@@ -121,20 +122,36 @@ def forward(
         if past_key_values is not None:
             cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
             k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)
-        attn_fn: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-        attn_output, attn_weights = attn_fn(
-            self,
-            q,
-            k,
-            v,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,
-            **kwargs,
-        )
+        if self.config._attn_implementation == "flex_attention":
+            kernel_options = dict(kwargs.pop("kernel_options", None) or {})
+            backend = flex_attention_backend()
+            if backend is not None:
+                kernel_options["BACKEND"] = backend
+
+            attn_output = compile_friendly_flex_attention(
+                q,
+                k,
+                v,
+                block_mask=attention_mask,
+                enable_gqa=True,
+                scale=self.scaling,
+                kernel_options=kernel_options or None,
+            ).transpose(1, 2)
+            attn_weights = None
+        else:
+            attn_fn: Callable = eager_attention_forward
+            if self.config._attn_implementation != "eager":
+                attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+            attn_output, attn_weights = attn_fn(
+                self,
+                q,
+                k,
+                v,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                **kwargs,
+            )
         attn_output = attn_output.reshape(bsz, q_len, -1)
         attn_output = self.o_proj(attn_output)
         return attn_output, attn_weights

specforge/modeling/draft/flex_attention_backend.py

@@ -0,0 +1,28 @@
+"""FlexAttention backend selection shared by DFlash training components."""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+from specforge.torch_compat import patch_inductor_cutedsl_lowerings
+
+_VALID_BACKENDS = {"AUTO", "TRITON", "FLASH", "TRITON_DECODE"}
+_BACKEND_ENV = "SPECFORGE_FLEX_ATTENTION_BACKEND"
+
+
+def flex_attention_backend() -> Optional[str]:
+    backend = os.environ.get(_BACKEND_ENV, "").upper()
+    if not backend:
+        return None
+    if backend not in _VALID_BACKENDS:
+        raise ValueError(
+            f"{_BACKEND_ENV} must be one of {sorted(_VALID_BACKENDS)}, "
+            f"got {backend!r}"
+        )
+    if backend == "FLASH":
+        patch_inductor_cutedsl_lowerings()
+    return backend
+
+
+__all__ = ["flex_attention_backend"]

specforge/torch_compat.py

@@ -0,0 +1,62 @@
+"""PyTorch compatibility shims used by optional fast paths."""
+
+from __future__ import annotations
+
+import importlib
+
+import sympy
+import torch
+from packaging.version import InvalidVersion, Version
+
+
+def patch_inductor_cutedsl_lowerings() -> bool:
+    """Backfill CuteDSL lowering needed by Torch 2.11 FLASH FlexAttention."""
+    try:
+        torch_version = Version(torch.__version__.split("+", 1)[0])
+    except InvalidVersion:
+        return False
+    if torch_version.major != 2 or torch_version.minor != 11:
+        return False
+
+    try:
+        module = importlib.import_module(
+            "torch._inductor.codegen.cutedsl.cutedsl_op_overrides"
+        )
+    except ImportError:
+        return False
+    from torch._inductor.utils import get_bounds_index_expr
+    from torch._inductor.virtualized import V
+
+    overrides = module.CuteDSLOpOverrides
+    if getattr(overrides, "_specforge_cutedsl_patch", False):
+        return True
+
+    def _minimum(a, b):
+        return overrides.where(overrides.lt(a, b), a, b)
+
+    def _maximum(a, b):
+        return overrides.where(overrides.gt(a, b), a, b)
+
+    def _index_expr(expr: sympy.Expr, dtype: torch.dtype):
+        if isinstance(expr, (int, sympy.Integer)):
+            return overrides.constant(int(expr), dtype)
+
+        idx_str = V.kernel.kexpr(V.kernel.rename_indexing(expr))
+        result = V.kernel.cse.generate(
+            V.kernel.body,
+            idx_str,
+            bounds=get_bounds_index_expr(expr),
+            dtype=dtype,
+        )
+        result.is_scalar_expr = True
+        result.index_expr = V.graph.sizevars.simplify(expr)
+        return result
+
+    overrides.minimum = staticmethod(_minimum)
+    overrides.maximum = staticmethod(_maximum)
+    overrides.index_expr = staticmethod(_index_expr)
+    overrides._specforge_cutedsl_patch = True
+    return True
+
+
+__all__ = ["patch_inductor_cutedsl_lowerings"]