# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # 1. Redistributions of source code must retain the above copyright notice, this # list of conditions and the following disclaimer. # 2. Redistributions in binary form must reproduce the above copyright notice, # this list of conditions and the following disclaimer in the documentation # and/or other materials provided with the distribution. # 3. Neither the name of the copyright holder nor the names of its # contributors may be used to endorse or promote products derived from # this software without specific prior written permission. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. from typing import Optional, Type, Tuple, Union import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync, tcgen05 import cutlass.utils as utils import cutlass.pipeline as pipeline import cutlass.cute.testing as testing import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack import cutlass.cute.math as math import inspect # Mathematical constant: log2(e) for converting exp(x) to exp2(x * log2(e)) LOG2_E = 1.4426950408889634 """ A high-performance persistent batched dense GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. - Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") - Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") - Matrix AB12 is MxNxL, L is batch dimension, AB12 can be row-major("N") or column-major("M") This GEMM kernel supports the following features: - Utilizes Tensor Memory Access (TMA) for efficient memory operations - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) - Implements TMA multicast with cluster to reduce L2 memory traffic - Support persistent tile scheduling to better overlap memory load/store with mma between tiles - Support warp specialization to avoid explicit pipelining between mainloop load and mma This GEMM works as follows: 1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. 2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. 3. EPILOGUE warp: - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. - Type convert C matrix to output type. - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) SM100 tcgen05.mma instructions operate as follows: - Read matrix A from SMEM - Read matrix B from SMEM - Write accumulator to TMEM The accumulator in TMEM must then be loaded to registers before writing back to GMEM. Constraints are same as dense_gemm.py: * Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), see detailed valid dtype combinations in below PersistentDenseGemmKernel class documentation * A/B tensor must have the same data type * Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) * Mma tiler N must be 32-256, step 32 * Cluster shape M/N must be positive and power of 2, total cluster size <= 16 * Cluster shape M must be multiple of 2 if use_2cta_instrs=True * The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, Float16/BFloat16, and Int8/Uint8/Float8, respectively. * OOB tiles are not allowed when TMA store is disabled """ class PersistentDenseGemmKernel: """This class implements batched matrix multiplication (C = A x B) with support for various data types and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. :param acc_dtype: Data type for accumulation during computation :type acc_dtype: type[cutlass.Numeric] :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation :type use_2cta_instrs: bool :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) :type mma_tiler_mn: Tuple[int, int] :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing :type cluster_shape_mn: Tuple[int, int] :note: This kernel always uses Tensor Memory Access (TMA) for storing results. :note: In current version, A and B tensor must have the same data type - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported :note: Supported A/B data types: - TFloat32 - Float16/BFloat16 - Int8/Uint8 - Float8E4M3FN/Float8E5M2 :note: Supported accumulator data types: - Float32 (for all floating point A/B data types) - Float16 (only for fp16 and fp8 A/B data types) - Int32 (only for uint8/int8 A/B data types) :note: Supported C data types: - Float32 (for float32 and int32 accumulator data types) - Int32 (for float32 and int32 accumulator data types) - Float16/BFloat16 (for fp16 and fp8 accumulator data types) - Int8/Uint8 (for uint8/int8 accumulator data types) - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) :note: Constraints: - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) - MMA tiler N must be 32-256, step 32 - Cluster shape M must be multiple of 2 if use_2cta_instrs=True - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 Example: >>> gemm = PersistentDenseGemmKernel( ... acc_dtype=cutlass.Float32, ... use_2cta_instrs=True, ... mma_tiler_mn=(128, 128), ... cluster_shape_mn=(2, 2) ... ) >>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream) """ def __init__( self, acc_dtype: Type[cutlass.Numeric], use_2cta_instrs: bool, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], ): """Initializes the configuration for a Blackwell dense GEMM kernel. This configuration includes several key aspects: 1. MMA Instruction Settings (tcgen05): - acc_dtype: Data types for MMA accumulator. - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant with cta_group=2 should be used. 2. Cluster Shape: - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. 3. Output C tensor store mode: - TMA store is always enabled for output tensors. :param acc_dtype: Data type of the accumulator. :type acc_dtype: type[cutlass.Numeric] :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. :type mma_tiler_mn: Tuple[int, int] :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. :type use_2cta_instrs: bool :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. :type cluster_shape_mn: Tuple[int, int] """ self.acc_dtype: Type[cutlass.Numeric] = acc_dtype self.use_2cta_instrs = mma_tiler_mn[0] == 256 self.cluster_shape_mn = cluster_shape_mn # K dimension is deferred in _setup_attributes self.mma_tiler = (*mma_tiler_mn, 1) self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE self.occupancy = 1 # Set specialized warp ids self.epilog_warp_id = ( 0, 1, 2, 3, ) self.mma_warp_id = 4 self.tma_warp_id = 5 self.threads_per_cta = 32 * len((self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)) # Set barrier id for cta sync, epilogue sync and tmem ptr sync self.cta_sync_bar_id = 0 self.epilog_sync_bar_id = 1 self.tmem_ptr_sync_bar_id = 2 self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs This method configures various attributes based on the input tensor properties (data types, leading dimensions) and kernel settings: - Configuring tiled MMA - Computing MMA/cluster/tile shapes - Computing cluster layout - Computing multicast CTAs for A/B - Computing epilogue subtile - Setting up A/B/C stage counts in shared memory - Computing A/B/C shared memory layout - Computing tensor memory allocation columns """ # Configure tiled mma tiled_mma = sm100_utils.make_trivial_tiled_mma( self.a_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler[:2], ) # Compute mma/cluster/tile shapes mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) mma_inst_tile_k = 4 self.mma_tiler = ( self.mma_tiler[0], self.mma_tiler[1], mma_inst_shape_k * mma_inst_tile_k, ) self.mma_tiler_c = ( self.mma_tiler[0], self.mma_tiler[1] // 2, # divide by 2 because Glu advnces by half on N dimension self.mma_tiler[2], ) self.cta_tile_shape_mnk = ( self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), self.mma_tiler[1], self.mma_tiler[2], ) self.cta_tile_shape_mnk_c = ( self.mma_tiler_c[0] // cute.size(tiled_mma.thr_id.shape), self.mma_tiler_c[1], self.mma_tiler_c[2], ) # Compute cluster layout self.cluster_layout_vmnk = cute.tiled_divide( cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,), ) # Compute number of multicast CTAs for A/B self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) self.is_a_mcast = self.num_mcast_ctas_a > 1 self.is_b_mcast = self.num_mcast_ctas_b > 1 # Compute epilogue subtile self.epi_tile = sm100_utils.compute_epilogue_tile_shape( self.cta_tile_shape_mnk, self.use_2cta_instrs, self.ab12_layout, self.ab12_dtype, ) self.epi_tile_c = sm100_utils.compute_epilogue_tile_shape( self.cta_tile_shape_mnk_c, self.use_2cta_instrs, self.c_layout, self.c_dtype, ) # Setup A/B/AB12 stage count in shared memory and ACC stage count in tensor memory self.num_acc_stage, self.num_ab_stage, self.num_ab12_stage, self.num_c_stage = self._compute_stages( tiled_mma, self.mma_tiler, self.a_dtype, self.b_dtype, self.epi_tile, self.epi_tile_c, self.ab12_dtype, self.ab12_layout, self.c_dtype, self.c_layout, self.smem_capacity, self.occupancy, ) # Compute A/B/AB12 shared memory layout self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage, ) self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage, ) self.ab12_smem_layout_staged = sm100_utils.make_smem_layout_epi( self.ab12_dtype, self.ab12_layout, self.epi_tile, self.num_ab12_stage, ) self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( self.c_dtype, self.c_layout, self.epi_tile_c, self.num_c_stage, ) # Compute the number of tensor memory allocation columns self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols(tiled_mma, self.mma_tiler, self.num_acc_stage) @cute.jit def __call__( self, a: cute.Tensor, b: cute.Tensor, ab12: cute.Tensor, c: cute.Tensor, alpha: cutlass.Float32, max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x / (1 + math.exp(-x, True)), ): """Execute the GEMM operation in steps: - Setup static attributes before smem/grid/tma computation - Setup TMA load/store atoms and tensors - Compute grid size with regard to hardware constraints - Define shared storage for kernel - Launch the kernel asynchronously :param a: Input tensor A :type a: cute.Tensor :param b: Input tensor B :type b: cute.Tensor :param ab12: Output tensor AB12 (full GEMM result) :type ab12: cute.Tensor :param c: Output tensor C (SwiGLU result) :type c: cute.Tensor :param max_active_clusters: Maximum number of active clusters :type max_active_clusters: cutlass.Constexpr :param stream: CUDA stream for asynchronous execution :type stream: cuda.CUstream :param epilogue_op: Optional elementwise lambda function to apply to the output tensor :type epilogue_op: cutlass.Constexpr :raises TypeError: If input data types are incompatible with the MMA instruction. :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. """ # Setup static attributes before smem/grid/tma computation self.a_dtype: Type[cutlass.Numeric] = a.element_type self.b_dtype: Type[cutlass.Numeric] = b.element_type self.ab12_dtype: Type[cutlass.Numeric] = ab12.element_type self.c_dtype: Type[cutlass.Numeric] = c.element_type self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() self.ab12_layout = utils.LayoutEnum.from_tensor(ab12) self.c_dtype: Type[cutlass.Numeric] = c.element_type self.c_layout = utils.LayoutEnum.from_tensor(c) # Check if input data types are compatible with MMA instruction if cutlass.const_expr(self.a_dtype != self.b_dtype): raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") # Setup attributes that dependent on gemm inputs self._setup_attributes() tiled_mma = sm100_utils.make_trivial_tiled_mma( self.a_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler[:2], ) atom_thr_size = cute.size(tiled_mma.thr_id.shape) # Setup TMA load for A a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( a_op, a, a_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape, internal_type=(cutlass.TFloat32 if a.element_type is cutlass.Float32 else None), ) # Setup TMA load for B b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( b_op, b, b_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape, internal_type=(cutlass.TFloat32 if b.element_type is cutlass.Float32 else None), ) a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size # Setup TMA store for AB12 and C ab12_cta_v_layout = cute.composition(cute.make_identity_layout(ab12.shape), self.epi_tile) c_cta_v_layout = cute.composition(cute.make_identity_layout(c.shape), self.epi_tile_c) epi_smem_layout = cute.slice_(self.ab12_smem_layout_staged, (None, None, 0)) epi_smem_layout_c = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) tma_atom_ab12, tma_tensor_ab12 = cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileS2GOp(), ab12, epi_smem_layout, ab12_cta_v_layout, ) tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout_c, c_cta_v_layout, ) # Compute grid size self.tile_sched_params, grid = self._compute_grid(ab12, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters) self.buffer_align_bytes = 1024 ab12_smem_size = cute.cosize(self.ab12_smem_layout_staged.outer) # ab12_smem_size: S<1,4,3> o 0 o ((8,16),(32,1),(1,8)):((32,256),(1,0),(0,4096)) c_smem_size = cute.cosize(self.c_smem_layout_staged.outer) # Define shared storage for kernel @cute.struct class SharedStorage: ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] tmem_dealloc_mbar_ptr: cutlass.Int64 tmem_holding_buf: cutlass.Int32 # (EPI_TILE_M, EPI_TILE_N, STAGE) sAB12: cute.struct.Align[ cute.struct.MemRange[ self.ab12_dtype, ab12_smem_size, ], self.buffer_align_bytes, ] # (EPI_TILE_M, EPI_TILE_N, STAGE) sC: cute.struct.Align[ cute.struct.MemRange[ self.c_dtype, c_smem_size, ], self.buffer_align_bytes, ] # c_smem_size: S<1,4,3> o 0 o ((8,16),(32,1),(1,8)):((32,256),(1,0),(0,4096)) # (MMA, MMA_M, MMA_K, STAGE) sA: cute.struct.Align[ cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], self.buffer_align_bytes, ] # (MMA, MMA_N, MMA_K, STAGE) sB: cute.struct.Align[ cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], self.buffer_align_bytes, ] self.shared_storage = SharedStorage # Launch the kernel asynchronously self.kernel( tiled_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, tma_atom_ab12, tma_atom_c, tma_tensor_ab12, tma_tensor_c, self.cluster_layout_vmnk, self.a_smem_layout_staged, self.b_smem_layout_staged, self.ab12_smem_layout_staged, self.c_smem_layout_staged, self.epi_tile, self.epi_tile_c, self.tile_sched_params, epilogue_op, alpha, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], cluster=(*self.cluster_shape_mn, 1), smem=self.shared_storage.size_in_bytes(), stream=stream, ) return # GPU device kernel @cute.kernel def kernel( self, tiled_mma: cute.TiledMma, tma_atom_a: cute.CopyAtom, mA_mkl: cute.Tensor, tma_atom_b: cute.CopyAtom, mB_nkl: cute.Tensor, tma_atom_ab12: Optional[cute.CopyAtom], tma_atom_c: Optional[cute.CopyAtom], mAB12_mnl: cute.Tensor, mC_mnl: cute.Tensor, cluster_layout_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, b_smem_layout_staged: cute.ComposedLayout, ab12_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], epi_tile: cute.Tile, epi_tile_c: cute.Tile, tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, alpha: cutlass.Float32, ): """ GPU device kernel performing the Persistent batched GEMM computation. """ warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) # # Prefetch tma desc # if warp_idx == self.tma_warp_id: cpasync.prefetch_descriptor(tma_atom_a) cpasync.prefetch_descriptor(tma_atom_b) cpasync.prefetch_descriptor(tma_atom_ab12) cpasync.prefetch_descriptor(tma_atom_c) use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 # # Setup cta/thread coordinates # # Coords inside cluster bidx, bidy, bidz = cute.arch.block_idx() mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) is_leader_cta = mma_tile_coord_v == 0 cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) # Coord inside cta tidx, _, _ = cute.arch.thread_idx() # # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier # smem = utils.SmemAllocator() storage = smem.allocate(self.shared_storage) tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr tmem_holding_buf = storage.tmem_holding_buf # Initialize mainloop ab_pipeline (barrier) and states ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) ab_pipeline = pipeline.PipelineTmaUmma.create( barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), num_stages=self.num_ab_stage, producer_group=ab_pipeline_producer_group, consumer_group=ab_pipeline_consumer_group, tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cluster_layout_vmnk, ) # Initialize acc_pipeline (barrier) and states acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) acc_pipeline = pipeline.PipelineUmmaAsync.create( barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), num_stages=self.num_acc_stage, producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, ) # Tensor memory dealloc barrier init if use_2cta_instrs: if warp_idx == self.tma_warp_id: num_tmem_dealloc_threads = 32 with cute.arch.elect_one(): cute.arch.mbarrier_init(tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads) cute.arch.mbarrier_init_fence() # Cluster arrive after barrier init if cute.size(self.cluster_shape_mn) > 1: cute.arch.cluster_arrive_relaxed() # # Setup smem tensor A/B/AB12/C # # (EPI_TILE_M, EPI_TILE_N, STAGE) sAB12 = storage.sAB12.get_tensor(ab12_smem_layout_staged.outer, swizzle=ab12_smem_layout_staged.inner) # (EPI_TILE_M, EPI_TILE_N, STAGE) sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) # (MMA, MMA_M, MMA_K, STAGE) sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) # (MMA, MMA_N, MMA_K, STAGE) sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) # # Compute multicast mask for A/B buffer full # a_full_mcast_mask = None b_full_mcast_mask = None if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1) # # Local_tile partition global tensors # # (bM, bK, RestM, RestK, RestL) gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) # (bN, bK, RestN, RestK, RestL) gB_nkl = cute.local_tile( mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None), # Half of the tile ) # (bM, bN, RestM, RestN, RestL) gAB12_mnl = cute.local_tile(mAB12_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) gC_mnl = cute.local_tile( mC_mnl, cute.slice_(self.mma_tiler_c, (None, None, 0)), (None, None, None), ) k_block_cnt = cute.size(gA_mkl, mode=[3]) # # Partition global tensor for TiledMMA_A/B/C # thr_mma = tiled_mma.get_slice(mma_tile_coord_v) # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) tCgA = thr_mma.partition_A(gA_mkl) # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) tCgB = thr_mma.partition_B(gB_nkl) # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) tCgAB12 = thr_mma.partition_C(gAB12_mnl) tCgC = thr_mma.partition_C(gC_mnl) tidx, _, _ = cute.arch.thread_idx() # # Partition global/shared tensor for TMA load A/B # # TMA load A partition_S/D a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) # ((atom_v, rest_v), STAGE) # ((atom_v, rest_v), RestM, RestK, RestL) tAsA, tAgA = cpasync.tma_partition( tma_atom_a, block_in_cluster_coord_vmnk[2], a_cta_layout, cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3), ) # TMA load B partition_S/D b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) # ((atom_v, rest_v), STAGE) # ((atom_v, rest_v), RestM, RestK, RestL) tBsB, tBgB = cpasync.tma_partition( tma_atom_b, block_in_cluster_coord_vmnk[1], b_cta_layout, cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3), ) # # Partition shared/tensor memory tensor for TiledMMA_A/B/C # # (MMA, MMA_M, MMA_K, STAGE) tCrA = tiled_mma.make_fragment_A(sA) # (MMA, MMA_N, MMA_K, STAGE) tCrB = tiled_mma.make_fragment_B(sB) # (MMA, MMA_M, MMA_N) acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) # (MMA, MMA_M, MMA_N, STAGE) tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) # # Cluster wait before tensor memory alloc # if cute.size(self.cluster_shape_mn) > 1: cute.arch.cluster_wait() else: cute.arch.barrier(barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta) # # Specialized TMA load warp # if warp_idx == self.tma_warp_id: # # Persistent tile scheduling loop # tile_sched = utils.StaticPersistentTileScheduler.create(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) work_tile = tile_sched.initial_work_tile_info() ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) while work_tile.is_valid_tile: # Get tile coord from tile scheduler cur_tile_coord = work_tile.tile_idx mma_tile_coord_mnl = ( cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), cur_tile_coord[1], cur_tile_coord[2], ) # # Slice to per mma tile index # # ((atom_v, rest_v), RestK) tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] # ((atom_v, rest_v), RestK) tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt ab_producer_state.reset_count() peek_ab_empty_status = cutlass.Boolean(1) if ab_producer_state.count < k_block_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) # # Tma load loop # for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1): # Conditionally wait for AB buffer empty ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) # TMA load A/B cute.copy( tma_atom_a, tAgA_slice[(None, ab_producer_state.count)], tAsA[(None, ab_producer_state.index)], tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), mcast_mask=a_full_mcast_mask, ) cute.copy( tma_atom_b, tBgB_slice[(None, ab_producer_state.count)], tBsB[(None, ab_producer_state.index)], tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), mcast_mask=b_full_mcast_mask, ) # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1 ab_producer_state.advance() peek_ab_empty_status = cutlass.Boolean(1) if ab_producer_state.count < k_block_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) # # Advance to next tile # tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() # # Wait A/B buffer empty # ab_pipeline.producer_tail(ab_producer_state) # # Specialized MMA warp # if warp_idx == self.mma_warp_id: # # Bar sync for retrieve tensor memory ptr from shared mem # tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) cute.arch.barrier( barrier_id=self.tmem_ptr_sync_bar_id, number_of_threads=tmem_ptr_read_threads, ) # # Retrieving tensor memory ptr and make accumulator tensor # tmem_ptr = cute.arch.retrieve_tmem_ptr( self.acc_dtype, alignment=16, ptr_to_buffer_holding_addr=tmem_holding_buf, ) # (MMA, MMA_M, MMA_N, STAGE) tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) # # Persistent tile scheduling loop # tile_sched = utils.StaticPersistentTileScheduler.create(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) work_tile = tile_sched.initial_work_tile_info() ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) while work_tile.is_valid_tile: # Get tile coord from tile scheduler cur_tile_coord = work_tile.tile_idx mma_tile_coord_mnl = ( cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), cur_tile_coord[1], cur_tile_coord[2], ) # Set tensor memory buffer for current tile # (MMA, MMA_M, MMA_N) tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] # Peek (try_wait) AB buffer full for k_block = 0 ab_consumer_state.reset_count() peek_ab_full_status = cutlass.Boolean(1) if ab_consumer_state.count < k_block_cnt and is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) # # Wait for accumulator buffer empty # if is_leader_cta: acc_pipeline.producer_acquire(acc_producer_state) # # Reset the ACCUMULATE field for each tile # tiled_mma.set(tcgen05.Field.ACCUMULATE, False) # # Mma mainloop # for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1): if is_leader_cta: # Conditionally wait for AB buffer full ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) # tCtAcc += tCrA * tCrB num_kphases = cute.size(tCrA, mode=[2]) for kphase_idx in cutlass.range(num_kphases, unroll_full=True): kphase_coord = ( None, None, kphase_idx, ab_consumer_state.index, ) cute.gemm( tiled_mma, tCtAcc, tCrA[kphase_coord], tCrB[kphase_coord], tCtAcc, ) # do something with tCtAcc1 and tCtAcc # Enable accumulate on tCtAcc after first kphase tiled_mma.set(tcgen05.Field.ACCUMULATE, True) # Async arrive AB buffer empty ab_pipeline.consumer_release(ab_consumer_state) # Peek (try_wait) AB buffer full for k_block = k_block + 1 ab_consumer_state.advance() peek_ab_full_status = cutlass.Boolean(1) if ab_consumer_state.count < k_block_cnt: if is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) # # Async arrive accumulator buffer full # if is_leader_cta: acc_pipeline.producer_commit(acc_producer_state) acc_producer_state.advance() # # Advance to next tile # tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() # # Wait for accumulator buffer empty # acc_pipeline.producer_tail(acc_producer_state) # # Specialized epilogue warps # if warp_idx < self.mma_warp_id: # # Alloc tensor memory buffer # if warp_idx == self.epilog_warp_id[0]: cute.arch.alloc_tmem( self.num_tmem_alloc_cols, tmem_holding_buf, is_two_cta=use_2cta_instrs, ) # # Bar sync for retrieve tensor memory ptr from shared memory # tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) cute.arch.barrier( barrier_id=self.tmem_ptr_sync_bar_id, number_of_threads=tmem_ptr_read_threads, ) # # Retrieving tensor memory ptr and make accumulator tensor # tmem_ptr = cute.arch.retrieve_tmem_ptr( self.acc_dtype, alignment=16, ptr_to_buffer_holding_addr=tmem_holding_buf, ) # (MMA, MMA_M, MMA_N, STAGE) tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) # # Partition for epilogue # epi_tidx = tidx ( tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc, tTR_rAcc1, ) = self.epilog_tmem_copy_and_partition( epi_tidx, tCtAcc_base, tCgAB12, tCgC, epi_tile, epi_tile_c, use_2cta_instrs, ) tTR_rAB12 = None tTR_rC = None tiled_copy_r2s = None tRS_rAB12 = None tRS_rC = None tRS_sAB12 = None tRS_sC = None bSG_sAB12 = None bSG_sC = None bSG_gAB12_partitioned = None bSG_gC_partitioned = None tTR_rAB12 = cute.make_rmem_tensor(tTR_rAcc.shape, self.ab12_dtype) tTR_rAB12_1 = cute.make_rmem_tensor(tTR_rAcc.shape, self.ab12_dtype) tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) tiled_copy_r2s, tRS_rAB12, tRS_rAB12_1, tRS_rC, tRS_sAB12, tRS_sC = self.epilog_smem_copy_and_partition( tiled_copy_t2r, tTR_rAB12, tTR_rAB12_1, tTR_rC, epi_tidx, sAB12, sC ) ( tma_atom_ab12, tma_atom_c, bSG_sAB12, bSG_sC, bSG_gAB12_partitioned, bSG_gC_partitioned, ) = self.epilog_gmem_copy_and_partition( epi_tidx, tma_atom_ab12, tma_atom_c, tCgAB12, tCgC, epi_tile, epi_tile_c, sAB12, sC, ) # # Persistent tile scheduling loop # tile_sched = utils.StaticPersistentTileScheduler.create(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) work_tile = tile_sched.initial_work_tile_info() acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) # Threads/warps participating in tma store pipeline c_producer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, 32 * len(self.epilog_warp_id), ) c_pipeline = pipeline.PipelineTmaStore.create( num_stages=self.num_ab12_stage, producer_group=c_producer_group, ) while work_tile.is_valid_tile: # Get tile coord from tile scheduler cur_tile_coord = work_tile.tile_idx mma_tile_coord_mnl = ( cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), cur_tile_coord[1], cur_tile_coord[2], ) # # Slice to per mma tile index # # ((ATOM_V, REST_V), EPI_M, EPI_N) bSG_gAB12 = bSG_gAB12_partitioned[ ( None, None, None, *mma_tile_coord_mnl, ) ] bSG_gC = bSG_gC_partitioned[ ( None, None, None, *mma_tile_coord_mnl, ) ] # Set tensor memory buffer for current tile # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] # # Wait for accumulator buffer full # acc_pipeline.consumer_wait(acc_consumer_state) # Get for the single CGA tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) bSG_gAB12 = cute.group_modes(bSG_gAB12, 1, cute.rank(bSG_gAB12)) bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) # # Store accumulator to global memory in subtiles # subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt for subtile_idx in cutlass.range(0, subtile_cnt, 2): # # Load accumulator from tensor memory buffer to register tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] # input tile0 tTR_tAcc_mn1 = tTR_tAcc[(None, None, None, subtile_idx + 1)] # input tile 1 cute.copy(tiled_copy_t2r, tTR_tAcc_mn1, tTR_rAcc1) # copy input tile 1 cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) # copy input tile 0 # Convert to C type acc_vec0 = tiled_copy_r2s.retile(tTR_rAcc).load() # copy input tile 0 acc_vec1 = tiled_copy_r2s.retile(tTR_rAcc1).load() # copy input tile 1 acc_vec0 = acc_vec0 * alpha acc_vec1 = acc_vec1 * alpha # Use exp2 with log2(e) conversion since cute.math.exp is not available # exp(x) = 2^(x * log2(e)) gate_rcp = (1 + cute.math.exp2(-1 * acc_vec1 * LOG2_E, True)).to(self.acc_dtype) res = cute.make_rmem_tensor(gate_rcp.shape, cutlass.Float32) res.store(gate_rcp) for i in cutlass.range_constexpr(cute.size(res.shape)): res[i] = cute.arch.rcp_approx(res[i]) gate = res.load() gate = gate * acc_vec1 acc_vec_c = (acc_vec0 * gate).to(self.c_dtype) acc_vec0 = (acc_vec0).to(self.ab12_dtype) acc_vec1 = (acc_vec1).to(self.ab12_dtype) tRS_rAB12.store(acc_vec0) # both of them are pure Gemm Output. tRS_rAB12_1.store(acc_vec1) tRS_rC.store(acc_vec_c) # Store AB12 and C to shared memory ab12_buffer0 = (num_prev_subtiles + subtile_idx) % self.num_ab12_stage ab12_buffer1 = (num_prev_subtiles + subtile_idx + 1) % self.num_ab12_stage c_buffer = (num_prev_subtiles + subtile_idx // 2) % self.num_c_stage cute.copy( tiled_copy_r2s, tRS_rAB12, tRS_sAB12[(None, None, None, ab12_buffer0)], ) # copy the gemm output for bprop to smem cute.copy( tiled_copy_r2s, tRS_rAB12_1, tRS_sAB12[(None, None, None, ab12_buffer1)], ) cute.copy( tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store cute.arch.fence_proxy( cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta, ) epilog_threads = 32 * len(self.epilog_warp_id) cute.arch.barrier( barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads, ) # TMA store AB12 and C to global memory if warp_idx == self.epilog_warp_id[0]: cute.copy( tma_atom_ab12, bSG_sAB12[(None, ab12_buffer0)], bSG_gAB12[(None, subtile_idx)], ) cute.copy( tma_atom_ab12, bSG_sAB12[(None, ab12_buffer1)], bSG_gAB12[(None, subtile_idx + 1)], ) cute.copy( tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, subtile_idx // 2)], ) # Fence and barrier to make sure shared memory store is visible to TMA store c_pipeline.producer_commit() c_pipeline.producer_acquire() cute.arch.barrier( barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads, ) # # Async arrive accumulator buffer empty # with cute.arch.elect_one(): acc_pipeline.consumer_release(acc_consumer_state) acc_consumer_state.advance() # # Advance to next tile # tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() # # Dealloc the tensor memory buffer # if warp_idx == self.epilog_warp_id[0]: cute.arch.relinquish_tmem_alloc_permit(is_two_cta=use_2cta_instrs) epilog_threads = 32 * len(self.epilog_warp_id) cute.arch.barrier(barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads) if warp_idx == self.epilog_warp_id[0]: if use_2cta_instrs: cute.arch.mbarrier_arrive(tmem_dealloc_mbar_ptr, cta_rank_in_cluster ^ 1) cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0) cute.arch.dealloc_tmem(tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=use_2cta_instrs) # # Wait for C store complete # c_pipeline.producer_tail() def epilog_tmem_copy_and_partition( self, tidx: cutlass.Int32, tAcc: cute.Tensor, gAB12_mnl: cute.Tensor, gC_mnl: cute.Tensor, epi_tile: cute.Tile, epi_tile_c: cute.Tile, use_2cta_instrs: Union[cutlass.Boolean, bool], ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]: """ Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). :param tidx: The thread index in epilogue warp groups :type tidx: cutlass.Int32 :param tAcc: The accumulator tensor to be copied and partitioned :type tAcc: cute.Tensor :param gAB12_mnl: The global tensor AB12 :type gAB12_mnl: cute.Tensor :param gC_mnl: The global tensor C :type gC_mnl: cute.Tensor :param epi_tile: The epilogue tiler :type epi_tile: cute.Tile :param epi_tile_c: The epilogue tiler for C :type epi_tile_c: cute.Tile :param use_2cta_instrs: Whether use_2cta_instrs is enabled :type use_2cta_instrs: bool :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) - tTR_tAcc: The partitioned accumulator tensor - tTR_rAcc: The accumulated tensor in register used to hold t2r results :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] """ # Make tiledCopy for tensor memory load copy_atom_t2r = sm100_utils.get_tmem_load_op( self.cta_tile_shape_mnk, self.ab12_layout, self.ab12_dtype, self.acc_dtype, epi_tile, use_2cta_instrs, ) # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) tAcc_epi = cute.flat_divide( tAcc[((None, None), 0, 0, None)], epi_tile, ) # (EPI_TILE_M, EPI_TILE_N) tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) gAB12_mnl_epi = cute.flat_divide(gAB12_mnl[((None, None), 0, 0, None, None, None)], epi_tile) # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) tTR_gAB12 = thr_copy_t2r.partition_D(gAB12_mnl_epi) # (T2R, T2R_M, T2R_N) tTR_rAcc = cute.make_rmem_tensor(tTR_gAB12[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) tTR_rAcc1 = cute.make_rmem_tensor(tTR_gAB12[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) return tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_rAcc1 def epilog_smem_copy_and_partition( self, tiled_copy_t2r: cute.TiledCopy, tTR_rAB12: cute.Tensor, tTR_rAB12_1: cute.Tensor, tTR_rC: cute.Tensor, tidx: cutlass.Int32, sAB12: cute.Tensor, sC: cute.Tensor, ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor]: """ Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) :type tiled_copy_t2r: cute.TiledCopy :param tTR_rAB12: The partitioned accumulator tensor for AB12 :type tTR_rAB12: cute.Tensor :param tTR_rAB12_1: The partitioned accumulator tensor for AB12 (second tile) :type tTR_rAB12_1: cute.Tensor :param tTR_rC: The partitioned accumulator tensor for C :type tTR_rC: cute.Tensor :param tidx: The thread index in epilogue warp groups :type tidx: cutlass.Int32 :param sAB12: The shared memory tensor for AB12 :type sAB12: cute.Tensor :param sC: The shared memory tensor for C :type sC: cute.Tensor :return: A tuple containing (tiled_copy_r2s, tRS_rAB12, tRS_rAB12_1, tRS_rC, tRS_sAB12, tRS_sC) where: - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) - tRS_rAB12: The partitioned tensor AB12 (register source) - tRS_sAB12: The partitioned tensor AB12 (smem destination) - tRS_rC: The partitioned tensor C (register source) - tRS_sC: The partitioned tensor C (smem destination) :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] """ copy_atom_r2s = sm100_utils.get_smem_store_op(self.ab12_layout, self.ab12_dtype, self.acc_dtype, tiled_copy_t2r) tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) # (R2S, R2S_M, R2S_N, PIPE_D) thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) tRS_sAB12 = thr_copy_r2s.partition_D(sAB12) tRS_sC = thr_copy_r2s.partition_D(sC) # (R2S, R2S_M, R2S_N) tRS_rAB12 = tiled_copy_r2s.retile(tTR_rAB12) tRS_rAB12_1 = tiled_copy_r2s.retile(tTR_rAB12_1) tRS_rC = tiled_copy_r2s.retile(tTR_rC) return tiled_copy_r2s, tRS_rAB12, tRS_rAB12_1, tRS_rC, tRS_sAB12, tRS_sC def epilog_gmem_copy_and_partition( self, tidx: cutlass.Int32, atom1: Union[cute.CopyAtom, cute.TiledCopy], atom2: Union[cute.CopyAtom, cute.TiledCopy], gAB12_mnl: cute.Tensor, gC_mnl: cute.Tensor, epi_tile: cute.Tile, epi_tile_c: cute.Tile, sAB12: cute.Tensor, sC: cute.Tensor, ) -> Tuple[cute.CopyAtom, cute.CopyAtom, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor]: """Make tiledCopy for global memory store, then use it to: - partition register array (source) and global memory (destination) for none TMA store version; - partition shared memory (source) and global memory (destination) for TMA store version. :param tidx: The thread index in epilogue warp groups :type tidx: cutlass.Int32 :param atom1: The copy_atom for AB12 TMA store :type atom1: cute.CopyAtom or cute.TiledCopy :param atom2: The copy_atom for C TMA store :type atom2: cute.CopyAtom or cute.TiledCopy :param gAB12_mnl: The global tensor AB12 :type gAB12_mnl: cute.Tensor :param gC_mnl: The global tensor C :type gC_mnl: cute.Tensor :param epi_tile: The epilogue tiler for AB12 :type epi_tile: cute.Tile :param epi_tile_c: The epilogue tiler for C :type epi_tile_c: cute.Tile :param sAB12: The shared memory tensor for AB12 :type sAB12: cute.Tensor :param sC: The shared memory tensor for C :type sC: cute.Tensor :return: A tuple containing: - tma_atom_ab12: The TMA copy atom for AB12 - tma_atom_c: The TMA copy atom for C - bSG_sAB12: The partitioned shared memory tensor AB12 - bSG_sC: The partitioned shared memory tensor C - bSG_gAB12: The partitioned global tensor AB12 - bSG_gC: The partitioned global tensor C :rtype: Tuple[cute.CopyAtom, cute.CopyAtom, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] """ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) gAB12_epi = cute.flat_divide(gAB12_mnl[((None, None), 0, 0, None, None, None)], epi_tile) gC_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile_c) tma_atom_ab12 = atom1 tma_atom_c = atom2 sAB12_for_tma_partition = cute.group_modes(sAB12, 0, 2) sC_for_tma_partition = cute.group_modes(sC, 0, 2) gAB12_for_tma_partition = cute.group_modes(gAB12_epi, 0, 2) gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) # ((ATOM_V, REST_V), EPI_M, EPI_N) # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) bSG_sAB12, bSG_gAB12 = cpasync.tma_partition( tma_atom_ab12, 0, cute.make_layout(1), sAB12_for_tma_partition, gAB12_for_tma_partition, ) bSG_sC, bSG_gC = cpasync.tma_partition( tma_atom_c, 0, cute.make_layout(1), sC_for_tma_partition, gC_for_tma_partition, ) return tma_atom_ab12, tma_atom_c, bSG_sAB12, bSG_sC, bSG_gAB12, bSG_gC @staticmethod def _compute_stages( tiled_mma: cute.TiledMma, mma_tiler_mnk: Tuple[int, int, int], a_dtype: Type[cutlass.Numeric], b_dtype: Type[cutlass.Numeric], epi_tile: cute.Tile, epi_tile_c: cute.Tile, ab12_dtype: Type[cutlass.Numeric], ab12_layout: utils.LayoutEnum, c_dtype: Type[cutlass.Numeric], c_layout: utils.LayoutEnum, smem_capacity: int, occupancy: int, ) -> Tuple[int, int, int]: """Computes the number of stages for A/B/AB12/C operands based on heuristics. :param tiled_mma: The tiled MMA object defining the core computation. :type tiled_mma: cute.TiledMma :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. :type mma_tiler_mnk: tuple[int, int, int] :param a_dtype: Data type of operand A. :type a_dtype: type[cutlass.Numeric] :param b_dtype: Data type of operand B. :type b_dtype: type[cutlass.Numeric] :param epi_tile: The epilogue tile shape for AB12. :type epi_tile: cute.Tile :param epi_tile_c: The epilogue tile shape for C. :type epi_tile_c: cute.Tile :param ab12_dtype: Data type of operand AB12 (full GEMM output). :type ab12_dtype: type[cutlass.Numeric] :param ab12_layout: Layout enum of operand AB12. :type ab12_layout: utils.LayoutEnum :param c_dtype: Data type of operand C (SwiGLU output). :type c_dtype: type[cutlass.Numeric] :param c_layout: Layout enum of operand C. :type c_layout: utils.LayoutEnum :param smem_capacity: Total available shared memory capacity in bytes. :type smem_capacity: int :param occupancy: Target number of CTAs per SM (occupancy). :type occupancy: int :return: A tuple containing the computed number of stages for: (ACC stages, A/B operand stages, AB12 stages, C stages) :rtype: tuple[int, int, int, int] """ # Default ACC stages num_acc_stage = 2 # Default epilogue stages (TMA store always enabled) num_ab12_stage = 4 num_c_stage = 2 # Calculate smem layout and size for one stage of A, B, AB12, and C a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( tiled_mma, mma_tiler_mnk, a_dtype, 1, # a tmp 1 stage is provided ) b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( tiled_mma, mma_tiler_mnk, b_dtype, 1, # a tmp 1 stage is provided ) ab12_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( ab12_dtype, ab12_layout, epi_tile, 1, ) c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( c_dtype, c_layout, epi_tile_c, 1, ) ab_bytes_per_stage = cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) mbar_helpers_bytes = 1024 ab12_bytes_per_stage = cute.size_in_bytes(ab12_dtype, ab12_smem_layout_staged_one) ab12_bytes = ab12_bytes_per_stage * num_ab12_stage c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) c_bytes = c_bytes_per_stage * num_c_stage # Calculate A/B stages: # Start with total smem per CTA (capacity / occupancy) # Subtract reserved bytes and initial AB12/C stages bytes # Divide remaining by bytes needed per A/B stage num_ab_stage = (smem_capacity // occupancy - (mbar_helpers_bytes + ab12_bytes + c_bytes)) // ab_bytes_per_stage # Refine epilogue stages: # Calculate remaining smem after allocating for A/B stages and reserved bytes # Add remaining unused smem to epilogue # num_ab12_stage += ( # smem_capacity # - occupancy * ab_bytes_per_stage * num_ab_stage # - occupancy * (mbar_helpers_bytes + ab12_bytes) # ) // (occupancy * ab12_bytes_per_stage) # Assert: Check total shared memory usage doesn't exceed capacity total_ab_smem = occupancy * ab_bytes_per_stage * num_ab_stage total_output_smem = occupancy * (ab12_bytes_per_stage * num_ab12_stage + c_bytes_per_stage * num_c_stage) total_smem_used = total_ab_smem + total_output_smem + occupancy * mbar_helpers_bytes return num_acc_stage, num_ab_stage, num_ab12_stage, num_c_stage @staticmethod def _compute_grid( ab12: cute.Tensor, cta_tile_shape_mnk: Tuple[int, int, int], cluster_shape_mn: Tuple[int, int], max_active_clusters: cutlass.Constexpr, ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: """Use persistent tile scheduler to compute the grid size for the output tensor AB12. :param ab12: The output tensor AB12 :type ab12: cute.Tensor :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. :type cta_tile_shape_mnk: tuple[int, int, int] :param cluster_shape_mn: Shape of each cluster in M, N dimensions. :type cluster_shape_mn: tuple[int, int] :param max_active_clusters: Maximum number of active clusters. :type max_active_clusters: cutlass.Constexpr :return: A tuple containing: - tile_sched_params: Parameters for the persistent tile scheduler. - grid: Grid shape for kernel launch. :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] """ ab12_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) gab12 = cute.zipped_divide(ab12, tiler=ab12_shape) num_ctas_mnl = gab12[(0, (None, None, None))].shape cluster_shape_mnl = (*cluster_shape_mn, 1) tile_sched_params = utils.PersistentTileSchedulerParams(num_ctas_mnl, cluster_shape_mnl) grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters) return tile_sched_params, grid @staticmethod def _compute_num_tmem_alloc_cols( tiled_mma: cute.TiledMma, mma_tiler: Tuple[int, int, int], num_acc_stage: int, ) -> int: """ Compute the number of tensor memory allocation columns. :param tiled_mma: The tiled MMA object defining the core computation. :type tiled_mma: cute.TiledMma :param mma_tiler: The shape (M, N, K) of the MMA tile. :type mma_tiler: tuple[int, int, int] :param num_acc_stage: The stage of the accumulator tensor. :type num_acc_stage: int :return: The number of tensor memory allocation columns. :rtype: int """ acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) return num_tmem_alloc_cols class PersistentDenseGemmKernelNoDlpack: """Wrapper around PersistentDenseGemmKernel that avoids DLPack. This wrapper constructs cute.Tensors directly from cute.Pointer, shapes, and explicit layout orders for operands A, B, AB12 and C. """ def __init__( self, acc_dtype: Type[cutlass.Numeric], use_2cta_instrs: bool, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], ): self.kernel = PersistentDenseGemmKernel( acc_dtype=acc_dtype, use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, ) @cute.jit def __call__( self, a_ptr: cute.Pointer, a_shape: cutlass.Constexpr[Tuple[int, int, int]], a_order: cutlass.Constexpr[Tuple[int, int, int]], b_ptr: cute.Pointer, b_shape: cutlass.Constexpr[Tuple[int, int, int]], b_order: cutlass.Constexpr[Tuple[int, int, int]], ab12_ptr: cute.Pointer, ab12_shape: cutlass.Constexpr[Tuple[int, int, int]], ab12_order: cutlass.Constexpr[Tuple[int, int, int]], c_cute: cute.Tensor, alpha: cutlass.Float32, max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x / (1 + math.exp(-x, True)), ): a_cute = cute.make_tensor(a_ptr, layout=cute.make_ordered_layout(a_shape, order=a_order)) b_cute = cute.make_tensor(b_ptr, layout=cute.make_ordered_layout(b_shape, order=b_order)) ab12_cute = cute.make_tensor(ab12_ptr, layout=cute.make_ordered_layout(ab12_shape, order=ab12_order)) self.kernel( a_cute, b_cute, ab12_cute, c_cute, alpha, max_active_clusters, stream, epilogue_op, )