quack-kernels 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

quack/gemm_config.py

@@ -1,6 +1,7 @@
 # Copyright (C) 2025, Fri Dao.
 import itertools
-from typing import Optional, List
+from typing import Optional, List, Literal
+from functools import partial
 from dataclasses import dataclass
 
 
@@ -13,57 +14,82 @@ class GemmConfig:
     cluster_n: int = 1
     swap_ab: bool = False
     # raster_order: int = 1
-    # max_swizzle_size: int = 8
+    max_swizzle_size: int = 8
 
 
 def get_all_configs(
+    device_capacity: Literal[9, 10] = 9,
     epilogue: Optional[str] = None,
     tune_coop: bool = True,
     # tune_raster_order=True,
 ) -> List[GemmConfig]:
-    tile_n_vals = [128, 144, 160, 176, 192, 208]
-    tile_mn_coop_vals = [(256, tile_n) for tile_n in tile_n_vals] + [
-        (128, 224),
-        (128, 256),
-        # (192, 256),  # Getting IOT instruction (core dumped) in the bwd
-    ]
-    tile_mn_pingpong_vals = [(128, tile_n) for tile_n in tile_n_vals] + [(192, 128)]
-    if epilogue in ["gated"]:
-        tile_mn_coop_vals = [(m, n) for m, n in tile_mn_coop_vals if n % 32 == 0 and m != 192]
-        tile_mn_pingpong_vals = [(m, n) for m, n in tile_mn_pingpong_vals if n % 32 == 0]
-    elif epilogue in ["lse"]:
-        tile_mn_coop_vals = [(m, n) for m, n in tile_mn_coop_vals if m != 192]
-    tile_mn_vals = []
-    if tune_coop:
-        tile_mn_vals += [(m, n, False) for m, n in tile_mn_coop_vals]
-    tile_mn_vals += [(m, n, True) for m, n in tile_mn_pingpong_vals]
-    cluster = [(1, 2), (2, 1)]
-    # cluster = [(1, 1), (1, 2), (2, 1)]
-    if epilogue in ["lse"]:
+    assert device_capacity in [9, 10]
+    if device_capacity == 9:
+        tile_n_vals = [128, 144, 160, 176, 192, 208]
+        tile_mn_coop_vals = [(256, tile_n) for tile_n in tile_n_vals] + [
+            (128, 224),
+            (128, 256),
+            # (192, 256),  # Getting IOT instruction (core dumped) in the bwd
+        ]
+        tile_mn_pingpong_vals = [(128, tile_n) for tile_n in tile_n_vals] + [(192, 128)]
+        if epilogue in ["gated"]:
+            tile_mn_coop_vals = [(m, n) for m, n in tile_mn_coop_vals if n % 32 == 0 and m != 192]
+            tile_mn_pingpong_vals = [(m, n) for m, n in tile_mn_pingpong_vals if n % 32 == 0]
+        elif epilogue in ["lse"]:
+            tile_mn_coop_vals = [(m, n) for m, n in tile_mn_coop_vals if m != 192]
+        tile_mn_vals = []
+        if tune_coop:
+            tile_mn_vals += [(m, n, False) for m, n in tile_mn_coop_vals]
+        tile_mn_vals += [(m, n, True) for m, n in tile_mn_pingpong_vals]
         cluster = [(1, 2), (2, 1)]
