[NVIDIA][CuTe,Fwd,sm120] Implement Pack-GQA on SM120 (+ graceful SplitKV fallback) (#2656)

* [CuTe,Fwd,sm120] Fix use_tma_O crash on SM120 (issue #2649)

On SM120 (Blackwell GeForce / RTX PRO 6000 / DGX Spark) the forward kernel set
`use_tma_O = self.arch >= Arch.sm_90`, enabling the TMA-based O-store epilogue.
But SM120 does not build the TMA store atom (tma_atom_O is None), so any forward
call crashes in cpasync.tma_partition with:

    AttributeError: 'NoneType' object has no attribute '_trait'

This makes the CuTe-DSL forward unusable on every SM120 GPU.

Restrict the TMA O-store to sm_90..sm_119, which is where the WGMMA-era epilogue
path is actually available:

    self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120

SM120 falls back to the non-TMA register->gmem O store (already used for the
SM80 path), which is correct and what the CpAsync SM120 kernel expects.

Verified on RTX PRO 6000 Blackwell (sm_120, cc 12.0), torch 2.12.0+cu130,
nvidia-cutlass-dsl 4.5.2: forward now runs and matches PyTorch SDPA reference
for hdim 64/96/128, causal and non-causal (max abs err <= 8e-3 in bf16). Before
this fix every SM120 forward call raised the AttributeError above.

* [CuTe,Fwd,sm120] Implement Pack-GQA on SM120; graceful SplitKV fallback

Pack-GQA was only half-wired in the SM80/SM120 CpAsync forward: the epilogue
referenced PackGQA.store_O/store_LSE, but the Q-load and head-indexing used the
plain (unpacked) path. So pack_gqa=True crashed in pack_gqa.store_O (crd2idx on a
packed (h_idx, m_idx) coordinate against an unpacked mO layout).

This implements Pack-GQA end to end on SM120 (and SM80), mirroring the SM90 path:
- Reshape mQ/mO (head_idx=2) and mLSE (head_idx=1) via pack_gqa_layout so
  qhead_per_kvhead folds into the seqlen mode ((qhead, seqlen)).
- Scheduler args use cute.size(mQ.shape[0]) (packed total rows) and seqlen_q_static
  = mQ.shape[0][1] (logical seqlen), so causal/mask q_idx stay correct.
- Kernel head-indexing: when pack_gqa, num_head from the scheduler already indexes
  the KV head (mQ/mK share nheads_kv); no division.
- Q-load: gather rows via PackGQA.load_Q (per-row (h_idx, m_idx) gmem pointers)
  instead of the contiguous local_tile path.

SplitKV (num_splits>1) is an SM100-only feature (SM80/SM90 also assert it
unsupported); SM120 has no forward+combine path. Fall back to num_splits=1, which
is numerically correct, instead of crashing in _check_type on the fp32 partials.

Verified on RTX PRO 6000 Blackwell (sm_120): pack_gqa=True matches PyTorch SDPA
GQA/MQA reference (err <= 8.4e-3 bf16) AND is bit-identical to the unpacked path
(max |packed - unpacked| = 0.0) across MHA/GQA/MQA, causal/non-causal, hd 64/128,
seqlen 512-2048. num_splits=3 falls back and matches reference (err 6.8e-4).
Stacked on the SM120 use_tma_O fix (#2649).

* re-enable SM120 pack-gqa after rebase

* clean up SM120 pack-gqa split handling

* fix SM120 varlen pack-gqa offset

---------

Co-authored-by: drisspg <drisspguessous@gmail.com>

Author: johnnynunez

flash_attn/cute/flash_fwd.py

@@ -30,7 +30,7 @@
 from flash_attn.cute.softmax import Softmax, apply_score_mod_inner
 from flash_attn.cute.seqlen_info import SeqlenInfoQK
 from flash_attn.cute.block_info import BlockInfo
-from flash_attn.cute.pack_gqa import PackGQA
+from flash_attn.cute.pack_gqa import PackGQA, pack_gqa_layout
 from flash_attn.cute.named_barrier import NamedBarrierFwd
 from flash_attn.cute.block_sparsity import BlockSparseTensors
 from flash_attn.cute.tile_scheduler import SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments
@@ -655,8 +655,7 @@ def __call__(
         self.num_producer_threads = self.num_threads
         self.num_Q_load_threads = self.num_threads
         self.num_epilogue_threads = self.num_threads
-        # self.use_tma_O = self.arch >= 90 and mCuSeqlensQ is None
-        self.use_tma_O = self.arch >= Arch.sm_90
+        self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120
         self._setup_attributes()
         SharedStorage = self._get_shared_storage_cls()
         mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)]
@@ -674,6 +673,12 @@ def __call__(
         if const_expr(mLSE is not None):
             LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
             mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose))
+        if const_expr(self.pack_gqa):
+            nheads_kv = mK.shape[2]
+            mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2)
+            mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2)
+            if const_expr(mLSE is not None):
+                mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1)
         # TileScheduler for varlen, simple grid for non-varlen
         if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
             TileScheduler = SingleTileVarlenScheduler
