[Cute,Sm90] Move gemm helper functions to hopper_helpers.py

Author: tridao

flash_attn/cute/copy_utils.py

@@ -119,11 +119,15 @@ def cpasync_reduce_bulk_add_f32(
     ip=None,
 ):
     smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    # cache_hint = cutlass.Int64(0x14F0000000000000)  # EVICT_LAST
     llvm.inline_asm(
         None,
         [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()],
         "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;",
         "l,r,r",
+        # [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()],
+        # "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;",
+        # "l,r,r,l",
         has_side_effects=True,
         is_align_stack=False,
         asm_dialect=llvm.AsmDialect.AD_ATT,

flash_attn/cute/flash_bwd_sm90.py

@@ -21,37 +21,6 @@
 from flash_attn.cute.named_barrier import NamedBarrierFwd, NamedBarrierBwd
 
 
-def mma_zero_init(
-    tiled_mma: cute.TiledMma,
-    shape: cute.Shape,
-    tCrA: cute.Tensor,
-    tCrB: cute.Tensor,
-    A_idx: Optional[Int32] = None,
-    B_idx: Optional[Int32] = None,
-    wg_wait: int = -1,
-) -> cute.Tensor:
-    acc = cute.make_fragment(tiled_mma.partition_shape_C(shape), Float32)
-    rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
-    rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
-    sm90_utils.gemm(tiled_mma, acc, rA, rB, zero_init=True, wg_wait=wg_wait)
-    return acc
-
-
-def mma_sm90(
-    tiled_mma: cute.TiledMma,
-    acc: cute.Tensor,
-    tCrA: cute.Tensor,
-    tCrB: cute.Tensor,
-    zero_init: Boolean,
-    A_idx: Optional[Int32] = None,
-    B_idx: Optional[Int32] = None,
-    wg_wait: int = -1,
-) -> None:
-    rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
-    rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
-    sm90_utils.gemm(tiled_mma, acc, rA, rB, zero_init=zero_init, wg_wait=wg_wait)
-
-
 class FlashAttentionBackwardSm90:
     arch = 90
 
@@ -153,7 +122,6 @@ def _setup_attributes(self):
                 ((self.tile_m, self.tile_n), self.dS_stage),
             ]
         ]