-    swap_ab_vals = [False, True]
-    if epilogue in ["lse", "gated"]:
-        swap_ab_vals = [False]
-    # raster_swizzle = (
-    #     [(0, 1)]
-    #     if not tune_raster_order
-    #     else [(1, 1), (1, 2), (1, 4), (1, 8), (2, 1), (2, 2), (2, 4), (2, 8)]
-    # )
-    return [
-        GemmConfig(
-            tile_m=tile_m,
-            tile_n=tile_n,
-            pingpong=pingpong,
-            cluster_m=cluster_m,
-            cluster_n=cluster_n,
-            swap_ab=swap_ab,
-            # raster_order=raster_order,
-            # max_swizzle_size=max_swizzle_size,
+        # cluster = [(1, 1), (1, 2), (2, 1)]
+        if epilogue in ["lse"]:
+            cluster = [(1, 2), (2, 1)]
+        swap_ab_vals = [False, True]
+        if epilogue in ["lse", "gated"]:
+            swap_ab_vals = [False]
+        # raster_swizzle = (
+        #     [(0, 1)]
+        #     if not tune_raster_order
+        #     else [(1, 1), (1, 2), (1, 4), (1, 8), (2, 1), (2, 2), (2, 4), (2, 8)]
+        # )
+        return [
+            GemmConfig(
+                tile_m=tile_m,
+                tile_n=tile_n,
+                pingpong=pingpong,
+                cluster_m=cluster_m,
+                cluster_n=cluster_n,
+                swap_ab=swap_ab,
+                # raster_order=raster_order,
+                # max_swizzle_size=max_swizzle_size,
+            )
+            for (tile_m, tile_n, pingpong), (cluster_m, cluster_n), swap_ab in itertools.product(
+                tile_mn_vals,
+                cluster,
+                swap_ab_vals,
+                # raster_swizzle,
+            )
+        ]
+    elif device_capacity == 10:
+        tile_n_vals = [128, 160, 192, 224, 256]
+        tile_n_64_vals = [128, 192, 256]
+        tile_mn_cluster_vals = (
+            [(128, tile_n, (1, 2)) for tile_n in tile_n_vals]
+            # + [(128, tile_n, (2, 1)) for tile_n in tile_n_64_vals]
+            + [(128, tile_n, (2, 1)) for tile_n in tile_n_vals]
+            + [(256, tile_n, (2, 1)) for tile_n in tile_n_vals]
         )
-        for (tile_m, tile_n, pingpong), (cluster_m, cluster_n), swap_ab in itertools.product(
-            tile_mn_vals,
-            cluster,
-            swap_ab_vals,
-            # raster_swizzle,
-        )
-    ]
+        swap_ab_vals = [False, True]
+        if epilogue in ["lse", "gated"]:
+            swap_ab_vals = [False]
+        max_swizzle_size_vals = [4, 8, 16]
+        GemmConfigCls = partial(GemmConfig, pingpong=False)  # There's no pingpong on Sm100
+        return [
+            GemmConfigCls(
+                tile_m=m, tile_n=n, cluster_m=cm, cluster_n=cn, swap_ab=sab, max_swizzle_size=ms
+            )
+            for (m, n, (cm, cn)), sab, ms in itertools.product(
+                tile_mn_cluster_vals, swap_ab_vals, max_swizzle_size_vals
+            )
+        ]

quack/{gemm_dact_sm90.py → gemm_dact.py}

@@ -1,40 +1,57 @@
 # Copyright (c) 2025, Tri Dao.
-from typing import Optional
+from typing import Optional, Tuple
+from functools import partial
 
 from torch import Tensor
 
 import cutlass
 import cutlass.cute as cute
-from cutlass import const_expr
+from cutlass import Float32, const_expr
 import cutlass.torch as cutlass_torch
 
-from quack.gemm_act_sm90 import GemmActSm90
-from quack.cute_dsl_utils import get_max_active_clusters
+from quack.gemm_sm90 import GemmSm90
+from quack.gemm_sm100 import GemmSm100
+from quack.gemm_default_epi import GemmDefaultEpiMixin
+from quack.gemm_act import GemmActMixin
+from quack.cute_dsl_utils import get_device_capacity, get_max_active_clusters
 from quack.gemm_wrapper_utils import GemmWrapperBase
 import quack.activation
 
 
-class GemmDActSm90(GemmActSm90):
+class GemmDActMixin(GemmActMixin):
     # Different from GemmActSm90, here act_bwd_fn must take in 2 arguments (x, dout)
     # and return 2 arguments (dx, out)
-    EpilogueArguments = GemmActSm90.EpilogueArguments
-    EpilogueParams = GemmActSm90.EpilogueParams
+    EpilogueArguments = GemmActMixin.EpilogueArguments
+    EpilogueParams = GemmActMixin.EpilogueParams
 
     @cute.jit
-    def epi_visit_acc_subtile(
+    def epi_visit_subtile(
         self,
         params: EpilogueParams,
+        epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
     ) -> Optional[cute.Tensor]:
         assert tRS_rC is not None
+        # We don't add C to the accumulator
+        GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC=None)
         tRS_rC_acc = cute.make_fragment_like(tRS_rC, self.acc_dtype)
         tRS_rC_acc.store(tRS_rC.load().to(self.acc_dtype))
         # If we don't have .shape here, the compiler generates local stores and loads
         if const_expr(params.act_fn is not None):
             tRS_rPostAct = cute.make_fragment(tRS_rD.layout.shape, self.acc_dtype)
-            for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True):
-                tRS_rD[i], tRS_rPostAct[i] = params.act_fn(tRS_rC_acc[i], tRS_rD[i])
+            if const_expr(self.arch < 100):
+                for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True):
+                    tRS_rD[i], tRS_rPostAct[i] = params.act_fn(tRS_rC_acc[i], tRS_rD[i])
+            else:
+                for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True):
+                    (
+                        (tRS_rD[2 * i], tRS_rD[2 * i + 1]),
+                        (tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1]),
+                    ) = params.act_fn(
+                        (tRS_rC_acc[2 * i], tRS_rC_acc[2 * i + 1]),
+                        (tRS_rD[2 * i], tRS_rD[2 * i + 1]),
+                    )
         else:
             tRS_rPostAct = tRS_rC_acc
         # Type conversion
