[Bwd,Sm90] Use quack.copy_utils

Author: tridao

flash_attn/cute/block_sparse_utils.py

@@ -1122,8 +1122,7 @@ def _load_q_do_block_sm90(
     else:
         pipeline_Q.producer_acquire(producer_state_Q)
     load_Q(m_block, producer_state=producer_state_Q)
-    with cute.arch.elect_one():
-        load_LSE(m_block, producer_state=producer_state_Q)
+    load_LSE(m_block, producer_state=producer_state_Q)
 
     producer_state_dO_cur = (
         producer_state_dO if const_expr(not Q_stage_eq_dO_stage) else producer_state_Q
@@ -1134,8 +1133,7 @@ def _load_q_do_block_sm90(
     else:
         pipeline_dO.producer_acquire(producer_state_dO_cur)
     load_dO(m_block, producer_state=producer_state_dO_cur)
-    with cute.arch.elect_one():
-        load_dPsum(m_block, producer_state=producer_state_dO_cur)
+    load_dPsum(m_block, producer_state=producer_state_dO_cur)
 
     producer_state_Q.advance()
     producer_state_dO.advance()

flash_attn/cute/flash_bwd_sm90.py

@@ -12,13 +12,13 @@
 from cutlass import Float32, Int32, Boolean, const_expr
 from cutlass.utils import LayoutEnum
 
+from quack import copy_utils
 from quack import layout_utils
 from quack import sm90_utils
 from quack.sm90_utils import gemm_zero_init, gemm_w_idx
 
 from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
 from flash_attn.cute import utils
-from flash_attn.cute import copy_utils
 from flash_attn.cute.mask import AttentionMask
 from flash_attn.cute.seqlen_info import SeqlenInfoQK
 from flash_attn.cute.block_info import BlockInfo
@@ -825,8 +825,7 @@ def load(
                         )
                         load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q))
                         load_Q(first_m_block, producer_state=producer_state_Q)
-                        with cute.arch.elect_one():
-                            load_LSE(first_m_block, producer_state=producer_state_Q)
+                        load_LSE(first_m_block, producer_state=producer_state_Q)
                         producer_state_dO_cur = (
                             producer_state_dO
                             if const_expr(self.Q_stage != self.dO_stage)
@@ -837,25 +836,22 @@ def load(
                         )
                         load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_cur))
                         load_dO(first_m_block, producer_state=producer_state_dO_cur)
-                        with cute.arch.elect_one():
-                            load_dPsum(first_m_block, producer_state=producer_state_dO_cur)
+                        load_dPsum(first_m_block, producer_state=producer_state_dO_cur)
                         producer_state_Q.advance()
                         producer_state_dO.advance()
 
                         for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1):
                             pipeline_Q.producer_acquire(producer_state_Q)
                             load_Q(m_block, producer_state=producer_state_Q)
-                            with cute.arch.elect_one():
-                                load_LSE(m_block, producer_state=producer_state_Q)
+                            load_LSE(m_block, producer_state=producer_state_Q)
                             producer_state_dO_cur = (
                                 producer_state_dO
                                 if const_expr(self.Q_stage != self.dO_stage)
                                 else producer_state_Q
                             )
                             pipeline_dO.producer_acquire(producer_state_dO_cur)
                             load_dO(m_block, producer_state=producer_state_dO_cur)
-                            with cute.arch.elect_one():
-                                load_dPsum(m_block, producer_state=producer_state_dO_cur)
+                            load_dPsum(m_block, producer_state=producer_state_dO_cur)
                             producer_state_Q.advance()
                             producer_state_dO.advance()
                     else: