[CuTe, SM120] Support automatic num_splits (#2758)

* [CuTe, SM120] Support automatic num_splits

* [CuTe, SM120] Cover negative automatic num_splits

Author: bvolpato

flash_attn/cute/interface.py

@@ -414,7 +414,8 @@ def _get_fwd_config(
     num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n
     num_SMs = None
     if arch // 10 == 12:
-        assert num_splits == 1, "SM120 forward only supports num_splits=1"
+        assert num_splits <= 1, "SM120 forward only supports num_splits=1"
+        num_splits = 1
     elif num_splits < 1:
         num_SMs = get_num_sms_for_selection(device.index, arch)
         num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128)

tests/cute/test_flash_attn.py

@@ -92,10 +92,14 @@ def test_flash_attn_auto_splits_more_tiles_than_sms():
 
 
 @pytest.mark.skipif(not IS_SM120, reason="SM120-only SplitKV unsupported behavior")
-def test_flash_attn_sm120_rejects_splitkv():
+@pytest.mark.parametrize("num_splits", [0, -1])
+def test_flash_attn_sm120_num_splits(num_splits):
     q = torch.randn(1, 16, 4, 64, device="cuda", dtype=torch.bfloat16)
     k = torch.randn(1, 16, 1, 64, device="cuda", dtype=torch.bfloat16)
     v = torch.randn(1, 16, 1, 64, device="cuda", dtype=torch.bfloat16)
+    out = flash_attn_func(q, k, v)
+    out_auto = flash_attn_func(q, k, v, num_splits=num_splits)
+    torch.testing.assert_close(out_auto, out)
     with pytest.raises(AssertionError, match="SM120 forward only supports num_splits=1"):
         flash_attn_func(q, k, v, num_splits=3)