[Bwd,Sm90] For dQ, move wait_group before TMA atomic add

Author: tridao

flash_attn/cute/block_sparse_utils.py

@@ -1352,6 +1352,12 @@ def _store_one_dQaccum_sm90(
     tma_copy_bytes_dQ,
 ):
     """Store dQaccum for a single m_block."""
+    for warp_group_idx in cutlass.range_constexpr(num_mma_warp_groups):
+        cute.arch.cp_async_bulk_wait_group(num_mma_warp_groups - 1 - warp_group_idx, read=True)
+        cute.arch.barrier_arrive(
+            barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
+            number_of_threads=num_threads_per_warp_group + cute.arch.WARP_SIZE,
+        )
     for warp_group_idx in cutlass.range_constexpr(num_mma_warp_groups):
         cute.arch.barrier(
             barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx,
@@ -1364,12 +1370,6 @@ def _store_one_dQaccum_sm90(
                 tma_copy_bytes_dQ,
             )
         cute.arch.cp_async_bulk_commit_group()
-    for warp_group_idx in cutlass.range_constexpr(num_mma_warp_groups):
-        cute.arch.cp_async_bulk_wait_group(num_mma_warp_groups - 1 - warp_group_idx, read=True)
-        cute.arch.barrier_arrive(
-            barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
-            number_of_threads=num_threads_per_warp_group + cute.arch.WARP_SIZE,
-        )
 
 
 @cute.jit

flash_attn/cute/flash_bwd_sm90.py

@@ -668,11 +668,6 @@ def kernel(
                     qhead_per_kvhead_divmod,
                 )
             if warp_idx == 1:
-                for warp_group_idx in cutlass.range(self.num_mma_warp_groups):
-                    cute.arch.barrier_arrive(
-                        barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
-                        number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE,
-                    )
                 self.dQaccum_store(
                     mdQaccum,
                     sdQaccum,
@@ -1605,6 +1600,16 @@ def dQaccum_store(
                         m_block = m_block_min + iter_idx
                         m_block_safe = m_block
 
+                        for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
+                            cute.arch.cp_async_bulk_wait_group(
+                                self.num_mma_warp_groups - 1 - warp_group_idx, read=True
+                            )
+                            cute.arch.barrier_arrive(
+                                barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
+                                number_of_threads=self.num_threads_per_warp_group
+                                + cute.arch.WARP_SIZE,
+                            )
+
                         for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
                             cute.arch.barrier(
                                 barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx,
@@ -1618,15 +1623,6 @@ def dQaccum_store(
                                     self.tma_copy_bytes["dQ"],
                                 )
                             cute.arch.cp_async_bulk_commit_group()
-                        for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
-                            cute.arch.cp_async_bulk_wait_group(
-                                self.num_mma_warp_groups - 1 - warp_group_idx, read=True
-                            )
-                            cute.arch.barrier_arrive(
-                                barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
-                                number_of_threads=self.num_threads_per_warp_group
-                                + cute.arch.WARP_SIZE,
-                            )
                 else:
                     dQaccum_store_block_sparse_bwd_sm90(
                         blocksparse_tensors,
@@ -1643,3 +1639,5 @@ def dQaccum_store(
                     )
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
+
+        cute.arch.cp_async_bulk_wait_group(0, read=True)