-
         self.sdQaccum_layout = cute.make_layout(self.tile_m * self.tile_hdim)
         # dQaccum R->S
         self.r2s_tiled_copy_dQaccum = copy_utils.tiled_copy_1d(
@@ -792,14 +760,16 @@ def mma(
             Float32,
         )
 
-        mma_qk_fn = partial(mma_zero_init, tiled_mma_SdP, (self.tile_m, self.tile_n), tSrQ, tSrK)
+        mma_qk_fn = partial(
+            sm90_utils.gemm_zero_init, tiled_mma_SdP, (self.tile_m, self.tile_n), tSrQ, tSrK
+        )
         mma_dov_fn = partial(
-            mma_zero_init, tiled_mma_SdP, (self.tile_m, self.tile_n), tdPrdO, tdPrV
+            sm90_utils.gemm_zero_init, tiled_mma_SdP, (self.tile_m, self.tile_n), tdPrdO, tdPrV
         )
-        mma_pdo_fn = partial(mma_sm90, tiled_mma_dV, acc_dV, tdVrPt, tdVrdOt)
-        mma_dsq_fn = partial(mma_sm90, tiled_mma_dK, acc_dK, tdKrdSt, tdKrQt)
+        mma_pdo_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_dV, acc_dV, tdVrPt, tdVrdOt)
+        mma_dsq_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_dK, acc_dK, tdKrdSt, tdKrQt)
         mma_dsk_fn = partial(
-            mma_zero_init, tiled_mma_dQ, (self.tile_m, self.tile_hdim), tdQrdS, tdQrKt
+            sm90_utils.gemm_zero_init, tiled_mma_dQ, (self.tile_m, self.tile_hdim), tdQrdS, tdQrKt
         )
 
         mma_one_m_block_all = partial(

flash_attn/cute/flash_fwd.py

@@ -35,22 +35,6 @@
 from flash_attn.cute.fast_math import FastDivmod
 
 
-def mma_qk(tiled_mma_qk: cute.TiledMma, shape: cute.Shape, tSrQ: cute.Tensor, tSrK: cute.Tensor, smem_idx: Int32, wg_wait: int = -1) -> cute.Tensor:
-    acc_S = cute.make_fragment(tiled_mma_qk.partition_shape_C(shape), Float32)
-    sm90_utils.gemm(
-        tiled_mma_qk, acc_S, tSrQ, tSrK[None, None, None, smem_idx], zero_init=True, wg_wait=wg_wait
-    )
-    return acc_S
-
-
-def mma_pv(tiled_mma_pv: cute.TiledMma, acc_O: cute.Tensor, tOrP: cute.Tensor, tOrVt: cute.Tensor, smem_idx: Int32, zero_init: Boolean, wg_wait: int = -1) -> None:
-    sm90_utils.gemm(
-        tiled_mma_pv, acc_O, tOrP,
-        tOrVt[None, None, None, smem_idx],
-        zero_init=zero_init, wg_wait=wg_wait
-    )
-
-
 class FlashAttentionForwardBase:
 
     arch: int = 80
@@ -1557,7 +1541,6 @@ def load(
                 work_tile = tile_scheduler.get_current_work()
                 # End of persistent scheduler loop
 
-
     @cute.jit
     def mma(
         self,
@@ -1627,8 +1610,10 @@ def mma(
         acc_O = cute.make_fragment(acc_shape_O, Float32)
         smem_copy_params = SimpleNamespace(smem_thr_copy_P=smem_thr_copy_P, tPsP=tPsP)
 
-        mma_qk_fn = partial(mma_qk, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK)
-        mma_pv_fn = partial(mma_pv, tiled_mma_pv, acc_O, tOrP, tOrVt)
+        mma_qk_fn = partial(
+            sm90_utils.gemm_zero_init, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK
+        )
+        mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt)
 
         mma_one_n_block_all = partial(
             self.mma_one_n_block_intrawg_overlap if const_expr(self.intra_wg_overlap) else self.mma_one_n_block,
@@ -1692,7 +1677,7 @@ def mma(
             # First iteration with seqlen masking
             if const_expr(self.intra_wg_overlap):
                 pipeline_k.consumer_wait(kv_consumer_state, pipeline_k.consumer_try_wait(kv_consumer_state))
-                acc_S = mma_qk_fn(kv_consumer_state.index, wg_wait=0)
+                acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=0)
                 pipeline_k.consumer_release(kv_consumer_state)
                 # Use vectorized score modification
                 if cutlass.const_expr(score_mod_fn is not None):
@@ -1767,7 +1752,7 @@ def mma(
             # Last "half" iteration
             if const_expr(self.intra_wg_overlap):
                 pipeline_v.consumer_wait(kv_consumer_state, pipeline_v.consumer_try_wait(kv_consumer_state))
-                mma_pv_fn(kv_consumer_state.index, zero_init=not O_should_accumulate, wg_wait=0)
+                mma_pv_fn(B_idx=kv_consumer_state.index, zero_init=not O_should_accumulate, wg_wait=0)
                 pipeline_v.consumer_release(kv_consumer_state)
                 kv_consumer_state.advance()
             else:
@@ -1821,7 +1806,8 @@ def mma_one_n_block(
         check_inf: cutlass.Constexpr = True,
     ):
         pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read))
-        acc_S = mma_qk_fn(smem_pipe_read.index, wg_wait=-1)
+        # S = Q @ K.T
+        acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1)
         self.warp_scheduler_barrier_arrive()
         warpgroup.wait_group(0)
         pipeline_k.consumer_release(smem_pipe_read)
@@ -1850,7 +1836,8 @@ def mma_one_n_block(
             cute.arch.sync_warp()  # Only need syncwarp since each warp is using its own P values for MmaPV
         pipeline_v.consumer_wait(smem_pipe_read, pipeline_v.consumer_try_wait(smem_pipe_read))
         self.warp_scheduler_barrier_sync()
-        mma_pv_fn(smem_pipe_read.index, wg_wait=0)
+        # O += P @ V
+        mma_pv_fn(B_idx=smem_pipe_read.index, wg_wait=0)
         pipeline_v.consumer_release(smem_pipe_read)
         smem_pipe_read.advance()
         return smem_pipe_read
@@ -1877,9 +1864,11 @@ def mma_one_n_block_intrawg_overlap(
         smem_pipe_read.advance()
         pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read))
         self.warp_scheduler_barrier_sync()
-        acc_S = mma_qk_fn(smem_pipe_read.index, wg_wait=-1)
+        # S = Q @ K.T
+        acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1)
         pipeline_v.consumer_wait(smem_pipe_read_v, pipeline_v.consumer_try_wait(smem_pipe_read_v))
-        mma_pv_fn(smem_pipe_read_v.index, wg_wait=-1)
+        # O += P @ V
+        mma_pv_fn(B_idx=smem_pipe_read_v.index, wg_wait=-1)
         self.warp_scheduler_barrier_arrive()
         warpgroup.wait_group(1)
         pipeline_k.consumer_release(smem_pipe_read)

flash_attn/cute/hopper_helpers.py

@@ -2,7 +2,7 @@
 from typing import Type, Union, Optional
 import cutlass
 import cutlass.cute as cute
-from cutlass import Int32, const_expr
+from cutlass import Int32, Float32, Boolean, const_expr
 from cutlass.cute.nvgpu import warpgroup
 from cutlass._mlir.dialects import llvm
 from cutlass.cutlass_dsl import Numeric, dsl_user_op
@@ -37,6 +37,37 @@ def gemm(
             warpgroup.wait_group(wg_wait)
 
 
+def gemm_zero_init(
+    tiled_mma: cute.TiledMma,
+    shape: cute.Shape,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    A_idx: Optional[Int32] = None,
+    B_idx: Optional[Int32] = None,
+    wg_wait: int = -1,
+) -> cute.Tensor:
+    acc = cute.make_fragment(tiled_mma.partition_shape_C(shape), Float32)
+    rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
+    rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
+    gemm(tiled_mma, acc, rA, rB, zero_init=True, wg_wait=wg_wait)
+    return acc
+
+
+def gemm_w_idx(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    zero_init: Boolean,
+    A_idx: Optional[Int32] = None,
+    B_idx: Optional[Int32] = None,
+    wg_wait: int = -1,
+) -> None:
+    rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
+    rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
+    gemm(tiled_mma, acc, rA, rB, zero_init=zero_init, wg_wait=wg_wait)
+
+
 @dsl_user_op
 def make_smem_layout(
     dtype: Type[Numeric],