@@ -682,10 +687,10 @@ def __call__(
         num_batch = (
             mCuSeqlensQ.shape[0] - 1
             if const_expr(mCuSeqlensQ is not None)
-            else mQ.shape[3]
+            else cute.size(mQ.shape[3])
         )
         tile_sched_args = TileSchedulerArguments(
-            num_block=cute.ceil_div(mQ.shape[0], self.tile_m),
+            num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m),
             num_head=cute.size(mQ.shape[2]),
             num_batch=num_batch,
             num_splits=1,
@@ -794,7 +799,7 @@ def kernel(
         )
         seqlen = SeqlenInfoQK.create(
             batch_idx=batch_size,
-            seqlen_q_static=mQ.shape[0],
+            seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1],
             seqlen_k_static=mK.shape[0],
             mCuSeqlensQ=mCuSeqlensQ,
             mCuSeqlensK=mCuSeqlensK,
@@ -814,18 +819,16 @@ def kernel(
         blkQ_shape = (self.tile_m, self.tile_hdim)
         blkK_shape = (self.tile_n, self.tile_hdim)
         blkV_shape = (self.tile_n, self.tile_hdimv)
-        num_head_kv = num_head // self.qhead_per_kvhead
-        if const_expr(not seqlen.has_cu_seqlens_q):
-            mQ_cur = mQ[None, None, num_head, batch_size]
-        else:
-            mQ_cur = cute.domain_offset((seqlen.offset_q, 0), mQ[None, None, num_head])
+        num_head_kv = num_head if const_expr(self.pack_gqa) else num_head // self.qhead_per_kvhead
+        mQ_cur = seqlen.offset_batch_Q(mQ, batch_size, dim=3)[None, None, num_head]
         if const_expr(not seqlen.has_cu_seqlens_k):
             mK_cur = mK[None, None, num_head_kv, batch_size]
             mV_cur = mV[None, None, num_head_kv, batch_size]
         else:
             mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, num_head_kv])
             mV_cur = cute.domain_offset((seqlen.offset_k, 0), mV[None, None, num_head_kv])
-        gQ = cute.local_tile(mQ_cur, blkQ_shape, (m_block, 0))
+        if const_expr(not self.pack_gqa):
+            gQ = cute.local_tile(mQ_cur, blkQ_shape, (m_block, 0))
         gK = cute.local_tile(mK_cur, blkK_shape, (None, 0))
         gV = cute.local_tile(mV_cur, blkV_shape, (None, 0))
 
@@ -957,7 +960,11 @@ def kernel(
         # ///////////////////////////////////////////////////////////////////////////////
         # Start async loads of the last mn-tile, where we take care of the mn residue
         gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx)
