[Sm90] Use functions from quack.sm90_utils

Author: tridao

flash_attn/cute/flash_bwd_postprocess.py

@@ -14,11 +14,12 @@
 from cutlass import Float32, const_expr
 from cutlass.utils import LayoutEnum
 
+import quack.sm90_utils as sm90_utils
+
 from flash_attn.cute import utils
 from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
 from flash_attn.cute import copy_utils
 from flash_attn.cute import ampere_helpers as sm80_utils
-from flash_attn.cute import hopper_helpers as sm90_utils
 from flash_attn.cute.seqlen_info import SeqlenInfoQK
 import cutlass.cute.nvgpu.tcgen05 as tcgen05
 from flash_attn.cute.tile_scheduler import (

flash_attn/cute/flash_bwd_sm90.py

@@ -12,11 +12,12 @@
 from cutlass import Float32, Int32, Boolean, const_expr
 from cutlass.utils import LayoutEnum
 
-from flash_attn.cute import hopper_helpers as sm90_utils
+import quack.sm90_utils as 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.hopper_helpers import gemm_zero_init, gemm_w_idx
 from flash_attn.cute.mask import AttentionMask
 from flash_attn.cute.seqlen_info import SeqlenInfoQK
 from flash_attn.cute.block_info import BlockInfo
@@ -33,21 +34,6 @@
 )
 
 
-def mma_partition_fragment_AB(
-    thr_mma: cute.core.ThrMma, sA: Optional[cute.Tensor], sB: Optional[cute.Tensor], swap_AB: bool
-):
-    if const_expr(not swap_AB):
-        return (
-            thr_mma.make_fragment_A(thr_mma.partition_A(sA)) if sA is not None else None,
-            thr_mma.make_fragment_B(thr_mma.partition_B(sB)) if sB is not None else None,
-        )
-    else:
-        return (
-            thr_mma.make_fragment_B(thr_mma.partition_B(sA)) if sA is not None else None,
-            thr_mma.make_fragment_A(thr_mma.partition_A(sB)) if sB is not None else None,
-        )
-
-
 class FlashAttentionBackwardSm90:
     arch = 90
 
