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