-        self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1])
+        if const_expr(not self.pack_gqa):
+            self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1])
+        else:
+            pack_gqa = PackGQA(self.tile_m, self.tile_hdim, self.check_hdim_oob, self.qhead_per_kvhead)
+            pack_gqa.load_Q(mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q)
         cute.arch.cp_async_commit_group()
 
         def preprocess_Q():

flash_attn/cute/interface.py

@@ -459,8 +459,6 @@ def _flash_attn_fwd(
     qhead_per_kvhead = num_head // num_head_kv
     if pack_gqa is None:
         pack_gqa = qhead_per_kvhead > 1
-    if arch // 10 == 12:
-        pack_gqa = False
 
     is_fp8 = v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
     requires_grad = any(t is not None and t.requires_grad for t in [q, k, v, qv])
@@ -566,7 +564,9 @@ def _flash_attn_fwd(
     total_mblocks = batch_size * num_head_kv * num_m_blocks
     num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n
     num_SMs = 132 if is_fake_mode() else torch.cuda.get_device_properties(device).multi_processor_count
-    if num_splits < 1:
+    if arch // 10 == 12:
+        assert num_splits == 1, "SM120 forward only supports num_splits=1"
+    elif num_splits < 1:
         num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128)
 
     # SplitKV uses float32 partial output, which doubles the O buffer size

tests/cute/test_flash_attn.py

@@ -79,12 +79,23 @@ def check_tensor_vs_ref(name, actual, ref, pt, rtol=2, atol=None):
 # When operating fake tensors, we cannot perform data-dependent operations (e.g., `tensor.max()`).
 USE_FAKE_TENSOR = int(os.getenv("FLASH_ATTENTION_FAKE_TENSOR", 0)) == 1
 DISABLE_SPLIT = os.getenv("FLASH_ATTENTION_DISABLE_SPLIT", "FALSE") == "TRUE"
-# SplitKV is not supported on SM90
+# SplitKV is not supported on SM90 or SM120
 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()[0] == 12
 TEST_BWD_ONLY = False
 VERBOSE = True
 
+
+@pytest.mark.skipif(not IS_SM120, reason="SM120-only SplitKV unsupported behavior")
+def test_flash_attn_sm120_rejects_splitkv():
+    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)
+    with pytest.raises(AssertionError, match="SM120 forward only supports num_splits=1"):
+        flash_attn_func(q, k, v, num_splits=3)
+
+
 # @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"])
@@ -321,8 +332,8 @@ 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 and not has_qv 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
             if IS_SM100 and (d >= 192 and dv >= 192) and not (d == 256 and dv == 256):
                 continue
@@ -854,8 +865,8 @@ 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(wangsiyu): SM100 head_dim=256 2CTA kernel does not support pack_gqa yet.
             # pack_gqa=None means auto-enable for GQA/MQA (qhead_per_kvhead > 1)
@@ -1477,8 +1488,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(
@@ -2559,8 +2570,8 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device):
         pack_gqa_vals = [True]
         num_splits_vals = [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
             out_unpad, lse = flash_attn_varlen_func(
                 q_unpad if unpad_q else q,

tests/cute/test_flash_attn_fast.py

@@ -25,6 +25,7 @@
 
 USE_FAKE_TENSOR = int(os.getenv("FLASH_ATTENTION_FAKE_TENSOR", 0)) == 1
 IS_SM90 = torch.cuda.get_device_capability()[0] == 9
+IS_SM120 = torch.cuda.get_device_capability()[0] == 12
 
 
 # ---------------------------------------------------------------------------
@@ -47,8 +48,8 @@
 )
 @maybe_fake_tensor_mode(USE_FAKE_TENSOR)
 def test_flash_attn_output(seqlen_q, seqlen_k, d, causal, num_splits, mha_type, dtype):
-    if IS_SM90 and num_splits > 1:
-        pytest.skip("SM90 fwd doens't support num_splits > 1")
+    if (IS_SM90 or IS_SM120) and num_splits > 1:
+        pytest.skip("SM90/SM120 fwd doesn't support num_splits > 1")
     device = "cuda"
     torch.random.manual_seed(0)
     random.seed(0)