@@ -1033,20 +1019,56 @@ def mma(
         wg_mma_dV = tiled_mma_dV.get_slice(warp_group_thread_layout(warp_group_idx))
         wg_mma_dQ = tiled_mma_dQ.get_slice(warp_group_thread_layout(warp_group_idx))
         # S = Q @ K.T
-        tSrQ, tSrK = mma_partition_fragment_AB(wg_mma_SdP, sQ, sK, self.SdP_swapAB)
+        shape_mnk_S = (self.tile_m, self.tile_n, self.tile_hdim)
+        _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC(
+            wg_mma_SdP, shape_mnk_S, sQ, sK, swap_AB=self.SdP_swapAB
+        )
+        mma_qk_fn = partial(
+            gemm_zero_init, tiled_mma_SdP, shape_mnk_S[:2], tSrQ, tSrK, swap_AB=self.SdP_swapAB
+        )
         # dP = dO @ V.T
-        tdPrdO, tdPrV = mma_partition_fragment_AB(wg_mma_SdP, sdO, sV, self.SdP_swapAB)
+        shape_mnk_dP = (self.tile_m, self.tile_n, self.tile_hdimv)
+        _, tdPrdO, tdPrV = sm90_utils.partition_fragment_ABC(
+            wg_mma_SdP, shape_mnk_dP, sdO, sV, swap_AB=self.SdP_swapAB
+        )
+        mma_dov_fn = partial(
+            gemm_zero_init, tiled_mma_SdP, shape_mnk_dP[:2], tdPrdO, tdPrV, swap_AB=self.SdP_swapAB
+        )
         # dV += P.T @ dO
         sPt = utils.transpose_view(sP) if sP is not None else None
         sdOt = utils.transpose_view(sdO)
-        tdVrPt, tdVrdOt = mma_partition_fragment_AB(wg_mma_dV, sPt, sdOt, self.dKV_swapAB)
+        shape_mnk_dV = (self.tile_n, self.tile_hdimv, self.tile_m)
+        acc_dV, tdVrPt, tdVrdOt = sm90_utils.partition_fragment_ABC(
+            wg_mma_dV, shape_mnk_dV, sPt, sdOt, swap_AB=self.dKV_swapAB
+        )
+        if const_expr(not self.mma_dkv_is_rs):
+            mma_pdo_fn = partial(
+                gemm_w_idx, tiled_mma_dV, acc_dV, tdVrPt, tdVrdOt, swap_AB=self.dKV_swapAB
+            )
+        else:
+            mma_pdo_fn = partial(gemm_w_idx, tiled_mma_dV, acc_dV, tCrB=tdVrdOt)
         # dK += dS.T @ Q
         sdSt = utils.transpose_view(sdS)
         sQt = utils.transpose_view(sQ)
-        tdKrdSt, tdKrQt = mma_partition_fragment_AB(wg_mma_dK, sdSt, sQt, self.dKV_swapAB)
+        shape_mnk_dK = (self.tile_n, self.tile_hdim, self.tile_m)
+        acc_dK, tdKrdSt, tdKrQt = sm90_utils.partition_fragment_ABC(
+            wg_mma_dK, shape_mnk_dK, sdSt, sQt, swap_AB=self.dKV_swapAB
+        )
+        if const_expr(not self.mma_dkv_is_rs):
+            mma_dsq_fn = partial(
+                gemm_w_idx, tiled_mma_dK, acc_dK, tdKrdSt, tdKrQt, swap_AB=self.dKV_swapAB
+            )
+        else:
+            mma_dsq_fn = partial(gemm_w_idx, tiled_mma_dK, acc_dK, tCrB=tdKrQt)
         # dQ = dS @ K
         sKt = utils.transpose_view(sK)
-        tdQrdS, tdQrKt = mma_partition_fragment_AB(wg_mma_dQ, sdS, sKt, self.dQ_swapAB)
+        shape_mnk_dQ = (self.tile_m, self.tile_hdim, self.tile_n)
+        _, tdQrdS, tdQrKt = sm90_utils.partition_fragment_ABC(
+            wg_mma_dQ, shape_mnk_dQ, sdS, sKt, swap_AB=self.dQ_swapAB
+        )
+        mma_dsk_fn = partial(
+            gemm_zero_init, tiled_mma_dQ, shape_mnk_dQ[:2], tdQrdS, tdQrKt, swap_AB=self.dQ_swapAB
+        )
 
         # Smem copy atom tiling
         smem_copy_atom_PdS = utils.get_smem_store_atom(
@@ -1084,53 +1106,6 @@ def mma(
         smem_thr_copy_dQaccum = r2s_tiled_copy_dQaccum.get_slice(tidx)
         tdQsdQaccum = smem_thr_copy_dQaccum.partition_D(sdQaccum)
 
-        dV_shape = (self.tile_n, self.tile_hdimv)
-        acc_dV = cute.make_fragment(
-            tiled_mma_dV.partition_shape_C(dV_shape if not self.dKV_swapAB else dV_shape[::-1]),
-            Float32,
-        )
-        dK_shape = (self.tile_n, self.tile_hdim)
-        acc_dK = cute.make_fragment(
-            tiled_mma_dK.partition_shape_C(dK_shape if not self.dKV_swapAB else dK_shape[::-1]),
-            Float32,
-        )
-
-        mma_qk_fn = partial(
-            gemm_zero_init,
-            tiled_mma_SdP,
-            (self.tile_m, self.tile_n),
-            tSrQ,
-            tSrK,
-            swap_AB=self.SdP_swapAB,
-        )
-        mma_dov_fn = partial(
-            gemm_zero_init,
-            tiled_mma_SdP,
-            (self.tile_m, self.tile_n),
-            tdPrdO,
-            tdPrV,
-            swap_AB=self.SdP_swapAB,
-        )
-        if const_expr(not self.mma_dkv_is_rs):
-            mma_pdo_fn = partial(
-                gemm_w_idx, tiled_mma_dV, acc_dV, tdVrPt, tdVrdOt, swap_AB=self.dKV_swapAB
-            )
-            mma_dsq_fn = partial(
-                gemm_w_idx, tiled_mma_dK, acc_dK, tdKrdSt, tdKrQt, swap_AB=self.dKV_swapAB
-            )
-        else:
-            assert not self.dKV_swapAB
-            mma_pdo_fn = partial(gemm_w_idx, tiled_mma_dV, acc_dV, tCrB=tdVrdOt)
-            mma_dsq_fn = partial(gemm_w_idx, tiled_mma_dK, acc_dK, tCrB=tdKrQt)
-        mma_dsk_fn = partial(
-            gemm_zero_init,
-            tiled_mma_dQ,
-            (self.tile_m, self.tile_hdim),
-            tdQrdS,
-            tdQrKt,
-            swap_AB=self.dQ_swapAB,
-        )
-
         mma_one_m_block_all = partial(
             self.mma_one_m_block,
             warp_group_idx=warp_group_idx,

flash_attn/cute/flash_fwd.py

@@ -25,7 +25,6 @@
 
 from flash_attn.cute import ampere_helpers as sm80_utils
 from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
-from flash_attn.cute import hopper_helpers as sm90_utils
 from flash_attn.cute import utils
 from flash_attn.cute.mask import AttentionMask
 from flash_attn.cute.softmax import Softmax, apply_score_mod_inner
@@ -1206,17 +1205,7 @@ def _get_tiled_mma(self):
             if self.mma_pv_is_rs
             else warpgroup.OperandSource.SMEM,
         )
-        tiled_mma_pv_rs = sm90_utils_basic.make_trivial_tiled_mma(
-            self.dtype,
-            self.dtype,
-            warpgroup.OperandMajorMode.K,
-            warpgroup.OperandMajorMode.MN,
-            Float32,
-            atom_layout_mnk=(self.tile_m // 64, 1, 1),  # Might need (1, 2, 1) for hdim 512
-            tiler_mn=(64, self.tile_hdimv),
-            a_source=warpgroup.OperandSource.RMEM,
-        )
-        return tiled_mma_qk, tiled_mma_pv, tiled_mma_pv_rs
+        return tiled_mma_qk, tiled_mma_pv
 
     def _get_shared_storage_cls(self):
         sQ_struct, sK_struct, sV_struct = [
@@ -1296,7 +1285,7 @@ def __call__(
         LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
         mLSE = utils.select(mLSE, LSE_layout_transpose) if const_expr(mLSE is not None) else None
 
-        tiled_mma_qk, tiled_mma_pv, tiled_mma_pv_rs = self._get_tiled_mma()
+        tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma()
         self.num_mma_threads = tiled_mma_qk.size
         self.num_threads_per_warp_group = 128
         self.num_mma_warp_groups = self.num_mma_threads // self.num_threads_per_warp_group
@@ -1342,7 +1331,7 @@ def __call__(
         self.sP_layout = None
         if const_expr(not self.mma_pv_is_rs):
             self.sP_layout = sm90_utils.make_smem_layout(
-                mV.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_n)
+                mV.element_type, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_n)
             )
 
         SharedStorage = self._get_shared_storage_cls()
@@ -1526,7 +1515,6 @@ def __call__(
             self.gmem_tiled_copy_O,
             tiled_mma_qk,
             tiled_mma_pv,
-            tiled_mma_pv_rs,
             tile_sched_params,
             TileScheduler,
             SharedStorage,
@@ -1572,7 +1560,6 @@ def kernel(
         gmem_tiled_copy_O: cute.TiledCopy,
         tiled_mma_qk: cute.TiledMma,
         tiled_mma_pv: cute.TiledMma,
-        tiled_mma_pv_rs: cute.TiledMma,
         tile_sched_params: ParamsBase,
         TileScheduler: cutlass.Constexpr[Callable],
         SharedStorage: cutlass.Constexpr[Callable],
@@ -1701,7 +1688,6 @@ def kernel(
             self.mma(
                 tiled_mma_qk,
                 tiled_mma_pv,
-                tiled_mma_pv_rs,
                 mQ,
                 mO,
                 mLSE,
@@ -1855,7 +1841,6 @@ def mma(
         self,
         tiled_mma_qk: cute.TiledMma,
         tiled_mma_pv: cute.TiledMma,
-        tiled_mma_pv_rs: cute.TiledMma,
         # softmax: Softmax,
         # acc_O: cute.Tensor,
         mQ: cute.Tensor,
@@ -1891,46 +1876,32 @@ def mma(
         thr_mma_qk = tiled_mma_qk.get_slice(tidx)
         wg_mma_qk = tiled_mma_qk.get_slice(warp_group_thread_layout(warp_group_idx))
         wg_mma_pv = tiled_mma_pv.get_slice(warp_group_thread_layout(warp_group_idx))
-        tSrQ = tiled_mma_qk.make_fragment_A(wg_mma_qk.partition_A(sQ))
-        tSrK = tiled_mma_qk.make_fragment_B(wg_mma_qk.partition_B(sK))
-        if const_expr(self.mma_pv_is_rs):
-            acc_S_shape = tiled_mma_qk.partition_shape_C((self.tile_m, self.tile_n))
-            tOrP = cute.make_fragment(
-                utils.convert_layout_acc_frgA(cute.make_layout(acc_S_shape)), self.dtype
-            )
-        else:
-            tOrP = tiled_mma_pv.make_fragment_A(wg_mma_pv.partition_A(sP))
-        tOrVt = tiled_mma_pv.make_fragment_B(wg_mma_pv.partition_B(sVt))
+        _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC(
+            wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ, sK
+        )
+        mma_qk_fn = partial(
+            sm90_utils.gemm_zero_init, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK
+        )
+        acc_O, tOrP, tOrVt = sm90_utils.partition_fragment_ABC(
+            wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt
+        )
+        mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt)
 
         # ///////////////////////////////////////////////////////////////////////////////
         # Smem copy atom tiling
         # ///////////////////////////////////////////////////////////////////////////////
         smem_copy_atom_P = utils.get_smem_store_atom(self.arch, self.dtype)
         smem_thr_copy_P = cute.make_tiled_copy_C(smem_copy_atom_P, tiled_mma_qk).get_slice(tidx)
-        # tPsP = smem_thr_copy_P.partition_D(sP_pi) if const_expr(sP_pi is not None) else None
         tPsP = smem_thr_copy_P.partition_D(sP) if const_expr(sP is not None) else None
-        # if cute.arch.thread_idx()[0] == 0:
-        #     cute.printf(sP_pi.layout, sP_pi.iterator)
-        #     cute.printf(sP.layout, sP.iterator)
-        #     cute.printf(tPsP.layout, tPsP.iterator)
-
-        self.mma_init()
-
-        acc_shape_O = tiled_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv))
-        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(
-            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)
+        self.mma_init()
 
         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,
             mma_qk_fn=mma_qk_fn,
-            tiled_mma_pv_rs=tiled_mma_pv_rs,
             pipeline_k=pipeline_k,
             pipeline_v=pipeline_v,
             acc_O=acc_O,
@@ -2273,7 +2244,6 @@ def mma_one_n_block(
         n_block: Int32,
         mma_qk_fn: Callable,
         mma_pv_fn: Callable,
-        tiled_mma_pv_rs: cute.TiledMma,
         pipeline_k: cutlass.pipeline.PipelineAsync,
         pipeline_v: cutlass.pipeline.PipelineAsync,
         acc_O: cute.Tensor,
@@ -2333,7 +2303,6 @@ def mma_one_n_block_intrawg_overlap(
         n_block: Int32,
         mma_qk_fn: Callable,
         mma_pv_fn: Callable,
-        tiled_mma_pv_rs: cute.TiledMma,
         pipeline_k: cutlass.pipeline.PipelineAsync,
         pipeline_v: cutlass.pipeline.PipelineAsync,
         acc_O: cute.Tensor,