From d3334eb55f79ebf1e7f37982c261adf990303b2d Mon Sep 17 00:00:00 2001 From: Ziming Date: Wed, 1 Apr 2026 05:34:55 -0400 Subject: [PATCH 1/2] [CuTe, SM120] Fix forward/backward bugs and add SM120 test guards Bug fixes: - flash_fwd_sm120.py: Override self.arch after parent __init__ to prevent SM80 epilogue from enabling TMA-O store on SM120 - interface.py: Initialize dQ_single_wg and num_stages_PdS in SM120 backward config (fixes NameError in compile_key) - interface.py: Fix hdim>128 forward tile sizes for SM120 99KB SMEM - utils.py: Replace nvvm.atomicrmw with inline PTX red.global.add.f32 for newer NVVM Python bindings compatibility - __init__.py: Graceful FA2 import fallback when flash_attn_2_cuda is not built Test guards (tests/cute/test_flash_attn.py): - Add IS_SM120 flag - Skip: SplitKV, deterministic backward, learnable_sink, paged KV - Skip: GQA/MQA (pack_gqa crd2idx incompatibility on newer CUTLASS DSL) - TODO markers reference tracking PRs (#2336, #2348, #2389) Tested on RTX 5090 Laptop GPU (SM 12.0): 240 test_flash_attn_output cases passed (MHA, bf16, hdim 64/128/192/256, causal/non-causal, multiple sequence lengths, forward + backward). Builds on #2404 (bug fixes by @NJX-njx). Co-Authored-By: NJX-njx --- flash_attn/__init__.py | 21 ++++++++++------- flash_attn/cute/flash_fwd_sm120.py | 7 ++++++ flash_attn/cute/interface.py | 18 +++++++++----- flash_attn/cute/utils.py | 26 +++++++------------- tests/cute/test_flash_attn.py | 38 +++++++++++++++++++++++++----- 5 files changed, 72 insertions(+), 38 deletions(-) diff --git a/flash_attn/__init__.py b/flash_attn/__init__.py index ebf8176e25f..fe7e416e0e2 100644 --- a/flash_attn/__init__.py +++ b/flash_attn/__init__.py @@ -5,12 +5,15 @@ __version__ = "2.8.4" -from flash_attn.flash_attn_interface import ( - flash_attn_func, - flash_attn_kvpacked_func, - flash_attn_qkvpacked_func, - flash_attn_varlen_func, - flash_attn_varlen_kvpacked_func, - flash_attn_varlen_qkvpacked_func, - flash_attn_with_kvcache, -) +try: + from flash_attn.flash_attn_interface import ( + flash_attn_func, + flash_attn_kvpacked_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_varlen_kvpacked_func, + flash_attn_varlen_qkvpacked_func, + flash_attn_with_kvcache, + ) +except ImportError: + pass diff --git a/flash_attn/cute/flash_fwd_sm120.py b/flash_attn/cute/flash_fwd_sm120.py index 08d219acfa8..828c588fe36 100644 --- a/flash_attn/cute/flash_fwd_sm120.py +++ b/flash_attn/cute/flash_fwd_sm120.py @@ -16,6 +16,13 @@ class FlashAttentionForwardSm120(FlashAttentionForwardSm80): # The compilation target is determined by the GPU at compile time, not this field. arch = 80 + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Override the runtime arch detection from parent __init__ so that the + # SM80 __call__ path correctly disables TMA for the output store. + from cutlass.base_dsl.arch import Arch + self.arch = Arch.sm_80 + @staticmethod def can_implement( dtype, diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index ef624677f01..97279d72238 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -470,12 +470,13 @@ def _flash_attn_fwd( if tile_mn is None: if arch // 10 == 12: # SM120 tile sizes tuned for 99 KB SMEM capacity: - # D<=64: 128x128 → 48 KB (good occupancy) - # D>64: 128x64 → 64 KB (128x128 would use 96 KB, hurting occupancy) + # SMEM = tile_m*D*2 + tile_n*D*stages*2 + tile_n*Dv*stages*2 if head_dim <= 64: - fwd_cfg = FwdConfig(128, 128, True, True) + fwd_cfg = FwdConfig(128, 128, True, True) # 48 KB + elif head_dim <= 128: + fwd_cfg = FwdConfig(128, 64, True, True) # 64 KB else: - fwd_cfg = FwdConfig(128, 64, True, True) + fwd_cfg = FwdConfig(64, 64, True, True) # 64 KB for D=256 elif arch // 10 == 8: fwd_cfg = FwdConfig(128, 64, True, True) # SM80, should tune elif arch // 10 == 9: @@ -747,8 +748,8 @@ def _flash_attn_fwd( elif arch // 10 == 12: # SM120 (Blackwell GeForce / DGX Spark): uses SM80 MMA with SM120 SMEM capacity assert not use_block_sparsity, "Block sparsity not supported on SM 12.0" - assert page_table is None, "Paged KV not supported on SM 12.0 in this PR" - assert not is_split_kv, "SplitKV not supported on SM 12.0 in this PR" + assert page_table is None, "Paged KV not supported on SM 12.0" + assert not is_split_kv, "SplitKV not supported on SM 12.0" fa_fwd = FlashAttentionForwardSm120( dtype, head_dim, @@ -1037,9 +1038,14 @@ def _flash_attn_bwd( AtomLayoutNdKV = 4 AtomLayoutMdQ = 4 V_in_regs = False + dQ_single_wg = False + num_stages_PdS = 1 cluster_size = 1 use_2cta_instrs = False num_threads = 128 + # TODO(SM120): These features require an SM90-style backward kernel. The SM80 + # backward kernel used by SM120 does not implement them. Forward block sparsity + # is tracked in #2389, SplitKV in #2336, paged KV in #2348. assert not (block_sparse_tensors is not None), "Block sparsity backward not supported on SM 12.0" assert score_mod is None and score_mod_bwd is None, "score_mod backward not supported on SM 12.0" assert mask_mod is None, "mask_mod backward not supported on SM 12.0" diff --git a/flash_attn/cute/utils.py b/flash_attn/cute/utils.py index 31186618569..f20f916605c 100644 --- a/flash_attn/cute/utils.py +++ b/flash_attn/cute/utils.py @@ -403,23 +403,15 @@ def fadd_reduce( @dsl_user_op def atomic_add_fp32(a: float | Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> None: - # gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value() - # # cache_hint = cutlass.Int64(0x12F0000000000000) - # llvm.inline_asm( - # None, - # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip)], - # # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()], - # "red.global.add.f32 [$0], $1;", - # # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;", - # # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;", - # "l,f", - # # "l,f,l", - # has_side_effects=True, - # is_align_stack=False, - # asm_dialect=llvm.AsmDialect.AD_ATT, - # ) - nvvm.atomicrmw( - res=T.f32(), op=nvvm.AtomicOpKind.FADD, ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value() + gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value() + llvm.inline_asm( + None, + [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip)], + "red.global.add.f32 [$0], $1;", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, ) diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index 69e6308fb60..c38b97e6ebf 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -37,6 +37,7 @@ # SplitKV is not supported on SM90 IS_SM90 = torch.cuda.get_device_capability()[0] == 9 IS_SM100 = torch.cuda.get_device_capability()[0] == 10 +IS_SM120 = torch.cuda.get_device_capability() == (12, 0) TEST_BWD_ONLY = False VERBOSE = True @@ -113,6 +114,11 @@ def test_flash_attn_output( local = local_enum > 0 if local and causal: pytest.skip() + if IS_SM120 and has_learnable_sink: + pytest.skip("learnable_sink not supported on SM120") + # TODO(SM120): GQA/MQA hits crd2idx incompatibility in pack_gqa on newer CUTLASS DSL + if IS_SM120 and mha_type != "mha": + pytest.skip("GQA/MQA not yet supported on SM120 (pack_gqa crd2idx issue)") device = "cuda" # set seed seed = 0 @@ -250,8 +256,11 @@ def test_flash_attn_output( # pack_gqa_vals = [False] num_splits_vals = [1, 3] if d < 192 and not DISABLE_SPLIT and not TEST_BWD_ONLY else [1] for pack_gqa, num_splits in itertools.product(pack_gqa_vals, num_splits_vals): - # SplitKV not supported on SM90 - skip this iteration - if IS_SM90 and num_splits > 1: + # SplitKV not supported on SM90/SM120 - skip this iteration + if (IS_SM90 or IS_SM120) and num_splits > 1: + continue + # TODO(SM120): pack_gqa hits crd2idx incompatibility on newer CUTLASS DSL + if IS_SM120 and pack_gqa: continue if IS_SM100 and (d >= 192 and dv >= 192): # hdim 192 and 256 not support on SM100 continue @@ -301,6 +310,8 @@ def test_flash_attn_output( pytest.xfail("hdim > 192 backward: SM90 not supported yet") if d != dv and mha_type != "mha" and IS_SM90: pytest.xfail("SM90 GQA bwd currently requires headdim == headdim_v") + if IS_SM120 and deterministic: + pytest.skip("deterministic backward not supported on SM120") g = torch.randn_like(out) # do_o = ((g.float() * out.float()).sum(-1)).transpose(1, 2) dq, dk, dv = torch.autograd.grad(out, (q, k, v), g) @@ -472,6 +483,10 @@ def test_flash_attn_varlen_output( local = local_enum > 0 if local and causal: pytest.skip() + if IS_SM120 and has_learnable_sink: + pytest.skip("learnable_sink not supported on SM120") + if IS_SM120 and mha_type != "mha": + pytest.skip("GQA/MQA not yet supported on SM120 (pack_gqa crd2idx issue)") if ( causal or local ): # Right now reference only supports causal attention with seqlen_k == seqlen_q @@ -682,8 +697,11 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): # SplitKV is not supported for hdim >= 192 num_splits_vals = [1, 3] if d < 192 and not DISABLE_SPLIT and not TEST_BWD_ONLY else [1] for pack_gqa, num_splits in itertools.product(pack_gqa_vals, num_splits_vals): - # SplitKV not supported on SM90 - skip this iteration - if IS_SM90 and num_splits > 1: + # SplitKV not supported on SM90/SM120 - skip this iteration + if (IS_SM90 or IS_SM120) and num_splits > 1: + continue + # TODO(SM120): pack_gqa hits crd2idx incompatibility on newer CUTLASS DSL + if IS_SM120 and pack_gqa: continue out_unpad, lse = flash_attn_varlen_func( q_unpad if unpad_q else q, @@ -749,6 +767,8 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): pytest.xfail("hdim > 192 backward: SM90 not supported yet") if d != dv and mha_type != "mha" and IS_SM90: pytest.xfail("SM90 GQA bwd currently requires headdim == headdim_v") + if IS_SM120 and deterministic: + pytest.skip("deterministic backward not supported on SM120") g_unpad = torch.randn_like(out_unpad) # do_o = ((g_unpad.float() * out_unpad.float()).sum(-1)).transpose(-1, -2) # import flash_attn_3_cuda @@ -949,6 +969,12 @@ def test_flash_attn_kvcache( ): if page_size is not None and seqlen_k % page_size != 0: pytest.skip() + if IS_SM120 and page_size is not None: + pytest.skip("Paged KV not supported on SM120") + if IS_SM120 and has_learnable_sink: + pytest.skip("learnable_sink not supported on SM120") + if IS_SM120 and mha_type != "mha": + pytest.skip("GQA/MQA not yet supported on SM120 (pack_gqa crd2idx issue)") if seqlen_q > seqlen_k and new_kv: pytest.skip() if not new_kv and rotary_fraction > 0.0: @@ -1291,8 +1317,8 @@ def test_flash_attn_kvcache( for num_splits, precompute_metadata in itertools.product( num_splits_vals, precompute_metadata_vals ): - # SplitKV not supported on SM90 - skip this iteration - if IS_SM90 and num_splits > 1: + # SplitKV not supported on SM90/SM120 - skip this iteration + if (IS_SM90 or IS_SM120) and num_splits > 1: continue # if precompute_metadata: # scheduler_metadata = get_scheduler_metadata( From 715919cb3a3a855269f4e8f2c245fce66641386c Mon Sep 17 00:00:00 2001 From: Ziming Date: Wed, 1 Apr 2026 05:43:41 -0400 Subject: [PATCH 2/2] fix: narrow except to ModuleNotFoundError in flash_attn __init__ Only suppress the missing flash_attn_2_cuda module error, not unexpected ImportError from ABI mismatch or other failures. --- flash_attn/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flash_attn/__init__.py b/flash_attn/__init__.py index fe7e416e0e2..952bf8c04c4 100644 --- a/flash_attn/__init__.py +++ b/flash_attn/__init__.py @@ -15,5 +15,6 @@ flash_attn_varlen_qkvpacked_func, flash_attn_with_kvcache, ) -except ImportError: +except ModuleNotFoundError: + # flash_attn_2_cuda not built — FA4 (flash_attn.cute) can still be used directly pass