From bb4503f453bcce938629393e5c022a1692c0f8b5 Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Mon, 20 Jul 2026 20:29:17 -0700 Subject: [PATCH] DFlash/Domino: Add support for the FlexAttention FLASH backend Adds an option to turn on FAv4 backend for FlexAttention. This is significantly faster on Blackwell than the default backend (which is triton). We unfortunately need to monkey-patch PyTorch (in 2.11, which is what SGLang is pinned to right now). Whenever SpecForge can upgrade to PyTorch >= 2.13 then we can get rid of the monkeypatches (and also the correctness test). This PR also cleans up the attention implementation a little. --- .../algorithms/common/dflash_family_model.py | 17 +- specforge/modeling/draft/dflash.py | 55 ++++-- .../modeling/draft/flex_attention_backend.py | 28 +++ specforge/torch_compat.py | 62 +++++++ .../test_flex_attention_backend.py | 173 ++++++++++++++++++ 5 files changed, 315 insertions(+), 20 deletions(-) create mode 100644 specforge/modeling/draft/flex_attention_backend.py create mode 100644 specforge/torch_compat.py create mode 100644 tests/test_modeling/test_flex_attention_backend.py diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index 20d6f0b7e..5ff973c44 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/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( diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 719b341ad..e48e41532 100644 --- a/specforge/modeling/draft/dflash.py +++ b/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 diff --git a/specforge/modeling/draft/flex_attention_backend.py b/specforge/modeling/draft/flex_attention_backend.py new file mode 100644 index 000000000..80c127425 --- /dev/null +++ b/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"] diff --git a/specforge/torch_compat.py b/specforge/torch_compat.py new file mode 100644 index 000000000..2bf868853 --- /dev/null +++ b/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"] diff --git a/tests/test_modeling/test_flex_attention_backend.py b/tests/test_modeling/test_flex_attention_backend.py new file mode 100644 index 000000000..13d77b8c9 --- /dev/null +++ b/tests/test_modeling/test_flex_attention_backend.py @@ -0,0 +1,173 @@ +import os +import unittest +from unittest import mock + +import torch +from torch.nn.attention.flex_attention import flex_attention +from transformers import Qwen3Config + +from specforge.algorithms.common.dflash_family_model import create_dflash_block_mask +from specforge.modeling.draft.dflash import Qwen3DFlashAttention +from specforge.modeling.draft.dflash_kernels import DEFAULT_DFLASH_KERNELS +from specforge.modeling.draft.flex_attention_backend import flex_attention_backend + + +class FlexAttentionBackendTest(unittest.TestCase): + # This correctness regression test can be deleted when we require + # torch>=2.13; it tests the Torch 2.11 Inductor monkeypatch for CuteDSL + # operations in patch_inductor_cutedsl_lowerings(). + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10, + "FLASH FlexAttention correctness requires a Blackwell CUDA device", + ) + def test_flash_matches_triton_forward_and_backward(self): + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + batch_size, num_query_heads, num_key_value_heads = 1, 3, 1 + context_len, head_dim = 256, 64 + num_blocks, draft_block_size = 4, 64 + query_len = num_blocks * draft_block_size + kv_len = context_len + query_len + anchors = torch.tensor([[64, 128, 192, 224]], device=device) + keep_blocks = torch.ones( + (batch_size, num_blocks), dtype=torch.bool, device=device + ) + + inputs = ( + torch.randn( + batch_size, + num_query_heads, + query_len, + head_dim, + device=device, + dtype=dtype, + ), + torch.randn( + batch_size, + num_key_value_heads, + kv_len, + head_dim, + device=device, + dtype=dtype, + ), + torch.randn( + batch_size, + num_key_value_heads, + kv_len, + head_dim, + device=device, + dtype=dtype, + ), + ) + + def run_backend(backend, flex_block_size=None): + block_mask = create_dflash_block_mask( + anchor_positions=anchors, + block_keep_mask=keep_blocks, + S=context_len, + block_size=draft_block_size, + device=device, + flex_block_size=flex_block_size, + ) + + compiled_attention = torch.compile( + lambda query, key, value, mask: flex_attention( + query, + key, + value, + block_mask=mask, + enable_gqa=True, + kernel_options={"BACKEND": backend}, + ), + fullgraph=True, + ) + + query, key, value = [ + tensor.detach().clone().requires_grad_(True) for tensor in inputs + ] + output = compiled_attention(query, key, value, block_mask) + output.float().square().mean().backward() + torch.cuda.synchronize() + return output.detach(), tuple( + tensor.grad.detach() for tensor in (query, key, value) + ) + + triton_output, triton_grads = run_backend("TRITON") + with mock.patch.dict(os.environ, {"SPECFORGE_FLEX_ATTENTION_BACKEND": "FLASH"}): + self.assertEqual(flex_attention_backend(), "FLASH") + flash_output, flash_grads = run_backend("FLASH", (256, 128)) + + self.assertTrue(torch.isfinite(flash_output).all()) + torch.testing.assert_close(flash_output, triton_output, atol=3e-3, rtol=2e-2) + for flash_grad, triton_grad in zip(flash_grads, triton_grads): + self.assertTrue(torch.isfinite(flash_grad).all()) + torch.testing.assert_close(flash_grad, triton_grad, atol=5e-6, rtol=2e-2) + + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10, + "FLASH FlexAttention correctness requires a Blackwell CUDA device", + ) + def test_dflash_flash_attention_forward_backward_smoke(self): + config = Qwen3Config( + hidden_size=256, + intermediate_size=512, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=1, + head_dim=64, + layer_types=["full_attention"], + attention_dropout=0.0, + ) + config._attn_implementation = "flex_attention" + attention = Qwen3DFlashAttention( + config, + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ).to(device="cuda", dtype=torch.bfloat16) + hidden_states = torch.randn( + 1, + 256, + config.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + target_hidden = torch.randn( + 1, + 256, + config.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + block_mask = create_dflash_block_mask( + anchor_positions=torch.tensor([[64, 128, 192, 224]], device="cuda"), + block_keep_mask=torch.ones(1, 4, dtype=torch.bool, device="cuda"), + S=256, + block_size=64, + device=torch.device("cuda"), + flex_block_size=(256, 128), + ) + cos = torch.ones(1, 512, config.head_dim, device="cuda", dtype=torch.bfloat16) + sin = torch.zeros_like(cos) + + with mock.patch.dict(os.environ, {"SPECFORGE_FLEX_ATTENTION_BACKEND": "FLASH"}): + output, weights = attention( + hidden_states=hidden_states, + target_hidden=target_hidden, + position_embeddings=(cos, sin), + attention_mask=block_mask, + ) + output.float().square().mean().backward() + torch.cuda.synchronize() + + self.assertIsNone(weights) + self.assertIsNotNone(hidden_states.grad) + self.assertIsNotNone(target_hidden.grad) + self.assertTrue(torch.isfinite(hidden_states.grad).all()) + self.assertTrue(torch.isfinite(target_hidden.grad).all()) + + +if __name__ == "__main__": + unittest.main()