@@ -43,6 +60,14 @@ class GemmDActSm90(GemmActSm90):
         return tRS_rPostAct_out
 
 
+class GemmDActSm90(GemmDActMixin, GemmSm90):
+    pass
+
+
+class GemmDActSm100(GemmDActMixin, GemmSm100):
+    pass
+
+
 dact_fn_map = {
     None: None,
     "relu": quack.activation.drelu,
@@ -51,7 +76,7 @@ dact_fn_map = {
 }
 
 
-def gemm_dact_sm90(
+def gemm_dact(
     A: Tensor,  # (l, m, k) or (total_m, k) if varlen_m or (whatever, k) if gather_A with varlen_m
     B: Tensor,  # (l, n, k)
     Out: Tensor,  # (l, m, n) or (total_m, n) if varlen_m
@@ -65,6 +90,7 @@ def gemm_dact_sm90(
     cluster_N: int,
     pingpong: bool = True,
     persistent: bool = True,
+    max_swizzle_size: int = 8,
     cu_seqlens_m: Optional[Tensor] = None,  # (l+1,) cumulative sum of m values for variable length
     A_idx: Optional[Tensor] = None,  # (total_m,) if gather_A with varlen_m
 ) -> None:
@@ -100,10 +126,14 @@ def gemm_dact_sm90(
     }
     GemmWrapperBase.determine_major_orders(tensor_infos, major_configs)
 
-    acc_dtype = cutlass.Float32
+    device_capacity = get_device_capacity(A.device)
+    assert device_capacity[0] in [9, 10], "Only SM90 and SM100 are supported"
+    GemmCls = GemmDActSm100 if device_capacity[0] > 9 else GemmDActSm90
+
+    acc_dtype = Float32
     tile_shape_mn = (tile_M, tile_N)
     cluster_shape_mnk = (cluster_M, cluster_N, 1)
-    if not GemmDActSm90.is_valid_dtypes(
+    if not GemmCls.is_valid_dtypes(
         tensor_infos["A"].dtype,
         tensor_infos["B"].dtype,
         acc_dtype,
@@ -116,9 +146,9 @@ def gemm_dact_sm90(
     max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     GemmWrapperBase.create_cute_tensors(tensor_infos, major_configs)
     act_fn = dact_fn_map[activation]
-    epi_args = GemmDActSm90.EpilogueArguments(tensor_infos["PostAct"].cute_tensor, act_fn)
+    epi_args = GemmCls.EpilogueArguments(tensor_infos["PostAct"].cute_tensor, act_fn)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
-        max_active_clusters, tile_count_semaphore
+        max_active_clusters, tile_count_semaphore, max_swizzle_size=max_swizzle_size
     )
 
     # Create varlen arguments if needed (assumes persistent=True when varlen_m)
@@ -129,7 +159,7 @@ def gemm_dact_sm90(
         max_active_clusters,
         cluster_shape_mnk,
         tensor_infos,
-        GemmDActSm90.num_epi_tensormaps,
+        GemmCls.num_epi_tensormaps,
         pingpong,
     )
 
@@ -142,19 +172,21 @@ def gemm_dact_sm90(
         pingpong,
         persistent,
         tile_count_semaphore is not None,
+        device_capacity,
+        max_swizzle_size,
         cu_seqlens_m is not None,
         A_idx is not None,
         key_tensor_names=("A", "B", "D", "PostAct", "C"),
     )
-    cache = gemm_dact_sm90.compile_cache
+    cache = gemm_dact.compile_cache
     if compile_key not in cache:
-        gemm = GemmDActSm90(
+        if device_capacity[0] == 9:
+            GemmCls = partial(GemmCls, pingpong=pingpong, is_persistent=persistent)
+        gemm = GemmCls(
             acc_dtype,
             tensor_infos["A"].dtype,
             tile_shape_mn,
             cluster_shape_mnk,
-            pingpong=pingpong,
-            is_persistent=persistent,
             gather_A=gather_A,
         )
         cache[compile_key] = cute.compile(
@@ -180,4 +212,4 @@ def gemm_dact_sm90(
     )
 
 
-gemm_dact_sm90.compile_cache = {}
+gemm_dact.compile_cache = {}

quack/gemm_default_epi.py

@@ -0,0 +1,259 @@
+# Copyright (c) 2025, Wentao Guo, Tri Dao.
+from typing import Optional, Tuple
+from functools import partial
+from dataclasses import dataclass
+
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, Float32, Boolean, const_expr
+
+from quack.cute_dsl_utils import ArgumentsBase, ParamsBase
+from quack.gemm_sm90 import GemmSm90
+from quack.gemm_sm100 import GemmSm100
+from quack.sm90_utils import partition_for_epilogue
+import quack.utils as utils
+import quack.copy_utils as copy_utils
+from quack.varlen_utils import VarlenManager
+
+
+class GemmDefaultEpiMixin:
+    num_epi_tensormaps: int = 0
+
+    @dataclass
+    class EpilogueArguments(ArgumentsBase):
+        alpha: Optional[Float32 | cute.Tensor] = None
+        beta: Optional[Float32 | cute.Tensor] = None
+        mRowVecBroadcast: Optional[cute.Tensor] = None
+        mColVecBroadcast: Optional[cute.Tensor] = None
+        add_to_output: bool = False
+
+    @dataclass
+    class EpilogueParams(ParamsBase):
+        alpha: Optional[Float32 | cute.Tensor] = None
+        beta: Optional[Float32 | cute.Tensor] = None
+        mRowVecBroadcast: Optional[cute.Tensor] = None
+        mColVecBroadcast: Optional[cute.Tensor] = None
+
+    def epi_to_underlying_arguments(
+        self, args: EpilogueArguments, *, loc=None, ip=None
+    ) -> EpilogueParams:
+        # Assume all strides are divisible by 32 bits except the last stride
+        new_stride = lambda t: tuple(
+            cute.assume(s, divby=32 // t.element_type.width) if not cute.is_static(s) else s
+            for s in t.stride
+        )
+        mRowVecBroadcast, mColVecBroadcast = [
+            cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t)))
+            if t is not None
+            else None
+            for t in (args.mRowVecBroadcast, args.mColVecBroadcast)
+        ]
+        return self.EpilogueParams(
+            alpha=args.alpha,
+            beta=args.beta,
+            mRowVecBroadcast=mRowVecBroadcast,
+            mColVecBroadcast=mColVecBroadcast,
+        )
+
+    @cute.jit
+    def epi_begin(
+        self,
+        params: EpilogueParams,
+        epi_smem_tensors: Tuple[cute.Tensor, ...],
+        epi_tile: cute.Tile,
+        tiled_copy_t2r: Optional[cute.TiledCopy],
+        tiled_copy_r2s: cute.TiledCopy,
+        tile_coord_mnkl: cute.Coord,
+        varlen_manager: VarlenManager,
+        epilogue_barrier: cutlass.pipeline.NamedBarrier,
+        tidx: Int32,
+    ):
+        alpha, beta = None, None
+        if const_expr(hasattr(params, "alpha") and params.alpha is not None):
+            alpha = utils.load_scalar_or_pointer(params.alpha)
+        if const_expr(hasattr(params, "beta") and params.beta is not None):
+            beta = utils.load_scalar_or_pointer(params.beta)
+        sRowVec, sColVec, *rest = epi_smem_tensors
+        tile_M, tile_N = self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1]
+        batch_idx = tile_coord_mnkl[3]
+        num_epi_threads = self.num_epi_warps * cute.arch.WARP_SIZE
+        # Don't need sync as we assume the previous epilogue has finished
+
+        partition_for_epilogue_fn = partial(
+            partition_for_epilogue,
+            epi_tile=epi_tile,
+            tiled_copy=tiled_copy_t2r if tiled_copy_t2r is not None else tiled_copy_r2s,
+            tidx=tidx,
+            reference_src=tiled_copy_t2r is None,
+        )
+
+        tDsRowVec = None
+        if const_expr(params.mRowVecBroadcast is not None):
+            rowvec_dtype = params.mRowVecBroadcast.element_type
+            num_copy_elems = const_expr(max(32, rowvec_dtype.width)) // rowvec_dtype.width
+            thr_copy_RV = copy_utils.tiled_copy_1d(
+                params.mRowVecBroadcast.element_type, num_epi_threads, num_copy_elems, is_async=True
+            ).get_slice(tidx)
+            mRowVec = params.mRowVecBroadcast[batch_idx, None]
+            gRowVec = cute.local_tile(mRowVec, (tile_N,), (tile_coord_mnkl[1],))
+            tRVgRV = thr_copy_RV.partition_S(gRowVec)
+            tRVsRV = thr_copy_RV.partition_D(sRowVec)
+            tRVcRV = thr_copy_RV.partition_S(cute.make_identity_tensor(tile_N))
+            limit_n = min(mRowVec.shape[0] - tile_coord_mnkl[1] * tile_N, tile_N)
+            tRVpRV = cute.make_fragment((1, cute.size(tRVsRV.shape[1])), Boolean)
+            for m in cutlass.range(cute.size(tRVsRV.shape[1]), unroll_full=True):
+                tRVpRV[0, m] = tRVcRV[0, m] < limit_n
+            cute.copy(thr_copy_RV, tRVgRV, tRVsRV, pred=tRVpRV)
+            # (CPY, CPY_M, CPY_N, EPI_M, EPI_N)
+            tDsRowVec = partition_for_epilogue_fn(
+                cute.make_tensor(
+                    sRowVec.iterator, cute.make_layout((tile_M, tile_N), stride=(0, 1))
+                )
+            )
+            if const_expr(tiled_copy_t2r is not None):
+                tDsRowVec = tiled_copy_r2s.retile(tDsRowVec)
+
+        tDsColVec = None
+        if const_expr(params.mColVecBroadcast is not None):
+            colvec_dtype = params.mColVecBroadcast.element_type
+            num_copy_elems = const_expr(max(32, colvec_dtype.width)) // colvec_dtype.width
+            thr_copy_CV = copy_utils.tiled_copy_1d(
+                params.mColVecBroadcast.element_type, num_epi_threads, num_copy_elems, is_async=True
+            ).get_slice(tidx)
+            if const_expr(not varlen_manager.varlen_m):
+                mColVec = params.mColVecBroadcast[batch_idx, None]
+            else:
+                mColVec = cute.domain_offset(
+                    (varlen_manager.params.cu_seqlens_m[batch_idx],), params.mColVecBroadcast
+                )
+            gColVec = cute.local_tile(mColVec, (tile_M,), (tile_coord_mnkl[0],))
+            tCVgCV = thr_copy_CV.partition_S(gColVec)
+            tCVsCV = thr_copy_CV.partition_D(sColVec)
+            tCVcCV = thr_copy_CV.partition_S(cute.make_identity_tensor(tile_M))
+            limit_m = min(varlen_manager.len_m(batch_idx) - tile_coord_mnkl[0] * tile_M, tile_M)
+            tCVpCV = cute.make_fragment((1, cute.size(tCVsCV.shape[1])), Boolean)
+            for m in cutlass.range(cute.size(tCVsCV.shape[1]), unroll_full=True):
+                tCVpCV[0, m] = tCVcCV[0, m] < limit_m
+            cute.copy(thr_copy_CV, tCVgCV, tCVsCV, pred=tCVpCV)
+            tDsColVec = partition_for_epilogue_fn(
+                cute.make_tensor(
+                    sColVec.iterator, cute.make_layout((tile_M, tile_N), stride=(1, 0))
+                )
+            )
+            if const_expr(tiled_copy_t2r is not None):
+                tDsColVec = tiled_copy_r2s.retile(tDsColVec)
+
+        if const_expr(params.mRowVecBroadcast is not None or params.mColVecBroadcast is not None):
+            cute.arch.cp_async_commit_group()
+            cute.arch.cp_async_wait_group(0)
+            epilogue_barrier.arrive_and_wait()
+        return alpha, beta, tDsRowVec, tDsColVec
+
+    def epi_begin_loop(self, params: EpilogueParams, epi_tensors, epi_coord: cute.Coord):
+        alpha, beta, tDsRowVec, tDsColVec = epi_tensors
+        tDrRowVec_cvt = None
+        if const_expr(tDsRowVec is not None):
+            tDsRowVec_cur = cute.group_modes(tDsRowVec, 3, cute.rank(tDsRowVec))[
+                None, None, None, epi_coord
+            ]
+            # tDrRowVec = cute.make_fragment_like(tDsRowVec_cur)
+            tDrRowVec = cute.make_fragment(tDsRowVec_cur.layout, tDsRowVec_cur.element_type)
+            cute.autovec_copy(cute.filter_zeros(tDsRowVec_cur), cute.filter_zeros(tDrRowVec))
+            tDrRowVec_cvt = cute.make_fragment_like(tDrRowVec, self.acc_dtype)
+            tDrRowVec_cvt.store(tDrRowVec.load().to(self.acc_dtype))
+        tDrColVec_cvt = None
+        if const_expr(tDsColVec is not None):
+            tDsColVec_cur = cute.group_modes(tDsColVec, 3, cute.rank(tDsColVec))[
+                None, None, None, epi_coord
+            ]
+            # This somehow doesn't work, some dim with stride 0 turns to non-zero stride
+            # tDrRowVec = cute.make_fragment_like(tDsRowVec_cur)
+            tDrColVec = cute.make_fragment(tDsColVec_cur.layout, tDsColVec_cur.element_type)
+            cute.autovec_copy(cute.filter_zeros(tDsColVec_cur), cute.filter_zeros(tDrColVec))
+            tDrColVec_cvt = cute.make_fragment_like(tDrColVec, self.acc_dtype)
+            tDrColVec_cvt.store(tDrColVec.load().to(self.acc_dtype))
+        return alpha, beta, tDrRowVec_cvt, tDrColVec_cvt
+
+    @cute.jit
+    def epi_visit_subtile(
+        self,
+        params: EpilogueParams,
+        epi_loop_tensors: Tuple[cute.Tensor, ...],
+        tRS_rD: cute.Tensor,
+        tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Optional[cute.Tensor]:
+        alpha, beta, tDrRowVec, tDrColVec = epi_loop_tensors
+        rD = tRS_rD.load()
+        # Apply alpha scaling to accumulator if alpha is provided (not None)
+        if const_expr(hasattr(params, "alpha") and params.alpha is not None):
+            alpha = utils.load_scalar_or_pointer(params.alpha)
+            rD *= alpha
+        # Apply C with beta scaling
+        if const_expr(tRS_rC is not None):
+            if const_expr(not hasattr(params, "beta") or params.beta is None):
+                # beta is None, default behavior: add C (beta=1.0)
+                rD += tRS_rC.load().to(tRS_rD.element_type)
+            else:
+                beta = utils.load_scalar_or_pointer(params.beta)
+                rD += beta * tRS_rC.load().to(tRS_rD.element_type)
+        tRS_rD.store(rD)
+        if const_expr(tDrRowVec is not None):
+            for i in cutlass.range(cute.size(tDrRowVec), unroll_full=True):
+                tRS_rD[i] += tDrRowVec[i]
+        if const_expr(tDrColVec is not None):
+            for i in cutlass.range(cute.size(tDrColVec), unroll_full=True):
+                tRS_rD[i] += tDrColVec[i]
+        return None
+
+    @staticmethod
+    def epi_smem_bytes_per_stage(
+        args: Optional[EpilogueArguments],
+        cta_tile_shape_mnk: Tuple[int, int, int],
+        epi_tile: cute.Tile,
+    ) -> int:
+        row_vec_smem_size = 0 if args.mRowVecBroadcast is None else cta_tile_shape_mnk[1]
+        col_vec_smem_size = 0 if args.mColVecBroadcast is None else cta_tile_shape_mnk[0]
+        row_vec_dtype = (
+            args.mRowVecBroadcast.element_type if args.mRowVecBroadcast is not None else Float32
+        )
+        col_vec_dtype = (
+            args.mColVecBroadcast.element_type if args.mColVecBroadcast is not None else Float32
+        )
+        return (
+            row_vec_smem_size * row_vec_dtype.width + col_vec_smem_size * col_vec_dtype.width
+        ) // 8
+
+    def epi_get_smem_struct(self, params: EpilogueParams):
+        row_vec_smem_size = 0 if params.mRowVecBroadcast is None else self.cta_tile_shape_mnk[1]
+        col_vec_smem_size = 0 if params.mColVecBroadcast is None else self.cta_tile_shape_mnk[0]
+        row_vec_dtype = (
+            params.mRowVecBroadcast.element_type if params.mRowVecBroadcast is not None else Float32
+        )
+        col_vec_dtype = (
+            params.mColVecBroadcast.element_type if params.mColVecBroadcast is not None else Float32
+        )
+
+        @cute.struct
+        class EpiSharedStorage:
+            sRowVec: cute.struct.Align[cute.struct.MemRange[row_vec_dtype, row_vec_smem_size], 16]
+            sColVec: cute.struct.Align[cute.struct.MemRange[col_vec_dtype, col_vec_smem_size], 16]
+
+        return EpiSharedStorage
+
+    def epi_get_smem_tensors(self, params: EpilogueParams, storage) -> Tuple[cute.Tensor, ...]:
+        sRowVec = None
+        if const_expr(params.mRowVecBroadcast is not None):
+            sRowVec = storage.epi.sRowVec.get_tensor(cute.make_layout(self.cta_tile_shape_mnk[1]))
+        sColVec = None
+        if const_expr(params.mColVecBroadcast is not None):
+            sColVec = storage.epi.sColVec.get_tensor(cute.make_layout(self.cta_tile_shape_mnk[0]))
+        return (sRowVec, sColVec)
+
+
+class GemmDefaultSm90(GemmDefaultEpiMixin, GemmSm90):
+    pass
+
+
+class GemmDefaultSm100(GemmDefaultEpiMixin, GemmSm100):
+    pass