[CuTe, SM120] Fix forward/backward bugs and add SM120 test guards

flash_attn/__init__.py

@@ -5,12 +5,16 @@
 
 __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 ModuleNotFoundError:
+    # flash_attn_2_cuda not built — FA4 (flash_attn.cute) can still be used directly
+    pass

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,

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"

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

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(