diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_all_gather_gemm_blackwell.py b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_all_gather_gemm_blackwell.py index c69b5135e..c7f03a8f3 100644 --- a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_all_gather_gemm_blackwell.py +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_all_gather_gemm_blackwell.py @@ -27,50 +27,17 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import argparse -import os from typing import Optional, Tuple, Type, Union - -import numpy as np -import torch -import torch.distributed._symmetric_memory as symm_mem -import torch.distributed as dist import cuda.bindings.driver as cuda -from cuda.bindings import driver -try: - from cuda.core import Device -except ImportError: - from cuda.core.experimental import Device -from cuda.pathfinder import load_nvidia_dynamic_lib import cutlass import cutlass.cute as cute -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline -import cutlass.utils.blackwell_helpers as sm100_utils -import cutlass.utils.distributed as dist_helpers +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack -from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait - -try: - import nvshmem.core -except ImportError as exc: - raise ImportError( - "nvshmem4py is required but not installed. Please install it using:\n" - " For CUDA 12: pip install nvshmem4py-cu12\n" - " For CUDA 13: pip install nvshmem4py-cu13\n" - "Note: nvshmem4py version >= 0.1.3 is recommended." - ) from None - -try: - load_nvidia_dynamic_lib("nvshmem_host") -except RuntimeError as exc: - raise ImportError( - "nvshmem lib is required but not installed. Please install it using:\n" - " For CUDA 12: pip install nvidia-nvshmem-cu12\n" - " For CUDA 13: pip install nvidia-nvshmem-cu13\n" - ) from None +from cuda.bindings import driver """ A high-performance All gather + dense GEMM example for the NVIDIA Blackwell SM100 architecture @@ -108,7 +75,7 @@ Input arguments to this example is same as dense_gemm_persistent.py. .. code-block:: bash - torchrun --nproc_per_node=8 examples/distributed/distributed_all_gather_gemm_blackwell.py \ + torchrun --nproc_per_node=8 examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_all_gather_gemm_blackwell.py \ --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ --mnkl 8192,8192,8192,1 \ @@ -118,7 +85,7 @@ To collect performance with NCU profiler: .. code-block:: bash - ncu torchrun --nproc_per_node=8 examples/distributed/distributed_all_gather_gemm_blackwell.py \ + ncu torchrun --nproc_per_node=8 examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_all_gather_gemm_blackwell.py \ --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ --mnkl 8192,8192,8192,1 \ @@ -197,34 +164,51 @@ class SyncNvlDevices: counter_ptr = cute.make_ptr( cutlass.Int32, device_arrival_counters[peer_rank], - cute.AddressSpace.gmem, + cutlass.AddressSpace.gmem, assumed_align=4, ) - _ = utils.distributed.atomicAdd(counter_ptr, val_one) + _ = cute.arch.atomic_add( + ptr=counter_ptr.llvm_ptr, + val=val_one, + sem="relaxed", + scope="sys", + ) # --- Phase 3: Wait for Arrivals (only thread 0) --- # Poll local arrival counter local_counter_ptr = cute.make_ptr( cutlass.Int32, device_arrival_counters[device_idx], - cute.AddressSpace.gmem, + cutlass.AddressSpace.gmem, assumed_align=4, ) - local_counter_tensor = cute.make_tensor( - local_counter_ptr, cute.make_layout(shape=(1,), stride=(1,)) + current_arrivals = cute.arch.load( + local_counter_ptr.llvm_ptr, + cutlass.Int32, + sem="relaxed", + scope="sys", ) - current_arrivals = utils.distributed.ld_bypass(local_counter_tensor)[0] # Wait until NP-1 peers have signaled # Add a small sleep/yield if supported and necessary to prevent excessive polling load while current_arrivals < max_arrivals: # Consider adding a small delay here if possible (__nanosleep equivalent) # For now, just poll: - current_arrivals = utils.distributed.ld_bypass(local_counter_tensor)[0] + current_arrivals = cute.arch.load( + local_counter_ptr.llvm_ptr, + cutlass.Int32, + sem="relaxed", + scope="sys", + ) # --- Phase 4: Reset Local Counter (only thread 0) --- # Atomically subtract (NP-1) from the local counter reset_val = cutlass.Int32(-(self.num_of_parallelism - 1)) - _ = utils.distributed.atomicAdd(local_counter_ptr, reset_val) + _ = cute.arch.atomic_add( + ptr=local_counter_ptr.llvm_ptr, + val=reset_val, + sem="relaxed", + scope="sys", + ) @cute.jit def __call__( @@ -250,83 +234,6 @@ class SyncNvlDevices: grid=grid, block=[1, 1, 1], cluster=(1, 1, 1), smem=0, stream=stream ) -def _compute_stages( - tiled_mma: cute.TiledMma, - mma_tiler_mnk: Tuple[int, int, int], - a_dtype: Type[cutlass.Numeric], - b_dtype: Type[cutlass.Numeric], - c_dtype: Type[cutlass.Numeric], - smem_capacity: int, - occupancy: int, - use_tma_store: bool, - c_smem_layout: Union[cute.Layout, None], -) -> Tuple[int, int, int]: - """Computes the number of stages for A/B/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 c_dtype: Data type of operand C (output). - :type c_dtype: type[cutlass.Numeric] - :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 - :param use_tma_store: Whether TMA store is enabled. - :type use_tma_store: bool - :param c_smem_layout: Layout of C operand in shared memory, or None if not using TMA store. - :type c_smem_layout: Union[cute.Layout, None] - - :return: A tuple containing the computed number of stages for: - (ACC stages, A/B operand stages, C stages) - :rtype: tuple[int, int, int] - """ - # Default ACC stages - num_acc_stage = 2 - - # Default C stages - num_c_stage = 2 if use_tma_store else 0 - - # Calculate smem layout and size for one stage of A, B, and C with 1-stage - a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( - tiled_mma, mma_tiler_mnk, a_dtype, 1 - ) - b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( - tiled_mma, mma_tiler_mnk, b_dtype, 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 - - c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) - 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 C stages bytes - # Divide remaining by bytes needed per A/B stage - num_ab_stage = ( - smem_capacity // occupancy - (mbar_helpers_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 - if use_tma_store: - num_c_stage += ( - smem_capacity - - occupancy * ab_bytes_per_stage * num_ab_stage - - occupancy * (mbar_helpers_bytes + c_bytes) - ) // (occupancy * c_bytes_per_stage) - return num_acc_stage, num_ab_stage, num_c_stage - class PersistentDenseGemmKernel: """This class implements batched matrix multiplication (C = A x B) with support for various data types @@ -434,16 +341,31 @@ class PersistentDenseGemmKernel: self.occupancy = 1 # Set specialized warp ids - self.epilog_warp_id = (0, 1, 2, 3) + 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.epilog_threads = 32 * len(self.epilog_warp_id) + # Barrier IDs: 0 reserved for sync_threads (pipeline internal), 1=epilog sync, + # 2=tmem alloc/retrieve (mma+epilog warps), 3=tmem dealloc (epilog warps only) self.epilog_sync_bar_id = 1 self.tmem_alloc_sync_bar_id = 2 self.tmem_dealloc_sync_bar_id = 3 + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=self.epilog_threads, + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") self.gated_a_load = gated_a_load @@ -462,8 +384,9 @@ class PersistentDenseGemmKernel: - Computing tensor memory allocation columns """ # Configure tiled mma - tiled_mma = sm100_utils.make_trivial_tiled_mma( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, + self.b_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, @@ -499,7 +422,7 @@ class PersistentDenseGemmKernel: # Compute epilogue subtile if cutlass.const_expr(self.use_tma_store): - self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, @@ -508,38 +431,43 @@ class PersistentDenseGemmKernel: else: self.epi_tile = self.cta_tile_shape_mnk[:2] - c_smem_layout = None - if cutlass.const_expr(self.use_tma_store): - c_smem_layout = sm100_utils.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, 1 - ) - # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory - self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( tiled_mma, self.mma_tiler, self.a_dtype, self.b_dtype, + self.epi_tile, self.c_dtype, + self.c_layout, self.smem_capacity, self.occupancy, self.use_tma_store, - c_smem_layout, ) # Compute A/B/C 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.a_smem_layout_staged = utils.sm100.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.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, ) - - self.c_smem_layout_staged = None - if self.use_tma_store: - self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( - self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + self.c_smem_layout_staged = ( + utils.sm100.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, ) + if self.use_tma_store + else None + ) # Compute the number of tensor memory allocation columns self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( @@ -554,7 +482,7 @@ class PersistentDenseGemmKernel: c: cute.Tensor, max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, - flag_offset: int, + flag_offset: cutlass.Constexpr, gate_a_flags: cute.Tensor, epilogue_op: cutlass.Constexpr = lambda x: x, ): @@ -604,8 +532,9 @@ class PersistentDenseGemmKernel: # Setup attributes that dependent on gemm inputs self._setup_attributes() - tiled_mma = sm100_utils.make_trivial_tiled_mma( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, + self.b_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, @@ -615,7 +544,7 @@ class PersistentDenseGemmKernel: 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( + a_op = utils.sm100.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)) @@ -632,7 +561,7 @@ class PersistentDenseGemmKernel: ) # Setup TMA load for B - b_op = sm100_utils.cluster_shape_to_tma_atom_B( + b_op = utils.sm100.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)) @@ -656,9 +585,12 @@ class PersistentDenseGemmKernel: tma_atom_c = None tma_tensor_c = None if cutlass.const_expr(self.use_tma_store): - epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + self.epi_tile, ) # Compute grid size @@ -666,6 +598,47 @@ class PersistentDenseGemmKernel: c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters ) + self.buffer_align_bytes = 1024 + + c_smem_size = ( + cute.cosize(self.c_smem_layout_staged.outer) if self.use_tma_store else 0 + ) + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + # full+empty contiguous so pipeline.create() can compute empty = base + num_stages + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + c_smem_size, + ], + self.buffer_align_bytes, + ] + # (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 synchronously gate_a_flag_tensor = cute.make_tensor( gate_a_flags.iterator + flag_offset, @@ -752,34 +725,22 @@ class PersistentDenseGemmKernel: # # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier # - # Define shared storage for kernel - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[ - cutlass.Int64, self.num_acc_stage * 2 - ] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) + storage = smem.allocate(self.shared_storage) # 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_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Warp) + 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, + enable_multicast_signaling=True, defer_sync=True, - ).make_participants() + ) # Initialize acc_pipeline (barrier) and states acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) @@ -797,44 +758,46 @@ class PersistentDenseGemmKernel: cta_layout_vmnk=cluster_layout_vmnk, defer_sync=True, ) - tmem_alloc_barrier = pipeline.NamedBarrier( - barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), - ) + + # tmem_dealloc_barrier: only epilog warps (128 threads) participate in dealloc sync. + # Not needed for use_tma_store (c_pipeline.producer_tail() provides ordering instead). tmem_dealloc_barrier = None if cutlass.const_expr(not self.use_tma_store): tmem_dealloc_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_dealloc_sync_bar_id, - num_threads=32 * len(self.epilog_warp_id), + num_threads=self.epilog_threads, ) + # Tensor memory dealloc barrier init tmem = utils.TmemAllocator( storage.tmem_holding_buf.ptr, - barrier_for_retrieve=tmem_alloc_barrier, + barrier_for_retrieve=self.tmem_alloc_barrier, allocator_warp_id=self.epilog_warp_id[0], is_two_cta=use_2cta_instrs, two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, ) - # Cluster arrive after barrier init - pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + # Cluster arrive after barrier init (deferred from pipeline creates above) + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = ( + storage.sC.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + if self.use_tma_store + else None + ) # (MMA, MMA_M, MMA_K, STAGE) - sA = smem.allocate_tensor( - element_type=self.a_dtype, - layout=a_smem_layout_staged.outer, - byte_alignment=128, - swizzle=a_smem_layout_staged.inner, + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner ) # (MMA, MMA_N, MMA_K, STAGE) - sB = smem.allocate_tensor( - element_type=self.b_dtype, - layout=b_smem_layout_staged.outer, - byte_alignment=128, - swizzle=b_smem_layout_staged.inner, + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner ) # @@ -925,7 +888,7 @@ class PersistentDenseGemmKernel: # # Cluster wait before tensor memory alloc # - pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) # # Specialized TMA load warp @@ -940,16 +903,30 @@ class PersistentDenseGemmKernel: ) work_tile = tile_sched.initial_work_tile_info() + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) ######################################################### if self.gated_a_load: # wait_on_flag if lane_idx == 0: - # Wait until input_ready_flags[iteration_i] is set - ready_flag = utils.distributed.ld_bypass(gate_a_flag)[0] - # Need to use volatile load to prevent compiler optimizing away the polling loop + # Wait until input_ready_flags[iteration_i] is set. + # Use cop="cv" (cache-volatile) so every iteration bypasses L1 + # and reads from L2/DRAM, guaranteeing visibility of the plain + # cuMemsetD32Async write issued by the copy stream. + # ld.relaxed.sys requires a matching st.release on the writer side; + # cuMemsetD32Async does not provide that, so cop="cv" is correct here. + ready_flag = cute.arch.load( + gate_a_flag.iterator, + cutlass.Int32, + cop="cv", + ) while ready_flag == 0: - # Keep polling until ready using volatile load - ready_flag = utils.distributed.ld_bypass(gate_a_flag)[0] + ready_flag = cute.arch.load( + gate_a_flag.iterator, + cutlass.Int32, + cop="cv", + ) cute.arch.sync_warp() ######################################################### while work_tile.is_valid_tile: @@ -974,36 +951,45 @@ class PersistentDenseGemmKernel: ] # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt - ab_producer.reset() - peek_ab_empty_status = ab_producer.try_acquire() + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) # # Tma load loop # for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): # Conditionally wait for AB buffer empty - handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) # TMA load A/B cute.copy( tma_atom_a, - tAgA_slice[(None, handle.count)], - tAsA[(None, handle.index)], - tma_bar_ptr=handle.barrier, + 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, handle.count)], - tBsB[(None, handle.index)], - tma_bar_ptr=handle.barrier, + 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_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() peek_ab_empty_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_empty_status = ab_producer.try_acquire() + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) # # Advance to next tile @@ -1014,16 +1000,20 @@ class PersistentDenseGemmKernel: # # Wait A/B buffer empty # - ab_producer.tail() + ab_pipeline.producer_tail(ab_producer_state) # # Specialized MMA warp # if warp_idx == self.mma_warp_id: # - # Retrieving tensor memory ptr and make accumulator tensor + # Bar sync for retrieve tensor memory ptr from shared mem # tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) # (MMA, MMA_M, MMA_N, STAGE) tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) @@ -1036,6 +1026,9 @@ class PersistentDenseGemmKernel: ) 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 ) @@ -1054,10 +1047,12 @@ class PersistentDenseGemmKernel: tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] # Peek (try_wait) AB buffer full for k_tile = 0 - ab_consumer.reset() + ab_consumer_state.reset_count() peek_ab_full_status = cutlass.Boolean(1) - if is_leader_cta: - peek_ab_full_status = ab_consumer.try_wait() + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) # # Wait for accumulator buffer empty @@ -1065,42 +1060,34 @@ class PersistentDenseGemmKernel: 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_tile in range(k_tile_cnt): if is_leader_cta: # Conditionally wait for AB buffer full - handle = ab_consumer.wait_and_advance(peek_ab_full_status) - + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) # tCtAcc += tCrA * tCrB - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - - cute.gemm( - tiled_mma, - tCtAcc, - tCrA[kblk_crd], - tCrB[kblk_crd], - tCtAcc, - ) - # Enable accumulate on tCtAcc after first kblock - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, ab_consumer_state.index) + cute.gemm( + tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc + ) # Async arrive AB buffer empty - handle.release() + ab_pipeline.consumer_release(ab_consumer_state) - # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 - peek_ab_full_status = cutlass.Boolean(1) - if handle.count + 1 < k_tile_cnt: - peek_ab_full_status = ab_consumer.try_wait() + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) # # Async arrive accumulator buffer full @@ -1119,16 +1106,6 @@ class PersistentDenseGemmKernel: # Wait for accumulator buffer empty # acc_pipeline.producer_tail(acc_producer_state) - - sC = None - if cutlass.const_expr(self.use_tma_store): - # (EPI_TILE_M, EPI_TILE_N, STAGE) - sC = smem.allocate_tensor( - element_type=self.c_dtype, - layout=c_smem_layout_staged.outer, - byte_alignment=128, - swizzle=c_smem_layout_staged.inner, - ) # # Specialized epilogue warps # @@ -1138,306 +1115,227 @@ class PersistentDenseGemmKernel: # tmem.allocate(self.num_tmem_alloc_cols) + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() # # Retrieving tensor memory ptr and make accumulator tensor # - - tmem.wait_for_alloc() tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) # (MMA, MMA_M, MMA_N, STAGE) tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - # Persistent tile scheduling loop for epilogue + # + # Partition for epilogue + # + epi_tidx = tidx + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = None + tiled_copy_r2s = None + simt_atom = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + tTR_gC_partitioned = None + if cutlass.const_expr(self.use_tma_store): + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + else: + ( + simt_atom, + tTR_rC, + tTR_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tiled_copy_t2r, tCgC, epi_tile, 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 + ) + + c_pipeline = None if cutlass.const_expr(self.use_tma_store): - - assert tma_atom_c is not None and sC is not None - self.epilogue_tma_store( - tidx, - warp_idx, - acc_pipeline, - tiled_mma, - tma_atom_c, - tCtAcc_base, - sC, - tCgC, - epi_tile, - tile_sched, - epilogue_op, + # 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_c_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 + # + bSG_gC = None + tTR_gC = None + if cutlass.const_expr(self.use_tma_store): + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + else: + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + 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) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + if cutlass.const_expr(self.use_tma_store): + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + else: + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_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(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + if cutlass.const_expr(self.use_tma_store): + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + 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( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # + # Convert to C type + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tTR_rC.store(acc_vec) + + # + # Store C to global memory + # + cute.copy( + simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)] + ) + + # + # 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() + + # + # Wait for C store complete / sync before TMEM dealloc + # + if cutlass.const_expr(self.use_tma_store): + c_pipeline.producer_tail() else: - self.epilogue( - tidx, - acc_pipeline, - tiled_mma, - tCtAcc_base, - tCgC, - epi_tile, - tile_sched, - epilogue_op, - tmem_dealloc_barrier, - ) - + tmem_dealloc_barrier.arrive_and_wait() # # Dealloc the tensor memory buffer # tmem.relinquish_alloc_permit() tmem.free(tmem_ptr) - @cute.jit - def epilogue_tma_store( - self, - epi_tidx: cutlass.Int32, - warp_idx: cutlass.Int32, - acc_pipeline: pipeline.PipelineAsync, - tiled_mma: cute.TiledMma, - tma_atom_c: cute.CopyAtom, - # Input of epilogue - tCtAcc_base: cute.Tensor, - # Staging of epilogue - sC: cute.Tensor, - # Output of epilogue - tCgC: cute.Tensor, - epi_tile: cute.Tile, - tile_sched: utils.StaticPersistentTileScheduler, - epilogue_op: cutlass.Constexpr, - ) -> None: - tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( - epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs - ) - - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( - tiled_copy_t2r, tTR_rC, epi_tidx, sC - ) - - # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) - tCgC_epi = cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - # ((ATOM_V, REST_V), EPI_M, EPI_N) - # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) - bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( - tma_atom_c, - 0, - cute.make_layout(1), - cute.group_modes(sC, 0, 2), - cute.group_modes(tCgC_epi, 0, 2), - ) - - 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_c_stage, producer_group=c_producer_group - ) - - epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=self.epilog_sync_bar_id, - num_threads=32 * len(self.epilog_warp_id), - ) - - work_tile = tile_sched.initial_work_tile_info() - 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_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) - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - 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(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # - # Convert to C type - # - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tRS_rC.store(acc_vec) - - # - # Store C to shared memory - # - c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage - 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("async.shared", space="cta") - epilog_sync_barrier.arrive_and_wait() - - # - # TMA store C to global memory - # - if warp_idx == self.epilog_warp_id[0]: - cute.copy( - tma_atom_c, - bSG_sC[(None, c_buffer)], - bSG_gC[(None, subtile_idx)], - ) - # Fence and barrier to make sure shared memory store is visible to TMA store - c_pipeline.producer_commit() - c_pipeline.producer_acquire() - epilog_sync_barrier.arrive_and_wait() - - epilog_sync_barrier.arrive_and_wait() - - # - # 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() - - # Wait for C store complete - c_pipeline.producer_tail() - - @cute.jit - def epilogue( - self, - epi_tidx: cutlass.Int32, - acc_pipeline: pipeline.PipelineAsync, - tiled_mma: cute.TiledMma, - tCtAcc_base: cute.Tensor, - tCgC: cute.Tensor, - epi_tile: cute.Tile, - tile_sched: utils.StaticPersistentTileScheduler, - epilogue_op: cutlass.Constexpr, - tmem_dealloc_barrier: pipeline.NamedBarrier, - ) -> None: - tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( - epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs - ) - - gC_epi = cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) - thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) - tTR_gC_partitioned = thr_copy_t2r.partition_D(gC_epi) - # (T2R, T2R_M, T2R_N) - tTR_rC = cute.make_rmem_tensor( - tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype - ) - simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - - work_tile = tile_sched.initial_work_tile_info() - 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 - # - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_gC = tTR_gC_partitioned[ - (None, None, None, None, None, *mma_tile_coord_mnl) - ] - - # Set tensor memory buffer for current tile - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_tAcc = tTR_tAcc_base[ - (None, None, None, None, None, acc_consumer_state.index) - ] - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) - - # - # Wait for accumulator buffer full - # - acc_pipeline.consumer_wait(acc_consumer_state) - - # - # Store accumulator to global memory in subtiles - # - subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) - for subtile_idx in cutlass.range(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # - # Convert to C type - # - acc_vec = tTR_rAcc.load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tTR_rC.store(acc_vec) - - # - # Store C to global memory - # - cute.copy(simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)]) - - # - # 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() - - # Synchronize before TMEM dealloc (done by the caller) - tmem_dealloc_barrier.arrive_and_wait() - def epilog_tmem_copy_and_partition( self, tidx: cutlass.Int32, @@ -1467,7 +1365,7 @@ class PersistentDenseGemmKernel: :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] """ # Make tiledCopy for tensor memory load - copy_atom_t2r = sm100_utils.get_tmem_load_op( + copy_atom_t2r = utils.sm100.get_tmem_load_op( self.cta_tile_shape_mnk, self.c_layout, self.c_dtype, @@ -1476,7 +1374,10 @@ class PersistentDenseGemmKernel: 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) + 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)] @@ -1524,7 +1425,7 @@ class PersistentDenseGemmKernel: - tRS_sC: The partitioned tensor C (smem destination) :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] """ - copy_atom_r2s = sm100_utils.get_smem_store_op( + copy_atom_r2s = utils.sm100.get_smem_store_op( self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r ) tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) @@ -1535,6 +1436,169 @@ class PersistentDenseGemmKernel: tRS_rC = tiled_copy_r2s.retile(tTR_rC) return tiled_copy_r2s, tRS_rC, tRS_sC + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> Tuple[cute.CopyAtom, 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 atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing either: + - For TMA store: (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + - For non-TMA store: (simt_atom, tTR_rC, tTR_gC) where: + - simt_atom: The SIMT copy atom + - tTR_rC: The register tensor C + - tTR_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + if cutlass.const_expr(self.use_tma_store): + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 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_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_c, bSG_sC, bSG_gC + else: + tiled_copy_t2r = atom + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_gC = thr_copy_t2r.partition_D(gC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype + ) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return simt_atom, tTR_rC, tTR_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, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + ) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/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. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (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 + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + c_smem_layout_staged_one = ( + utils.sm100.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + if use_tma_store + else None + ) + 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 + c_bytes_per_stage = ( + cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + if use_tma_store + else 0 + ) + 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 C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_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 + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + @staticmethod def _compute_grid( c: cute.Tensor, @@ -1598,7 +1662,9 @@ class PersistentDenseGemmKernel: return num_tmem_alloc_cols def is_valid_dtypes( - self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] + self, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], ) -> bool: """ Check if the dtypes are valid @@ -1748,22 +1814,25 @@ class PersistentDenseGemmKernel: """ is_valid = True - # TODO: move to utils - def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): major_mode_idx = 0 if is_mode0_major else 1 num_major_elements = tensor_shape[major_mode_idx] num_contiguous_elements = 16 * 8 // dtype.width return num_major_elements % num_contiguous_elements == 0 if ( - not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) - or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) - or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) ): is_valid = False return is_valid - def is_valid_epilog_store_option(self, m: int, n: int) -> bool: + def is_valid_epilog_store_option( + self, + m: int, + n: int, + ) -> bool: """ Check if the epilogue store option is valid @@ -1801,7 +1870,7 @@ class PersistentDenseGemmKernel: :rtype: bool """ m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] - + print(f"m: {m}, n: {n}, k: {k}, l: {l}") # infer a_major, b_major, c_major is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() @@ -1829,123 +1898,17 @@ class PersistentDenseGemmKernel: return can_implement -def create_tensors( - l, - m_per_iteration, - n, - k, - a_major, - b_major, - c_major, - ab_dtype, - c_dtype, - num_steps, - local_rank, -): - torch.manual_seed(1111) - a_torch_cpu = cutlass_torch.matrix(l, m_per_iteration, k, a_major == "m", ab_dtype) - b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype) - c_torch_cpu = cutlass_torch.matrix(l, m_per_iteration, n, c_major == "m", c_dtype) - - # create local buffers - a_local_list = [] - a_local_torch_list = [] - c_local_list = [] - c_local_torch_list = [] - for m_dim_offset in range(num_steps): - a_temp = a_torch_cpu.clone() - a_temp.zero_() - if m_dim_offset == local_rank: - a_temp.copy_(a_torch_cpu) - a_local, a_local_torch = cutlass_torch.cute_tensor_like( - a_temp, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - a_local_list.append(a_local) - a_local_torch_list.append(a_local_torch) - c_temp = c_torch_cpu.clone().zero_() - c_local, c_local_torch = cutlass_torch.cute_tensor_like( - c_temp, c_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_local_list.append(c_local) - c_local_torch_list.append(c_local_torch) - a_torch_unique = torch.empty( - a_torch_cpu.shape, device="cuda", dtype=a_torch_cpu.dtype - ) - a_torch_unique.copy_(a_torch_cpu) - a_tensor_symm = nvshmem.core.tensor(a_torch_cpu.shape, dtype=a_torch_cpu.dtype) - a_tensor_symm.copy_(a_local_torch_list[local_rank]) - a_tensor_peers = [nvshmem.core.get_peer_tensor(a_tensor_symm, rank) for rank in range(world_size)] - - b_tensor, b_torch = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - - gate_a_flags = torch.zeros(num_steps, device="cuda", dtype=torch.int32) - - device_arrival_counters = torch.tensor([0], dtype=torch.int32, device="cuda") - device_arrival_counters_symm = nvshmem.core.tensor(1, dtype=torch.int32) - device_arrival_counters_symm.copy_(device_arrival_counters) - device_arrival_counters_peers = [nvshmem.core.get_peer_tensor(device_arrival_counters_symm, rank) for rank in range(world_size)] - - return ( - a_local_list, - a_local_torch_list, - a_torch_unique, - a_tensor_symm, - a_tensor_peers, - b_tensor, - b_torch, - c_local_list, - c_local_torch_list, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - gate_a_flags, - device_arrival_counters_symm, - device_arrival_counters_peers, +def create_tensors(*args, **kwargs): + raise RuntimeError( + "create_tensors is only available inside run(). Use run() to execute." ) -def compare( - a_torch_unique, - b_torch_cpu, - c_torch_local_list, - c_dtype, - tolerance, - rank, - world_size, -): - # Copy gpu result back - kernel_result = torch.cat(c_torch_local_list, dim=0).cpu() - - all_gather_a = [torch.zeros_like(a_torch_unique) for _ in range(world_size)] - dist.all_gather(all_gather_a, a_torch_unique) - a = torch.cat(all_gather_a, dim=0) - - # Convert GPU tensors to CPU before reference computation - # This ensures the same computation path as dense_gemm_persistent.py - a_cpu = a.cpu() - - # Compute reference result using CPU tensors (like dense_gemm_persistent.py) - num_rows = a_cpu.shape[0] - ref = torch.einsum( - "mkl,nkl->mnl", - a_cpu.to(dtype=torch.float32), - b_torch_cpu.to(dtype=torch.float32), - ) - - # Convert ref to c_dtype - _, ref_torch_gpu = cutlass_torch.cute_tensor_like( - ref, c_dtype, is_dynamic_layout=True, assumed_align=16 - ) - ref_result = ref_torch_gpu.cpu() - # Assert close results - torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05) +def compare(*args, **kwargs): + raise RuntimeError("compare is only available inside run(). Use run() to execute.") def run( - rank, - world_size, mnkl: Tuple[int, int, int, int], ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], @@ -1964,79 +1927,155 @@ def run( use_cold_l2: bool = False, **kwargs, ): - """Execute a patthen of All gather + persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + import os as _os + import torch + import torch.distributed as dist + import torch.distributed._symmetric_memory as symm_mem + import cutlass.torch as cutlass_torch - This function prepares input tensors and flags needed for sync between communicaton and gemm, configures and launches the persistent GEMM kernel and Peer mem copys and memset, - and construt them into a cuda graph. - optionally performs reference validation, and benchmarks the execution performance. + # helper functions localized here to keep torch usage inside run + def _create_tensors( + l, + m_per_iteration, + n, + k, + a_major, + b_major, + c_major, + ab_dtype, + c_dtype, + num_steps, + local_rank, + ): + torch.manual_seed(1111) + a_torch_cpu = cutlass_torch.matrix( + l, m_per_iteration, k, a_major == "m", ab_dtype + ) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype) + c_torch_cpu = cutlass_torch.matrix( + l, m_per_iteration, n, c_major == "m", c_dtype + ) - :param mnkl: Problem size (M, N, K, L) - :type mnkl: Tuple[int, int, int, int] - :param ab_dtype: Data type for input tensors A and B - :type ab_dtype: Type[cutlass.Numeric] - :param c_dtype: Data type for output tensor C - :type c_dtype: Type[cutlass.Numeric] - :param acc_dtype: Data type for accumulation during matrix multiplication - :type acc_dtype: Type[cutlass.Numeric] - :param a_major/b_major/c_major: Memory layout of tensor A/B/C - :type a_major/b_major/c_major: str - :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the - default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters. - :type mma_tiler_mn: Tuple[int, int], optional - :param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the - default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters. - :type cluster_shape_mn: Tuple[int, int], optional - :param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner - will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. - :type use_2cta_instrs: bool, optional - :param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use - the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. - :type use_tma_store: bool, optional - :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 - :type tolerance: float, optional - :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 - :type warmup_iterations: int, optional - :param iterations: Number of benchmark iterations to run, defaults to 1 - :type iterations: int, optional - :param skip_ref_check: Whether to skip reference result validation, defaults to False - :type skip_ref_check: bool, optional - :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False - :type use_cold_l2: bool, optional - :raises RuntimeError: If CUDA GPU is not available - :raises ValueError: If the configuration is invalid or unsupported by the kernel - :return: Execution time of the GEMM kernel - :rtype: float - """ - print("Running Blackwell Persistent Dense GEMM test with:") - print(f"per GPU mnkl: {mnkl}") - print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") - print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") - print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") - print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") - print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") - print(f"Tolerance: {tolerance}") - print(f"Warmup iterations: {warmup_iterations}") - print(f"Iterations: {iterations}") - print(f"Skip reference checking: {skip_ref_check}") - print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + a_local_list = [] + a_local_torch_list = [] + c_local_list = [] + c_local_torch_list = [] + for m_dim_offset in range(num_steps): + a_temp = a_torch_cpu.clone() + a_temp.zero_() + if m_dim_offset == local_rank: + a_temp.copy_(a_torch_cpu) + a_local, a_local_torch = cutlass_torch.cute_tensor_like( + a_temp, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + a_local_list.append(a_local) + a_local_torch_list.append(a_local_torch) + c_temp = c_torch_cpu.clone().zero_() + c_local, c_local_torch = cutlass_torch.cute_tensor_like( + c_temp, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_local_list.append(c_local) + c_local_torch_list.append(c_local_torch) + a_torch_unique = torch.empty( + a_torch_cpu.shape, device="cuda", dtype=a_torch_cpu.dtype + ) + a_torch_unique.copy_(a_torch_cpu) + a_tensor_symm = symm_mem.empty( + a_torch_cpu.shape, device="cuda", dtype=a_torch_cpu.dtype + ) + a_tensor_symm.copy_(a_local_torch_list[local_rank]) + a_tensor_symm_hdl = symm_mem.rendezvous(a_tensor_symm, dist.group.WORLD) + a_tensor_symm_hdl.barrier(channel=0) + + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + gate_a_flags = torch.zeros(num_steps, device="cuda", dtype=torch.int32) + + device_arrival_counters = torch.tensor([0], dtype=torch.int32, device="cuda") + device_arrival_counters_symm = symm_mem.empty( + 1, device="cuda", dtype=torch.int32 + ) + device_arrival_counters_symm.copy_(device_arrival_counters) + device_arrival_counters_hdl = symm_mem.rendezvous( + device_arrival_counters_symm, dist.group.WORLD + ) + device_arrival_counters_hdl.barrier(channel=0) + + return ( + a_local_list, + a_local_torch_list, + a_torch_unique, + a_tensor_symm_hdl, + b_tensor, + b_torch, + c_local_list, + c_local_torch_list, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + gate_a_flags, + device_arrival_counters_hdl, + ) + + def _compare( + a_torch_unique, + b_torch_cpu, + c_torch_local_list, + c_dtype, + tolerance, + rank, + world_size, + ): + kernel_result = torch.cat(c_torch_local_list, dim=0).cpu() + all_gather_a = [torch.zeros_like(a_torch_unique) for _ in range(world_size)] + dist.all_gather(all_gather_a, a_torch_unique) + a = torch.cat(all_gather_a, dim=0) + a_cpu = a.cpu() + ref = torch.einsum( + "mkl,nkl->mnl", + a_cpu.to(dtype=torch.float32), + b_torch_cpu.to(dtype=torch.float32), + ) + _, ref_torch_gpu = cutlass_torch.cute_tensor_like( + ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + ref_result = ref_torch_gpu.cpu() + torch.testing.assert_close( + kernel_result, ref_result, atol=tolerance, rtol=1e-05 + ) + + # init torch distributed here + local_rank_env = _os.environ.get("LOCAL_RANK") + if local_rank_env is not None: + import torch as _t + + _t.cuda.set_device(int(local_rank_env)) + if "dist" not in globals(): + pass + if "torch" not in globals(): + pass + + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() # Unpack parameters m, n, k, l = mnkl m_per_step = m // world_size - if not torch.cuda.is_available(): - raise RuntimeError("GPU is required to run this example!") - # Get current CUDA stream from PyTorch torch_stream = torch.cuda.current_stream() # Get the raw stream pointer as a CUstream current_stream = cuda.CUstream(torch_stream.cuda_stream) + ( a_local_list, a_local_torch_list, a_torch_unique, - a_tensor_symm, - a_tensor_peers, + a_tensor_symm_hdl, b_tensor, b_torch, c_local_list, @@ -2045,9 +2084,8 @@ def run( b_torch_cpu, c_torch_cpu, gate_a_flags, - device_arrival_counters_symm, - device_arrival_counters_peers, - ) = create_tensors( + device_arrival_counters_hdl, + ) = _create_tensors( l, m_per_step, n, @@ -2090,29 +2128,42 @@ def run( max_active_clusters = utils.HardwareInfo().get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) - compiled_gemm = cute.compile( - gemm, - a_local_list[0], - b_tensor, - c_local_list[0], - max_active_clusters, - current_stream, - flag_offset=0, - gate_a_flags=from_dlpack(gate_a_flags), - ) - compiled_gemm_with_gated_a_load = cute.compile( - gemm_with_gated_a_load, - a_local_list[0], - b_tensor, - c_local_list[0], - max_active_clusters, - current_stream, - flag_offset=0, - gate_a_flags=from_dlpack(gate_a_flags), - ) + num_steps = world_size + # flag_offset is a Constexpr, so each m_dim_offset value gets its own + # compiled kernel instance; compile one per step so each polls the + # correct flag slot. + compiled_gemm_list = [] + compiled_gemm_with_gated_a_load_list = [] + for j in range(num_steps): + m_dim_offset = (rank + j) % num_steps + compiled_gemm_list.append( + cute.compile( + gemm, + a_local_list[m_dim_offset], + b_tensor, + c_local_list[m_dim_offset], + max_active_clusters, + current_stream, + flag_offset=m_dim_offset, + gate_a_flags=from_dlpack(gate_a_flags), + ) + ) + if j > 0: + compiled_gemm_with_gated_a_load_list.append( + cute.compile( + gemm_with_gated_a_load, + a_local_list[m_dim_offset], + b_tensor, + c_local_list[m_dim_offset], + max_active_clusters, + current_stream, + flag_offset=m_dim_offset, + gate_a_flags=from_dlpack(gate_a_flags), + ) + ) device_arrival_counters_ptrs = torch.tensor( - [tensor.data_ptr() for tensor in device_arrival_counters_peers], device="cuda", dtype=torch.int64 + device_arrival_counters_hdl.buffer_ptrs, device="cuda", dtype=torch.int64 ) sync_nvl_devices_instance = SyncNvlDevices(world_size) compiled_sync_nvl_devices = cute.compile( @@ -2128,7 +2179,6 @@ def run( copy_stream = torch.cuda.Stream(local_rank) gemm_stream = torch.cuda.Stream(local_rank) - num_steps = world_size with torch.cuda.graph(g, stream=capture_stream): compiled_sync_nvl_devices( rank, @@ -2140,7 +2190,7 @@ def run( for j in range(num_steps): m_dim_offset = (local_rank + j) % num_steps if j == 0: # gemm process local data directly - compiled_gemm( + compiled_gemm_list[j]( a_local_list[m_dim_offset], b_tensor, c_local_list[m_dim_offset], @@ -2149,7 +2199,7 @@ def run( gate_a_flags=from_dlpack(gate_a_flags), ) else: # gemm process remote data with gated a load - compiled_gemm_with_gated_a_load( + compiled_gemm_with_gated_a_load_list[j - 1]( a_local_list[m_dim_offset], b_tensor, c_local_list[m_dim_offset], @@ -2167,7 +2217,7 @@ def run( * a_local_torch_list[m_dim_offset].numel() ) dst_address = a_local_torch_list[m_dim_offset].data_ptr() - src_address = a_tensor_peers[m_dim_offset].data_ptr() + src_address = a_tensor_symm_hdl.buffer_ptrs[m_dim_offset] driver.cuMemcpyDtoDAsync( dst_address, src_address, @@ -2221,7 +2271,7 @@ def run( # reference check if not skip_ref_check: # compiled_gemm(a_tensor, b_tensor, c_tensor, current_stream) - compare( + _compare( a_torch_unique, b_torch_cpu, c_local_torch_list, @@ -2231,54 +2281,8 @@ def run( world_size, ) - for i in range(world_size): - if i != dist.get_rank(): - nvshmem.core.free_tensor(a_tensor_peers[i]) - nvshmem.core.free_tensor(device_arrival_counters_peers[i]) - nvshmem.core.free_tensor(a_tensor_symm) - nvshmem.core.free_tensor(device_arrival_counters_symm) - - return exec_time # Return execution time in microseconds - - -def torchrun_uid_init_bcast(): - """ - Initialize NVSHMEM using UniqueID with `torchrun` as the launcher - - It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting - """ - # Set Torch device - local_rank = int(os.environ['LOCAL_RANK']) - torch.cuda.set_device(local_rank) - - # nvshmem4py requires a cuda.core Device at init time - dev = Device(local_rank) - dev.set_current() - global stream - stream = dev.create_stream() - - # Initialize torch.distributed process group - dist.init_process_group( - backend="cpu:gloo,cuda:nccl", - ) - - # Extract rank, nranks from process group - num_ranks = dist.get_world_size() - - # Create an empty uniqueid for all ranks - uid = nvshmem.core.get_unique_id(empty=(local_rank != 0)) - uid_bytes = uid._data.view(np.uint8).copy() - uid_tensor = torch.from_numpy(uid_bytes).cuda() - dist.broadcast(uid_tensor, src=0) - dist.barrier() - uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype) - - nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid") - - -def torchrun_finalize(): - nvshmem.core.finalize() dist.destroy_process_group() + return exec_time # Return execution time in microseconds if __name__ == "__main__": @@ -2360,15 +2364,16 @@ if __name__ == "__main__": if len(args.cluster_shape_mn) != 2: parser.error("--cluster_shape_mn must contain exactly 2 values") - torchrun_uid_init_bcast() + import torch + import torch.distributed as dist + + dist.init_process_group(backend="nccl") local_rank = dist.get_rank() world_size = dist.get_world_size() print(f"World size: {world_size}, local rank: {local_rank}") torch.cuda.set_device(local_rank) run( - local_rank, - world_size, args.mnkl, args.ab_dtype, args.c_dtype, @@ -2386,5 +2391,4 @@ if __name__ == "__main__": args.skip_ref_check, args.use_cold_l2, ) - torchrun_finalize() print("PASS") diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_blackwell.py b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_blackwell.py index b0f8e0ff7..7821a2f55 100644 --- a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_blackwell.py +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_blackwell.py @@ -28,11 +28,16 @@ import os import argparse -from typing import Optional, Tuple, Type, Union +import csv +import glob +import json +import socket +import time +from pathlib import Path + +from math import gcd +from typing import Optional, Tuple, Type, Union, Literal -import numpy as np -import torch -import torch.distributed as dist import cuda.bindings.driver as cuda try: from cuda.core import Device @@ -40,299 +45,410 @@ except ImportError: from cuda.core.experimental import Device from cuda.pathfinder import load_nvidia_dynamic_lib +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.torch as cutlass_torch +from cutlass.torch import dtype as torch_dtype +from cutlass import testing import cutlass.utils as utils import cutlass.pipeline as pipeline -import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cutlass_dsl import BaseDSL from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack -from cutlass.cute.typing import Int32, Float16, BFloat16, Float32, Float8E5M2, Float8E4M3FN -from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait - -try: - import nvshmem.core -except ImportError as exc: - raise ImportError( - "nvshmem4py is required but not installed. Please install it using:\n" - " For CUDA 12: pip install nvshmem4py-cu12\n" - " For CUDA 13: pip install nvshmem4py-cu13\n" - "Note: nvshmem4py version >= 0.1.3 is recommended." - ) from None - -try: - load_nvidia_dynamic_lib("nvshmem_host") -except RuntimeError as exc: - raise ImportError( - "nvshmem lib is required but not installed. Please install it using:\n" - " For CUDA 12: pip install nvidia-nvshmem-cu12\n" - " For CUDA 13: pip install nvidia-nvshmem-cu13\n" - ) from None """ -A high-performance persistent batched dense GEMM AR 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 C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") -- Matrix C_mc is a multicast C matrix that changes can be broadcasted to all GPUs by multimem instructions. +SM100 GEMM + AllReduce (LDMCxSTMC) Fusion Kernel -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 - - Support all-reduce epilogue with multimem instructions to distribute the workload to all GPUs +* Test (correctness only — single launch, no warmup/iteration loop) +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs --all_reduce LDMCxSTMC \ + --benchmark_or_test test -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)) -4. All reduce epilogue: - - Load and reduce the 128bit data from all ranks by multimem instructions. - - Broadcast the reduced data to all ranks by multimem instructions. - - current implementation only supports LDMCxSTMC all-reduce which means each rank only computes a portion of - the output tensor and broadcast the result to all ranks. - - the all-reduce epilogue is only supported when use_tma_store is True. - - the all-reduce epilogue is only supported when c_dtype is Float16, Float32, BFloat16, Float8E4M3FN, Float8E5M2. +* Benchmark One +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --use_2cta_instrs --all_reduce LDMCxSTMC \ + --benchmark_or_test benchmark --use_cuda_graph --csv_report ./csv_report.csv -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. +* Benchmark All +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --all_reduce LDMCxSTMC \ + --benchmark_or_test benchmark_all --use_cuda_graph --csv_report ./csv_report.csv -Input arguments to this example is same as dense_gemm.py. +NOTE: +* multimem.ld_reduce does NOT handle NaN issue for f8. -.. code-block:: bash - - torchrun --nproc-per-node 8 examples/distributed/distributed_dense_gemm_persistent_all_reduce.py \ - --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ - --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ - --mnkl 1024,1000,1024,1 --warmup_iterations 3 --iterations 10 \ - --use_tma_store --use_2cta_instrs --all_reduce LDMCxSTMC - -To collect performance with NSYS profiler: - -.. code-block:: bash - - nsys profile --gpu-metrics-devices=cuda-visible \ - torchrun --nproc-per-node 8 examples/distributed/distributed_dense_gemm_persistent_all_reduce.py \ - --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ - --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ - --mnkl 8192,8192,8192,1 \ - --use_tma_store --use_2cta_instrs --warmup_iterations 3 --iterations 10 \ - --skip_ref_check --all_reduce LDMCxSTMC - - -Constraints are same as dense_gemm_persistent.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 -* when all_reduce is "two_shot", M and N must be multiple of 128, world_size must be 8 +ALGO: +-> Producer +for each output tile + * store gemm output [data] (UC) to local rank global memory with .weak semantics (tma is default .weak semantic) + * [flag] (MC) update with .release semantics and .gpu scope (prior store is to local rank global memory thus .gpu scope) on multicast flag ptr (thus flag update broadcast to all rank) +-> Consumer +for each output tile + * spin lock wait [flag] (UC) with .relaxed / .acquire and .gpu scope (all rank have write their data to local rank gmem) + * multimem.ld_reduce send [data] (MC) reduce request to switch, switch send request to all rank's gmem, reduce on switch, return to requester rank + * multimem.st broadcast reduced [data] (MC) to all rank as each output tile is handled collectively by P rank +before kernel exit + * [flag] (MC) update with .release semantics and .sys scope (prior store broadcast data to all rank in sys) + * spin lock wait [flag] (UC) with .relaxed / .acquire and .sys scope """ +##################################################################### +#### KERNEL IMPLEMENTATION START HERE #### +##################################################################### -def _compute_stages( + +@cute.jit +def LDMCxSTMC( + kernel, + cur_tile_coord, tiled_mma: cute.TiledMma, - mma_tiler_mnk: Tuple[int, int, int], - a_dtype: Type[cutlass.Numeric], - b_dtype: Type[cutlass.Numeric], - c_dtype: Type[cutlass.Numeric], - smem_capacity: int, - occupancy: int, - use_tma_store: bool, - c_smem_layout: Union[cute.Layout, None], -) -> Tuple[int, int, int]: - """Computes the number of stages for A/B/C operands based on heuristics. + comm_tidx: cutlass.Int32, + tCgCommInMC: cute.Tensor, + tCgCommOutMC: cute.Tensor, +) -> None: + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) - :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 c_dtype: Data type of operand C (output). - :type c_dtype: type[cutlass.Numeric] - :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 - :param use_tma_store: Whether TMA store is enabled. - :type use_tma_store: bool - :param c_smem_layout: Layout of C operand in shared memory, or None if not using TMA store. - :type c_smem_layout: Union[cute.Layout, None] + # partition and slice at tile level + tCgCommInMC_slice = tCgCommInMC[((None, None), 0, 0, *mma_tile_coord_mnl)] + tCgCommOutMC_slice = tCgCommOutMC[((None, None), 0, 0, *mma_tile_coord_mnl)] - :return: A tuple containing the computed number of stages for: - (ACC stages, A/B operand stages, C stages) - :rtype: tuple[int, int, int] + # + # Partition s.t. each rank handle 1/P of this tile + # For row-major (N-contiguous): split along M to keep memory contiguous per rank + # For col-major (M-contiguous): split along N to keep memory contiguous per rank + # + cta_mma_tile_m = kernel.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + cta_tile_n = kernel.mma_tiler[1] + total_comm_threads = len(kernel.comm_warp_id) * cute.arch.WARP_SIZE + + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + # Row-major (N-contiguous): split along M + m_local_rank = int(cta_mma_tile_m / kernel.num_ranks) + tCgCommInMC_slice_partitioned = cute.zipped_divide( + tCgCommInMC_slice, (m_local_rank, cta_tile_n) + ) + tCgCommOutMC_slice_partitioned = cute.zipped_divide( + tCgCommOutMC_slice, (m_local_rank, cta_tile_n) + ) + # Data processed by local rank: select M chunk + tCgCommInMC_local_rank = cute.slice_( + tCgCommInMC_slice_partitioned, ((None, None), (kernel.rank_id, 0)) + ) + tCgCommOutMC_local_rank = cute.slice_( + tCgCommOutMC_slice_partitioned, ((None, None), (kernel.rank_id, 0)) + ) + local_tile_m = m_local_rank + local_tile_n = cta_tile_n + else: + # Col-major (M-contiguous): split along N + n_local_rank = int(cta_tile_n / kernel.num_ranks) + tCgCommInMC_slice_partitioned = cute.zipped_divide( + tCgCommInMC_slice, (cta_mma_tile_m, n_local_rank) + ) + tCgCommOutMC_slice_partitioned = cute.zipped_divide( + tCgCommOutMC_slice, (cta_mma_tile_m, n_local_rank) + ) + # Data processed by local rank: select N chunk + tCgCommInMC_local_rank = cute.slice_( + tCgCommInMC_slice_partitioned, ((None, None), (0, kernel.rank_id)) + ) + tCgCommOutMC_local_rank = cute.slice_( + tCgCommOutMC_slice_partitioned, ((None, None), (0, kernel.rank_id)) + ) + local_tile_m = cta_mma_tile_m + local_tile_n = n_local_rank + + # + # Determine vectorization length based on data size per thread + # Always use 128-bit instructions (guaranteed by _pick_num_comm_warp_for_128b) + # + elements_per_thread = (local_tile_m * local_tile_n) // total_comm_threads + + # Always use 128-bit atom size + # Guaranteed by _pick_num_comm_warp_for_128b which selects T to ensure + # elements_per_thread is divisible by 128-bit atom (16B = 128 bits) + atom_num_elements_128 = 128 // kernel.c_dtype.width + atom_num_elements = atom_num_elements_128 + + # COMPILE-TIME ASSERTION: Verify 128-bit PTX instruction usage + # _pick_num_comm_warp_for_128b selects num_warps such that every thread + # in the AllReduce warp group handles exactly atom_num_elements_128 elements, + # which means exactly one 128-bit LDMC/STMC instruction per AR slab iteration. + assert elements_per_thread % atom_num_elements_128 == 0, ( + f"KERNEL ASSERTION: 128-bit PTX instruction invariant violated. " + f"elements_per_thread={elements_per_thread} must be divisible by " + f"atom_num_elements_128={atom_num_elements_128}. " + f"This ensures every thread executes exactly one 128-bit instruction. " + f"Configuration: local_tile=({local_tile_m},{local_tile_n}), " + f"total_comm_threads={total_comm_threads}, c_dtype={kernel.c_dtype}" + ) + + # + # Create thr_copy_fake with dynamically selected atom size + # For row-major: threads spread along N (contiguous), load along N + # For col-major: threads spread along M (contiguous), load along M + # + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + # Row-major: N is contiguous, spread threads along N. + # Use GCD to handle non-power-of-2 tile widths (e.g. cta_n=192) + # that don't naively divide the thread count. + max_thr_n = local_tile_n // atom_num_elements + atom_thr_n = gcd(max_thr_n, total_comm_threads) + atom_thr_m = total_comm_threads // atom_thr_n + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) + val_layout = cute.make_layout( + (1, atom_num_elements), stride=(atom_num_elements, 1) + ) + else: + # Col-major: M is contiguous, spread threads along M. + max_thr_m = local_tile_m // atom_num_elements + atom_thr_m = gcd(max_thr_m, total_comm_threads) + atom_thr_n = total_comm_threads // atom_thr_m + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(1, atom_thr_m)) + val_layout = cute.make_layout( + (atom_num_elements, 1), stride=(1, atom_num_elements) + ) + + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), kernel.c_dtype) + tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + thr_copy_fake = tiled_copy_fake.get_slice(comm_tidx - kernel.comm_warp_id[0] * 32) + + # + # Partition per thread + # + tTR_gCommInMC = thr_copy_fake.partition_S(tCgCommInMC_local_rank) + tTR_gCommOutMC = thr_copy_fake.partition_S(tCgCommOutMC_local_rank) + _atom, loop_m, loop_n = tTR_gCommInMC.shape + + # + # All Reduce LDMCxSTMC — adaptive variant based on ldmcxstmc_default_inflight_depth: + # - ldmcxstmc_default_inflight_depth=2: pair-issue (2-element chunks) + # - ldmcxstmc_default_inflight_depth=4: quad-issue (4-element chunks) + # The chunk loop is staged as scf.for with explicit `unroll=1` hint. + # + total_iters = loop_m * loop_n + + if cutlass.const_expr(kernel.ldmcxstmc_default_inflight_depth == 2): + # Pair-issue variant (inflight2) + num_pairs = total_iters // 2 + for k in cutlass.range(num_pairs, unroll=1): + i0 = (2 * k) // loop_n + j0 = (2 * k) % loop_n + i1 = (2 * k + 1) // loop_n + j1 = (2 * k + 1) % loop_n + + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + + # Trailing odd iteration (singleton) when total_iters is odd. + if cutlass.const_expr(total_iters % 2 == 1): + last = total_iters - 1 + i = last // loop_n + j = last % loop_n + in_ = tTR_gCommInMC[None, i, j].iterator + out_ = tTR_gCommOutMC[None, i, j].iterator + regs = utils.distributed.multimem_ld_reduce( + in_, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out_, *regs) + else: + # Quad-issue variant (inflight4, default) + num_quads = total_iters // 4 + for k in cutlass.range(num_quads, unroll=1): + i0 = (4 * k) // loop_n + j0 = (4 * k) % loop_n + i1 = (4 * k + 1) // loop_n + j1 = (4 * k + 1) % loop_n + i2 = (4 * k + 2) // loop_n + j2 = (4 * k + 2) % loop_n + i3 = (4 * k + 3) // loop_n + j3 = (4 * k + 3) % loop_n + + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + in2 = tTR_gCommInMC[None, i2, j2].iterator + out2 = tTR_gCommOutMC[None, i2, j2].iterator + in3 = tTR_gCommInMC[None, i3, j3].iterator + out3 = tTR_gCommOutMC[None, i3, j3].iterator + + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs2 = utils.distributed.multimem_ld_reduce( + in2, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs3 = utils.distributed.multimem_ld_reduce( + in3, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + utils.distributed.multimem_st(out2, *regs2) + utils.distributed.multimem_st(out3, *regs3) + + # Tail: total_iters % 4 == 1 → 1 trailing LDMC + 1 STMC. + if cutlass.const_expr(total_iters % 4 == 1): + last = total_iters - 1 + i = last // loop_n + j = last % loop_n + in_ = tTR_gCommInMC[None, i, j].iterator + out_ = tTR_gCommOutMC[None, i, j].iterator + regs = utils.distributed.multimem_ld_reduce( + in_, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out_, *regs) + # Tail: total_iters % 4 == 2 → 2 trailing LDMCs + 2 STMCs. + if cutlass.const_expr(total_iters % 4 == 2): + flat0 = total_iters - 2 + flat1 = total_iters - 1 + i0 = flat0 // loop_n + j0 = flat0 % loop_n + i1 = flat1 // loop_n + j1 = flat1 % loop_n + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + # Tail: total_iters % 4 == 3 → 3 trailing LDMCs + 3 STMCs. + if cutlass.const_expr(total_iters % 4 == 3): + flat0 = total_iters - 3 + flat1 = total_iters - 2 + flat2 = total_iters - 1 + i0 = flat0 // loop_n + j0 = flat0 % loop_n + i1 = flat1 // loop_n + j1 = flat1 % loop_n + i2 = flat2 // loop_n + j2 = flat2 % loop_n + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + in2 = tTR_gCommInMC[None, i2, j2].iterator + out2 = tTR_gCommOutMC[None, i2, j2].iterator + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs2 = utils.distributed.multimem_ld_reduce( + in2, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + utils.distributed.multimem_st(out2, *regs2) + + +class Sm100PersistentDenseGemmAllReduceLDMCxSTMCKernel: """ - # Default ACC stages - num_acc_stage = 2 - - # Default C stages - num_c_stage = 2 if use_tma_store else 0 - - # Calculate smem layout and size for one stage of A, B, and C with 1-stage - a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( - tiled_mma, mma_tiler_mnk, a_dtype, 1 - ) - b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( - tiled_mma, mma_tiler_mnk, b_dtype, 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 - - c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) - 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 C stages bytes - # Divide remaining by bytes needed per A/B stage - num_ab_stage = ( - smem_capacity // occupancy - (mbar_helpers_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 - if use_tma_store: - num_c_stage += ( - smem_capacity - - occupancy * ab_bytes_per_stage * num_ab_stage - - occupancy * (mbar_helpers_bytes + c_bytes) - ) // (occupancy * c_bytes_per_stage) - return num_acc_stage, num_ab_stage, num_c_stage - - -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] - :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results - :type use_tma_store: bool - :param all_reduce: All-reduce mode, can be "none", "LDMCxSTMC" - :type all_reduce: str - - :note: only support two_shot all-reduce mode in a single node right now - - :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) + **Example:** + gemm = Sm100PersistentDenseGemmAllReduceLDMCxSTMCKernel( + acc_dtype=cutlass.Float32, + c_dtype=cutlass.BFloat16, + use_2cta_instrs=True, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + use_tma_store=True, + rank_id=rank_id, + num_ranks=num_ranks, + all_reduce="LDMCxSTMC", + swizzle_size=1, + raster_order="m", + ) + compiled_gemm = cute.compile( + gemm, + a=a, b=b, c=c, + comm_in_multicast_tensor=comm_in_mc, + comm_out_multicast_tensor=comm_out_mc, + barrier_flag_unicast=flag_unicast, + barrier_flag_multicast=flag_multicast, + stream=stream, + max_active_clusters=max_active_clusters, + ) + compiled_gemm( + a=a, b=b, c=c, + comm_in_multicast_tensor=comm_in_mc, + comm_out_multicast_tensor=comm_out_mc, + barrier_flag_unicast=flag_unicast, + barrier_flag_multicast=flag_multicast, + stream=stream, + ) """ def __init__( self, acc_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], use_2cta_instrs: bool, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], use_tma_store: bool, rank_id: int, - world_size: int, - all_reduce="two_shot", + num_ranks: int, + all_reduce: str = "LDMCxSTMC", + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", ): - """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: - - use_tma_store: Boolean indicating whether to use Tensor Memory Access (TMA) for storing results. - - :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] - :param use_tma_store: Use Tensor Memory Access (TMA) or normal store for output C tensor. - :type use_tma_store: bool - :param all_reduce: All-reduce mode, can be "none", "LDMCxSTMC" - :type all_reduce: str - """ - self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.c_dtype: Type[cutlass.Numeric] = c_dtype self.use_2cta_instrs = use_2cta_instrs self.cluster_shape_mn = cluster_shape_mn + self.swizzle_size = swizzle_size + self.raster_order = raster_order # K dimension is deferred in _setup_attributes self.mma_tiler_mn = mma_tiler_mn self.mma_tiler = (*mma_tiler_mn, 1) self.use_tma_store = use_tma_store + # Capture the JIT target arch so SMEM/TMEM sizing tracks the actual + # hardware (sm_100 / sm_103 / sm_107 / ...). Matches the FMHA pattern. + arch_enum = BaseDSL._get_dsl().get_arch_enum() + self.arch = f"sm_{arch_enum.major}{arch_enum.minor}" self.cta_group = ( tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE @@ -343,46 +459,74 @@ class PersistentDenseGemmKernel: self.epilogue_warp_id = (0, 1, 2, 3) self.mma_warp_id = 4 self.tma_warp_id = 5 - self.all_reduce_warp_id = (6, 7, 8, 9) + self.all_reduce = all_reduce + if all_reduce == "LDMCxSTMC": + self.rank_id = rank_id + self.num_ranks = num_ranks + num_comm_warps = self._pick_num_comm_warp_for_128b(c_dtype) + self.comm_warp_id = tuple(range(6, 6 + num_comm_warps)) + # Determine inflight depth: use inflight2 for special case (128x256 P=8), else inflight4 + # For 2SM kernel: cta_tile_m = mma_tile_m / 2, for 1SM: cta_tile_m = mma_tile_m + mma_tile_m, mma_tile_n = mma_tiler_mn + cta_tile_m = mma_tile_m // 2 if use_2cta_instrs else mma_tile_m + cta_tile_n = mma_tile_n + if cta_tile_m == 128 and cta_tile_n == 256 and num_ranks == 8: + # Special case for largest tile + rank = 8 prefer in-flight = 2 + self.ldmcxstmc_default_inflight_depth = 2 + else: + # For most cases, in-flight depth = 4 is best as it fully utilize overlapping ability + self.ldmcxstmc_default_inflight_depth = 4 + else: + # No AR mode — no comm warps. + self.comm_warp_id = () + self.rank_id = 0 + self.num_ranks = 1 + self.threads_per_cta = 32 * len( ( self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id, - *self.all_reduce_warp_id, + *self.comm_warp_id, ) ) # Set barrier id for cta sync, epilogue sync and tmem ptr sync - self.epilogue_sync_bar_id = 1 + self.epilog_sync_bar_id = 1 self.tmem_alloc_sync_bar_id = 2 self.tmem_dealloc_sync_bar_id = 3 - self.all_reduce_sync_bar_id = 4 - self.all_reduce_sync_barrier = pipeline.NamedBarrier( - barrier_id=self.all_reduce_sync_bar_id, - num_threads=32 * len(self.all_reduce_warp_id), + self.comm_sync_bar_id = 4 + + def _pick_num_comm_warp_for_128b(self, c_dtype): + """Pick the LARGEST comm-warp count W ∈ {4, 3, 2, 1} (T = W·32 ∈ + {128, 96, 64, 32}) such that the AR slab is evenly distributed + across T threads AND every thread gets a full 128 b LDMC payload + (= ``atom_128`` elements, where ``atom_128 = 16 / sizeof(c_dtype)``). + + Iterates W = 4 → 1 and returns the first match, so the largest + viable W is returned. Smaller W is only chosen when strictly + necessary, because larger W means more parallel AR threads and + proportionally fewer per-thread PTX instructions. + """ + # cta-tile dims (the per-CTA M·N output the AR walks). 2cta halves M. + cta_mma_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + slab_elements = (cta_mma_tile_m * cta_tile_n) // self.num_ranks + atom_128 = 16 // (c_dtype.width // 8) + for num_warps in (4, 3, 2, 1): + T = num_warps * 32 + if slab_elements % T: + continue + if (slab_elements // T) % atom_128 == 0: + return num_warps + raise ValueError( + f"no comm-warp count enables 128 b PTX for c_dtype={c_dtype}, " + f"cta=({cta_mma_tile_m},{cta_tile_n}), P={self.num_ranks}" ) - self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - self.all_reduce = all_reduce - self.rank_id = rank_id - self.num_ranks = world_size 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( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, + self.b_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, @@ -390,7 +534,6 @@ class PersistentDenseGemmKernel: 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 = ( @@ -404,21 +547,18 @@ class PersistentDenseGemmKernel: self.mma_tiler[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 if cutlass.const_expr(self.use_tma_store): - self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, @@ -429,12 +569,13 @@ class PersistentDenseGemmKernel: c_smem_layout = None if cutlass.const_expr(self.use_tma_store): - c_smem_layout = sm100_utils.make_smem_layout_epi( + c_smem_layout = utils.sm100.make_smem_layout_epi( self.c_dtype, self.c_layout, self.epi_tile, 1 ) - # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory - self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( tiled_mma, self.mma_tiler, self.a_dtype, @@ -446,23 +587,21 @@ class PersistentDenseGemmKernel: c_smem_layout, ) - # Compute A/B/C shared memory layout - self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + self.a_smem_layout_staged = utils.sm100.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( + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage ) self.c_smem_layout_staged = None if self.use_tma_store: - self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( self.c_dtype, self.c_layout, self.epi_tile, 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 + tiled_mma, self.mma_tiler, self.num_acc_stage, self.arch ) @cute.jit @@ -471,45 +610,24 @@ class PersistentDenseGemmKernel: a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - c_mc: cute.Tensor, - barrier_flag: cute.Tensor, - barrier_flag_mc: cute.Tensor, - max_active_clusters: cutlass.Constexpr, + comm_in_multicast_tensor: cute.Tensor, + comm_out_multicast_tensor: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, stream: cuda.CUstream, + max_active_clusters: cutlass.Constexpr, epilogue_op: cutlass.Constexpr = lambda x: x, ): - """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 synchronously - - :param a: Input tensor A - :type a: cute.Tensor - :param b: Input tensor B - :type b: cute.Tensor - :param c: Output tensor C - :type c: cute.Tensor - :param c_mc: Output symmetric tensor C_mc, any write or read to a multicast tensor will be broadcasted to all GPUs - :type c_mc: cute.Tensor - :param barrier_flag: flag tensor acting as a barrier for GPUs to synchronize - :type barrier_flag: cute.Tensor - :param barrier_flag_mc: the multicast tensor of barrier_flag - :type barrier_flag_mc: 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 + # Setup static attributes before smem/grid/tma computation. + # NOTE: self.c_dtype was set in __init__ from the explicit c_dtype + # arg. We only sanity-check here that the C tensor matches that hint. self.a_dtype: Type[cutlass.Numeric] = a.element_type self.b_dtype: Type[cutlass.Numeric] = b.element_type - self.c_dtype: Type[cutlass.Numeric] = c.element_type + if cutlass.const_expr(c.element_type != self.c_dtype): + raise TypeError( + f"C tensor element_type ({c.element_type}) does not match " + f"the c_dtype declared at __init__ ({self.c_dtype})" + ) 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.c_layout = utils.LayoutEnum.from_tensor(c) @@ -518,21 +636,23 @@ class PersistentDenseGemmKernel: 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( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, + self.b_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler[:2], ) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + 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( + a_op = utils.sm100.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)) @@ -549,7 +669,7 @@ class PersistentDenseGemmKernel: ) # Setup TMA load for B - b_op = sm100_utils.cluster_shape_to_tma_atom_B( + b_op = utils.sm100.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)) @@ -580,28 +700,46 @@ class PersistentDenseGemmKernel: # Compute grid size self.tile_sched_params, grid = self._compute_grid( - c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + self.swizzle_size, + self.raster_order, + max_active_clusters, ) + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + self.shared_storage = SharedStorage + # Launch the kernel synchronously self.kernel( - tiled_mma, - tma_atom_a, - tma_tensor_a, - tma_atom_b, - tma_tensor_b, - tma_atom_c, - tma_tensor_c if self.use_tma_store else c, - c_mc, - barrier_flag, - barrier_flag_mc, - self.cluster_layout_vmnk, - self.a_smem_layout_staged, - self.b_smem_layout_staged, - self.c_smem_layout_staged, - self.epi_tile, - self.tile_sched_params, - epilogue_op, + tiled_mma=tiled_mma, + tma_atom_a=tma_atom_a, + mA_mkl=tma_tensor_a, + tma_atom_b=tma_atom_b, + mB_nkl=tma_tensor_b, + tma_atom_c=tma_atom_c, + mC_mnl=(tma_tensor_c if self.use_tma_store else c), + mCommInMC_mnl=comm_in_multicast_tensor, + mCommOutMC_mnl=comm_out_multicast_tensor, + barrier_flag_unicast=barrier_flag_unicast, + barrier_flag_multicast=barrier_flag_multicast, + cluster_layout_vmnk=self.cluster_layout_vmnk, + a_smem_layout_staged=self.a_smem_layout_staged, + b_smem_layout_staged=self.b_smem_layout_staged, + c_smem_layout_staged=self.c_smem_layout_staged, + epi_tile=self.epi_tile, + tile_sched_params=self.tile_sched_params, + epilogue_op=epilogue_op, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -621,9 +759,10 @@ class PersistentDenseGemmKernel: mB_nkl: cute.Tensor, tma_atom_c: Optional[cute.CopyAtom], mC_mnl: cute.Tensor, - c_mc: cute.Tensor, - barrier_flag: cute.Tensor, - barrier_flag_mc: cute.Tensor, + mCommInMC_mnl: cute.Tensor, + mCommOutMC_mnl: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, cluster_layout_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, b_smem_layout_staged: cute.ComposedLayout, @@ -632,9 +771,6 @@ class PersistentDenseGemmKernel: tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, ): - """ - GPU device kernel performing the Persistent batched GEMM computation. - """ warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) @@ -668,25 +804,12 @@ class PersistentDenseGemmKernel: # # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier # - # Define shared storage for kernel - @cute.struct - class SharedStorage: - ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] - acc_full_mbar_ptr: cute.struct.MemRange[ - cutlass.Int64, self.num_acc_stage * 2 - ] - tmem_dealloc_mbar: cutlass.Int64 - tmem_holding_buf: cutlass.Int32 - smem = utils.SmemAllocator() - storage = smem.allocate(SharedStorage) + storage = smem.allocate(self.shared_storage) # 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_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Warp) ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), num_stages=self.num_ab_stage, @@ -694,6 +817,7 @@ class PersistentDenseGemmKernel: consumer_group=ab_pipeline_consumer_group, tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cluster_layout_vmnk, + enable_multicast_signaling=True, defer_sync=True, ).make_participants() @@ -734,7 +858,7 @@ class PersistentDenseGemmKernel: ) # Cluster arrive after barrier init - pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C @@ -842,7 +966,17 @@ class PersistentDenseGemmKernel: # # Cluster wait before tensor memory alloc # - pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() # # Specialized TMA load warp @@ -852,11 +986,6 @@ class PersistentDenseGemmKernel: # # 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() - while work_tile.is_valid_tile: # Get tile coord from tile scheduler cur_tile_coord = work_tile.tile_idx @@ -936,11 +1065,6 @@ class PersistentDenseGemmKernel: # # 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_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.num_acc_stage ) @@ -970,11 +1094,6 @@ class PersistentDenseGemmKernel: 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 # @@ -984,19 +1103,11 @@ class PersistentDenseGemmKernel: handle = ab_consumer.wait_and_advance(peek_ab_full_status) # tCtAcc += tCrA * tCrB - num_kblocks = cute.size(tCrA, mode=[2]) - for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): - kblk_crd = (None, None, kblk_idx, handle.index) - - cute.gemm( - tiled_mma, - tCtAcc, - tCrA[kblk_crd], - tCrB[kblk_crd], - tCtAcc, - ) - # Enable accumulate on tCtAcc after first kblock - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm( + tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc + ) # Async arrive AB buffer empty handle.release() @@ -1051,45 +1162,134 @@ class PersistentDenseGemmKernel: # (MMA, MMA_M, MMA_N, STAGE) tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + m_cta_tiles_in_total = gC_mnl.shape[2] * cute.size(tiled_mma.thr_id.shape) + n_cta_tiles_in_total = gC_mnl.shape[3] + # # Persistent tile scheduling loop for epilogue # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + + epilogue_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), ) - if cutlass.const_expr(self.use_tma_store): + 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], + ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() - assert tma_atom_c is not None and sC is not None - self.epilogue_tma_store_release_flag( - tidx, - warp_idx, - acc_pipeline, - tiled_mma, - tma_atom_c, - tCtAcc_base, - sC, - tCgC, - epi_tile, - tile_sched, - epilogue_op, - flag_base=barrier_flag_mc, - flag_mem_scope="gpu", - ) + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + + # Post-epilogue: arrive flag for distributed communication + if cutlass.const_expr(self.all_reduce == "LDMCxSTMC"): + # 1D linear index of current output tile + cta_tile_id_linear = cute.Int32( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + + # Wait for C store complete + # Unlike regular epilogue where we only wait C store complete once at end of each kernel. + # Here we need to wait for C store complete for each output tile before we set the release flag. + c_pipeline.producer_tail() + + # Update flag with release semantic with GPU scope + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = ( + barrier_flag_multicast.iterator + cta_tile_id_linear + ) + # Release flag MUST be `release` semantic + # Refer to CUDA PTX doc Memory Consistency Model for more detail + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + order="release", + scope="gpu", + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + # Post-epilogue: arrive flag for distributed communication + if cutlass.const_expr(self.all_reduce == "LDMCxSTMC"): + # 1D linear index of current output tile + cta_tile_id_linear = cute.Int32( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + + # Sync all thread in epilogue warp + # Ensure all thread in epilogue warp have issue store operation + epilogue_sync_barrier.arrive_and_wait() + + # Update flag with release semantic with GPU scope + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = ( + barrier_flag_multicast.iterator + cta_tile_id_linear + ) + # Release flag MUST be `release` semantic + # Refer to CUDA PTX doc Memory Consistency Model for more detail + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + order="release", + scope="gpu", + ) + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete (only if not using LDMCxSTMC which does per-tile tail) + if cutlass.const_expr(self.all_reduce != "LDMCxSTMC"): + c_pipeline.producer_tail() else: - self.epilogue_release_flag( - tidx, - acc_pipeline, - tiled_mma, - tCtAcc_base, - tCgC, - epi_tile, - tile_sched, - epilogue_op, - tmem_dealloc_barrier, - flag_base=barrier_flag_mc, - flag_mem_scope="gpu", - ) + # Synchronize before TMEM dealloc + if cutlass.const_expr(self.all_reduce != "LDMCxSTMC"): + tmem_dealloc_barrier.arrive_and_wait() # # Dealloc the tensor memory buffer @@ -1098,146 +1298,99 @@ class PersistentDenseGemmKernel: tmem.free(tmem_ptr) # /////////////////////////////////////////////////////////////////////////////// - # Allreduce warps + # AllReduce warps # /////////////////////////////////////////////////////////////////////////////// if cutlass.const_expr(self.all_reduce == "LDMCxSTMC"): - if warp_idx >= self.all_reduce_warp_id[0]: - # /////////////////////////////////////////////////////////////////////////////// - # Add persistent tile loop - # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.comm_warp_id[0]: + # (bM, bN, RestM, RestN, RestL) with multicast memory as comm input + gCommInMC_mnl = cute.local_tile( + mCommInMC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + # (bM, bN, RestM, RestN, RestL) with multicast memory as comm output + gCommOutMC_mnl = cute.local_tile( + mCommOutMC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) with multicast memory + tCgCommInMC = thr_mma.partition_C(gCommInMC_mnl) + tCgCommOutMC = thr_mma.partition_C(gCommOutMC_mnl) - rank_id = self.rank_id - num_ranks = Int32(self.num_ranks) - lane_id = cute.arch.lane_idx() + m_cta_tiles_in_total = gCommInMC_mnl.shape[2] * cute.size( + tiled_mma.thr_id.shape + ) + n_cta_tiles_in_total = gCommInMC_mnl.shape[3] - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + # + # Persistent tile scheduling loop for all reduce + # + comm_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.comm_sync_bar_id, + num_threads=32 * len(self.comm_warp_id), ) - work_tile = tile_sched.initial_work_tile_info() - # we want 128bit ld/st for better performance - atom_val = 128 // c_mc.element_type.width - atom_thr_n = self.mma_tiler[1] // atom_val - atom_thr_m = len(self.all_reduce_warp_id) * cute.arch.WARP_SIZE // atom_thr_n - thr_layout = cute.make_layout( - (atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1) - ) - val_layout = cute.make_layout((1, atom_val), stride=(atom_val, 1)) - - copy_atom_load = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), c_mc.element_type - ) - tiled_copy_fake = cute.make_tiled_copy_tv( - copy_atom_load, thr_layout, val_layout - ) - thr_copy_fake = tiled_copy_fake.get_slice( - tidx - self.all_reduce_warp_id[0] * 32 - ) - # predicate tensor - idC = cute.make_identity_tensor(c_mc.shape) + # AllReduce communication pattern (w/ rank4 as example): + # + # [rank0 A] [rank1 A] [rank2 A] [rank3 A] -> RS -> [rank0+1+2+3 A] [rank0+1+2+3 A] [rank0+1+2+3 A] [rank0+1+2+3 A] + # [rank0 B] [rank1 B] [rank2 B] [rank3 B] -> RS -> [rank0+1+2+3 B] [rank0+1+2+3 B] [rank0+1+2+3 B] [rank0+1+2+3 B] + # [rank0 C] [rank1 C] [rank2 C] [rank3 C] -> RS -> [rank0+1+2+3 C] [rank0+1+2+3 C] [rank0+1+2+3 C] [rank0+1+2+3 C] + # [rank0 D] [rank1 D] [rank2 D] [rank3 D] -> RS -> [rank0+1+2+3 D] [rank0+1+2+3 D] [rank0+1+2+3 D] [rank0+1+2+3 D] + # rank0 rank1 rank2 rank3 -> RS -> rank0 rank1 rank2 rank3 + # + # Each rank's AR warp only handle 1/P of overall output tiles, and relies on other rank to broadcast data to it + # rank0: compute each tile's [0:1/4] and broadcast it to all rank + # rank1: compute each tile's [1/4:2:4] and broadcast it to all rank while work_tile.is_valid_tile: cur_tile_coord = work_tile.tile_idx - tile_id = Int32( - tile_sched._current_work_linear_idx - * cute.size(self.cluster_shape_mn) - + cute.arch.block_idx_in_cluster() - ) - mma_tile_coord_mnl = ( - cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), - cur_tile_coord[1], - cur_tile_coord[2], - ) - # System barrier to make sure that data from each GPU is in memory before allreduce - if warp_idx == self.all_reduce_warp_id[0]: + # (1) Per-tile flag wait. + cta_tile_id_linear = ( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + if warp_idx == self.comm_warp_id[0]: with cute.arch.elect_one(): - flag = barrier_flag.iterator + tile_id - utils.distributed.spin_lock_atom_cas_relaxed_wait( - lock_ptr=flag, - expected_val=num_ranks, + utils.distributed.spin_lock_atom_cas_acquire_wait( + lock_ptr=barrier_flag_unicast.iterator + + cta_tile_id_linear, + expected_val=self.num_ranks, reset_val=0, scope="gpu", ) - self.all_reduce_sync_barrier.arrive_and_wait() - # partition and slice at tile level - gC_mc = cute.local_tile( - c_mc, - cute.slice_(self.mma_tiler, (None, None, 0)), - (None, None, None), - ) - cC = cute.local_tile( - idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) - ) - - tCgC_mc = thr_mma.partition_C(gC_mc) - tCpC = thr_mma.partition_C(cC) - tCgC_mc_slice = tCgC_mc[((None, None), 0, 0, *mma_tile_coord_mnl)] - tCpC_slice = tCpC[((None, None), 0, 0, *mma_tile_coord_mnl)] + comm_sync_barrier.arrive_and_wait() - # partition based on the number of GPUs - cta_mma_tile_m = self.mma_tiler[0] // cute.size( - tiled_mma.thr_id.shape + # (2) Per-tile data movement (multimem.ld_reduce + multimem.st). + LDMCxSTMC( + self, + cur_tile_coord=cur_tile_coord, + tiled_mma=tiled_mma, + comm_tidx=tidx, + tCgCommInMC=tCgCommInMC, + tCgCommOutMC=tCgCommOutMC, ) - m_local_rank = int(cta_mma_tile_m / self.num_ranks) - tCgC_mc_slice_partitioned = cute.zipped_divide( - tCgC_mc_slice, (m_local_rank, self.mma_tiler[1]) - ) - tCpC_slice_partitioned = cute.zipped_divide(tCpC_slice, (m_local_rank, self.mma_tiler[1])) - tCgC_mc_local_rank = cute.slice_( - tCgC_mc_slice_partitioned, ((None, None), (rank_id, 0)) - ) - tCpC_local_rank = cute.slice_(tCpC_slice_partitioned, ((None, None), (rank_id, 0))) - # partition at thread level - frgC_mc = thr_copy_fake.partition_S(tCgC_mc_local_rank) - frpC = thr_copy_fake.partition_S(tCpC_local_rank) - atom, loop_m, loop_n = frgC_mc.shape - for i in cutlass.range_constexpr(loop_m): - for j in cutlass.range_constexpr(loop_n): - if cute.elem_less(frpC[0, i, j], c_mc.shape): - mc_ptr = frgC_mc[None, i, j].iterator - x, y, z, w = 0, 0, 0, 0 - if cutlass.const_expr(self.c_dtype == Float16): - x, y, z, w = utils.distributed.multimem_ld_reduce_8xf16( - mc_ptr - ) - elif cutlass.const_expr(self.c_dtype == Float32): - x, y, z, w = utils.distributed.multimem_ld_reduce_4xf32( - mc_ptr - ) - elif cutlass.const_expr(self.c_dtype == BFloat16): - x, y, z, w = ( - utils.distributed.multimem_ld_reduce_8xbf16(mc_ptr) - ) - elif cutlass.const_expr(self.c_dtype == Float8E4M3FN): - x, y, z, w = ( - utils.distributed.multimem_ld_reduce_16xe4m3(mc_ptr) - ) - elif cutlass.const_expr(self.c_dtype == Float8E5M2): - x, y, z, w = ( - utils.distributed.multimem_ld_reduce_16xe5m2(mc_ptr) - ) - utils.distributed.multimem_st_4xb32(mc_ptr, x, y, z, w) + # # Advance to next tile + # tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() - self.all_reduce_sync_barrier.arrive_and_wait() + # Ensure all threads in allreduce warp group have complete issuing all reduce operations. + comm_sync_barrier.arrive_and_wait() # # Set Per SM Flag with Release # # This ensure # 1. no rank early exit while other ranks are still issuing multimem.ld_reduce - # 2. each rank's prior multimem.st have become visiable to all other ranks in the system (w/ .SYS scope) - if warp_idx == self.all_reduce_warp_id[0]: + # 2. each rank's prior multimem.st have become visible to all other ranks in the system (w/ .SYS scope) + if warp_idx == self.comm_warp_id[0]: with cute.arch.elect_one(): # Offset to last tile flag idx - last_tile_id_linear = cute.size( - tile_sched.params.problem_layout_ncluster_mnl - ) * cute.size(self.cluster_shape_mn) + total_num_cta_tile = m_cta_tiles_in_total * n_cta_tiles_in_total # Linear id of current SM. sm_id_linear = ( cute.arch.block_idx()[0] @@ -1248,445 +1401,91 @@ class PersistentDenseGemmKernel: ) # Release flag with sys scope utils.distributed.multimem_red_add1( - lock_ptr=barrier_flag_mc.iterator - + last_tile_id_linear + lock_ptr=barrier_flag_multicast.iterator + + total_num_cta_tile + sm_id_linear, scope="sys", order="release", ) - # Relaxed spin-lock wait flag with sys scope - utils.distributed.spin_lock_atom_cas_relaxed_wait( - lock_ptr=barrier_flag.iterator - + last_tile_id_linear + # Acquire spin-lock wait flag with sys scope + utils.distributed.spin_lock_atom_cas_acquire_wait( + lock_ptr=barrier_flag_unicast.iterator + + total_num_cta_tile + sm_id_linear, - expected_val=num_ranks, + expected_val=self.num_ranks, reset_val=0, scope="sys", ) - @cute.jit - def epilogue_tma_store_release_flag( - self, - epi_tidx: cutlass.Int32, - warp_idx: cutlass.Int32, - acc_pipeline: pipeline.PipelineAsync, + @staticmethod + def _compute_stages( tiled_mma: cute.TiledMma, - tma_atom_c: cute.CopyAtom, - # Input of epilogue - tCtAcc_base: cute.Tensor, - # Staging of epilogue - sC: cute.Tensor, - # Output of epilogue - tCgC: cute.Tensor, - epi_tile: cute.Tile, - tile_sched: utils.StaticPersistentTileScheduler, - epilogue_op: cutlass.Constexpr, - flag_base: cute.Tensor, - flag_mem_scope: str, - ) -> None: - tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilogue_tmem_copy_and_partition( - epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + c_smem_layout: Union[cute.Layout, None], + ) -> Tuple[int, int, int]: + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 ) - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - tiled_copy_r2s, tRS_rC, tRS_sC = self.epilogue_smem_copy_and_partition( - tiled_copy_t2r, tTR_rC, epi_tidx, sC - ) + 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 - # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) - tCgC_epi = cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - # ((ATOM_V, REST_V), EPI_M, EPI_N) - # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) - bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( - tma_atom_c, - 0, - cute.make_layout(1), - cute.group_modes(sC, 0, 2), - cute.group_modes(tCgC_epi, 0, 2), - ) + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage - # Threads/warps participating in tma store pipeline - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, - 32 * len(self.epilogue_warp_id), - ) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group - ) - - epilogue_sync_barrier = pipeline.NamedBarrier( - barrier_id=self.epilogue_sync_bar_id, - num_threads=32 * len(self.epilogue_warp_id), - ) - - work_tile = tile_sched.initial_work_tile_info() - 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_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) - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - 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(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # - # Convert to C type - # - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tRS_rC.store(acc_vec) - - # - # Store C to shared memory - # - c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage - 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("async.shared", space="cta") - epilogue_sync_barrier.arrive_and_wait() - - # - # TMA store C to global memory - # - if warp_idx == self.epilogue_warp_id[0]: - cute.copy( - tma_atom_c, - bSG_sC[(None, c_buffer)], - bSG_gC[(None, subtile_idx)], - ) - # Fence and barrier to make sure shared memory store is visible to TMA store - c_pipeline.producer_commit() - c_pipeline.producer_acquire() - epilogue_sync_barrier.arrive_and_wait() - - epilogue_sync_barrier.arrive_and_wait() - - # - # Async arrive accumulator buffer empty - # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - - # - # Set Per Output Tile Flag with Release - # - # 1D linear index of current output tile - tile_id_linear = Int32( - tile_sched._current_work_linear_idx * cute.size(self.cluster_shape_mn) - + cute.arch.block_idx_in_cluster() - ) - # Wait for C store complete - # Unlike regular epilogue where we only wait C store complete once at end of each kernel. - # Here we need to wait for C store complete for each output tile before we set the release flag. - c_pipeline.producer_tail() - # Update flag with release semantic with GPU scope - if warp_idx == self.epilogue_warp_id[0]: - with cute.arch.elect_one(): - flag_curr_tile = flag_base.iterator + tile_id_linear - utils.distributed.multimem_red_add1( - lock_ptr=flag_curr_tile, - scope=flag_mem_scope, - order="release", - ) - - # - # Advance to next tile - # - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - - @cute.jit - def epilogue_release_flag( - self, - epi_tidx: cutlass.Int32, - acc_pipeline: pipeline.PipelineAsync, - tiled_mma: cute.TiledMma, - tCtAcc_base: cute.Tensor, - tCgC: cute.Tensor, - epi_tile: cute.Tile, - tile_sched: utils.StaticPersistentTileScheduler, - epilogue_op: cutlass.Constexpr, - tmem_dealloc_barrier: pipeline.NamedBarrier, - flag_base: cute.Tensor, - flag_mem_scope: str, - ) -> None: - tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilogue_tmem_copy_and_partition( - epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs - ) - - gC_epi = cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) - thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) - tTR_gC_partitioned = thr_copy_t2r.partition_D(gC_epi) - # (T2R, T2R_M, T2R_N) - tTR_rC = cute.make_rmem_tensor( - tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype - ) - simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - - work_tile = tile_sched.initial_work_tile_info() - 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 - # - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_gC = tTR_gC_partitioned[ - (None, None, None, None, None, *mma_tile_coord_mnl) - ] - - # Set tensor memory buffer for current tile - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_tAcc = tTR_tAcc_base[ - (None, None, None, None, None, acc_consumer_state.index) - ] - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) - - # - # Wait for accumulator buffer full - # - acc_pipeline.consumer_wait(acc_consumer_state) - - # - # Store accumulator to global memory in subtiles - # - subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) - for subtile_idx in cutlass.range(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # - # Convert to C type - # - acc_vec = tTR_rAcc.load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tTR_rC.store(acc_vec) - - # - # Store C to global memory - # - cute.copy(simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)]) - - # - # Async arrive accumulator buffer empty - # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - - # - # Set Per Output Tile Flag with Release - # - # 1D linear index of current output tile - tile_id_linear = Int32( - tile_sched._current_work_linear_idx * cute.size(self.cluster_shape_mn) - + cute.arch.block_idx_in_cluster() - ) - # Wait for C store complete - # Unlike regular epilogue where we only wait C store complete once at end of each kernel. - # Here we need to wait for C store complete for each output tile before we set the release flag. - c_pipeline.producer_tail() - # Update flag with release semantic with GPU scope - if warp_idx == self.epilogue_warp_id[0]: - with cute.arch.elect_one(): - flag_curr_tile = flag_base.iterator + tile_id_linear - utils.distributed.multimem_red_add1( - lock_ptr=flag_curr_tile, - scope=flag_mem_scope, - order="release", - ) - - # Advance to next tile - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - - # Synchronize before TMEM dealloc (done by the caller) - tmem_dealloc_barrier.arrive_and_wait() - - def epilogue_tmem_copy_and_partition( - self, - tidx: cutlass.Int32, - tAcc: cute.Tensor, - gC_mnl: cute.Tensor, - epi_tile: cute.Tile, - use_2cta_instrs: Union[cutlass.Boolean, bool], - ) -> Tuple[cute.TiledCopy, 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 gC_mnl: The global tensor C - :type gC_mnl: cute.Tensor - :param epi_tile: The epilogue tiler - :type epi_tile: 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.c_layout, - self.c_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) - gC_mnl_epi = cute.flat_divide( - gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile - ) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) - tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) - # (T2R, T2R_M, T2R_N) - tTR_rAcc = cute.make_rmem_tensor( - tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype - ) - return tiled_copy_t2r, tTR_tAcc, tTR_rAcc - - def epilogue_smem_copy_and_partition( - self, - tiled_copy_t2r: cute.TiledCopy, - tTR_rC: cute.Tensor, - tidx: cutlass.Int32, - sC: cute.Tensor, - ) -> Tuple[cute.TiledCopy, 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_rC: The partitioned accumulator tensor - :type tTR_rC: cute.Tensor - :param tidx: The thread index in epilogue warp groups - :type tidx: cutlass.Int32 - :param sC: The shared memory tensor to be copied and partitioned - :type sC: cute.Tensor - :type sepi: cute.Tensor - - :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: - - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) - - tRS_rC: The partitioned tensor C (register source) - - tRS_sC: The partitioned tensor C (smem destination) - :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] - """ - copy_atom_r2s = sm100_utils.get_smem_store_op( - self.c_layout, self.c_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_sC = thr_copy_r2s.partition_D(sC) - # (R2S, R2S_M, R2S_N) - tRS_rC = tiled_copy_r2s.retile(tTR_rC) - return tiled_copy_r2s, tRS_rC, tRS_sC + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage @staticmethod def _compute_grid( c: cute.Tensor, cta_tile_shape_mnk: Tuple[int, int, int], cluster_shape_mn: Tuple[int, int], + swizzle_size: int, + raster_order: Literal["m", "n"], 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 C. - - :param c: The output tensor C - :type c: 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]] - """ c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) gc = cute.zipped_divide(c, tiler=c_shape) num_ctas_mnl = gc[(0, (None, None, None))].shape cluster_shape_mnl = (*cluster_shape_mn, 1) tile_sched_params = utils.PersistentTileSchedulerParams( - num_ctas_mnl, cluster_shape_mnl + num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_order == "m" ) grid = utils.StaticPersistentTileScheduler.get_grid_shape( tile_sched_params, max_active_clusters @@ -1699,40 +1498,17 @@ class PersistentDenseGemmKernel: tiled_mma: cute.TiledMma, mma_tiler: Tuple[int, int, int], num_acc_stage: int, + arch: str, ) -> 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) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) return num_tmem_alloc_cols - def is_valid_dtypes( + def check_supported_dtypes( self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] ) -> bool: - """ - Check if the dtypes are valid - - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] - :param c_dtype: The data type of the output tensor - :type c_dtype: Type[cutlass.Numeric] - - :return: True if the dtypes are valid, False otherwise - :rtype: bool - """ valid_ab_dtypes = { cutlass.Float16, cutlass.BFloat16, @@ -1743,10 +1519,15 @@ class PersistentDenseGemmKernel: cutlass.Float8E5M2, } if ab_dtype not in valid_ab_dtypes: - return False + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype: {ab_dtype}" + ) - if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: - return False + valid_acc_dtypes = {cutlass.Float32, cutlass.Float16, cutlass.Int32} + if self.acc_dtype not in valid_acc_dtypes: + raise testing.CantImplementError( + f"[dtype] Unsupported accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and AB type acc_ab_compatibility = { @@ -1766,7 +1547,9 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and AB type if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype {ab_dtype} for accumulator {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and C type acc_c_compatibility = { @@ -1795,34 +1578,46 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and C type if c_dtype not in acc_c_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for accumulator {self.acc_dtype}" + ) # check if c_dtype is supported by multimem all-reduce if cutlass.const_expr( - c_dtype not in {cutlass.Float16, cutlass.Float32, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2} + c_dtype + not in { + cutlass.Float16, + cutlass.Float32, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } ): - return False + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for multimem all-reduce" + ) return True - def is_valid_mma_tiler_and_cluster_shape(self) -> bool: - """Check if the mma tiler and cluster shape are valid. - - :return: True if the mma tiler and cluster shape are valid, False otherwise - :rtype: bool - """ - is_valid = True + def check_mma_tiler_and_cluster_shape(self) -> bool: # Skip invalid mma tile shape - if not ( - (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) - or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) - ): - is_valid = False - if self.mma_tiler_mn[1] not in range(32, 257, 32): - is_valid = False + # 1SM: MMA tile {64, 128} x {64, 128, 192, 256} + # 2SM: MMA tile {128, 256} x {64, 128, 192, 256} + valid_mma_m = [64, 128] if not self.use_2cta_instrs else [128, 256] + valid_mma_n = [64, 128, 192, 256] + if self.mma_tiler_mn[0] not in valid_mma_m: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler M: {self.mma_tiler_mn[0]}, expected {valid_mma_m} for {'1SM' if not self.use_2cta_instrs else '2SM'}" + ) + if self.mma_tiler_mn[1] not in valid_mma_n: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler N: {self.mma_tiler_mn[1]}, expected {valid_mma_n}" + ) # Skip illegal cluster shape if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: - is_valid = False + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape M: {self.cluster_shape_mn[0]}, must be multiple of {2 if self.use_2cta_instrs else 1} for {'2SM' if self.use_2cta_instrs else '1SM'}" + ) # Skip invalid cluster shape is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 if ( @@ -1832,10 +1627,12 @@ class PersistentDenseGemmKernel: or not is_power_of_2(self.cluster_shape_mn[0]) or not is_power_of_2(self.cluster_shape_mn[1]) ): - is_valid = False - return is_valid + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape: {self.cluster_shape_mn}, must be positive power of 2 with total <= 16" + ) + return True - def is_valid_tensor_alignment( + def check_tensor_alignment( self, m: int, n: int, @@ -1843,226 +1640,600 @@ class PersistentDenseGemmKernel: l: int, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], ) -> bool: - """ - Check if the tensor alignment is valid - - :param m: The number of rows in the A tensor - :type m: int - :param n: The number of columns in the B tensor - :type n: int - :param k: The number of columns in the A tensor - :type k: int - :param l: The number of columns in the C tensor - :type l: int - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] - :param c_dtype: The data type of the output tensor - :type c_dtype: Type[cutlass.Numeric] - :param a_major: The major axis of the A tensor - :type a_major: str - :param b_major: The major axis of the B tensor - :type b_major: str - :param c_major: The major axis of the C tensor - :type c_major: str - - :return: True if the problem shape is valid, False otherwise - :rtype: bool - """ - is_valid = True - - def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): major_mode_idx = 0 if is_mode0_major else 1 num_major_elements = tensor_shape[major_mode_idx] num_contiguous_elements = 16 * 8 // dtype.width return num_major_elements % num_contiguous_elements == 0 if ( - not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) - or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) - or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) ): - is_valid = False + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: contiguous dimension must be 16B aligned for (M={m}, K={k}, N={n})" + ) - return is_valid + if m % 128 != 0 and n % 128 != 0: + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: at least one of M={m} or N={n} must be divisible by 128" + ) - def is_valid_epilogue_store_option(self, m: int, n: int) -> bool: - """ - Check if the epilogue store option is valid + if a_major not in ["m", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid a_major: {a_major}, expected 'm' or 'k'" + ) + if b_major not in ["n", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid b_major: {b_major}, expected 'n' or 'k'" + ) + if c_major not in ["m", "n"]: + raise testing.CantImplementError( + f"[alignment] Invalid c_major: {c_major}, expected 'm' or 'n'" + ) - :param m: The number of rows in the A tensor - :type m: int - :param n: The number of columns in the B tensor - :type n: int + return True - :return: True if the epilogue store option is valid, False otherwise - :rtype: bool - """ - - is_valid = True - # None TMA store version does not have predication, can not support OOB tiles + def check_epilogue_store_option(self, m: int, n: int) -> bool: + # Non TMA store version does not have predication, can not support OOB tiles cta_tile_shape_mn = ( self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), self.mma_tiler_mn[1], ) if not self.use_tma_store: if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): - is_valid = False - return is_valid + raise testing.CantImplementError( + f"[epilogue] Problem shape (M={m}, N={n}) must be divisible by cta tile shape {cta_tile_shape_mn} for non TMA store" + ) + # CTA swizzling improves the L2 cache utilization and reduces the number of cache misses. + # Make sure the swizzle size divides the cta/cluster count since non TMA epilogue don't support OOB tiles. + # Swizzle only applies to the dimension orthogonal to the raster direction. + m_per_swizzle = (m // cta_tile_shape_mn[0]) // self.cluster_shape_mn[0] + n_per_swizzle = (n // cta_tile_shape_mn[1]) // self.cluster_shape_mn[1] + if self.raster_order == "m": + # Swizzle applies to N dimension only + if n_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape n={n} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if m_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape m={m} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_n" + ) + return True - def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool: - """Check if the given tensors can be implemented by this kernel. + def check_valid_rank_config(self) -> bool: + if self.num_ranks not in [2, 4, 8] and self.all_reduce != "none": + raise testing.CantImplementError( + f"[rank_config] Invalid rank config: num_ranks={self.num_ranks}, must be 2, 4, or 8 for all-reduce" + ) + return True - :param a: Input tensor A - :type a: cute.Tensor - :param b: Input tensor B - :type b: cute.Tensor - :param c: Output tensor C - :type c: cute.Tensor + def check_valid_all_reduce_mode(self) -> bool: + if self.all_reduce not in ["LDMCxSTMC", "none"]: + raise testing.CantImplementError( + f"[all_reduce_mode] Invalid all-reduce mode: {self.all_reduce}, expected 'LDMCxSTMC' or 'none'" + ) + return True - :return: True if the gemm supports the given config, False otherwise - :rtype: bool - """ - m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] + def check_even_divisiable_by_cluster_tilesize( + self, m: int, n: int, k: int, l: int + ) -> bool: + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] - # infer a_major, b_major, c_major - is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() - is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() - is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c() - a_major = "m" if is_m_major_a else "k" - b_major = "n" if is_n_major_b else "k" - c_major = "m" if is_m_major_c else "n" + cluster_shape_m = self.cluster_shape_mn[0] + cluster_shape_n = self.cluster_shape_mn[1] - can_implement = True + cluster_tile_m = cta_tile_m * cluster_shape_m + cluster_tile_n = cta_tile_n * cluster_shape_n + + if m % cluster_tile_m != 0 or n % cluster_tile_n != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_cluster_tile_size] Problem shape (M={m}, N={n}) not divisible by cluster tile size ({cluster_tile_m}, {cluster_tile_n})" + ) + return True + + def check_even_divisible_by_swizzle(self, m: int, n: int) -> bool: + # Only check for LDMCxSTMC mode which uses multimem operations + if self.all_reduce != "LDMCxSTMC": + return True + + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + + # Number of clusters in each dimension + num_clusters_m = m // (cta_tile_m * self.cluster_shape_mn[0]) + num_clusters_n = n // (cta_tile_n * self.cluster_shape_mn[1]) + + if self.raster_order == "m": + # Swizzle applies to N dimension only + if num_clusters_n % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in N dimension ({num_clusters_n}) " + f"must be divisible by swizzle_size={self.swizzle_size} for LDMCxSTMC with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if num_clusters_m % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in M dimension ({num_clusters_m}) " + f"must be divisible by swizzle_size={self.swizzle_size} for LDMCxSTMC with raster_order_n" + ) + + return True + + def can_implement( + self, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + ) -> bool: # Skip unsupported types - if not self.is_valid_dtypes(a.element_type, c.element_type): - can_implement = False + self.check_supported_dtypes(ab_dtype, c_dtype) + # Skip invalid mma tile shape and cluster shape - if not self.is_valid_mma_tiler_and_cluster_shape(): - can_implement = False + self.check_mma_tiler_and_cluster_shape() + + # Unpack mnkl for clarity in calling the epilog check + m, n, k, l = mnkl + + # Skip if GEMM not even divisiable by tilesize + # Right now, AR Fusion **ONLY** support even divisiable by tilesize + self.check_even_divisiable_by_cluster_tilesize(m, n, k, l) + + # Skip if swizzle doesn't evenly divide the problem (for LDMCxSTMC communication) + self.check_even_divisible_by_swizzle(m, n) + # Skip illegal problem shape for load/store alignment - if not self.is_valid_tensor_alignment( - m, n, k, l, a.element_type, c.element_type, a_major, b_major, c_major - ): - can_implement = False + self.check_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ) + # Skip invalid epilogue store option - if not self.is_valid_epilogue_store_option(m, n): - can_implement = False - # currently only support world_size = 2, 4, 8 - if dist.get_world_size() not in [2, 4, 8]: - can_implement = False - return can_implement + self.check_epilogue_store_option(m, n) + + # Skip invalid all-reduce rank config + self.check_valid_rank_config() + + # Skip invalid all-reduce mode + self.check_valid_all_reduce_mode() + + return True -def create_mc_tensor(torch_tensor_cpu, dtype, leading_dim, is_dynamic_layout=True): - torch_tensor_gpu_local = nvshmem.core.tensor(torch_tensor_cpu.shape, dtype=torch_tensor_cpu.dtype) - torch_tensor_gpu_local.copy_(torch_tensor_cpu) - torch_tensor_gpu_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, torch_tensor_gpu_local) +##################################################################### +#### KERNEL IMPLEMENTATION END HERE #### +##################################################################### - # create MC tensor memref - cute_tensor_mc = from_dlpack(torch_tensor_gpu_mc, assumed_align=16) - cute_tensor = from_dlpack(torch_tensor_gpu_local, assumed_align=16) +_DISTRIBUTED_INITIALIZED = False - cute_tensor.element_type = dtype - cute_tensor_mc.element_type = dtype - if is_dynamic_layout: - cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) +# Check for device p2p access +def check_cuda_p2p_access(): + num_devices = torch.cuda.device_count() + for i in range(num_devices - 1): + for j in range(i + 1, num_devices): + if not torch.cuda.can_device_access_peer(i, j): + return False + return True + + +# Check for multicast support +def check_multicast_support(): + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + for device_id in range(torch.cuda.device_count()): + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device_id): + return False + return True + except ImportError: + return False + + +def init_distributed(): + global _DISTRIBUTED_INITIALIZED + + if not dist.is_initialized() and not _DISTRIBUTED_INITIALIZED: + # * Check before running * + if not torch.cuda.is_available(): + print("CUDA is not available") + exit(-1) + if torch.cuda.get_device_capability()[0] != 10: + print("Running on non-sm100 device") + exit(-1) + if torch.cuda.device_count() < 2: + print("Less than two devices") + exit(-1) + if torch.cuda.device_count() % 2 != 0: + print("Odd world size") + exit(-1) + for var in ("LOCAL_RANK", "RANK", "WORLD_SIZE"): + if var not in os.environ: + print(f"{var} not in environment variables") + exit(-1) + + # LOCAL_RANK = which of this host's GPUs we bind to. + # RANK = global identity across all hosts; used for NCCL init + # and for kernel rank_id (slab ownership in the AR). + local_rank = int(os.environ["LOCAL_RANK"]) + global_rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + world_size=world_size, + rank=global_rank, + device_id=device, + ) + + _DISTRIBUTED_INITIALIZED = True + + print( + f"[init] hostname={socket.gethostname()} global_rank={global_rank} " + f"local_rank={local_rank} world_size={world_size} device={device}" + ) + + global_rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", global_rank)) + device = torch.device("cuda", local_rank) + + # Return the GLOBAL rank — callers feed this into kernel rank_id. + return global_rank, world_size, device + + +def clean_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def test_simple_torch_allreduce(device): + try: + tensor = symm_mem.empty((16384,), dtype=torch.float32, device=device) + symm_mem.rendezvous(tensor, group=dist.group.WORLD) + torch.ops.symm_mem.multimem_all_reduce_( + tensor, "sum", dist.group.WORLD.group_name + ) + except Exception: + return False + return True + + +# --------------------------------------------------------------------------- +# Dump / Load helpers for determinism verification +# --------------------------------------------------------------------------- + + +def _make_tensor_filename( + local_rank, global_rank, world_size, hostname, name, tensor +) -> str: + dtype_str = str(tensor.dtype).replace("torch.", "") + shape_str = "x".join(str(d) for d in tensor.shape) + return ( + f"rank{local_rank}_global{global_rank}_world{world_size}" + f"_{hostname}_{name}_{dtype_str}_{shape_str}.pt" + ) + + +def _dump_tensors( + dump_dir, local_rank, global_rank, world_size, hostname, named_tensors: dict +): + os.makedirs(dump_dir, exist_ok=True) + for name, tensor in named_tensors.items(): + fname = _make_tensor_filename( + local_rank, global_rank, world_size, hostname, name, tensor + ) + torch.save(tensor, os.path.join(dump_dir, fname)) + + +def _write_metadata_json( + dump_dir, + global_rank, + world_size, + mnkl, + ab_dtype, + c_dtype, + acc_dtype, + a_major, + b_major, + c_major, +): + """ALL ranks call this (for all_gather_object). Only rank 0 writes the file.""" + hostnames = [None] * world_size + dist.all_gather_object(hostnames, socket.gethostname()) + if global_rank == 0: + os.makedirs(dump_dir, exist_ok=True) + meta = { + "mnkl": list(mnkl), + "world_size": world_size, + "ab_dtype": ab_dtype.__name__, + "c_dtype": c_dtype.__name__, + "acc_dtype": acc_dtype.__name__, + "a_major": a_major, + "b_major": b_major, + "c_major": c_major, + "hostnames": hostnames, + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), + } + with open(os.path.join(dump_dir, "metadata.json"), "w") as f: + json.dump(meta, f, indent=2) + + +def _load_input_tensors( + load_dir, + local_rank, + global_rank, + world_size, + ab_dtype, + expected_shape_a, + expected_shape_b, +): + """Load A and B for this rank from dump dir. Validates dtype and shape.""" + expected_torch_dtype = torch_dtype(ab_dtype) + + def _load_one(name, expected_shape): + pattern = os.path.join( + load_dir, + f"rank*_global{global_rank}_world{world_size}_*_{name}_*.pt", + ) + matches = glob.glob(pattern) + if len(matches) == 0: + raise FileNotFoundError(f"[load_input] No file for {name}: {pattern}") + if len(matches) > 1: + raise RuntimeError( + f"[load_input] Multiple files matched for {name}: {matches}" + ) + tensor = torch.load(matches[0], weights_only=True) + if tensor.dtype != expected_torch_dtype: + raise ValueError( + f"[load_input] dtype mismatch for {name}: got {tensor.dtype}, " + f"expected {expected_torch_dtype}" + ) + if tuple(tensor.shape) != tuple(expected_shape): + raise ValueError( + f"[load_input] shape mismatch for {name}: got {tuple(tensor.shape)}, " + f"expected {tuple(expected_shape)}" + ) + return tensor + + torch_tensor_a_cpu = _load_one("a", expected_shape_a) + torch_tensor_b_cpu = _load_one("b", expected_shape_b) + return torch_tensor_a_cpu, torch_tensor_b_cpu + + +def allocate_tensors( + *, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + num_workspace: int, + device: int, + slot_init_mode: Literal["test", "benchmark"], + global_rank: int = 0, + local_rank: int = 0, + world_size: int = 1, + load_input_dir: Optional[str] = None, +): + _init_lo, _init_hi = (-5, 5) if slot_init_mode == "test" else (-20, 20) + + def _create_multicast_tensor(*, torch_tensor_ref, dtype, leading_dim): + is_col_major = leading_dim == 0 + shape = torch_tensor_ref.shape + ref_dtype = torch_tensor_ref.dtype + + if is_col_major: + reversed_shape = tuple(reversed(shape)) + perm = tuple(range(len(shape) - 1, -1, -1)) + torch_symm_base = symm_mem.empty( + reversed_shape, device=device, dtype=ref_dtype + ) + torch_tensor_uc = torch_symm_base.permute(perm) + else: + torch_symm_base = symm_mem.empty(shape, device=device, dtype=ref_dtype) + torch_tensor_uc = torch_symm_base + + torch_tensor_uc.copy_(torch_tensor_ref) + + torch_tensor_hdl = symm_mem.rendezvous(torch_symm_base, group=dist.group.WORLD) + + cute_tensor_uc = from_dlpack(torch_tensor_uc, assumed_align=16) + cute_tensor_uc.element_type = dtype + cute_tensor_uc = cute_tensor_uc.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor_uc = cutlass_torch.convert_cute_tensor( + torch_tensor_uc, + cute_tensor_uc, + dtype, + is_dynamic_layout=True, + ) + + if is_col_major: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, reversed_shape, ref_dtype + ).permute(perm) + else: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, shape, ref_dtype + ) + cute_tensor_mc = from_dlpack(mc_torch, assumed_align=16) + cute_tensor_mc.element_type = dtype cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim) - return cute_tensor, cute_tensor_mc, torch_tensor_gpu_local, torch_tensor_gpu_mc + return cute_tensor_uc, cute_tensor_mc, torch_tensor_uc + m, n, k, l = mnkl -def create_tensors( - l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype, mma_tiler_mn, use_2cta_instrs -): - torch.manual_seed(1111) - cta_tile_shape_mn = ( - mma_tiler_mn[0] // (2 if use_2cta_instrs else 1), - mma_tiler_mn[1], - ) - problem_shape_ntile_mn = (m // cta_tile_shape_mn[0], n // cta_tile_shape_mn[1]) - num_tiles = problem_shape_ntile_mn[0] * problem_shape_ntile_mn[1] - num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + if load_input_dir is not None: + # k-major A stored as (m,k,l); m-major A stored as (l,m,k) — same as cutlass_torch.matrix permute + expected_shape_a = (m, k, l) if a_major == "k" else (l, m, k) + expected_shape_b = (n, k, l) if b_major == "k" else (l, n, k) + torch_tensor_a_cpu, torch_tensor_b_cpu = _load_input_tensors( + load_dir=load_input_dir, + local_rank=local_rank, + global_rank=global_rank, + world_size=world_size, + ab_dtype=ab_dtype, + expected_shape_a=expected_shape_a, + expected_shape_b=expected_shape_b, + ) + else: + torch.manual_seed(42 + global_rank) + torch_tensor_a_int_cpu = cutlass_torch.matrix( + l, + m, + k, + a_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=_init_lo, max_val=_init_hi + ), + ) + torch_tensor_b_int_cpu = cutlass_torch.matrix( + l, + n, + k, + b_major == "n", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=_init_lo, max_val=_init_hi + ), + ) + torch_tensor_a_cpu = torch_tensor_a_int_cpu.to(dtype=torch_dtype(ab_dtype)) + torch_tensor_b_cpu = torch_tensor_b_int_cpu.to(dtype=torch_dtype(ab_dtype)) - a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) - b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) - c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype) - barrier_flag_torch_cpu = torch.zeros(num_tiles+num_sms, device="cuda", dtype=torch.int32) - - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + torch_tensor_c_int_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc = create_mc_tensor( - c_torch_cpu, c_dtype, (1 if c_major == "n" else 0), is_dynamic_layout=True - ) - barrier_flag, barrier_flag_mc, barrier_flag_torch_gpu, barrier_flag_torch_gpu_mc = create_mc_tensor( - barrier_flag_torch_cpu, cutlass.Int32, 0, is_dynamic_layout=True + torch_tensor_comm_out_c_int_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), ) - return ( - a_tensor, - b_tensor, - c_tensor, - c_tensor_mc, - barrier_flag, - barrier_flag_mc, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - barrier_flag_torch_cpu, - c_torch_gpu, - c_torch_gpu_mc, - barrier_flag_torch_gpu, - barrier_flag_torch_gpu_mc, + torch_tensor_c_cpu = torch_tensor_c_int_cpu.to(dtype=torch_dtype(c_dtype)) + torch_tensor_comm_out_c_cpu = torch_tensor_comm_out_c_int_cpu.to( + dtype=torch_dtype(c_dtype) ) + cute_tensor_a_list = [] + cute_tensor_b_list = [] + torch_tensor_a_gpu_list = [] + torch_tensor_b_gpu_list = [] + for _ in range(num_workspace): + cute_tensor_a_i, torch_tensor_a_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_a_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_b_i, torch_tensor_b_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_b_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_a_list.append(cute_tensor_a_i) + cute_tensor_b_list.append(cute_tensor_b_i) + torch_tensor_a_gpu_list.append(torch_tensor_a_gpu_i) + torch_tensor_b_gpu_list.append(torch_tensor_b_gpu_i) -def compare( - a_torch_cpu, b_torch_cpu, c_torch_gpu, ab_dtype, c_dtype, tolerance -): + assert not torch.isnan(torch_tensor_a_gpu_list[0]).any() + assert not torch.isnan(torch_tensor_b_gpu_list[0]).any() - # Compute reference result - ref = torch.einsum( - "mkl,nkl->mnl", - a_torch_cpu.to(dtype=torch.float32), - b_torch_cpu.to(dtype=torch.float32), + ( + cute_tensor_c, + cute_tensor_comm_in_mc, + torch_tensor_c_gpu, + ) = _create_multicast_tensor( + torch_tensor_ref=torch_tensor_c_cpu, + dtype=c_dtype, + leading_dim=(1 if c_major == "n" else 0), + ) + ( + _cute_tensor_comm_out_uc, + cute_tensor_comm_out_mc, + torch_tensor_comm_out_gpu, + ) = _create_multicast_tensor( + torch_tensor_ref=torch_tensor_comm_out_c_cpu, + dtype=c_dtype, + leading_dim=(1 if c_major == "n" else 0), ) - acc_type = cutlass.Float32 - if c_dtype.width == 16: - acc_dtype = cutlass.Float32 - elif c_dtype.width == 8: - acc_dtype = cutlass.Float16 - - torch_acc_dtype = cutlass_torch.dtype(acc_dtype) - torch_c_dtype = cutlass_torch.dtype(c_dtype) + # Flag tensor sized for the worst-case kernel config: smallest cta_tile is + # 64x64, so max #tiles = (m/64) * (n/64). Pad by 160 for per-SM slots + # (B200 has 148; 160 is a safe upper bound). The kernel resets these flags + # internally on each launch, so a single allocation can be reused across + # (mma_tiler, use_2cta) candidates in a benchmark sweep. + num_flags = (m // 64) * (n // 64) + 160 - # note that in this example, f8 reduce sum is done by f16 acc precision and f16 by f32 - ref = ref.to(torch_c_dtype).to(torch_acc_dtype).cuda() - torch.distributed.all_reduce(ref, op=torch.distributed.ReduceOp.SUM) - ref = ref.to(torch_c_dtype) + torch_tensor_flag_unicast = symm_mem.empty( + (num_flags,), dtype=torch.int32, device=device + ) + torch_tensor_flag_unicast.fill_(0) + torch_flag_hdl = symm_mem.rendezvous( + torch_tensor_flag_unicast, group=dist.group.WORLD + ) + torch_tensor_flag_multicast = cutlass_torch.as_tensor( + torch_flag_hdl.multicast_ptr, + torch_tensor_flag_unicast.shape, + torch_tensor_flag_unicast.dtype, + ) + cute_tensor_flag_unicast = from_dlpack(torch_tensor_flag_unicast, assumed_align=16) + cute_tensor_flag_unicast.element_type = cutlass.Int32 + cute_tensor_flag_unicast = cute_tensor_flag_unicast.mark_layout_dynamic() + cute_tensor_flag_multicast = from_dlpack( + torch_tensor_flag_multicast, assumed_align=16 + ) + cute_tensor_flag_multicast.element_type = cutlass.Int32 + cute_tensor_flag_multicast = cute_tensor_flag_multicast.mark_layout_dynamic() - # convert to higher precision to satisfy assert_close function - ref_result = ref.cpu().to(torch_acc_dtype) - kernel_result = c_torch_gpu.view(torch_c_dtype).cpu().to(torch_acc_dtype) - - max_val = torch.finfo(kernel_result.dtype).max - min_val = torch.finfo(kernel_result.dtype).min - kernel_result = torch.nan_to_num(kernel_result, nan=max_val, posinf=max_val, neginf=min_val) - ref_result = torch.nan_to_num(ref_result, nan=max_val, posinf=max_val, neginf=min_val) - - # Assert close results - if dist.get_rank() == 0: - print(f"kernel_result: {kernel_result[:8,:8,0]}") - print(f"ref_result: {ref_result[:8,:8,0]}") - torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05) + return { + "cute_tensor_a_list": cute_tensor_a_list, + "cute_tensor_b_list": cute_tensor_b_list, + "cute_tensor_c": cute_tensor_c, + "cute_tensor_comm_in_mc": cute_tensor_comm_in_mc, + "cute_tensor_comm_out_mc": cute_tensor_comm_out_mc, + "cute_tensor_flag_unicast": cute_tensor_flag_unicast, + "cute_tensor_flag_multicast": cute_tensor_flag_multicast, + "torch_tensor_a_cpu": torch_tensor_a_cpu, + "torch_tensor_b_cpu": torch_tensor_b_cpu, + "torch_tensor_c_cpu": torch_tensor_c_cpu, + "torch_tensor_c_gpu": torch_tensor_c_gpu, + "torch_tensor_comm_out_gpu": torch_tensor_comm_out_gpu, + "torch_tensor_flag_unicast": torch_tensor_flag_unicast, + "_anchors": ( + torch_tensor_a_gpu_list, + torch_tensor_b_gpu_list, + torch_tensor_c_gpu, + torch_tensor_comm_out_gpu, + torch_tensor_flag_unicast, + torch_tensor_flag_multicast, + torch_flag_hdl, + ), + } def run( @@ -2070,273 +2241,866 @@ def run( ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], acc_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", use_2cta_instrs: bool = True, use_tma_store: bool = True, - tolerance: float = 1e-01, warmup_iterations: int = 0, iterations: int = 1, - skip_ref_check: bool = False, - use_cold_l2: bool = False, - all_reduce: str = "two_shot", - **kwargs, + all_reduce: str = "LDMCxSTMC", + benchmark_or_test: str = "test", + use_cuda_graph: bool = False, + csv_report: str = "./csv_report.csv", + dump_input: Optional[str] = None, + dump_output: Optional[str] = None, + load_input: Optional[str] = None, ): - """Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + # * Init Distributed * + global_rank, world_size, device = init_distributed() - This function prepares input tensors, configures and launches the persistent GEMM kernel, - optionally performs reference validation, and benchmarks the execution performance. + # * Validate dump/load flags * + if (dump_input or dump_output or load_input) and benchmark_or_test != "test": + raise ValueError( + "--dump_input / --dump_output / --load_input are only valid in test mode" + ) - :param mnkl: Problem size (M, N, K, L) - :type mnkl: Tuple[int, int, int, int] - :param ab_dtype: Data type for input tensors A and B - :type ab_dtype: Type[cutlass.Numeric] - :param c_dtype: Data type for output tensor C - :type c_dtype: Type[cutlass.Numeric] - :param acc_dtype: Data type for accumulation during matrix multiplication - :type acc_dtype: Type[cutlass.Numeric] - :param a_major/b_major/c_major: Memory layout of tensor A/B/C - :type a_major/b_major/c_major: str - :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the - default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters. - :type mma_tiler_mn: Tuple[int, int], optional - :param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the - default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters. - :type cluster_shape_mn: Tuple[int, int], optional - :param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner - will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. - :type use_2cta_instrs: bool, optional - :param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use - the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. - :type use_tma_store: bool, optional - :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 - :type tolerance: float, optional - :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 - :type warmup_iterations: int, optional - :param iterations: Number of benchmark iterations to run, defaults to 1 - :type iterations: int, optional - :param skip_ref_check: Whether to skip reference result validation, defaults to False - :type skip_ref_check: bool, optional - :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False - :type use_cold_l2: bool, optional - :param all_reduce: All-reduce mode, can be "none", "LDMCxSTMC" - :type all_reduce: str, optional - :raises RuntimeError: If CUDA GPU is not available - :raises ValueError: If the configuration is invalid or unsupported by the kernel - :return: Execution time of the GEMM kernel - :rtype: float - """ - print("Running Blackwell Persistent Dense GEMM AR test with:") - print(f"mnkl: {mnkl}") - print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") - print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") - print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") - print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") - print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") - print(f"Tolerance: {tolerance}") - print(f"Warmup iterations: {warmup_iterations}") - print(f"Iterations: {iterations}") - print(f"Skip reference checking: {skip_ref_check}") - print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") - print(f"Fused AllReduce Op: {all_reduce}") + # * Check based on AR mode * + if all_reduce == "LDMCxSTMC": + if not check_cuda_p2p_access(): + print("CUDA P2P access not available") + exit(-1) + if not check_multicast_support(): + print("Multicast support not available") + exit(-1) + if not test_simple_torch_allreduce(device): + print("Simple torch allreduce test failed") + exit(-1) - # Unpack parameters - m, n, k, l = mnkl + # * Verbose * + def _verbose(): + print("Running Blackwell SM100 Persistent Dense GEMM AR test with:") + print(f"mnkl: {mnkl}") + print(f"global_rank: {global_rank}, world_size: {world_size}, device: {device}") + print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}") + print(f"Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") + print(f"Swizzle Size: {swizzle_size}") + print(f"Raster Order: {raster_order}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Use CUDA Graph: {'True' if use_cuda_graph else 'False'}") + print(f"Mode: {benchmark_or_test}") + print(f"Fused AllReduce Op: {all_reduce}") - if not torch.cuda.is_available(): - raise RuntimeError("GPU is required to run this example!") + if global_rank == 0: + _verbose() - # Get current CUDA stream from PyTorch - torch_stream = torch.cuda.current_stream() - # Get the raw stream pointer as a CUstream - current_stream = cuda.CUstream(torch_stream.cuda_stream) - ( - a_tensor, - b_tensor, - c_tensor, - c_tensor_mc, - barrier_flag, - barrier_flag_mc, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - barrier_flag_torch_cpu, - c_torch_gpu, - c_torch_gpu_mc, - barrier_flag_torch_gpu, - barrier_flag_torch_gpu_mc, - ) = create_tensors( - l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype, mma_tiler_mn, use_2cta_instrs - ) + # * Run Distributed * + if benchmark_or_test in ("benchmark", "benchmark_all"): + run_distributed_benchmark( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + warmup_iterations=warmup_iterations, + iterations=iterations, + all_reduce=all_reduce, + use_cuda_graph=use_cuda_graph, + benchmark_all=(benchmark_or_test == "benchmark_all"), + csv_report=csv_report, + ) + else: + run_distributed_test( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + all_reduce=all_reduce, + dump_input=dump_input, + dump_output=dump_output, + load_input=load_input, + ) - # Build GEMM object - gemm = PersistentDenseGemmKernel( + clean_distributed() + + +def run_distributed_test( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + all_reduce: str = "LDMCxSTMC", + dump_input: Optional[str] = None, + dump_output: Optional[str] = None, + load_input: Optional[str] = None, +): + # * CUDA * + # Non-default stream as we'll use CUDA graphs + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + local_rank = int(os.environ.get("LOCAL_RANK", global_rank)) + hostname = socket.gethostname() + + # * Build GEMM AllReduce Object * + gemm_all_reduce = Sm100PersistentDenseGemmAllReduceLDMCxSTMCKernel( acc_dtype, + c_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store, - dist.get_rank(), - dist.get_world_size(), + rank_id=global_rank, + num_ranks=world_size, all_reduce=all_reduce, + swizzle_size=swizzle_size, + raster_order=raster_order, ) - # Check if configuration can be implemented - can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor) - if not can_implement: - raise ValueError( - f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " - f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " - f"use_tma_store = {use_tma_store}" + # * Set Filename & Funcname for different rank * + # Use hash to keep filename short and avoid "File name too long" error + try: + import hashlib + + param_str = f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}_{'2Sm' if use_2cta_instrs else '1Sm'}_{'EpiTma' if use_tma_store else 'EpiStg'}" + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_ar_{param_hash}_{all_reduce}_r{global_rank}_w{world_size}" ) + except Exception: + unique_name = f"sm100_gemm_ar_{all_reduce}_r{global_rank}_w{world_size}" + gemm_all_reduce.__call__.set_name_prefix(unique_name) + + # * can_implement * + can_implement = gemm_all_reduce.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, all_reduce={all_reduce}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + # * Allocate tensors * + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + num_workspace=1, + device=device, + slot_init_mode="test", + global_rank=global_rank, + local_rank=local_rank, + world_size=world_size, + load_input_dir=load_input, + ) + cute_tensor_a = tensors["cute_tensor_a_list"][0] + cute_tensor_b = tensors["cute_tensor_b_list"][0] + cute_tensor_c = tensors["cute_tensor_c"] + cute_tensor_comm_in_mc = tensors["cute_tensor_comm_in_mc"] + cute_tensor_comm_out_mc = tensors["cute_tensor_comm_out_mc"] + cute_tensor_flag_unicast = tensors["cute_tensor_flag_unicast"] + cute_tensor_flag_multicast = tensors["cute_tensor_flag_multicast"] + torch_tensor_a_cpu = tensors["torch_tensor_a_cpu"] + torch_tensor_b_cpu = tensors["torch_tensor_b_cpu"] + torch_tensor_c = tensors["torch_tensor_c_gpu"] + torch_tensor_comm_out = tensors["torch_tensor_comm_out_gpu"] + torch_tensor_flag_unicast = tensors["torch_tensor_flag_unicast"] + + # * Dump input (before compile/run) * + if dump_input is not None: + _dump_tensors( + dump_input, + local_rank, + global_rank, + world_size, + hostname, + {"a": torch_tensor_a_cpu, "b": torch_tensor_b_cpu}, + ) + _write_metadata_json( + dump_input, + global_rank, + world_size, + mnkl, + ab_dtype, + c_dtype, + acc_dtype, + a_major, + b_major, + c_major, + ) + + # * Compile * + compiled_gemm = cute.compile( + gemm_all_reduce, + a=cute_tensor_a, + b=cute_tensor_b, + c=cute_tensor_c, + comm_in_multicast_tensor=cute_tensor_comm_in_mc, + comm_out_multicast_tensor=cute_tensor_comm_out_mc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + stream=stream, + max_active_clusters=utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ), + ) + + # * Run * + compiled_gemm( + a=cute_tensor_a, + b=cute_tensor_b, + c=cute_tensor_c, + comm_in_multicast_tensor=cute_tensor_comm_in_mc, + comm_out_multicast_tensor=cute_tensor_comm_out_mc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + stream=stream, + ) + + # * Dump output (before compare — _compare_comm mutates C in-place for n-major) * + torch_stream.synchronize() + if dump_output is not None: + _dump_tensors( + dump_output, + local_rank, + global_rank, + world_size, + hostname, + {"c": torch_tensor_c.cpu(), "comm_out": torch_tensor_comm_out.cpu()}, + ) + _write_metadata_json( + dump_output, + global_rank, + world_size, + mnkl, + ab_dtype, + c_dtype, + acc_dtype, + a_major, + b_major, + c_major, + ) + + # * Compare * + # FP8 nan to saturate + def _fp8_nan_to_saturate(tensor_fp8: torch.Tensor) -> torch.Tensor: + if tensor_fp8.dtype == torch.float8_e4m3fn: + # Map to u8 raw tensor for easier byte handle + tensor_u8 = tensor_fp8.view(torch.uint8) + # pos & neg nan + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + # Map pos nan to max, neg nan to min + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 448 + tensor_f32[is_neg_nan] = -448 + # Cast back to f8 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + elif tensor_fp8.dtype == torch.float8_e5m2: + # Map to u8 raw tensor for easier byte handle + tensor_u8 = tensor_fp8.view(torch.uint8) + # pos & neg nan + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + # Map pos nan to max, neg nan to min + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 57344 + tensor_f32[is_neg_nan] = -57344 + # Cast back to f8 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + else: + # Not supported dtype + tensor_f8_out = tensor_fp8 + return tensor_f8_out + + def _compare_gemm( + torch_tensor_a_cpu: torch.Tensor, + torch_tensor_b_cpu: torch.Tensor, + torch_tensor_c: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + # GEMM Ref + gemm_ref_f32 = torch.einsum( + "mkl,nkl->mnl", + torch_tensor_a_cpu.to(dtype=torch.float32), + torch_tensor_b_cpu.to(dtype=torch.float32), + ) + + # For fp8_e4m3, clamp value to [-448, 448] to avoid nan and then cast to fp8_e4m3 + if c_dtype == cutlass.Float8E4M3FN: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-448, max=448) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + # For fp8_e5m2, clamp value to [-57344, 57344] to avoid nan and then cast to fp8_e5m2 + elif c_dtype == cutlass.Float8E5M2: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-57344, max=57344) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + else: + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + + # Cast ref to fp32 for comparison + gemm_ref_f32_c_f32 = gemm_ref_f32_c.to(dtype=torch.float32) + + # Cast kernel output to fp32 for comparison + gemm_out_f32 = torch_tensor_c.cpu().to(dtype=torch.float32) + + # Compare + torch.testing.assert_close( + gemm_out_f32, + gemm_ref_f32_c_f32, + atol=atol, + rtol=rtol, + ) + + def _compare_comm( + torch_tensor_comm_in: torch.Tensor, + torch_tensor_comm_out: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + comm_ref = torch_tensor_comm_in + + # Sync across all rank + dist.barrier() + + # Reference AR — NCCL dist.all_reduce for all supported dtypes. + if c_dtype in [cutlass.Float32, cutlass.BFloat16, cutlass.Float16]: + comm_ref = comm_ref.contiguous() + dist.all_reduce(comm_ref, op=torch.distributed.ReduceOp.SUM) + comm_ref_f32 = comm_ref.to(dtype=torch.float32).cpu() + # For fp8, cast to higher precision, do AR, then cast back to fp8 and handle nan + elif c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_ref_fp32_acc = comm_ref.to(dtype=torch.float32).contiguous() + dist.all_reduce(comm_ref_fp32_acc, op=torch.distributed.ReduceOp.SUM) + # Add small epsilon to mimic multimem.ld_reduce round up behavior where torch cast round down + epsilon = 1e-2 + comm_ref_f32_acc_epsilon = comm_ref_fp32_acc + epsilon + comm_ref_c = comm_ref_f32_acc_epsilon.to(dtype=torch_dtype(c_dtype)) + comm_ref_nan2sat = _fp8_nan_to_saturate(comm_ref_c) + comm_ref_f32 = comm_ref_nan2sat.to(dtype=torch.float32).cpu() + + # Fusion Kernel Result + # multimem.ld_reduce on fp8 will return nan if val is out of range + # Need to handle nan the same way as comm ref + comm_out = torch_tensor_comm_out.cpu() + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + # Handle nan + comm_out_nan2sat = _fp8_nan_to_saturate(comm_out) + # Cast to fp32 for comparison + comm_out_f32 = comm_out_nan2sat.to(dtype=torch.float32) + else: + # Cast to fp32 for comparison + comm_out_f32 = comm_out.to(dtype=torch.float32) + + # Compare + torch.testing.assert_close( + comm_ref_f32, + comm_out_f32, + atol=atol, + rtol=rtol, + ) + + # compare + # Impose a strick check on GEMM as input is uniform random int + gemm_atol = 1e-05 + gemm_rtol = 1e-05 + _compare_gemm( + torch_tensor_a_cpu=torch_tensor_a_cpu, + torch_tensor_b_cpu=torch_tensor_b_cpu, + torch_tensor_c=torch_tensor_c, + c_dtype=c_dtype, + atol=gemm_atol, + rtol=gemm_rtol, + ) + + if all_reduce == "LDMCxSTMC": + # Relaxed tolerance — reference uses NCCL dist.all_reduce, which + # reduces in a different order than the fused multimem kernel. + comm_atol = 5 + comm_rtol = 1e-1 + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_rtol = 1e-1 + comm_atol = 0.0 + _compare_comm( + torch_tensor_comm_in=torch_tensor_c, + torch_tensor_comm_out=torch_tensor_comm_out, + c_dtype=c_dtype, + atol=comm_atol, + rtol=comm_rtol, + ) + + # * Check flags are reset * + # The kernel uses the flag tensor for cross-CTA / cross-rank synchronization + # and is expected to leave every entry at 0 once it returns. If any entry is + # non-zero, a tile didn't drain its slot — re-running the kernel would block + # forever on that stale flag. + torch.cuda.synchronize() + flag_after = torch_tensor_flag_unicast.cpu() + nonzero_count = int(torch.count_nonzero(flag_after).item()) + assert nonzero_count == 0, ( + f"Kernel did not reset all barrier flags to zero: " + f"{nonzero_count}/{flag_after.numel()} entries non-zero " + f"(first 16 values: {flag_after[:16].tolist()})" + ) + + +def compile_and_benchmark( + *, + mnkl, + ab_dtype, + c_dtype, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + swizzle_size, + raster_order, + use_2cta_instrs, + use_tma_store, + all_reduce, + global_rank, + world_size, + num_workspace, + tensors, + torch_stream, + stream, + warmup_iterations, + iterations, + use_cuda_graph, +): + """Instantiate one LDMCxSTMC kernel candidate, compile it, and time the + warmup + benchmark loop. + + Reads cute views out of ``tensors`` (produced by :func:`allocate_tensors`) + and uses the caller-supplied ``torch_stream`` / ``stream`` pair so every + candidate in a sweep shares a single CUDA stream + event pool. + Returns the mean per-iteration time in microseconds over the benchmark + iterations (warmup timings are recorded but not included in the average). + """ + cute_tensor_a_list = tensors["cute_tensor_a_list"] + cute_tensor_b_list = tensors["cute_tensor_b_list"] + cute_tensor_c = tensors["cute_tensor_c"] + cute_tensor_comm_in_mc = tensors["cute_tensor_comm_in_mc"] + cute_tensor_comm_out_mc = tensors["cute_tensor_comm_out_mc"] + cute_tensor_flag_unicast = tensors["cute_tensor_flag_unicast"] + cute_tensor_flag_multicast = tensors["cute_tensor_flag_multicast"] + + # Build kernel + can_implement first — before any cross-rank op, so a + # deterministic skip keeps all ranks in lockstep for the next candidate. + gemm_all_reduce = Sm100PersistentDenseGemmAllReduceLDMCxSTMCKernel( + acc_dtype, + c_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + rank_id=global_rank, + num_ranks=world_size, + all_reduce=all_reduce, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + try: + import hashlib + + param_str = ( + f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}" + f"_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}" + f"_{'2Sm' if use_2cta_instrs else '1Sm'}" + f"_{'EpiTma' if use_tma_store else 'EpiStg'}" + ) + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_ar_{param_hash}_{all_reduce}_r{global_rank}_w{world_size}" + ) + except Exception: + unique_name = f"sm100_gemm_ar_{all_reduce}_r{global_rank}_w{world_size}" + gemm_all_reduce.__call__.set_name_prefix(unique_name) + + can_implement = gemm_all_reduce.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, all_reduce={all_reduce}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + def make_kernel_kwargs(ws): + return dict( + a=cute_tensor_a_list[ws], + b=cute_tensor_b_list[ws], + c=cute_tensor_c, + comm_in_multicast_tensor=cute_tensor_comm_in_mc, + comm_out_multicast_tensor=cute_tensor_comm_out_mc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + stream=stream, + ) + + kernel_kwargs_list = [make_kernel_kwargs(i) for i in range(num_workspace)] + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) - compiled_gemm = cute.compile( - gemm, - a_tensor, - b_tensor, - c_tensor, - c_tensor_mc, - barrier_flag, - barrier_flag_mc, - max_active_clusters, - current_stream, + gemm_all_reduce, + **kernel_kwargs_list[0], + max_active_clusters=max_active_clusters, ) - if not skip_ref_check: - compiled_gemm( - a_tensor, - b_tensor, - c_tensor, - c_tensor_mc, - barrier_flag, - barrier_flag_mc, - current_stream, - ) - compare( - a_torch_cpu, - b_torch_cpu, - c_torch_gpu, - ab_dtype, - c_dtype, - tolerance, - ) + # Per-iteration runnable. Under --use_cuda_graph we capture one graph per + # workspace (pointers are baked in at capture time, so N_workspace graphs + # are required to rotate A/B). Without graphs we dispatch the compiled + # kernel directly with the appropriate kernel_kwargs entry. + if use_cuda_graph: + graphs = [] + for ws in range(num_workspace): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch_stream): + compiled_gemm(**kernel_kwargs_list[ws]) + graphs.append(g) - # always free the multicast tensors first - nvshmem.core.free_tensor(c_torch_gpu_mc) - nvshmem.core.free_tensor(c_torch_gpu) - nvshmem.core.free_tensor(barrier_flag_torch_gpu_mc) - nvshmem.core.free_tensor(barrier_flag_torch_gpu) + def run_one_iter(i): + graphs[i % num_workspace].replay() + else: - free_func_and_tensor_pairs = [] - def add_free_func_and_tensor(free_func, tensor): - free_func_and_tensor_pairs.append((free_func, tensor)) + def run_one_iter(i): + compiled_gemm(**kernel_kwargs_list[i % num_workspace]) - def generate_tensors(): - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc = create_mc_tensor( - c_torch_cpu, c_dtype, (1 if c_major == "n" else 0), is_dynamic_layout=True - ) - barrier_flag, barrier_flag_mc, barrier_flag_torch_gpu, barrier_flag_torch_gpu_mc = create_mc_tensor( - barrier_flag_torch_cpu, cutlass.Int32, 0, is_dynamic_layout=True - ) + # One loop for warmup + benchmark. Every iteration is bracketed by its own + # event pair; only the trailing `iterations` event pairs contribute to the + # reported average so warmup transients don't skew the result. - ja = testing.JitArguments( - a_tensor, - b_tensor, - c_tensor, - c_tensor_mc, - barrier_flag, - barrier_flag_mc, - current_stream, - ) - tensors_to_free = [c_torch_gpu_mc, c_torch_gpu, barrier_flag_torch_gpu_mc, barrier_flag_torch_gpu] - for tensor in tensors_to_free: - add_free_func_and_tensor(nvshmem.core.free_tensor, tensor) - return ja + # Cross-rank sync before timing. torch.cuda.synchronize() drains any + # pending GPU work (graph captures land on the default stream; timing + # runs on torch_stream) so dist.barrier()'s host-side rendezvous is + # not racing in-flight kernels from the previous candidate. Without + # this, the slowest rank's leftover work would skew the measured + # window of every other rank in a benchmark_all sweep. + torch.cuda.synchronize() + dist.barrier() - workspace_count = 1 - if use_cold_l2: - one_workspace_bytes = ( - a_torch_cpu.numel() * a_torch_cpu.element_size() - + b_torch_cpu.numel() * b_torch_cpu.element_size() - + c_torch_cpu.numel() * c_torch_cpu.element_size() - + barrier_flag_torch_cpu.numel() * barrier_flag_torch_cpu.element_size() - ) - workspace_count = testing.get_workspace_count( - one_workspace_bytes, warmup_iterations, iterations - ) + total_iters = warmup_iterations + iterations + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + with torch.cuda.stream(torch_stream): + for i in range(total_iters): + start_events[i].record() + run_one_iter(i) + end_events[i].record() + end_events[-1].synchronize() - exec_time = testing.benchmark( - compiled_gemm, - workspace_generator=generate_tensors, - workspace_count=workspace_count, - stream=current_stream, + per_iter_us = [s.elapsed_time(e) * 1000.0 for s, e in zip(start_events, end_events)] + bench_times = per_iter_us[warmup_iterations:] + avg_us = sum(bench_times) / iterations + print( + f"Rank {global_rank} Execution time: {avg_us} us " + f"(avg of {iterations} iters, after {warmup_iterations} warmup)\n", + flush=True, + ) + return avg_us + + +def run_distributed_benchmark( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + warmup_iterations: int = 0, + iterations: int = 1, + all_reduce: str = "LDMCxSTMC", + use_cuda_graph: bool = False, + benchmark_all: bool = False, + csv_report: str = "./csv_report.csv", +): + # Allocate every kernel-side tensor (A/B/C/comm_out/flags) up front — + # shared across all candidates in the --benchmark_all loop. + num_workspace = 10 + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + num_workspace=num_workspace, + device=device, + slot_init_mode="benchmark", + global_rank=global_rank, + ) + + # One non-default stream shared across every candidate. `torch_stream` is + # kept alive because we capture it into a CUDA graph; creating a fresh + # stream per candidate would churn events + graph context for no reason. + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # Build the candidate list of args for compile_and_benchmark. For a single + # run the list has one entry; --benchmark_all sweeps: + # raster_order ∈ {m, n} + # swizzle_size ∈ {1, 2, 4, 8} + # use_2cta_instrs, cluster_shape_mn, mma_tiler_mn jointly, with the + # constraint that cluster_m ∈ {2,4,8} requires 2cta (smaller tiles and + # single-cta clusters are only valid without 2cta). + base_kwargs = dict( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + use_tma_store=use_tma_store, + all_reduce=all_reduce, + global_rank=global_rank, + world_size=world_size, + num_workspace=num_workspace, + tensors=tensors, + torch_stream=torch_stream, + stream=stream, warmup_iterations=warmup_iterations, iterations=iterations, + use_cuda_graph=use_cuda_graph, ) - if dist.get_rank() == 0: - print(f"exec_time: {exec_time}") - - for free_func, tensor in free_func_and_tensor_pairs: - free_func(tensor) + if benchmark_all: + raster_candidates = ["m", "n"] + swizzle_candidates = [1, 2, 4, 8] + # All cluster shapes from the allowed set; each is filtered per 2cta + # option below by the "cluster_m ∈ {2,4,8} needs 2cta" rule. + cluster_all = [ + (1, 1), + (1, 2), + (2, 1), + (2, 2), + ] + two_cta_options = [ + { + "use_2cta_instrs": True, + "clusters": [c for c in cluster_all if c[0] in (2, 4, 8)], + "mma_tilers": [ + (256, 256), + (256, 192), + (256, 128), + (256, 64), + (128, 256), + (128, 192), + (128, 128), + (128, 64), + ], + }, + { + "use_2cta_instrs": False, + "clusters": cluster_all, + "mma_tilers": [ + (128, 256), + (128, 192), + (128, 128), + (128, 64), + (64, 256), + (64, 192), + (64, 128), + (64, 64), + ], + }, + ] + candidates = [] + for opt in two_cta_options: + for cluster in opt["clusters"]: + for tile in opt["mma_tilers"]: + for r in raster_candidates: + for s in swizzle_candidates: + candidates.append( + { + **base_kwargs, + "use_2cta_instrs": opt["use_2cta_instrs"], + "cluster_shape_mn": cluster, + "mma_tiler_mn": tile, + "raster_order": r, + "swizzle_size": s, + } + ) + else: + candidates = [ + { + **base_kwargs, + "use_2cta_instrs": use_2cta_instrs, + "cluster_shape_mn": cluster_shape_mn, + "mma_tiler_mn": mma_tiler_mn, + "raster_order": raster_order, + "swizzle_size": swizzle_size, + } + ] - return exec_time # Return execution time in microseconds + if global_rank == 0: + print(f"Total number of candidate configs: {len(candidates)}", flush=True) + + # Run each candidate. Tensors (A/B/C/comm_out/flags) are reused across + # candidates; the kernel resets flags internally. Skip any combo the kernel + # can't implement — can_implement is deterministic so all ranks skip the + # same ones, keeping rendezvous collective operations lockstep. + # ~2s sleep between candidates to let the GPU cool down. + results = [] + for i, cand in enumerate(candidates): + if i > 0: + time.sleep(2) + try: + avg_us = compile_and_benchmark(**cand) + except testing.CantImplementError: + continue + results.append((cand, avg_us)) + + # Each rank writes its own CSV — one row per candidate, one config field + # per column. The path comes from --csv_report; per-rank suffix is appended + # before the extension so concurrent ranks don't clobber the same file. + global_rank = dist.get_rank() + m, n, k, l = mnkl + fieldnames = [ + "gemm_m", + "gemm_n", + "gemm_k", + "gemm_l", + "ab_dtype", + "c_dtype", + "acc_dtype", + "a_major", + "b_major", + "c_major", + "mma_tiler_m", + "mma_tiler_n", + "cluster_m", + "cluster_n", + "use_2cta_instrs", + "use_tma_store", + "all_reduce", + "use_cuda_graph", + "raster_order", + "swizzle_size", + "num_workspace", + "warmup_iterations", + "iterations", + "world_size", + "global_rank", + "benchmark_time_us", + ] + csv_path = Path(csv_report) + csv_path = csv_path.with_name( + f"{csv_path.stem}_rank{global_rank}_world{world_size}{csv_path.suffix}" + ) + csv_path.parent.mkdir(parents=True, exist_ok=True) + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for cand, avg_us in results: + writer.writerow( + { + "gemm_m": m, + "gemm_n": n, + "gemm_k": k, + "gemm_l": l, + "ab_dtype": str(ab_dtype), + "c_dtype": str(c_dtype), + "acc_dtype": str(acc_dtype), + "a_major": a_major, + "b_major": b_major, + "c_major": c_major, + "mma_tiler_m": cand["mma_tiler_mn"][0], + "mma_tiler_n": cand["mma_tiler_mn"][1], + "cluster_m": cand["cluster_shape_mn"][0], + "cluster_n": cand["cluster_shape_mn"][1], + "use_2cta_instrs": cand["use_2cta_instrs"], + "use_tma_store": use_tma_store, + "all_reduce": all_reduce, + "use_cuda_graph": use_cuda_graph, + "raster_order": cand["raster_order"], + "swizzle_size": cand["swizzle_size"], + "num_workspace": num_workspace, + "warmup_iterations": warmup_iterations, + "iterations": iterations, + "world_size": world_size, + "global_rank": global_rank, + "benchmark_time_us": avg_us, + } + ) + + return results[0][1] if len(results) == 1 else results -def torchrun_uid_init_bcast(): - """ - Initialize NVSHMEM using UniqueID with `torchrun` as the launcher - - It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting - """ - # Set Torch device - local_rank = int(os.environ['LOCAL_RANK']) - torch.cuda.set_device(local_rank) - - # nvshmem4py requires a cuda.core Device at init time - dev = Device(local_rank) - dev.set_current() - global stream - stream = dev.create_stream() - - # Initialize torch.distributed process group - dist.init_process_group( - backend="cpu:gloo,cuda:nccl", +def prepare_parser(): + parser = argparse.ArgumentParser( + description="Example of SM100 Persistent Dense GEMM AR on Blackwell." ) - # Extract rank, nranks from process group - num_ranks = dist.get_world_size() - - # Create an empty uniqueid for all ranks - uid = nvshmem.core.get_unique_id(empty=(local_rank != 0)) - uid_bytes = uid._data.view(np.uint8).copy() - uid_tensor = torch.from_numpy(uid_bytes).cuda() - dist.broadcast(uid_tensor, src=0) - dist.barrier() - uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype) - - nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid") - - -def torchrun_finalize(): - nvshmem.core.finalize() - dist.destroy_process_group() - - -if __name__ == "__main__": - - def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: try: return tuple(int(x.strip()) for x in s.split(",")) except ValueError: @@ -2344,25 +3108,21 @@ if __name__ == "__main__": "Invalid format. Expected comma-separated integers." ) - parser = argparse.ArgumentParser( - description="Example of Dense Persistent GEMM on Blackwell." - ) - parser.add_argument( "--mnkl", - type=parse_comma_separated_ints, + type=_parse_comma_separated_ints, default=(256, 256, 512, 1), help="mnkl dimensions (comma-separated)", ) parser.add_argument( "--mma_tiler_mn", - type=parse_comma_separated_ints, + type=_parse_comma_separated_ints, default=(128, 128), help="Mma tile shape (comma-separated)", ) parser.add_argument( "--cluster_shape_mn", - type=parse_comma_separated_ints, + type=_parse_comma_separated_ints, default=(1, 1), help="Cluster shape (comma-separated)", ) @@ -2381,7 +3141,13 @@ if __name__ == "__main__": "--use_tma_store", action="store_true", help="Use tma store or not" ) parser.add_argument( - "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + "--benchmark_or_test", + type=str, + choices=["test", "benchmark", "benchmark_all"], + default="test", + help="Select run mode: 'test' for a single-launch correctness check, " + "'benchmark' for cudaEvent timing + per-rank CSV on the given config, " + "'benchmark_all' for the same with a sweep across kernel-config variants.", ) parser.add_argument( "--warmup_iterations", type=int, default=0, help="Warmup iterations" @@ -2393,38 +3159,69 @@ if __name__ == "__main__": help="Number of iterations to run the kernel", ) parser.add_argument( - "--skip_ref_check", action="store_true", help="Skip reference checking" - ) - parser.add_argument( - "--use_cold_l2", + "--use_cuda_graph", action="store_true", default=False, - help="Use circular buffer tensor sets to ensure L2 cold cache", + help="Capture one iteration into a CUDA graph and replay per iteration", ) parser.add_argument( "--all_reduce", - choices=["none", "LDMCxSTMC"], + choices=["LDMCxSTMC", "none"], type=str, - default="two_shot", + default="LDMCxSTMC", help="Allreduce algorithm to fuse with gemm", ) + parser.add_argument( + "--swizzle_size", + type=int, + default=1, + help="Swizzling size in the unit of cluster for improving L2 cache hit rate", + ) + parser.add_argument( + "--raster_order", + type=str, + choices=["m", "n"], + default="m", + help="Rasterization order of clusters", + ) + parser.add_argument( + "--csv_report", + type=str, + default="./csv_report.csv", + help="Path for the per-rank benchmark CSV. Each rank appends " + "_rank{R}_world{W} before the extension to avoid collisions.", + ) + parser.add_argument( + "--dump_input", + type=str, + default=None, + help="(test mode only) Dump A and B tensors to this directory after init.", + ) + parser.add_argument( + "--dump_output", + type=str, + default=None, + help="(test mode only) Dump C and comm_out tensors to this directory after kernel.", + ) + parser.add_argument( + "--load_input", + type=str, + default=None, + help="(test mode only) Load A and B from this directory instead of random init.", + ) + return parser + +if __name__ == "__main__": + parser = prepare_parser() args = parser.parse_args() if len(args.mnkl) != 4: parser.error("--mnkl must contain exactly 4 values") - if len(args.mma_tiler_mn) != 2: parser.error("--mma_tiler_mn must contain exactly 2 values") - if len(args.cluster_shape_mn) != 2: parser.error("--cluster_shape_mn must contain exactly 2 values") - - if args.c_dtype.width == 8: - import warnings - warnings.warn("f8 output is easy to overflow and can be nan here") - - torchrun_uid_init_bcast() run( args.mnkl, @@ -2436,14 +3233,23 @@ if __name__ == "__main__": args.c_major, args.mma_tiler_mn, args.cluster_shape_mn, + args.swizzle_size, + args.raster_order, args.use_2cta_instrs, args.use_tma_store, - args.tolerance, args.warmup_iterations, args.iterations, - args.skip_ref_check, - args.use_cold_l2, - all_reduce=args.all_reduce, + args.all_reduce, + args.benchmark_or_test, + args.use_cuda_graph, + args.csv_report, + dump_input=args.dump_input, + dump_output=args.dump_output, + load_input=args.load_input, + ) + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + print( + f"PASS (hostname={socket.gethostname()}, global_rank={rank}, local_rank={local_rank})", + flush=True, ) - - torchrun_finalize() diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_lamport_blackwell.py b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_lamport_blackwell.py new file mode 100644 index 000000000..665ff4212 --- /dev/null +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_lamport_blackwell.py @@ -0,0 +1,3656 @@ +# Copyright (c) 2025 - 2026 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. + +import os +import argparse +import csv +import socket +import time + +from math import gcd +from pathlib import Path +from typing import List, Optional, Tuple, Type, Union, Literal + +import cuda.bindings.driver as cuda + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +import cutlass.torch as cutlass_torch +from cutlass.torch import dtype as torch_dtype +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cutlass_dsl import BaseDSL +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Pointer, Int32, Float32 +from cutlass.cutlass_dsl import T, dsl_user_op, while_generate, yield_out +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm + +""" +SM100 GEMM + AllReduce (Lamport) Fusion Kernel + +* Test (correctness only — single launch, no warmup/iteration loop) +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_lamport_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs --all_reduce Lamport \ + --benchmark_or_test test + +* Benchmark One +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_lamport_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --use_2cta_instrs --all_reduce Lamport \ + --benchmark_or_test benchmark --use_cuda_graph --csv_report ./csv_report.csv + +* Benchmark All +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_lamport_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --all_reduce Lamport \ + --benchmark_or_test benchmark_all --use_cuda_graph --csv_report ./csv_report.csv + +NOTE: +* Float8E4M3FN / Float8E5M2 are rejected at allocate_tensors — the -0.0 → +0.0 + epilogue scrub is implemented for Float32 / Float16 / BFloat16 only. +* use_tma_store=True is required: multicast TMA store is the only producer + path that fans the gemm-epi store out to peers atomically per 16B. +* C (workspace buffer) is a 3-slot grouped allocation (ping / pong / cooling). + +ALGO: +-> Producer (each rank, per output tile) + * Clear() kernel scrubs slot[pong] back to the per-byte -0.0 sentinel so the next-iter producer has an armed slot to overwrite + * epilogue scrub: acc_vec + (+0.0) — flips any -0.0 lanes in gemm output to +0.0 (IEEE RTNE) so peers don't observe the sentinel as data (for f32/bf16/f16) + * TMA multicast store [data] to slot[ping][local_rank] — NVSwitch fans out to every peer's mirror of slot[ping][local_rank] atomically per 16B + * NO separate flag tensor — data IS the flag + +-> Consumer (each rank, per output tile) + * for each peer r: load with .relaxed .gpu from this rank's UC view of slot[ping][r] (direct P2P, no multimem) + * spin while loaded data == -0.0 sentinel; + * once data arrive, load data, convert to f32, reduce on f32, convert back. + each cta now handle entire output tile, no collaboration with other rank + * store data with .weak .gpu to local rank's global memory +""" + +##################################################################### +#### KERNEL IMPLEMENTATION START HERE #### +##################################################################### + + +@dsl_user_op +def ld_weak_global_4xu32( + ptr: Pointer, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Int32, Int32, Int32, Int32]: + ptr_int = ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + return_struct = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(i32,i32,i32,i32)>"), + [ptr_int], + "ld.weak.global.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return_regs = [llvm.extractvalue(T.i32(), return_struct, [i]) for i in range(4)] + return return_regs[0], return_regs[1], return_regs[2], return_regs[3] + + +@dsl_user_op +def ld_relaxed_gpu_global_4xu32( + ptr: Pointer, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Int32, Int32, Int32, Int32]: + ptr_int = ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + return_struct = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(i32,i32,i32,i32)>"), + [ptr_int], + "ld.relaxed.gpu.global.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return_regs = [llvm.extractvalue(T.i32(), return_struct, [i]) for i in range(4)] + return return_regs[0], return_regs[1], return_regs[2], return_regs[3] + + +@dsl_user_op +def st_weak_global_4xu32( + ptr: Pointer, + x: Int32, + y: Int32, + z: Int32, + w: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + ptr_int = ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + x_ir = Int32(x).ir_value(loc=loc, ip=ip) + y_ir = Int32(y).ir_value(loc=loc, ip=ip) + z_ir = Int32(z).ir_value(loc=loc, ip=ip) + w_ir = Int32(w).ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ptr_int, x_ir, y_ir, z_ir, w_ir], + "st.weak.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_bf16x2_to_2xf32( + bf16x2_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32)>"), + [bf16x2_b32], + "{ .reg .b16 hl, hh;\n" + " mov.b32 {hl, hh}, $2;\n" + " cvt.f32.bf16 $0, hl;\n" + " cvt.f32.bf16 $1, hh; }", + "=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + f_lo = Float32(llvm.extractvalue(T.f32(), ret, [0])) + f_hi = Float32(llvm.extractvalue(T.f32(), ret, [1])) + return f_lo, f_hi + + +@dsl_user_op +def cvt_f16x2_to_2xf32( + f16x2_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32)>"), + [f16x2_b32], + "{ .reg .b16 hl, hh;\n" + " mov.b32 {hl, hh}, $2;\n" + " cvt.f32.f16 $0, hl;\n" + " cvt.f32.f16 $1, hh; }", + "=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + f_lo = Float32(llvm.extractvalue(T.f32(), ret, [0])) + f_hi = Float32(llvm.extractvalue(T.f32(), ret, [1])) + return f_lo, f_hi + + +@dsl_user_op +def cvt_2xf32_to_bf16x2( + f_lo: Float32, + f_hi: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f_lo).ir_value(loc=loc, ip=ip), + Float32(f_hi).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_2xf32_to_f16x2( + f_lo: Float32, + f_hi: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f_lo).ir_value(loc=loc, ip=ip), + Float32(f_hi).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.f16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def add_rn_f32( + a: Float32, + b: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def add_rn_f32x2_pair( + a_lo: Float32, + a_hi: Float32, + b_lo: Float32, + b_hi: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32)>"), + [ + Float32(a_lo).ir_value(loc=loc, ip=ip), + Float32(a_hi).ir_value(loc=loc, ip=ip), + Float32(b_lo).ir_value(loc=loc, ip=ip), + Float32(b_hi).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b64 ab, bb, db;\n" + " mov.b64 ab, {$2, $3};\n" + " mov.b64 bb, {$4, $5};\n" + " add.rn.f32x2 db, ab, bb;\n" + " mov.b64 {$0, $1}, db; }", + "=f,=f,f,f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + sum_lo = Float32(llvm.extractvalue(T.f32(), ret, [0])) + sum_hi = Float32(llvm.extractvalue(T.f32(), ret, [1])) + return sum_lo, sum_hi + + +@dsl_user_op +def cvt_e4m3x4_to_4xf32( + e4m3x4_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32, Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [e4m3x4_b32], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " .reg .b32 f16x2_lo, f16x2_hi;\n" + " .reg .b16 h0, h1, h2, h3;\n" + " mov.b32 {e8x2_lo, e8x2_hi}, $4;\n" + " cvt.rn.f16x2.e4m3x2 f16x2_lo, e8x2_lo;\n" + " cvt.rn.f16x2.e4m3x2 f16x2_hi, e8x2_hi;\n" + " mov.b32 {h0, h1}, f16x2_lo;\n" + " mov.b32 {h2, h3}, f16x2_hi;\n" + " cvt.f32.f16 $0, h0;\n" + " cvt.f32.f16 $1, h1;\n" + " cvt.f32.f16 $2, h2;\n" + " cvt.f32.f16 $3, h3; }", + "=f,=f,=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return tuple(Float32(llvm.extractvalue(T.f32(), ret, [i])) for i in range(4)) + + +@dsl_user_op +def cvt_e5m2x4_to_4xf32( + e5m2x4_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32, Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [e5m2x4_b32], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " .reg .b32 f16x2_lo, f16x2_hi;\n" + " .reg .b16 h0, h1, h2, h3;\n" + " mov.b32 {e8x2_lo, e8x2_hi}, $4;\n" + " cvt.rn.f16x2.e5m2x2 f16x2_lo, e8x2_lo;\n" + " cvt.rn.f16x2.e5m2x2 f16x2_hi, e8x2_hi;\n" + " mov.b32 {h0, h1}, f16x2_lo;\n" + " mov.b32 {h2, h3}, f16x2_hi;\n" + " cvt.f32.f16 $0, h0;\n" + " cvt.f32.f16 $1, h1;\n" + " cvt.f32.f16 $2, h2;\n" + " cvt.f32.f16 $3, h3; }", + "=f,=f,=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return tuple(Float32(llvm.extractvalue(T.f32(), ret, [i])) for i in range(4)) + + +@dsl_user_op +def cvt_4xf32_to_e4m3x4( + f0: Float32, + f1: Float32, + f2: Float32, + f3: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f0).ir_value(loc=loc, ip=ip), + Float32(f1).ir_value(loc=loc, ip=ip), + Float32(f2).ir_value(loc=loc, ip=ip), + Float32(f3).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " cvt.rn.satfinite.e4m3x2.f32 e8x2_lo, $2, $1;\n" + " cvt.rn.satfinite.e4m3x2.f32 e8x2_hi, $4, $3;\n" + " mov.b32 $0, {e8x2_lo, e8x2_hi}; }", + "=r,f,f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_4xf32_to_e5m2x4( + f0: Float32, + f1: Float32, + f2: Float32, + f3: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f0).ir_value(loc=loc, ip=ip), + Float32(f1).ir_value(loc=loc, ip=ip), + Float32(f2).ir_value(loc=loc, ip=ip), + Float32(f3).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " cvt.rn.satfinite.e5m2x2.f32 e8x2_lo, $2, $1;\n" + " cvt.rn.satfinite.e5m2x2.f32 e8x2_hi, $4, $3;\n" + " mov.b32 $0, {e8x2_lo, e8x2_hi}; }", + "=r,f,f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_b32_to_1xf32( + b32_in: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32]: + return (Float32(llvm.bitcast(T.f32(), b32_in, loc=loc, ip=ip)),) + + +@dsl_user_op +def cvt_1xf32_to_b32( + f: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.bitcast(T.i32(), Float32(f).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + + +@dsl_user_op +def elem_wise_add_f32( + a: Tuple[Float32, ...], + b: Tuple[Float32, ...], + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, ...]: + assert len(a) == len(b), f"Mismatched tuple sizes: {len(a)} vs {len(b)}" + n = len(a) + if n == 1: + return (add_rn_f32(a[0], b[0], loc=loc, ip=ip),) + assert n % 2 == 0, f"elem_wise_add_f32: tuple length {n} must be 1 or even" + # Pair consecutive f32 values into FADD2 invocations. + out: Tuple[Float32, ...] = () + for i in range(0, n, 2): + s_lo, s_hi = add_rn_f32x2_pair(a[i], a[i + 1], b[i], b[i + 1], loc=loc, ip=ip) + out = out + (s_lo, s_hi) + return out + + +@cute.jit +def Lamport( + kernel, + cur_tile_coord, + tiled_mma: cute.TiledMma, + comm_tidx: cutlass.Int32, + tCgCommInUC_per_peer: Tuple[cute.Tensor, ...], + tCgCommOutUC: cute.Tensor, +) -> None: + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Per-peer UC views share shape/strides with the local UC tensor (only the + # base pointer differs), so one coord pattern slices every peer the same way. + tCgCommInUC_slice_per_peer = tuple( + t[((None, None), 0, 0, *mma_tile_coord_mnl)] for t in tCgCommInUC_per_peer + ) + tCgCommOutUC_slice = tCgCommOutUC[((None, None), 0, 0, *mma_tile_coord_mnl)] + + cta_mma_tile_m = kernel.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + cta_tile_n = kernel.mma_tiler[1] + total_comm_threads = len(kernel.comm_warp_id) * cute.arch.WARP_SIZE + + local_tile_m = cta_mma_tile_m + local_tile_n = cta_tile_n + + tCgCommInUC_local_rank_per_peer = tCgCommInUC_slice_per_peer + tCgCommOutUC_local_rank = tCgCommOutUC_slice + + elements_per_thread = (local_tile_m * local_tile_n) // total_comm_threads + atom_num_elements_128 = 128 // kernel.c_dtype.width + atom_num_elements = atom_num_elements_128 + + # 128-bit ld/st is mandatory; the cvt wrappers assume 32-bit packed atoms. + # _pick_num_comm_warp_for_128b ensures divisibility; this assert fires at + # @cute.jit time if that invariant ever breaks (a `raise` here would be + # rejected by the AST preprocessor). + assert cutlass.const_expr(elements_per_thread % atom_num_elements_128 == 0), ( + f"[Lamport] 128-bit instruction required but not viable: " + f"elements_per_thread={elements_per_thread} is not divisible by " + f"atom_num_elements_128={atom_num_elements_128}. " + f"This should never happen if _pick_num_comm_warp_for_128b is correct." + ) + + # Row-major: threads contiguous along N; col-major: along M. GCD handles + # non-pow2 widths (e.g. cta_n=192) that don't naively divide thread count. + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + max_thr_n = local_tile_n // atom_num_elements + atom_thr_n = gcd(max_thr_n, total_comm_threads) + atom_thr_m = total_comm_threads // atom_thr_n + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) + val_layout = cute.make_layout( + (1, atom_num_elements), stride=(atom_num_elements, 1) + ) + else: + max_thr_m = local_tile_m // atom_num_elements + atom_thr_m = gcd(max_thr_m, total_comm_threads) + atom_thr_n = total_comm_threads // atom_thr_m + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(1, atom_thr_m)) + val_layout = cute.make_layout( + (atom_num_elements, 1), stride=(1, atom_num_elements) + ) + + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), kernel.c_dtype) + tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + thr_copy_fake = tiled_copy_fake.get_slice(comm_tidx - kernel.comm_warp_id[0] * 32) + + tTR_gCommInUC_per_peer = tuple( + thr_copy_fake.partition_S(t) for t in tCgCommInUC_local_rank_per_peer + ) + tTR_gCommOutUC = thr_copy_fake.partition_S(tCgCommOutUC_local_rank) + _atom, loop_m, loop_n = tTR_gCommOutUC.shape + + # Dtype dispatch — cvt_to_xf32: i32 → n_per_b32 f32; cvt_back: f32… → i32. + # neg_zero_dtype_32bit packs the dtype-specific -0.0 across all lanes of + # one i32 so the spin's break test is a single setp.eq, no bit-extract. + if cutlass.const_expr(kernel.c_dtype is cutlass.BFloat16): + cvt_to_xf32 = cvt_bf16x2_to_2xf32 + cvt_back = cvt_2xf32_to_bf16x2 + n_per_b32 = 2 + neg_zero_dtype_32bit = 0x80008000 # bf16 -0.0 × 2 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float16): + cvt_to_xf32 = cvt_f16x2_to_2xf32 + cvt_back = cvt_2xf32_to_f16x2 + n_per_b32 = 2 + neg_zero_dtype_32bit = 0x80008000 # f16 -0.0 × 2 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float8E4M3FN): + cvt_to_xf32 = cvt_e4m3x4_to_4xf32 + cvt_back = cvt_4xf32_to_e4m3x4 + n_per_b32 = 4 + neg_zero_dtype_32bit = 0x80808080 # e4m3 -0.0 × 4 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float8E5M2): + cvt_to_xf32 = cvt_e5m2x4_to_4xf32 + cvt_back = cvt_4xf32_to_e5m2x4 + n_per_b32 = 4 + neg_zero_dtype_32bit = 0x80808080 # e5m2 -0.0 × 4 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float32): + cvt_to_xf32 = cvt_b32_to_1xf32 + cvt_back = cvt_1xf32_to_b32 + n_per_b32 = 1 + neg_zero_dtype_32bit = 0x80000000 # f32 -0.0 × 1 lane + else: + raise ValueError( + f"Lamport: unsupported c_dtype {kernel.c_dtype} " + "(supported: BFloat16, Float16, Float8E4M3FN, Float8E5M2, Float32)" + ) + + # Accumulator stays in f32 across the full peer chain — one cvt_back at + # the end preserves the same precision a multimem.ld_reduce path would. + # Flat-tuple length n_atom * n_per_b32 is always even (4/8/16), so + # elem_wise_add_f32 always lands on FADD2. + num_peers = len(tTR_gCommInUC_per_peer) + n_atom = 4 # 128b / 32b per i32 + + # `unroll=1` on the outer (i, j) lattice forces the CuTe preprocessor to + # emit ONE shared body. Full unroll explodes register pressure + # (num_peers × n_atom × cvt-tuple f32s live per (i, j)) and the allocator + # spills to local memory. Inner peer / atom loops stay range_constexpr + # so their unrolling helps instruction scheduling within one (i, j). + for i in cutlass.range(loop_m, unroll=1): + for j in cutlass.range(loop_n, unroll=1): + # (1) Per-peer load with flag-via-data spin: ld.relaxed.gpu on the + # peer's UC view, break when g3 (last i32 of the 16B chunk) leaves + # the sentinel. g3-only is sound because the producer is a TMA + # multicast store, which lands at peer mirrors atomically per 16B + # — no torn writes where g3 could update before g0/g1/g2. Random + # GEMM bits hitting the sentinel are ~2^-32 / lane and the + # producer-side scrub (epilogue_tma_store_handle_neg_zero) flips + # -0.0 → +0.0 to remove the only systematic source. If a non-TMA + # producer is ever wired in, widen the predicate. + peer_loaded = [] + for r in cutlass.range_constexpr(num_peers): + ptr = tTR_gCommInUC_per_peer[r][None, i, j].iterator + init_g0, init_g1, init_g2, init_g3 = ld_relaxed_gpu_global_4xu32(ptr) + spin = while_generate( + [init_g0, init_g1, init_g2, init_g3], + lambda _g0, _g1, _g2, g3: g3 == neg_zero_dtype_32bit, + ) + with spin: + new_regs = ld_relaxed_gpu_global_4xu32(ptr) + yield_out(list(new_regs)) + peer_loaded.append(tuple(spin.results)) + + # (2) build flat f32 tuple from each peer; init accumulator from + # peer 0, then vectorized-add the rest. + acc = () + for r in cutlass.range_constexpr(num_peers): + p_acc = () + for k in cutlass.range_constexpr(n_atom): + p_acc = p_acc + cvt_to_xf32(peer_loaded[r][k]) + acc = p_acc if r == 0 else elem_wise_add_f32(acc, p_acc) + + # (3) cvt accumulator back to source dtype, once per atom. + out_regs = [] + for k in cutlass.range_constexpr(n_atom): + chunk = acc[k * n_per_b32 : (k + 1) * n_per_b32] + out_regs.append(cvt_back(*chunk)) + + # (4) plain st.weak.global.v4.u32 to this rank's local UC slot. + out_uc_ptr = tTR_gCommOutUC[None, i, j].iterator + st_weak_global_4xu32(out_uc_ptr, *out_regs) + + +@cute.jit +def Clear( + kernel, + cur_tile_coord, + tiled_mma: cute.TiledMma, + comm_tidx: cutlass.Int32, + tCgCommClearUC_per_peer: Tuple[cute.Tensor, ...], +) -> None: + # Per-tile scrub of slot[pong] back to the dtype-specific -0.0 sentinel. + + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tCgCommClearUC_slice_per_peer = tuple( + t[((None, None), 0, 0, *mma_tile_coord_mnl)] for t in tCgCommClearUC_per_peer + ) + + cta_mma_tile_m = kernel.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + cta_tile_n = kernel.mma_tiler[1] + # Partitioned over EPILOGUE warps (vs Lamport's comm warps). + total_threads = len(kernel.epilogue_warp_id) * cute.arch.WARP_SIZE + + local_tile_m = cta_mma_tile_m + local_tile_n = cta_tile_n + + elements_per_thread = (local_tile_m * local_tile_n) // total_threads + atom_num_elements_128 = 128 // kernel.c_dtype.width + atom_num_elements = atom_num_elements_128 + + assert cutlass.const_expr(elements_per_thread % atom_num_elements_128 == 0), ( + f"[Clear] 128-bit instruction required but not viable: " + f"elements_per_thread={elements_per_thread} is not divisible by " + f"atom_num_elements_128={atom_num_elements_128}." + ) + + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + max_thr_n = local_tile_n // atom_num_elements + atom_thr_n = gcd(max_thr_n, total_threads) + atom_thr_m = total_threads // atom_thr_n + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) + val_layout = cute.make_layout( + (1, atom_num_elements), stride=(atom_num_elements, 1) + ) + else: + max_thr_m = local_tile_m // atom_num_elements + atom_thr_m = gcd(max_thr_m, total_threads) + atom_thr_n = total_threads // atom_thr_m + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(1, atom_thr_m)) + val_layout = cute.make_layout( + (atom_num_elements, 1), stride=(1, atom_num_elements) + ) + + copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), kernel.c_dtype) + tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout) + thr_copy_fake = tiled_copy_fake.get_slice( + comm_tidx - kernel.epilogue_warp_id[0] * 32 + ) + + tTR_gCommClearUC_per_peer = tuple( + thr_copy_fake.partition_S(t) for t in tCgCommClearUC_slice_per_peer + ) + _atom, loop_m, loop_n = tTR_gCommClearUC_per_peer[0].shape + + # Same dtype → -0.0 byte pattern as Lamport's neg_zero_dtype_32bit table. + if cutlass.const_expr(kernel.c_dtype is cutlass.BFloat16): + neg_zero_dtype_32bit = 0x80008000 # bf16 -0.0 × 2 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float16): + neg_zero_dtype_32bit = 0x80008000 # f16 -0.0 × 2 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float8E4M3FN): + neg_zero_dtype_32bit = 0x80808080 # e4m3 -0.0 × 4 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float8E5M2): + neg_zero_dtype_32bit = 0x80808080 # e5m2 -0.0 × 4 lanes + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float32): + neg_zero_dtype_32bit = 0x80000000 # f32 -0.0 × 1 lane + else: + raise ValueError( + f"Clear: unsupported c_dtype {kernel.c_dtype} " + "(supported: BFloat16, Float16, Float8E4M3FN, Float8E5M2, Float32)" + ) + + num_peers = len(tTR_gCommClearUC_per_peer) + + # `unroll=1` outer (same reason as Lamport); inner peer loop is short + # const_expr so it unrolls into back-to-back STG.E.128. + for i in cutlass.range(loop_m, unroll=1): + for j in cutlass.range(loop_n, unroll=1): + for r in cutlass.range_constexpr(num_peers): + ptr = tTR_gCommClearUC_per_peer[r][None, i, j].iterator + st_weak_global_4xu32( + ptr, + neg_zero_dtype_32bit, + neg_zero_dtype_32bit, + neg_zero_dtype_32bit, + neg_zero_dtype_32bit, + ) + + +@cute.jit +def epilogue_tma_store_handle_neg_zero( + gemm_kernel, + epi_tidx: cutlass.Int32, + warp_idx: cutlass.Int32, + tma_atom_c: cute.CopyAtom, + # Input of epilogue + tCtAcc_base: cute.Tensor, + # Staging of epilogue + sC: cute.Tensor, + # Output of epilogue + tCgC_base: cute.Tensor, + epi_tile: cute.Tile, + num_tiles_executed: cutlass.Int32, + epilogue_op: cutlass.Constexpr, + mma_tile_coord_mnl: Tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32], + acc_consumer_state: pipeline.PipelineState, + acc_pipeline: pipeline.PipelineAsync, + c_pipeline: pipeline.PipelineTmaStore, +) -> pipeline.PipelineState: + # Fork of cutlass.utils.gemm.sm100.epilogue_tma_store with the Lamport + # -0.0 → +0.0 scrub spliced in below; otherwise behaviour-identical to + # the upstream helper. + + # Layout transformation for tCgC_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, TILE_M, TILE_N, TILE_K) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), TILE_M, TILE_N, TILE_K) + tCgC = utils.gemm.sm100.transform_partitioned_tensor_layout(tCgC_base) + + # Layout transformation for tCtAcc_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, STAGE) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), STAGE) + tCtAcc = utils.gemm.sm100.transform_partitioned_tensor_layout(tCtAcc_base) + + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = ( + utils.gemm.sm100.epilogue_tmem_copy_and_partition( + gemm_kernel, epi_tidx, tCtAcc, tCgC, epi_tile, gemm_kernel.use_2cta_instrs + ) + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = utils.gemm.sm100.epilogue_smem_copy_and_partition( + gemm_kernel, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide(tCgC, epi_tile) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=gemm_kernel.epilog_sync_bar_id, + num_threads=32 * len(gemm_kernel.epilogue_warp_id), + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + 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) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + 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]) # type: ignore[union-attr] + num_prev_subtiles = num_tiles_executed * subtile_cnt + for subtile_idx in range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] # type: ignore[call-overload] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) # type: ignore[operator] + + # Lamport scrub: peers spin on -0.0 as the "not arrived" sentinel, + # so a legitimate -0.0 in GEMM output must not reach C. IEEE RTNE + # gives `-0.0 + +0.0 = +0.0` and is bit-exact pass-through for every + # other value, so a +0 lane-tensor add flips only the sentinel + # lanes. Not constant-folded: cute doesn't enable nsz globally. + if cutlass.const_expr(gemm_kernel.all_reduce == "Lamport"): + acc_vec = acc_vec + cute.full_like(acc_vec, 0.0) + + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % gemm_kernel.num_c_stage + 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( + "async.shared", + space="cta", + ) + epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + return acc_consumer_state + + +class Sm100PersistentDenseGemmAllReduceLamportKernel: + """ + **Example:** + gemm = Sm100PersistentDenseGemmAllReduceLamportKernel( + acc_dtype=cutlass.Float32, + c_dtype=cutlass.BFloat16, + use_2cta_instrs=True, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + use_tma_store=True, + rank_id=rank_id, + num_ranks=num_ranks, + all_reduce="Lamport", + swizzle_size=1, + raster_order="m", + ) + compiled_gemm = cute.compile( + gemm, + a=a, b=b, + c_multicast_tensor=c_mc_slot_ping_local_rank, + comm_in_unicast_tensor_per_peer=c_uc_per_peer_slot_ping, + comm_clear_unicast_tensor_per_peer=c_uc_per_peer_slot_pong, + comm_out_unicast_tensor=comm_out_uc, + stream=stream, + max_active_clusters=max_active_clusters, + ) + compiled_gemm( + a=a, b=b, + c_multicast_tensor=c_mc_slot_ping_local_rank, + comm_in_unicast_tensor_per_peer=c_uc_per_peer_slot_ping, + comm_clear_unicast_tensor_per_peer=c_uc_per_peer_slot_pong, + comm_out_unicast_tensor=comm_out_uc, + stream=stream, + ) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_tma_store: bool, + rank_id: int, + num_ranks: int, + all_reduce: str = "Lamport", + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + ): + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.c_dtype: Type[cutlass.Numeric] = c_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.swizzle_size = swizzle_size + self.raster_order = raster_order + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + # Capture the JIT target arch so SMEM/TMEM sizing tracks the actual + # hardware (sm_100 / sm_103 / sm_107 / ...). Matches the FMHA pattern. + arch_enum = BaseDSL._get_dsl().get_arch_enum() + self.arch = f"sm_{arch_enum.major}{arch_enum.minor}" + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.all_reduce = all_reduce + if all_reduce == "Lamport": + self.rank_id = rank_id + self.num_ranks = num_ranks + num_comm_warps = self._pick_num_comm_warp_for_128b(c_dtype) + self.comm_warp_id = tuple(range(6, 6 + num_comm_warps)) + else: + # Other AR mode not support by now. + self.comm_warp_id = () + self.rank_id = 0 + self.num_ranks = 1 + + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilogue_warp_id, + *self.comm_warp_id, + ) + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.comm_sync_bar_id = 4 + + def _pick_num_comm_warp_for_128b(self, c_dtype): + """Pick the LARGEST comm-warp count W ∈ {4, 3, 2, 1} (T = W·32 ∈ + {128, 96, 64, 32}) such that the AR slab is evenly distributed + across T threads AND every thread gets a full 128 b LD/ST payload + (= ``atom_128`` elements, where ``atom_128 = 16 / sizeof(c_dtype)``). + + Iterates W = 4 → 1 and returns the first match, so the largest + viable W is returned. Smaller W is only chosen when strictly + necessary, because larger W means more parallel AR threads and + proportionally fewer per-thread PTX instructions. + """ + # cta-tile dims (the per-CTA M·N output the AR walks). 2cta halves M. + # Lamport: each rank handles the entire CTA tile (no // num_ranks split), + # so the slab a single rank's comm warps must cover is the full M·N. + cta_mma_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + slab_elements = cta_mma_tile_m * cta_tile_n + atom_128 = 16 // (c_dtype.width // 8) + for num_warps in (4, 3, 2, 1): + T = num_warps * 32 + if slab_elements % T: + continue + if (slab_elements // T) % atom_128 == 0: + return num_warps + raise ValueError( + f"no comm-warp count enables 128 b PTX for c_dtype={c_dtype}, " + f"cta=({cta_mma_tile_m},{cta_tile_n}), P={self.num_ranks}" + ) + + 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 + """ + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + 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.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + 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 + + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + c_smem_layout = None + if cutlass.const_expr(self.use_tma_store): + c_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + c_smem_layout, + ) + + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage, self.arch + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c_multicast_tensor: cute.Tensor, + comm_in_unicast_tensor_per_peer: Tuple[cute.Tensor, ...], + comm_clear_unicast_tensor_per_peer: Tuple[cute.Tensor, ...], + comm_out_unicast_tensor: cute.Tensor, + stream: cuda.CUstream, + max_active_clusters: cutlass.Constexpr, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # Setup static attributes before smem/grid/tma computation. + # NOTE: self.c_dtype was set in __init__ from the explicit c_dtype + # arg. We only sanity-check here that the C tensor matches that hint. + # Lamport mandates use_tma_store=True + if cutlass.const_expr(self.all_reduce == "Lamport"): + assert self.use_tma_store, ( + "Lamport requires use_tma_store=True (the C tensor is now a " + "multicast view; STG on a multicast address does not fan out)." + ) + # f8 is temporarily rejected + assert self.c_dtype not in ( + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ), ( + f"Lamport currently does not support c_dtype={self.c_dtype}: the " + "epilogue -0.0 → +0.0 scrub is implemented for Float32/Float16/" + "BFloat16 only. The f8 path will land once the bitwise haszero " + "scrub is verified." + ) + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + if cutlass.const_expr(c_multicast_tensor.element_type != self.c_dtype): + raise TypeError( + f"C tensor element_type ({c_multicast_tensor.element_type}) does not match " + f"the c_dtype declared at __init__ ({self.c_dtype})" + ) + 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.c_layout = utils.LayoutEnum.from_tensor(c_multicast_tensor) + + # 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}") + + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.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 = utils.sm100.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 C + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c_multicast_tensor, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c_multicast_tensor, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + self.swizzle_size, + self.raster_order, + max_active_clusters, + ) + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma=tiled_mma, + tma_atom_a=tma_atom_a, + mA_mkl=tma_tensor_a, + tma_atom_b=tma_atom_b, + mB_nkl=tma_tensor_b, + tma_atom_c=tma_atom_c, + mC_mc_mnl=tma_tensor_c, + mCommInUC_mnl_per_peer=comm_in_unicast_tensor_per_peer, + mCommClearUC_mnl_per_peer=comm_clear_unicast_tensor_per_peer, + mCommOutUC_mnl=comm_out_unicast_tensor, + cluster_layout_vmnk=self.cluster_layout_vmnk, + a_smem_layout_staged=self.a_smem_layout_staged, + b_smem_layout_staged=self.b_smem_layout_staged, + c_smem_layout_staged=self.c_smem_layout_staged, + epi_tile=self.epi_tile, + tile_sched_params=self.tile_sched_params, + epilogue_op=epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + 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_c: Optional[cute.CopyAtom], + mC_mc_mnl: cute.Tensor, + mCommInUC_mnl_per_peer: Tuple[cute.Tensor, ...], + mCommClearUC_mnl_per_peer: Tuple[cute.Tensor, ...], + mCommOutUC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + 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) + if cutlass.const_expr(self.use_tma_store): + 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) + + # 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_producer, ab_consumer = 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, + defer_sync=True, + ).make_participants() + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_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, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + 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) + ) + # (bM, bN, RestM, RestN, RestL). For Lamport this is the multicast view + # of THIS rank's C_tuple[rank_id]; the TMA store on it broadcasts to + # all peers' C_tuple[rank_id].uc. + gC_mnl = cute.local_tile( + mC_mc_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + k_tile_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) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # 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 + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + 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_tile = prefetch_k_tile_cnt + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_producer.tail() + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + 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_tile = 0 + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm( + tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc + ) + + # Async arrive AB buffer empty + handle.release() + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # + # 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) + + sC = None + if cutlass.const_expr(self.use_tma_store): + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + m_cta_tiles_in_total = gC_mnl.shape[2] * cute.size(tiled_mma.thr_id.shape) + n_cta_tiles_in_total = gC_mnl.shape[3] + + # + # Persistent tile scheduling loop for epilogue + # + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + + # Build the per-peer Clear partition + if cutlass.const_expr(self.all_reduce == "Lamport"): + gCommClearUC_mnl_per_peer = tuple( + cute.local_tile( + t, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + for t in mCommClearUC_mnl_per_peer + ) + tCgCommClearUC_per_peer = tuple( + thr_mma.partition_C(g) for g in gCommClearUC_mnl_per_peer + ) + + 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], + ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + + # Per-tile scrub of slot[pong]'s `world_size` local UC slabs + # back to the -0.0 sentinel, in preparation for the next-iter + # ping write. Runs in the epilogue warp BEFORE the TMA store + if cutlass.const_expr(self.all_reduce == "Lamport"): + Clear( + self, + cur_tile_coord=cur_tile_coord, + tiled_mma=tiled_mma, + comm_tidx=tidx, + tCgCommClearUC_per_peer=tCgCommClearUC_per_peer, + ) + + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = epilogue_tma_store_handle_neg_zero( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + # Lamport: no per-tile flag release. The data store itself is + # the flag — peers spin on the per-i32 -0.0 sentinel inside + # their AR warp and observe arrival when the broadcast lands. + + if cutlass.const_expr(self.use_tma_store): + # Drain all in-flight TMA stores once at end of the loop. For + # Lamport this is the only producer-side wait: peers detect + # individual tile arrival via data polling, not via this tail. + c_pipeline.producer_tail() + else: + # Synchronize before TMEM dealloc + if cutlass.const_expr(self.all_reduce != "Lamport"): + tmem_dealloc_barrier.arrive_and_wait() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # AllReduce warps + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.all_reduce == "Lamport"): + if warp_idx >= self.comm_warp_id[0]: + # Per peer: (bM, bN, RestM, RestN, RestL) with that peer's local + # UC view (direct P2P address). Same tile bounds as the MC view — + # only the base pointer differs across the tuple. Used by Lamport + # to issue per-peer ld.weak.global.v4.u32 loads. + gCommInUC_mnl_per_peer = tuple( + cute.local_tile( + t, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + for t in mCommInUC_mnl_per_peer + ) + # (bM, bN, RestM, RestN, RestL) with local UC memory as comm output + gCommOutUC_mnl = cute.local_tile( + mCommOutUC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + # Per-peer thread-MMA partitioning of the input UC views. + tCgCommInUC_per_peer = tuple( + thr_mma.partition_C(g) for g in gCommInUC_mnl_per_peer + ) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) with local UC memory + tCgCommOutUC = thr_mma.partition_C(gCommOutUC_mnl) + + # + # Persistent tile scheduling loop for all reduce + # + + # AllReduce communication pattern (w/ rank4 as example): + # + # [rank0 A] [rank1 A] [rank2 A] [rank3 A] -> RS -> [rank0+1+2+3 A] [rank0+1+2+3 A] [rank0+1+2+3 A] [rank0+1+2+3 A] + # [rank0 B] [rank1 B] [rank2 B] [rank3 B] -> RS -> [rank0+1+2+3 B] [rank0+1+2+3 B] [rank0+1+2+3 B] [rank0+1+2+3 B] + # [rank0 C] [rank1 C] [rank2 C] [rank3 C] -> RS -> [rank0+1+2+3 C] [rank0+1+2+3 C] [rank0+1+2+3 C] [rank0+1+2+3 C] + # [rank0 D] [rank1 D] [rank2 D] [rank3 D] -> RS -> [rank0+1+2+3 D] [rank0+1+2+3 D] [rank0+1+2+3 D] [rank0+1+2+3 D] + # rank0 rank1 rank2 rank3 -> RS -> rank0 rank1 rank2 rank3 + # + # Each rank's AR warp only handle 1/P of overall output tiles, and relies on other rank to broadcast data to it + # rank0: compute each tile's [0:1/4] and broadcast it to all rank + # rank1: compute each tile's [1/4:2:4] and broadcast it to all rank + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + Lamport( + self, + cur_tile_coord=cur_tile_coord, + tiled_mma=tiled_mma, + comm_tidx=tidx, + tCgCommInUC_per_peer=tCgCommInUC_per_peer, + tCgCommOutUC=tCgCommOutUC, + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + @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], + c_dtype: Type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + c_smem_layout: Union[cute.Layout, None], + ) -> Tuple[int, int, int]: + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 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 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + 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 C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_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 + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + swizzle_size: int, + raster_order: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_order == "m" + ) + 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, + arch: str, + ) -> 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, arch=arch) + + return num_tmem_alloc_cols + + def check_supported_dtypes( + self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] + ) -> bool: + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if ab_dtype not in valid_ab_dtypes: + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype: {ab_dtype}" + ) + + valid_acc_dtypes = {cutlass.Float32, cutlass.Float16, cutlass.Int32} + if self.acc_dtype not in valid_acc_dtypes: + raise testing.CantImplementError( + f"[dtype] Unsupported accumulator dtype: {self.acc_dtype}" + ) + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, # Float32 accumulator supports floating point AB types only + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and AB type + if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype {ab_dtype} for accumulator {self.acc_dtype}" + ) + + # Define compatibility mapping between accumulator type and C type + acc_c_compatibility = { + cutlass.Float32: { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + cutlass.Float16: { + cutlass.BFloat16, + cutlass.Float16, + }, + cutlass.Int32: { + cutlass.BFloat16, + cutlass.Float16, + cutlass.Float32, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + } + # Check compatibility between accumulator type and C type + if c_dtype not in acc_c_compatibility[self.acc_dtype]: + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for accumulator {self.acc_dtype}" + ) + + # check if c_dtype is supported by multimem all-reduce + if cutlass.const_expr( + c_dtype + not in { + cutlass.Float16, + cutlass.Float32, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + ): + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for multimem all-reduce" + ) + + return True + + def check_mma_tiler_and_cluster_shape(self) -> bool: + # Skip invalid mma tile shape + # 1SM: MMA tile {64, 128} x {64, 128, 192, 256} + # 2SM: MMA tile {128, 256} x {64, 128, 192, 256} + valid_mma_m = [64, 128] if not self.use_2cta_instrs else [128, 256] + valid_mma_n = [64, 128, 192, 256] + if self.mma_tiler_mn[0] not in valid_mma_m: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler M: {self.mma_tiler_mn[0]}, expected {valid_mma_m} for {'1SM' if not self.use_2cta_instrs else '2SM'}" + ) + if self.mma_tiler_mn[1] not in valid_mma_n: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler N: {self.mma_tiler_mn[1]}, expected {valid_mma_n}" + ) + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape M: {self.cluster_shape_mn[0]}, must be multiple of {2 if self.use_2cta_instrs else 1} for {'2SM' if self.use_2cta_instrs else '1SM'}" + ) + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16 + or self.cluster_shape_mn[0] <= 0 + or self.cluster_shape_mn[1] <= 0 + or not is_power_of_2(self.cluster_shape_mn[0]) + or not is_power_of_2(self.cluster_shape_mn[1]) + ): + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape: {self.cluster_shape_mn}, must be positive power of 2 with total <= 16" + ) + return True + + def check_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + ) -> bool: + def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: contiguous dimension must be 16B aligned for (M={m}, K={k}, N={n})" + ) + + if m % 128 != 0 and n % 128 != 0: + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: at least one of M={m} or N={n} must be divisible by 128" + ) + + if a_major not in ["m", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid a_major: {a_major}, expected 'm' or 'k'" + ) + if b_major not in ["n", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid b_major: {b_major}, expected 'n' or 'k'" + ) + if c_major not in ["m", "n"]: + raise testing.CantImplementError( + f"[alignment] Invalid c_major: {c_major}, expected 'm' or 'n'" + ) + + return True + + def check_epilogue_store_option(self, m: int, n: int) -> bool: + # Non TMA store version does not have predication, can not support OOB tiles + cta_tile_shape_mn = ( + self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_mn[1], + ) + if not self.use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + raise testing.CantImplementError( + f"[epilogue] Problem shape (M={m}, N={n}) must be divisible by cta tile shape {cta_tile_shape_mn} for non TMA store" + ) + # CTA swizzling improves the L2 cache utilization and reduces the number of cache misses. + # Make sure the swizzle size divides the cta/cga count since non TMA epilogue don't support OOB tiles. + # Swizzle only applies to the dimension orthogonal to the raster direction. + m_per_swizzle = (m // cta_tile_shape_mn[0]) // self.cluster_shape_mn[0] + n_per_swizzle = (n // cta_tile_shape_mn[1]) // self.cluster_shape_mn[1] + if self.raster_order == "m": + # Swizzle applies to N dimension only + if n_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape n={n} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if m_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape m={m} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_n" + ) + return True + + def check_valid_rank_config(self) -> bool: + if self.num_ranks not in [2, 4, 8] and self.all_reduce != "none": + raise testing.CantImplementError( + f"[rank_config] Invalid rank config: num_ranks={self.num_ranks}, must be 2, 4, or 8 for all-reduce" + ) + return True + + def check_valid_all_reduce_mode(self) -> bool: + if self.all_reduce not in ["Lamport", "none"]: + raise testing.CantImplementError( + f"[all_reduce_mode] Invalid all-reduce mode: {self.all_reduce}, expected 'Lamport' or 'none'" + ) + return True + + def check_even_divisiable_by_cluster_tilesize( + self, m: int, n: int, k: int, l: int + ) -> bool: + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + + cluster_shape_m = self.cluster_shape_mn[0] + cluster_shape_n = self.cluster_shape_mn[1] + + cluster_tile_m = cta_tile_m * cluster_shape_m + cluster_tile_n = cta_tile_n * cluster_shape_n + + if m % cluster_tile_m != 0 or n % cluster_tile_n != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_cluster_tile_size] Problem shape (M={m}, N={n}) not divisible by cluster tile size ({cluster_tile_m}, {cluster_tile_n})" + ) + return True + + def check_even_divisible_by_swizzle(self, m: int, n: int) -> bool: + # Only check for Lamport mode which uses multimem operations + if self.all_reduce != "Lamport": + return True + + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + + # Number of clusters in each dimension + num_clusters_m = m // (cta_tile_m * self.cluster_shape_mn[0]) + num_clusters_n = n // (cta_tile_n * self.cluster_shape_mn[1]) + + if self.raster_order == "m": + # Swizzle applies to N dimension only + if num_clusters_n % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in N dimension ({num_clusters_n}) " + f"must be divisible by swizzle_size={self.swizzle_size} for Lamport with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if num_clusters_m % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in M dimension ({num_clusters_m}) " + f"must be divisible by swizzle_size={self.swizzle_size} for Lamport with raster_order_n" + ) + + return True + + def can_implement( + self, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + ) -> bool: + # Skip unsupported types + self.check_supported_dtypes(ab_dtype, c_dtype) + + # Skip invalid mma tile shape and cluster shape + self.check_mma_tiler_and_cluster_shape() + + # Unpack mnkl for clarity in calling the epilog check + m, n, k, l = mnkl + + # Skip if GEMM not even divisiable by tilesize + # Right now, AR Fusion **ONLY** support even divisiable by tilesize + self.check_even_divisiable_by_cluster_tilesize(m, n, k, l) + + # Skip if swizzle doesn't evenly divide the problem (for Lamport communication) + self.check_even_divisible_by_swizzle(m, n) + + # Skip illegal problem shape for load/store alignment + self.check_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ) + + # Skip invalid epilogue store option + self.check_epilogue_store_option(m, n) + + # Skip invalid all-reduce rank config + self.check_valid_rank_config() + + # Skip invalid all-reduce mode + self.check_valid_all_reduce_mode() + + return True + + +##################################################################### +#### KERNEL IMPLEMENTATION END HERE #### +##################################################################### + +_DISTRIBUTED_INITIALIZED = False + + +# Check for device p2p access +def check_cuda_p2p_access(): + import torch + + num_devices = torch.cuda.device_count() + cuda_p2p_access_available = True + for i in range(num_devices - 1): + for j in range(i + 1, num_devices): + if not torch.cuda.can_device_access_peer(i, j): + cuda_p2p_access_available = False + break + if not cuda_p2p_access_available: + break + return cuda_p2p_access_available + + +# Check for multicast support +def check_multicast_support(): + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + device_count = torch.cuda.device_count() + for device_id in range(device_count): + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device_id): + return False + return True + except ImportError: + return False + + +def init_distributed(): + global _DISTRIBUTED_INITIALIZED + + if not dist.is_initialized() and not _DISTRIBUTED_INITIALIZED: + # * Check before running * + if not torch.cuda.is_available(): + print("CUDA is not available") + exit(-1) + if torch.cuda.get_device_capability()[0] != 10: + print("Running on non-sm100 device") + exit(-1) + if torch.cuda.device_count() < 2: + print("Less than two devices") + exit(-1) + if torch.cuda.device_count() % 2 != 0: + print("Odd world size") + exit(-1) + for var in ("LOCAL_RANK", "RANK", "WORLD_SIZE"): + if var not in os.environ: + print(f"{var} not in environment variables") + exit(-1) + + # LOCAL_RANK = which of this host's GPUs we bind to. + # RANK = global identity across all hosts; used for NCCL init + # and for kernel rank_id (slab ownership in the AR). + local_rank = int(os.environ["LOCAL_RANK"]) + global_rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + world_size=world_size, + rank=global_rank, + device_id=device, + ) + + _DISTRIBUTED_INITIALIZED = True + + print( + f"[init] hostname={socket.gethostname()} global_rank={global_rank} " + f"local_rank={local_rank} world_size={world_size} device={device}" + ) + + global_rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", global_rank)) + device = torch.device("cuda", local_rank) + + # Return the GLOBAL rank — callers feed this into kernel rank_id. + return global_rank, world_size, device + + +def clean_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def test_simple_torch_allreduce(device): + def _simple_torch_allreduce(): + tensor = symm_mem.empty((16384), dtype=torch.float32, device=device) + symm_mem.rendezvous(tensor, group=dist.group.WORLD) + torch.ops.symm_mem.multimem_all_reduce_( + tensor, "sum", dist.group.WORLD.group_name + ) + + try: + _simple_torch_allreduce() + except Exception: + return False + return True + + +def allocate_tensors( + *, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + num_workspace: int, + world_size: int, + global_rank: int, + device: int, + slot_init_mode: Literal["test", "benchmark"] = "test", +): + """Allocate every tensor the Lamport kernel touches. + + Lamport requires a 3-slot grouped C buffer (ping / pong / cooling). + Each outer slot holds + ``world_size`` symmetric-memory allocations (one peer-owned slab apiece), + and each allocation exposes a unicast (UC) view, a multicast (MC) view, + and a tuple of per-peer direct-P2P UC views built from + ``buffer_ptrs[r]`` of the symm_mem rendezvous handle. + + Slot initialization depends on ``slot_init_mode``: + + - ``"test"``: slot[0] is pre-seeded with the per-element Lamport ``-0.0`` + sentinel; slots [1] and [2] are seeded with random non-sentinel data so + the in-kernel clear of slot[1] (back to ``-0.0``) is observable as a + bit-pattern change in the post-run verification. The correctness path + uses ping=0, pong=1. + - ``"benchmark"``: every outer slot is pre-seeded with the ``-0.0`` + sentinel because the benchmark rotates through all three slots in the + ping role, so every slot must look "not-yet-arrived" the first time + it is read. + + A and B are allocated as ``num_workspace`` per-device copies (L2-cold + ring buffer), all bit-exact equal. ``comm_out`` is a plain per-device + tensor — Lamport only writes it locally via ``st.weak.global``; no peer + reads, no multicast. + + Returns a dict keyed for ``run_distributed_test`` / ``compile_and_benchmark``. + The ``_anchors`` tuple holds every torch handle that keeps the underlying + GPU allocations alive for the lifetime of the caller. + """ + NUM_C_BUFFERS = 3 + + if c_dtype in (cutlass.Float8E4M3FN, cutlass.Float8E5M2): + # Lamport's -0.0 sentinel epilogue is implemented for f32/f16/bf16 + # only (see the file header). Reject f8 here so the failure surfaces + # before any GPU allocation. + raise testing.CantImplementError( + f"Lamport currently does not support c_dtype={c_dtype}: the " + "epilogue -0.0 → +0.0 scrub is implemented for Float32 / Float16 / " + "BFloat16 only." + ) + + def _create_multicast_tensor( + *, + torch_tensor_ref: torch.Tensor, + dtype: Type[cutlass.Numeric], + leading_dim: int, + ): + # PyTorch symmetric-memory allocation + multicast view, mirroring the + # pattern used in distributed_gemm_blockscaled_all_reduce_ldmcxstmc_blackwell.py. + # + # After the collective ``symm_mem.rendezvous(t, group=...)`` the + # returned handle exposes: + # - ``multicast_ptr``: VA of the multicast view onto the same + # symmetric allocation. Writes/reads through it hit every rank's + # copy via multimem.*. + # - ``buffer_ptrs[r]``: per-peer VAs onto peer r's slab. Lamport + # issues plain ld.weak.global vector loads through these. + is_col_major = leading_dim == 0 + shape = torch_tensor_ref.shape + ref_dtype = torch_tensor_ref.dtype + + if is_col_major: + reversed_shape = tuple(reversed(shape)) + perm = tuple(range(len(shape) - 1, -1, -1)) + torch_symm_base = symm_mem.empty( + reversed_shape, device=device, dtype=ref_dtype + ) + torch_tensor_uc = torch_symm_base.permute(perm) + else: + torch_symm_base = symm_mem.empty(shape, device=device, dtype=ref_dtype) + torch_tensor_uc = torch_symm_base + + torch_tensor_uc.copy_(torch_tensor_ref) + + torch_tensor_hdl = symm_mem.rendezvous(torch_symm_base, group=dist.group.WORLD) + + if is_col_major: + torch_tensor_mc = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, reversed_shape, ref_dtype + ).permute(perm) + else: + torch_tensor_mc = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, shape, ref_dtype + ) + + def _make_peer_torch(r: int) -> torch.Tensor: + if is_col_major: + return cutlass_torch.as_tensor( + torch_tensor_hdl.buffer_ptrs[r], reversed_shape, ref_dtype + ).permute(perm) + return cutlass_torch.as_tensor( + torch_tensor_hdl.buffer_ptrs[r], shape, ref_dtype + ) + + torch_tensors_uc_per_peer = tuple( + _make_peer_torch(r) for r in range(world_size) + ) + + cute_tensor_uc = from_dlpack(torch_tensor_uc, assumed_align=16) + cute_tensor_uc.element_type = dtype + cute_tensor_uc = cute_tensor_uc.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor_uc = cutlass_torch.convert_cute_tensor( + torch_tensor_uc, + cute_tensor_uc, + dtype, + is_dynamic_layout=True, + ) + + cute_tensor_mc = from_dlpack(torch_tensor_mc, assumed_align=16) + cute_tensor_mc.element_type = dtype + cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim) + + def _wrap_peer_cute(peer_torch: torch.Tensor): + peer_cute = from_dlpack(peer_torch, assumed_align=16) + peer_cute.element_type = dtype + peer_cute = peer_cute.mark_layout_dynamic(leading_dim=leading_dim) + return cutlass_torch.convert_cute_tensor( + peer_torch, peer_cute, dtype, is_dynamic_layout=True + ) + + cute_tensors_uc_per_peer = tuple( + _wrap_peer_cute(peer_torch) for peer_torch in torch_tensors_uc_per_peer + ) + + return ( + cute_tensor_uc, + cute_tensor_mc, + cute_tensors_uc_per_peer, + torch_tensor_uc, + torch_tensor_mc, + torch_tensors_uc_per_peer, + ) + + # Per-rank seed: distinct A/B per rank so AR sees real cross-rank diversity + # (matches the original run_distributed init). fp8 is rejected above, so + # the fp8-specific same-seed branch is no longer needed. + torch.manual_seed(42 + global_rank) + + m, n, k, l = mnkl + + # CPU random Int32 in [-1, 1] then cast — Lamport uses this narrower + # range than ldmcxstmc ([-2, 2]) because the all-reduce sum across N + # ranks must not overflow the f16 / bf16 dynamic range. + _init_lo, _init_hi = (-1, 1) if slot_init_mode == "test" else (-20, 20) + torch_tensor_a_int_cpu = cutlass_torch.matrix( + l, + m, + k, + a_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + torch_tensor_b_int_cpu = cutlass_torch.matrix( + l, + n, + k, + b_major == "n", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + torch_tensor_c_int_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + torch_tensor_comm_out_c_int_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + + torch_tensor_a_cpu = torch_tensor_a_int_cpu.to(dtype=torch_dtype(ab_dtype)) + torch_tensor_b_cpu = torch_tensor_b_int_cpu.to(dtype=torch_dtype(ab_dtype)) + torch_tensor_c_cpu = torch_tensor_c_int_cpu.to(dtype=torch_dtype(c_dtype)) + torch_tensor_comm_out_c_cpu = torch_tensor_comm_out_c_int_cpu.to( + dtype=torch_dtype(c_dtype) + ) + + # The Lamport sentinel-vs-random init pattern depends on slot_init_mode. + # In test mode, slot[0] enters as the -0.0 sentinel and slots [1, 2] hold + # random data so the in-kernel clear of slot[1] is observable in + # _compare_buffer_clear. In benchmark mode every slot enters as the + # sentinel because the rotation cycles all three through the ping role. + torch_tensor_c_random_cpu = torch_tensor_c_cpu.clone() + torch_tensor_c_cpu.fill_(-0.0) + + # A, B GPU: num_workspace copies (L2-cold ring buffer), bit-exact equal. + cute_tensor_a_list = [] + cute_tensor_b_list = [] + torch_tensor_a_gpu_list = [] + torch_tensor_b_gpu_list = [] + for _ in range(num_workspace): + cute_tensor_a_i, torch_tensor_a_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_a_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_b_i, torch_tensor_b_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_b_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_a_list.append(cute_tensor_a_i) + cute_tensor_b_list.append(cute_tensor_b_i) + torch_tensor_a_gpu_list.append(torch_tensor_a_gpu_i) + torch_tensor_b_gpu_list.append(torch_tensor_b_gpu_i) + + assert not torch.isnan(torch_tensor_a_gpu_list[0]).any() + assert not torch.isnan(torch_tensor_b_gpu_list[0]).any() + + # Grouped C: NUM_C_BUFFERS outer × world_size inner symmetric allocations. + # Each rank "owns" allocation index `global_rank` for scatter — its TMA + # multicast store goes through the MC view of slot[s][global_rank], which + # NVSwitch fans out to slot[s][global_rank] on every peer. Every rank then + # reads slot[s][r] (its own UC view of allocation r) to find peer r's + # contribution. + leading_dim_c = 1 if c_major == "n" else 0 + cute_uc_grouped: List[Tuple[cute.Tensor, ...]] = [] + cute_mc_grouped: List[Tuple[cute.Tensor, ...]] = [] + torch_uc_grouped: List[Tuple[torch.Tensor, ...]] = [] + torch_mc_grouped: List[Tuple[torch.Tensor, ...]] = [] + for buf_idx in range(NUM_C_BUFFERS): + cute_uc_list: List[cute.Tensor] = [] + cute_mc_list: List[cute.Tensor] = [] + torch_uc_list: List[torch.Tensor] = [] + torch_mc_list: List[torch.Tensor] = [] + if slot_init_mode == "test": + # Test scaffold: slot[0] enters as the Lamport -0.0 sentinel + # (live ping this iteration); slots [1, 2] init with random + # non-sentinel data so the kernel-side clear of slot[1] is + # observable as a bit-pattern change in the post-run verify. + init_ref = torch_tensor_c_cpu if buf_idx == 0 else torch_tensor_c_random_cpu + else: + # Benchmark: every slot enters as -0.0 because the rotation + # cycles all three through the ping role. + init_ref = torch_tensor_c_cpu + for _ in range(world_size): + ( + _c_uc, + _c_mc, + _c_uc_per_peer, + _t_uc, + _t_mc, + _t_uc_per_peer, + ) = _create_multicast_tensor( + torch_tensor_ref=init_ref, + dtype=c_dtype, + leading_dim=leading_dim_c, + ) + cute_uc_list.append(_c_uc) + cute_mc_list.append(_c_mc) + torch_uc_list.append(_t_uc) + torch_mc_list.append(_t_mc) + cute_uc_grouped.append(tuple(cute_uc_list)) + cute_mc_grouped.append(tuple(cute_mc_list)) + torch_uc_grouped.append(tuple(torch_uc_list)) + torch_mc_grouped.append(tuple(torch_mc_list)) + + cute_tensors_c_uc_per_peer_grouped = tuple(cute_uc_grouped) + cute_tensors_c_mc_per_peer_grouped = tuple(cute_mc_grouped) + torch_tensors_c_uc_per_peer_grouped = tuple(torch_uc_grouped) + torch_tensors_c_mc_per_peer_grouped = tuple(torch_mc_grouped) + + # comm_out: plain per-device CUDA tensor (Lamport only writes it locally + # via st.weak.global — no multicast, no peer reads). + cute_tensor_comm_out_uc, torch_tensor_comm_out_gpu = cutlass_torch.cute_tensor_like( + torch_tensor_comm_out_c_cpu, + c_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + return { + "cute_tensor_a_list": cute_tensor_a_list, + "cute_tensor_b_list": cute_tensor_b_list, + "cute_tensors_c_uc_per_peer_grouped": cute_tensors_c_uc_per_peer_grouped, + "cute_tensors_c_mc_per_peer_grouped": cute_tensors_c_mc_per_peer_grouped, + "cute_tensor_comm_out_uc": cute_tensor_comm_out_uc, + "torch_tensor_a_cpu": torch_tensor_a_cpu, + "torch_tensor_b_cpu": torch_tensor_b_cpu, + "torch_tensor_c_cpu": torch_tensor_c_cpu, + "torch_tensors_c_uc_per_peer_grouped": torch_tensors_c_uc_per_peer_grouped, + "torch_tensor_comm_out_gpu": torch_tensor_comm_out_gpu, + "_anchors": ( + torch_tensor_a_gpu_list, + torch_tensor_b_gpu_list, + torch_tensors_c_uc_per_peer_grouped, + torch_tensors_c_mc_per_peer_grouped, + torch_tensor_comm_out_gpu, + ), + } + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + warmup_iterations: int = 0, + iterations: int = 1, + all_reduce: str = "Lamport", + benchmark_or_test: str = "test", + use_cuda_graph: bool = False, + csv_report: str = "./csv_report.csv", +): + """Top-level Lamport entry point. Initializes distributed, dispatches to + the test or benchmark driver, and tears distributed down on exit. + + ``benchmark_or_test`` selects the mode: + - ``"test"``: single-launch correctness check via run_distributed_test. + - ``"benchmark"``: cudaEvent timing + per-rank CSV on the supplied config. + - ``"benchmark_all"``: sweep across kernel-config variants, one row per + candidate in the per-rank CSV. + """ + global_rank, world_size, device = init_distributed() + + if all_reduce == "Lamport": + if not check_cuda_p2p_access(): + print("CUDA P2P access not available") + exit(-1) + if not check_multicast_support(): + print("Multicast support not available") + exit(-1) + if not test_simple_torch_allreduce(device): + print("Simple torch allreduce test failed") + exit(-1) + + def _verbose(): + print("Running Blackwell SM100 Persistent Dense GEMM AR test with:") + print(f"mnkl: {mnkl}") + print(f"global_rank: {global_rank}, world_size: {world_size}, device: {device}") + print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}") + print(f"Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") + print(f"Swizzle Size: {swizzle_size}") + print(f"Raster Order: {raster_order}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Use CUDA Graph: {'True' if use_cuda_graph else 'False'}") + print(f"Mode: {benchmark_or_test}") + print(f"Fused AllReduce Op: {all_reduce}") + + if global_rank == 0: + _verbose() + + if benchmark_or_test in ("benchmark", "benchmark_all"): + run_distributed_benchmark( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + warmup_iterations=warmup_iterations, + iterations=iterations, + all_reduce=all_reduce, + use_cuda_graph=use_cuda_graph, + benchmark_all=(benchmark_or_test == "benchmark_all"), + csv_report=csv_report, + ) + else: + run_distributed_test( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + all_reduce=all_reduce, + ) + + clean_distributed() + + +def run_distributed_test( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + all_reduce: str = "Lamport", +): + # * CUDA * + # Non-default stream as we'll use CUDA graphs + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # * Build GEMM AllReduce Object * + gemm_all_reduce = Sm100PersistentDenseGemmAllReduceLamportKernel( + acc_dtype, + c_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + rank_id=global_rank, + num_ranks=world_size, + all_reduce=all_reduce, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + + # * Set Filename & Funcname for different rank * + # Use hash to keep filename short and avoid "File name too long" error + try: + import hashlib + + param_str = f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}_{'2Sm' if use_2cta_instrs else '1Sm'}_{'EpiTma' if use_tma_store else 'EpiStg'}" + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_ar_{param_hash}_{all_reduce}_r{global_rank}_w{world_size}" + ) + except Exception: + unique_name = f"sm100_gemm_ar_{all_reduce}_r{global_rank}_w{world_size}" + gemm_all_reduce.__call__.set_name_prefix(unique_name) + + # * can_implement * + can_implement = gemm_all_reduce.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, all_reduce={all_reduce}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + # * Allocate tensors * + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + num_workspace=1, + world_size=world_size, + global_rank=global_rank, + device=device, + slot_init_mode="test", + ) + cute_tensor_a = tensors["cute_tensor_a_list"][0] + cute_tensor_b = tensors["cute_tensor_b_list"][0] + cute_tensors_c_uc_per_peer_grouped = tensors["cute_tensors_c_uc_per_peer_grouped"] + cute_tensors_c_mc_per_peer_grouped = tensors["cute_tensors_c_mc_per_peer_grouped"] + cute_tensor_comm_out_uc = tensors["cute_tensor_comm_out_uc"] + torch_tensor_a_cpu = tensors["torch_tensor_a_cpu"] + torch_tensor_b_cpu = tensors["torch_tensor_b_cpu"] + torch_tensors_c_uc_per_peer_grouped = tensors["torch_tensors_c_uc_per_peer_grouped"] + torch_tensor_comm_out = tensors["torch_tensor_comm_out_gpu"] + + # Test mode: ping=0, pong=1. Slot[0]/slot[1] indexing is inlined at the + # cute.compile / compiled_gemm call sites rather than hoisted into locals + # so the slot semantics stay visible at the kernel boundary. + compiled_gemm = cute.compile( + gemm_all_reduce, + a=cute_tensor_a, + b=cute_tensor_b, + c_multicast_tensor=cute_tensors_c_mc_per_peer_grouped[0][global_rank], + comm_in_unicast_tensor_per_peer=cute_tensors_c_uc_per_peer_grouped[0], + comm_clear_unicast_tensor_per_peer=cute_tensors_c_uc_per_peer_grouped[1], + comm_out_unicast_tensor=cute_tensor_comm_out_uc, + stream=stream, + max_active_clusters=utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ), + ) + compiled_gemm( + a=cute_tensor_a, + b=cute_tensor_b, + c_multicast_tensor=cute_tensors_c_mc_per_peer_grouped[0][global_rank], + comm_in_unicast_tensor_per_peer=cute_tensors_c_uc_per_peer_grouped[0], + comm_clear_unicast_tensor_per_peer=cute_tensors_c_uc_per_peer_grouped[1], + comm_out_unicast_tensor=cute_tensor_comm_out_uc, + stream=stream, + ) + + # * Compare * + # FP8 nan to saturate + def _fp8_nan_to_saturate(tensor_fp8: torch.Tensor) -> torch.Tensor: + if tensor_fp8.dtype == torch.float8_e4m3fn: + tensor_u8 = tensor_fp8.view(torch.uint8) + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 448 + tensor_f32[is_neg_nan] = -448 + return tensor_f32.to(dtype=tensor_fp8.dtype) + if tensor_fp8.dtype == torch.float8_e5m2: + tensor_u8 = tensor_fp8.view(torch.uint8) + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 57344 + tensor_f32[is_neg_nan] = -57344 + return tensor_f32.to(dtype=tensor_fp8.dtype) + return tensor_fp8 + + def _compare_gemm( + torch_tensor_a_cpu: torch.Tensor, + torch_tensor_b_cpu: torch.Tensor, + torch_tensor_c: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + gemm_ref_f32 = torch.einsum( + "mkl,nkl->mnl", + torch_tensor_a_cpu.to(dtype=torch.float32), + torch_tensor_b_cpu.to(dtype=torch.float32), + ) + if c_dtype == cutlass.Float8E4M3FN: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-448, max=448) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + elif c_dtype == cutlass.Float8E5M2: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-57344, max=57344) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + else: + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + gemm_ref_f32_c_f32 = gemm_ref_f32_c.to(dtype=torch.float32) + gemm_out_f32 = torch_tensor_c.cpu().to(dtype=torch.float32) + torch.testing.assert_close( + gemm_out_f32, gemm_ref_f32_c_f32, atol=atol, rtol=rtol + ) + + def _compare_comm( + torch_tensor_comm_in: torch.Tensor, + torch_tensor_comm_out: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + comm_ref = torch_tensor_comm_in + dist.barrier() + if c_dtype in [cutlass.Float32, cutlass.BFloat16, cutlass.Float16]: + comm_ref = comm_ref.contiguous() + dist.all_reduce(comm_ref, op=torch.distributed.ReduceOp.SUM) + comm_ref_f32 = comm_ref.to(dtype=torch.float32).cpu() + elif c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_ref_fp32_acc = comm_ref.to(dtype=torch.float32).contiguous() + dist.all_reduce(comm_ref_fp32_acc, op=torch.distributed.ReduceOp.SUM) + epsilon = 1e-2 + comm_ref_c = (comm_ref_fp32_acc + epsilon).to(dtype=torch_dtype(c_dtype)) + comm_ref_nan2sat = _fp8_nan_to_saturate(comm_ref_c) + comm_ref_f32 = comm_ref_nan2sat.to(dtype=torch.float32).cpu() + + comm_out = torch_tensor_comm_out.cpu() + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_out = _fp8_nan_to_saturate(comm_out) + comm_out_f32 = comm_out.to(dtype=torch.float32) + torch.testing.assert_close(comm_ref_f32, comm_out_f32, atol=atol, rtol=rtol) + + def _compare_buffer_clear( + torch_tensors_c_uc_per_peer_for_slot: Tuple[torch.Tensor, ...], + slot_idx: int, + c_dtype, + ): + # Lamport litmus test: after the kernel consumes slot[ping], slot[pong] + # must be byte-exact -0.0 across every per-peer slab (sentinel pattern + # mirrors the AR kernel's neg_zero_dtype_32bit table). + if c_dtype is cutlass.Float32: + view_dtype, sentinel, hex_width = torch.int32, -0x80000000, 8 + elif c_dtype is cutlass.Float16 or c_dtype is cutlass.BFloat16: + view_dtype, sentinel, hex_width = torch.int16, -0x8000, 4 + elif c_dtype is cutlass.Float8E4M3FN or c_dtype is cutlass.Float8E5M2: + view_dtype, sentinel, hex_width = torch.uint8, 0x80, 2 + else: + raise AssertionError( + f"_compare_buffer_clear: unsupported c_dtype {c_dtype}" + ) + mask = (1 << (hex_width * 4)) - 1 + sentinel_hex = f"0x{sentinel & mask:0{hex_width}x}" + + for p, slab in enumerate(torch_tensors_c_uc_per_peer_for_slot): + elements = slab.contiguous().cpu().view(view_dtype) + if not bool((elements == sentinel).all().item()): + bad_idx = (elements != sentinel).nonzero(as_tuple=False)[0].tolist() + bad_val = int(elements[tuple(bad_idx)].item()) & mask + raise AssertionError( + f"comm_in slot[{slot_idx}][{p}] not cleared to " + f"-0.0 sentinel: element at {bad_idx} = " + f"0x{bad_val:0{hex_width}x} (expected {sentinel_hex}). " + f"Kernel-side buffer-clear pass not yet wired." + ) + + torch_tensor_c_local_slot = torch_tensors_c_uc_per_peer_grouped[0][global_rank] + + # compare + # Impose a strict check on GEMM as input is uniform random int + gemm_atol = 1e-05 + gemm_rtol = 1e-05 + _compare_gemm( + torch_tensor_a_cpu=torch_tensor_a_cpu, + torch_tensor_b_cpu=torch_tensor_b_cpu, + torch_tensor_c=torch_tensor_c_local_slot, + c_dtype=c_dtype, + atol=gemm_atol, + rtol=gemm_rtol, + ) + + if all_reduce != "none": + # Relaxed tolerance — _compare_comm uses NCCL dist.all_reduce, which + # reduces in a different order than the fused multimem kernel. + comm_rtol = 1e-1 + comm_atol = 5 + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_rtol = 1e-1 + comm_atol = 0.0 + _compare_comm( + torch_tensor_comm_in=torch_tensor_c_local_slot, + torch_tensor_comm_out=torch_tensor_comm_out, + c_dtype=c_dtype, + atol=comm_atol, + rtol=comm_rtol, + ) + + # * Check slot[pong] cleared * + # Lamport scrubs slot[1] back to -0.0 after consuming slot[0]; verify here. + if all_reduce == "Lamport": + _compare_buffer_clear( + torch_tensors_c_uc_per_peer_grouped[1], + slot_idx=1, + c_dtype=c_dtype, + ) + + +def compile_and_benchmark( + *, + mnkl, + ab_dtype, + c_dtype, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + swizzle_size, + raster_order, + use_2cta_instrs, + use_tma_store, + all_reduce, + global_rank, + world_size, + num_workspace, + tensors, + torch_stream, + stream, + warmup_iterations, + iterations, + use_cuda_graph, +): + """Compile one Lamport candidate, time warmup + benchmark with slot rotation. + + Slot rotation is: + + signal = i (per-iteration counter) + ping = signal % NUM_C_BUFFERS # this iter's read+write+epi target + pong = (signal + 1) % NUM_C_BUFFERS # this iter's clear target + cooling = (signal + 2) % NUM_C_BUFFERS # untouched (draining) + + ``num_workspace`` is locked to ``len(cute_tensors_c_uc_per_peer_grouped)`` + so the A/B ring buffer rotates in lockstep with the slot rotation — + a single set of NUM_C_BUFFERS captured CUDA graphs covers both. + + Returns the mean per-iteration time in microseconds (warmup excluded). + """ + cute_tensor_a_list = tensors["cute_tensor_a_list"] + cute_tensor_b_list = tensors["cute_tensor_b_list"] + cute_tensors_c_uc_per_peer_grouped = tensors["cute_tensors_c_uc_per_peer_grouped"] + cute_tensors_c_mc_per_peer_grouped = tensors["cute_tensors_c_mc_per_peer_grouped"] + cute_tensor_comm_out_uc = tensors["cute_tensor_comm_out_uc"] + + num_c_buffers = len(cute_tensors_c_uc_per_peer_grouped) + assert num_workspace == num_c_buffers, ( + f"compile_and_benchmark requires num_workspace == NUM_C_BUFFERS " + f"({num_c_buffers}) so the A/B ring buffer rotates in lockstep " + f"with the Lamport slot rotation; got num_workspace={num_workspace}." + ) + + # Build kernel + can_implement first — before any cross-rank op, so a + # deterministic skip keeps all ranks in lockstep for the next candidate. + gemm_all_reduce = Sm100PersistentDenseGemmAllReduceLamportKernel( + acc_dtype, + c_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + rank_id=global_rank, + num_ranks=world_size, + all_reduce=all_reduce, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + try: + import hashlib + + param_str = ( + f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}" + f"_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}" + f"_{'2Sm' if use_2cta_instrs else '1Sm'}" + f"_{'EpiTma' if use_tma_store else 'EpiStg'}" + ) + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_ar_{param_hash}_{all_reduce}_r{global_rank}_w{world_size}" + ) + except Exception: + unique_name = f"sm100_gemm_ar_{all_reduce}_r{global_rank}_w{world_size}" + gemm_all_reduce.__call__.set_name_prefix(unique_name) + + can_implement = gemm_all_reduce.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, all_reduce={all_reduce}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + def make_kernel_kwargs(i): + ping = i % num_c_buffers + pong = (i + 1) % num_c_buffers + return dict( + a=cute_tensor_a_list[i % num_workspace], + b=cute_tensor_b_list[i % num_workspace], + c_multicast_tensor=cute_tensors_c_mc_per_peer_grouped[ping][global_rank], + comm_in_unicast_tensor_per_peer=cute_tensors_c_uc_per_peer_grouped[ping], + comm_clear_unicast_tensor_per_peer=cute_tensors_c_uc_per_peer_grouped[pong], + comm_out_unicast_tensor=cute_tensor_comm_out_uc, + stream=stream, + ) + + kernel_kwargs_list = [make_kernel_kwargs(i) for i in range(num_c_buffers)] + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + compiled_gemm = cute.compile( + gemm_all_reduce, + **kernel_kwargs_list[0], + max_active_clusters=max_active_clusters, + ) + + # Under --use_cuda_graph capture NUM_C_BUFFERS graphs (one per slot + # rotation — pointers are baked at capture time). Without it dispatch + # the compiled kernel directly with the rotated kwargs. + if use_cuda_graph: + graphs = [] + for r in range(num_c_buffers): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch_stream): + compiled_gemm(**kernel_kwargs_list[r]) + graphs.append(g) + + def run_one_iter(i): + graphs[i % num_c_buffers].replay() + else: + + def run_one_iter(i): + compiled_gemm(**kernel_kwargs_list[i % num_c_buffers]) + + # One loop for warmup + benchmark. Every iteration is bracketed by its own + # event pair; only the trailing `iterations` event pairs contribute to the + # reported average so warmup transients don't skew the result. + + # Cross-rank sync before timing. torch.cuda.synchronize() drains any + # pending GPU work (graph captures land on the default stream; timing + # runs on torch_stream) so dist.barrier()'s host-side rendezvous is + # not racing in-flight kernels from the previous candidate. Without + # this, the slowest rank's leftover work would skew the measured + # window of every other rank in a benchmark_all sweep. + torch.cuda.synchronize() + dist.barrier() + + total_iters = warmup_iterations + iterations + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + with torch.cuda.stream(torch_stream): + for i in range(total_iters): + start_events[i].record() + run_one_iter(i) + end_events[i].record() + end_events[-1].synchronize() + + per_iter_us = [s.elapsed_time(e) * 1000.0 for s, e in zip(start_events, end_events)] + bench_times = per_iter_us[warmup_iterations:] + avg_us = sum(bench_times) / iterations + print( + f"Rank {global_rank} Execution time: {avg_us} us " + f"(avg of {iterations} iters, after {warmup_iterations} warmup)\n", + flush=True, + ) + return avg_us + + +def run_distributed_benchmark( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + warmup_iterations: int = 0, + iterations: int = 1, + all_reduce: str = "Lamport", + use_cuda_graph: bool = False, + benchmark_all: bool = False, + csv_report: str = "./csv_report.csv", +): + """Lamport benchmark driver — single config or full sweep + per-rank CSV. + + All allocations happen once up front; the A/B ring buffer length is + locked to ``len(cute_tensors_c_uc_per_peer_grouped)`` so a single set + of captured CUDA graphs covers both A/B and slot rotations. + + --benchmark_all sweeps raster ∈ {m, n}, swizzle ∈ {1, 2, 4, 8}, and the + (use_2cta_instrs, cluster_shape_mn, mma_tiler_mn) joint variants from + the LDMCxSTMC sweep list (cluster_m ∈ {2,4,8} requires 2cta). + + Lamport rejects use_tma_store=False at the kernel layer, so the + candidate list keeps the user-supplied value rather than sweeping it. + """ + # NUM_C_BUFFERS = 3 is fixed inside allocate_tensors. num_workspace is + # locked to the same value here so the A/B ring rotates in lockstep + # with the slot rotation. The compile_and_benchmark assert below + # enforces the invariant against whatever allocate_tensors returns. + num_workspace = 3 + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + num_workspace=num_workspace, + world_size=world_size, + global_rank=global_rank, + device=device, + slot_init_mode="benchmark", + ) + + # One non-default stream shared across every candidate. `torch_stream` is + # kept alive because we capture it into a CUDA graph; creating a fresh + # stream per candidate would churn events + graph context for no reason. + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # Build the candidate list of args for compile_and_benchmark. For a single + # run the list has one entry; --benchmark_all sweeps: + # raster_order ∈ {m, n} + # swizzle_size ∈ {1, 2, 4, 8} + # use_2cta_instrs, cluster_shape_mn, mma_tiler_mn jointly, with the + # constraint that cluster_m ∈ {2,4,8} requires 2cta (smaller tiles and + # single-cta clusters are only valid without 2cta). + base_kwargs = dict( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + use_tma_store=use_tma_store, + all_reduce=all_reduce, + global_rank=global_rank, + world_size=world_size, + num_workspace=num_workspace, + tensors=tensors, + torch_stream=torch_stream, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cuda_graph=use_cuda_graph, + ) + if benchmark_all: + raster_candidates = ["m", "n"] + swizzle_candidates = [1, 2, 4, 8] + # All cluster shapes from the allowed set; each is filtered per 2cta + # option below by the "cluster_m ∈ {2,4,8} needs 2cta" rule. + cluster_all = [ + (1, 1), + (1, 2), + (2, 1), + (2, 2), + ] + two_cta_options = [ + { + "use_2cta_instrs": True, + "clusters": [c for c in cluster_all if c[0] in (2, 4, 8)], + "mma_tilers": [ + (256, 256), + (256, 192), + (256, 128), + (256, 64), + (128, 256), + (128, 192), + (128, 128), + (128, 64), + ], + }, + { + "use_2cta_instrs": False, + "clusters": cluster_all, + "mma_tilers": [ + (128, 256), + (128, 192), + (128, 128), + (128, 64), + (64, 256), + (64, 192), + (64, 128), + (64, 64), + ], + }, + ] + candidates = [] + for opt in two_cta_options: + for cluster in opt["clusters"]: + for tile in opt["mma_tilers"]: + for r in raster_candidates: + for s in swizzle_candidates: + candidates.append( + { + **base_kwargs, + "use_2cta_instrs": opt["use_2cta_instrs"], + "cluster_shape_mn": cluster, + "mma_tiler_mn": tile, + "raster_order": r, + "swizzle_size": s, + } + ) + else: + candidates = [ + { + **base_kwargs, + "use_2cta_instrs": use_2cta_instrs, + "cluster_shape_mn": cluster_shape_mn, + "mma_tiler_mn": mma_tiler_mn, + "raster_order": raster_order, + "swizzle_size": swizzle_size, + } + ] + + if global_rank == 0: + print(f"Total number of candidate configs: {len(candidates)}", flush=True) + + # Run each candidate. Tensors (A/B/C/comm_out/flags) are reused across + # candidates; the kernel resets flags internally. Skip any combo the kernel + # can't implement — can_implement is deterministic so all ranks skip the + # same ones, keeping rendezvous collective operations lockstep. + # ~2s sleep between candidates to let the GPU cool down. + results = [] + for i, cand in enumerate(candidates): + if i > 0: + time.sleep(2) + try: + avg_us = compile_and_benchmark(**cand) + except testing.CantImplementError: + continue + results.append((cand, avg_us)) + + # Each rank writes its own CSV — one row per candidate, one config field + # per column. The path comes from --csv_report; per-rank suffix is appended + # before the extension so concurrent ranks don't clobber the same file. + global_rank = dist.get_rank() + m, n, k, l = mnkl + fieldnames = [ + "gemm_m", + "gemm_n", + "gemm_k", + "gemm_l", + "ab_dtype", + "c_dtype", + "acc_dtype", + "a_major", + "b_major", + "c_major", + "mma_tiler_m", + "mma_tiler_n", + "cluster_m", + "cluster_n", + "use_2cta_instrs", + "use_tma_store", + "all_reduce", + "use_cuda_graph", + "raster_order", + "swizzle_size", + "num_workspace", + "warmup_iterations", + "iterations", + "world_size", + "global_rank", + "benchmark_time_us", + ] + csv_path = Path(csv_report) + csv_path = csv_path.with_name( + f"{csv_path.stem}_rank{global_rank}_world{world_size}{csv_path.suffix}" + ) + csv_path.parent.mkdir(parents=True, exist_ok=True) + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for cand, avg_us in results: + writer.writerow( + { + "gemm_m": m, + "gemm_n": n, + "gemm_k": k, + "gemm_l": l, + "ab_dtype": str(ab_dtype), + "c_dtype": str(c_dtype), + "acc_dtype": str(acc_dtype), + "a_major": a_major, + "b_major": b_major, + "c_major": c_major, + "mma_tiler_m": cand["mma_tiler_mn"][0], + "mma_tiler_n": cand["mma_tiler_mn"][1], + "cluster_m": cand["cluster_shape_mn"][0], + "cluster_n": cand["cluster_shape_mn"][1], + "use_2cta_instrs": cand["use_2cta_instrs"], + "use_tma_store": use_tma_store, + "all_reduce": all_reduce, + "use_cuda_graph": use_cuda_graph, + "raster_order": cand["raster_order"], + "swizzle_size": cand["swizzle_size"], + "num_workspace": num_workspace, + "warmup_iterations": warmup_iterations, + "iterations": iterations, + "world_size": world_size, + "global_rank": global_rank, + "benchmark_time_us": avg_us, + } + ) + + return results[0][1] if len(results) == 1 else results + + +def prepare_parser(): + parser = argparse.ArgumentParser( + description="Example of SM100 Persistent Dense GEMM AR on Blackwell." + ) + + def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser.add_argument( + "--mnkl", + type=_parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=_parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=_parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--use_tma_store", action="store_true", help="Use tma store or not" + ) + parser.add_argument( + "--benchmark_or_test", + type=str, + choices=["test", "benchmark", "benchmark_all"], + default="test", + help="Select run mode: 'test' for a single-launch correctness check, " + "'benchmark' for cudaEvent timing + per-rank CSV on the given config, " + "'benchmark_all' for the same with a sweep across kernel-config variants.", + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--use_cuda_graph", + action="store_true", + default=False, + help="Capture one CUDA graph per slot rotation (3) and replay per iteration", + ) + parser.add_argument( + "--all_reduce", + choices=["Lamport", "none"], + type=str, + default="Lamport", + help="Allreduce algorithm to fuse with gemm", + ) + parser.add_argument( + "--swizzle_size", + type=int, + default=1, + help="Swizzling size in the unit of cluster for improving L2 cache hit rate", + ) + parser.add_argument( + "--raster_order", + type=str, + choices=["m", "n"], + default="m", + help="Rasterization order of clusters", + ) + parser.add_argument( + "--csv_report", + type=str, + default="./csv_report.csv", + help="Path for the per-rank benchmark CSV. Each rank appends " + "_rank{R}_world{W} before the extension to avoid collisions.", + ) + + return parser + + +if __name__ == "__main__": + parser = prepare_parser() + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + run( + args.mnkl, + args.ab_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.swizzle_size, + args.raster_order, + args.use_2cta_instrs, + args.use_tma_store, + args.warmup_iterations, + args.iterations, + args.all_reduce, + args.benchmark_or_test, + args.use_cuda_graph, + args.csv_report, + ) + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + print( + f"PASS (hostname={socket.gethostname()}, global_rank={rank}, local_rank={local_rank})", + flush=True, + ) diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_ldxstmc_blackwell.py b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_ldxstmc_blackwell.py new file mode 100644 index 000000000..3b4a0803c --- /dev/null +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_ldxstmc_blackwell.py @@ -0,0 +1,3354 @@ +# Copyright (c) 2025 - 2026 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. + +import os +import argparse +import csv +import socket +import time +from pathlib import Path + +from math import gcd +from typing import Optional, Tuple, Type, Union, Literal + +import cuda.bindings.driver as cuda + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +import cutlass.torch as cutlass_torch +from cutlass.torch import dtype as torch_dtype +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cutlass_dsl import BaseDSL +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Pointer, Int32, Float32 +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm + +""" +SM100 GEMM + AllReduce (LDxSTMC) Fusion Kernel + +* Test (correctness only — single launch, no warmup/iteration loop) +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_ldxstmc_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs --all_reduce LDxSTMC \ + --benchmark_or_test test + +* Benchmark One +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_ldxstmc_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --use_2cta_instrs --all_reduce LDxSTMC \ + --benchmark_or_test benchmark --use_cuda_graph --csv_report ./csv_report.csv + +* Benchmark All +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_all_reduce_ldxstmc_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --all_reduce LDxSTMC \ + --benchmark_or_test benchmark_all --use_cuda_graph --csv_report ./csv_report.csv + +NOTE: +* multimem.ld_reduce does NOT handle NaN issue for f8. + +ALGO: +-> Producer +for each output tile + * store gemm output [data] (UC) to local rank global memory with .weak semantics (tma is default .weak semantic) + * [flag] (MC) update with .release semantics and .gpu scope (prior store is to local rank global memory thus .gpu scope) on multicast flag ptr (thus flag update broadcast to all rank) + +-> Consumer +for each output tile + * spin lock wait [flag] (UC) with .relaxed / .acquire and .gpu scope (all rank have write their data to local rank gmem) on local unicast flag ptr + * for each peer r in [0, P): + - load [data] with .weak semantic from all peer rank's global memory + - convert to f32, reduce, convert back to original dtype + * multimem.st broadcast reduced [data] (MC) to all rank as each output tile is handled collectively by P rank +before kernel exit + * [flag] (MC) update with .release semantics and .sys scope (prior store broadcast data to all rank in sys) + * spin lock wait [flag] (UC) with .relaxed / .acquire and .sys scope +""" + +##################################################################### +#### KERNEL IMPLEMENTATION START HERE #### +##################################################################### + + +@dsl_user_op +def ld_weak_global_4xu32( + ptr: Pointer, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Int32, Int32, Int32, Int32]: + ptr_int = ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + return_struct = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(i32,i32,i32,i32)>"), + [ptr_int], + "ld.weak.global.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return_regs = [llvm.extractvalue(T.i32(), return_struct, [i]) for i in range(4)] + return return_regs[0], return_regs[1], return_regs[2], return_regs[3] + + +@dsl_user_op +def st_weak_global_4xu32( + ptr: Pointer, + x: Int32, + y: Int32, + z: Int32, + w: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + ptr_int = ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ptr_int, x, y, z, w], + "st.weak.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_bf16x2_to_2xf32( + bf16x2_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32)>"), + [bf16x2_b32], + "{ .reg .b16 hl, hh;\n" + " mov.b32 {hl, hh}, $2;\n" + " cvt.f32.bf16 $0, hl;\n" + " cvt.f32.bf16 $1, hh; }", + "=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + f_lo = Float32(llvm.extractvalue(T.f32(), ret, [0])) + f_hi = Float32(llvm.extractvalue(T.f32(), ret, [1])) + return f_lo, f_hi + + +@dsl_user_op +def cvt_f16x2_to_2xf32( + f16x2_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32)>"), + [f16x2_b32], + "{ .reg .b16 hl, hh;\n" + " mov.b32 {hl, hh}, $2;\n" + " cvt.f32.f16 $0, hl;\n" + " cvt.f32.f16 $1, hh; }", + "=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + f_lo = Float32(llvm.extractvalue(T.f32(), ret, [0])) + f_hi = Float32(llvm.extractvalue(T.f32(), ret, [1])) + return f_lo, f_hi + + +@dsl_user_op +def cvt_2xf32_to_bf16x2( + f_lo: Float32, + f_hi: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f_lo).ir_value(loc=loc, ip=ip), + Float32(f_hi).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_2xf32_to_f16x2( + f_lo: Float32, + f_hi: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f_lo).ir_value(loc=loc, ip=ip), + Float32(f_hi).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.f16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def add_rn_f32( + a: Float32, + b: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def add_rn_f32x2_pair( + a_lo: Float32, + a_hi: Float32, + b_lo: Float32, + b_hi: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32)>"), + [ + Float32(a_lo).ir_value(loc=loc, ip=ip), + Float32(a_hi).ir_value(loc=loc, ip=ip), + Float32(b_lo).ir_value(loc=loc, ip=ip), + Float32(b_hi).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b64 ab, bb, db;\n" + " mov.b64 ab, {$2, $3};\n" + " mov.b64 bb, {$4, $5};\n" + " add.rn.f32x2 db, ab, bb;\n" + " mov.b64 {$0, $1}, db; }", + "=f,=f,f,f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + sum_lo = Float32(llvm.extractvalue(T.f32(), ret, [0])) + sum_hi = Float32(llvm.extractvalue(T.f32(), ret, [1])) + return sum_lo, sum_hi + + +@dsl_user_op +def cvt_e4m3x4_to_4xf32( + e4m3x4_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32, Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [e4m3x4_b32], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " .reg .b32 f16x2_lo, f16x2_hi;\n" + " .reg .b16 h0, h1, h2, h3;\n" + " mov.b32 {e8x2_lo, e8x2_hi}, $4;\n" + " cvt.rn.f16x2.e4m3x2 f16x2_lo, e8x2_lo;\n" + " cvt.rn.f16x2.e4m3x2 f16x2_hi, e8x2_hi;\n" + " mov.b32 {h0, h1}, f16x2_lo;\n" + " mov.b32 {h2, h3}, f16x2_hi;\n" + " cvt.f32.f16 $0, h0;\n" + " cvt.f32.f16 $1, h1;\n" + " cvt.f32.f16 $2, h2;\n" + " cvt.f32.f16 $3, h3; }", + "=f,=f,=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return tuple(Float32(llvm.extractvalue(T.f32(), ret, [i])) for i in range(4)) + + +@dsl_user_op +def cvt_e5m2x4_to_4xf32( + e5m2x4_b32: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, Float32, Float32, Float32]: + ret = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [e5m2x4_b32], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " .reg .b32 f16x2_lo, f16x2_hi;\n" + " .reg .b16 h0, h1, h2, h3;\n" + " mov.b32 {e8x2_lo, e8x2_hi}, $4;\n" + " cvt.rn.f16x2.e5m2x2 f16x2_lo, e8x2_lo;\n" + " cvt.rn.f16x2.e5m2x2 f16x2_hi, e8x2_hi;\n" + " mov.b32 {h0, h1}, f16x2_lo;\n" + " mov.b32 {h2, h3}, f16x2_hi;\n" + " cvt.f32.f16 $0, h0;\n" + " cvt.f32.f16 $1, h1;\n" + " cvt.f32.f16 $2, h2;\n" + " cvt.f32.f16 $3, h3; }", + "=f,=f,=f,=f,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return tuple(Float32(llvm.extractvalue(T.f32(), ret, [i])) for i in range(4)) + + +@dsl_user_op +def cvt_4xf32_to_e4m3x4( + f0: Float32, + f1: Float32, + f2: Float32, + f3: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f0).ir_value(loc=loc, ip=ip), + Float32(f1).ir_value(loc=loc, ip=ip), + Float32(f2).ir_value(loc=loc, ip=ip), + Float32(f3).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " cvt.rn.satfinite.e4m3x2.f32 e8x2_lo, $2, $1;\n" + " cvt.rn.satfinite.e4m3x2.f32 e8x2_hi, $4, $3;\n" + " mov.b32 $0, {e8x2_lo, e8x2_hi}; }", + "=r,f,f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_4xf32_to_e5m2x4( + f0: Float32, + f1: Float32, + f2: Float32, + f3: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.inline_asm( + T.i32(), + [ + Float32(f0).ir_value(loc=loc, ip=ip), + Float32(f1).ir_value(loc=loc, ip=ip), + Float32(f2).ir_value(loc=loc, ip=ip), + Float32(f3).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b16 e8x2_lo, e8x2_hi;\n" + " cvt.rn.satfinite.e5m2x2.f32 e8x2_lo, $2, $1;\n" + " cvt.rn.satfinite.e5m2x2.f32 e8x2_hi, $4, $3;\n" + " mov.b32 $0, {e8x2_lo, e8x2_hi}; }", + "=r,f,f,f,f", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cvt_b32_to_1xf32( + b32_in: Int32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32]: + return (Float32(llvm.bitcast(T.f32(), b32_in, loc=loc, ip=ip)),) + + +@dsl_user_op +def cvt_1xf32_to_b32( + f: Float32, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + return llvm.bitcast(T.i32(), Float32(f).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + + +@dsl_user_op +def elem_wise_add_f32( + a: Tuple[Float32, ...], + b: Tuple[Float32, ...], + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Tuple[Float32, ...]: + assert len(a) == len(b), f"Mismatched tuple sizes: {len(a)} vs {len(b)}" + n = len(a) + if n == 1: + return (add_rn_f32(a[0], b[0], loc=loc, ip=ip),) + assert n % 2 == 0, f"elem_wise_add_f32: tuple length {n} must be 1 or even" + # Pair consecutive f32 values into FADD2 invocations. + out: Tuple[Float32, ...] = () + for i in range(0, n, 2): + s_lo, s_hi = add_rn_f32x2_pair(a[i], a[i + 1], b[i], b[i + 1], loc=loc, ip=ip) + out = out + (s_lo, s_hi) + return out + + +@cute.jit +def LDxSTMC( + kernel, + cur_tile_coord, + tiled_mma: cute.TiledMma, + comm_tidx: cutlass.Int32, + tCgCommInUC_per_peer: Tuple[cute.Tensor, ...], + tCgCommOutMC: cute.Tensor, +) -> None: + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Per-peer views share layout with the local UC; only the base pointer differs. + tCgCommInUC_slice_per_peer = tuple( + t[((None, None), 0, 0, *mma_tile_coord_mnl)] for t in tCgCommInUC_per_peer + ) + tCgCommOutMC_slice = tCgCommOutMC[((None, None), 0, 0, *mma_tile_coord_mnl)] + + cta_mma_tile_m = kernel.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + cta_tile_n = kernel.mma_tiler[1] + total_comm_threads = len(kernel.comm_warp_id) * cute.arch.WARP_SIZE + + # Split this rank's 1/P chunk along the contiguous axis so each rank's + # slice stays contiguous: split M for row-major C, split N for col-major C. + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + local_tile_m = int(cta_mma_tile_m / kernel.num_ranks) + local_tile_n = cta_tile_n + rank_slice_coord = (kernel.rank_id, 0) + else: + local_tile_m = cta_mma_tile_m + local_tile_n = int(cta_tile_n / kernel.num_ranks) + rank_slice_coord = (0, kernel.rank_id) + div_shape = (local_tile_m, local_tile_n) + + def _partition_to_local_rank(t): + return cute.slice_( + cute.zipped_divide(t, div_shape), ((None, None), rank_slice_coord) + ) + + tCgCommInUC_local_rank_per_peer = tuple( + _partition_to_local_rank(s) for s in tCgCommInUC_slice_per_peer + ) + tCgCommOutMC_local_rank = _partition_to_local_rank(tCgCommOutMC_slice) + + # _pick_num_comm_warp_for_128b sizes the comm warp group so each thread + # owns exactly one 128-bit atom per iteration. The assert guards that + # invariant; violation means the picker is wrong, not the input. + elements_per_thread = (local_tile_m * local_tile_n) // total_comm_threads + atom_num_elements_128 = 128 // kernel.c_dtype.width + atom_num_elements = atom_num_elements_128 + assert cutlass.const_expr(elements_per_thread % atom_num_elements_128 == 0), ( + f"[LDxSTMC] 128-bit invariant violated: elements_per_thread=" + f"{elements_per_thread} not divisible by atom={atom_num_elements_128}" + ) + + # Spread threads along the contiguous axis. GCD handles non-power-of-2 + # tile widths (e.g. cta_n=192) that don't naively divide thread count. + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + max_thr_n = local_tile_n // atom_num_elements + atom_thr_n = gcd(max_thr_n, total_comm_threads) + atom_thr_m = total_comm_threads // atom_thr_n + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) + val_layout = cute.make_layout( + (1, atom_num_elements), stride=(atom_num_elements, 1) + ) + else: + max_thr_m = local_tile_m // atom_num_elements + atom_thr_m = gcd(max_thr_m, total_comm_threads) + atom_thr_n = total_comm_threads // atom_thr_m + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(1, atom_thr_m)) + val_layout = cute.make_layout( + (atom_num_elements, 1), stride=(1, atom_num_elements) + ) + + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), kernel.c_dtype) + tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + thr_copy_fake = tiled_copy_fake.get_slice(comm_tidx - kernel.comm_warp_id[0] * 32) + + tTR_gCommInUC_per_peer = tuple( + thr_copy_fake.partition_S(t) for t in tCgCommInUC_local_rank_per_peer + ) + tTR_gCommOutMC = thr_copy_fake.partition_S(tCgCommOutMC_local_rank) + _atom, loop_m, loop_n = tTR_gCommOutMC.shape + + # cvt_to_xf32: i32 → tuple of n_per_b32 f32 ; cvt_back: inverse. + if cutlass.const_expr(kernel.c_dtype is cutlass.BFloat16): + cvt_to_xf32, cvt_back, n_per_b32 = cvt_bf16x2_to_2xf32, cvt_2xf32_to_bf16x2, 2 + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float16): + cvt_to_xf32, cvt_back, n_per_b32 = cvt_f16x2_to_2xf32, cvt_2xf32_to_f16x2, 2 + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float8E4M3FN): + cvt_to_xf32, cvt_back, n_per_b32 = cvt_e4m3x4_to_4xf32, cvt_4xf32_to_e4m3x4, 4 + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float8E5M2): + cvt_to_xf32, cvt_back, n_per_b32 = cvt_e5m2x4_to_4xf32, cvt_4xf32_to_e5m2x4, 4 + elif cutlass.const_expr(kernel.c_dtype is cutlass.Float32): + cvt_to_xf32, cvt_back, n_per_b32 = cvt_b32_to_1xf32, cvt_1xf32_to_b32, 1 + else: + raise ValueError(f"LDxSTMC: unsupported c_dtype {kernel.c_dtype}") + + # Accumulate in f32 across all peers; cvt_back only once at the end so + # partial sums keep f32 precision (matches multimem.ld_reduce semantics). + num_peers = len(tTR_gCommInUC_per_peer) + n_atom = 4 # 128 bits / 32 bits per i32 + + # Outer (loop_m, loop_n) is `unroll=1` so the DSL emits one shared body + # instead of loop_m * loop_n copies — keeps the (num_peers * n_atom * + # n_per_b32) live f32 register footprint from spilling. Inner peer/atom + # loops stay constexpr (short, benefit from unrolling for ILP). + for i in cutlass.range(loop_m, unroll=1): + for j in cutlass.range(loop_n, unroll=1): + # Explicit `.append` because the AST preprocessor doesn't unroll + # range_constexpr inside generator expressions. + peer_loaded = [] + for r in cutlass.range_constexpr(num_peers): + peer_loaded.append( + ld_weak_global_4xu32(tTR_gCommInUC_per_peer[r][None, i, j].iterator) + ) + + acc = () + for r in cutlass.range_constexpr(num_peers): + p_acc = () + for k in cutlass.range_constexpr(n_atom): + p_acc = p_acc + cvt_to_xf32(peer_loaded[r][k]) + acc = p_acc if r == 0 else elem_wise_add_f32(acc, p_acc) + + out_regs = [] + for k in cutlass.range_constexpr(n_atom): + chunk = acc[k * n_per_b32 : (k + 1) * n_per_b32] + out_regs.append(cvt_back(*chunk)) + + utils.distributed.multimem_st( + tTR_gCommOutMC[None, i, j].iterator, *out_regs + ) + + +class Sm100PersistentDenseGemmAllReduceLDxSTMCKernel: + """ + **Example:** + gemm = Sm100PersistentDenseGemmAllReduceLDxSTMCKernel( + acc_dtype=cutlass.Float32, + c_dtype=cutlass.BFloat16, + use_2cta_instrs=True, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + use_tma_store=True, + rank_id=rank_id, + num_ranks=num_ranks, + all_reduce="LDxSTMC", + swizzle_size=1, + raster_order="m", + ) + compiled_gemm = cute.compile( + gemm, + a=a, b=b, c=c, + comm_in_unicast_tensors_per_peer=comm_in_uc_per_peer, + comm_out_multicast_tensor=comm_out_mc, + barrier_flag_unicast=flag_unicast, + barrier_flag_multicast=flag_multicast, + stream=stream, + max_active_clusters=max_active_clusters, + ) + compiled_gemm( + a=a, b=b, c=c, + comm_in_unicast_tensors_per_peer=comm_in_uc_per_peer, + comm_out_multicast_tensor=comm_out_mc, + barrier_flag_unicast=flag_unicast, + barrier_flag_multicast=flag_multicast, + stream=stream, + ) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_tma_store: bool, + rank_id: int, + num_ranks: int, + all_reduce: str = "LDxSTMC", + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + ): + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.c_dtype: Type[cutlass.Numeric] = c_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.swizzle_size = swizzle_size + self.raster_order = raster_order + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + # Capture the JIT target arch so SMEM/TMEM sizing tracks the actual + # hardware (sm_100 / sm_103 / sm_107 / ...). Matches the FMHA pattern. + arch_enum = BaseDSL._get_dsl().get_arch_enum() + self.arch = f"sm_{arch_enum.major}{arch_enum.minor}" + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.all_reduce = all_reduce + if all_reduce == "LDxSTMC": + self.rank_id = rank_id + self.num_ranks = num_ranks + num_comm_warps = self._pick_num_comm_warp_for_128b(c_dtype) + self.comm_warp_id = tuple(range(6, 6 + num_comm_warps)) + else: + # No AR mode — no comm warps. + self.comm_warp_id = () + self.rank_id = 0 + self.num_ranks = 1 + + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilogue_warp_id, + *self.comm_warp_id, + ) + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.comm_sync_bar_id = 4 + + def _pick_num_comm_warp_for_128b(self, c_dtype): + """Pick the LARGEST comm-warp count W ∈ {4, 3, 2, 1} (T = W·32 ∈ + {128, 96, 64, 32}) such that the AR slab is evenly distributed + across T threads AND every thread gets a full 128 b LD/ST payload + (= ``atom_128`` elements, where ``atom_128 = 16 / sizeof(c_dtype)``). + + Iterates W = 4 → 1 and returns the first match, so the largest + viable W is returned. Smaller W is only chosen when strictly + necessary, because larger W means more parallel AR threads and + proportionally fewer per-thread PTX instructions. + """ + # cta-tile dims (the per-CTA M·N output the AR walks). 2cta halves M. + cta_mma_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + slab_elements = (cta_mma_tile_m * cta_tile_n) // self.num_ranks + atom_128 = 16 // (c_dtype.width // 8) + for num_warps in (4, 3, 2, 1): + T = num_warps * 32 + if slab_elements % T: + continue + if (slab_elements // T) % atom_128 == 0: + return num_warps + raise ValueError( + f"no comm-warp count enables 128 b PTX for c_dtype={c_dtype}, " + f"cta=({cta_mma_tile_m},{cta_tile_n}), P={self.num_ranks}" + ) + + 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 + """ + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + 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.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + 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 + + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + c_smem_layout = None + if cutlass.const_expr(self.use_tma_store): + c_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + c_smem_layout, + ) + + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage, self.arch + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + comm_in_unicast_tensors_per_peer: Tuple[cute.Tensor, ...], + comm_out_multicast_tensor: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, + stream: cuda.CUstream, + max_active_clusters: cutlass.Constexpr, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # Setup static attributes before smem/grid/tma computation. + # NOTE: self.c_dtype was set in __init__ from the explicit c_dtype + # arg. We only sanity-check here that the C tensor matches that hint. + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + if cutlass.const_expr(c.element_type != self.c_dtype): + raise TypeError( + f"C tensor element_type ({c.element_type}) does not match " + f"the c_dtype declared at __init__ ({self.c_dtype})" + ) + 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.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}") + + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.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 = utils.sm100.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 C + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + self.swizzle_size, + self.raster_order, + max_active_clusters, + ) + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma=tiled_mma, + tma_atom_a=tma_atom_a, + mA_mkl=tma_tensor_a, + tma_atom_b=tma_atom_b, + mB_nkl=tma_tensor_b, + tma_atom_c=tma_atom_c, + mC_mnl=(tma_tensor_c if self.use_tma_store else c), + mCommInUC_mnl_per_peer=comm_in_unicast_tensors_per_peer, + mCommOutMC_mnl=comm_out_multicast_tensor, + barrier_flag_unicast=barrier_flag_unicast, + barrier_flag_multicast=barrier_flag_multicast, + cluster_layout_vmnk=self.cluster_layout_vmnk, + a_smem_layout_staged=self.a_smem_layout_staged, + b_smem_layout_staged=self.b_smem_layout_staged, + c_smem_layout_staged=self.c_smem_layout_staged, + epi_tile=self.epi_tile, + tile_sched_params=self.tile_sched_params, + epilogue_op=epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + 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_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + mCommInUC_mnl_per_peer: Tuple[cute.Tensor, ...], + mCommOutMC_mnl: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + 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) + if cutlass.const_expr(self.use_tma_store): + 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) + + # 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_producer, ab_consumer = 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, + defer_sync=True, + ).make_participants() + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_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, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + 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) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_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) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # 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 + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + 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_tile = prefetch_k_tile_cnt + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_producer.tail() + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + 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_tile = 0 + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm( + tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc + ) + + # Async arrive AB buffer empty + handle.release() + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # + # 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) + + sC = None + if cutlass.const_expr(self.use_tma_store): + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + m_cta_tiles_in_total = gC_mnl.shape[2] * cute.size(tiled_mma.thr_id.shape) + n_cta_tiles_in_total = gC_mnl.shape[3] + + # + # Persistent tile scheduling loop for epilogue + # + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + + epilogue_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + + 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], + ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + + # Post-epilogue: arrive flag for distributed communication + if cutlass.const_expr(self.all_reduce == "LDxSTMC"): + # 1D linear index of current output tile + cta_tile_id_linear = cute.Int32( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + + # Wait for C store complete + # Unlike regular epilogue where we only wait C store complete once at end of each kernel. + # Here we need to wait for C store complete for each output tile before we set the release flag. + c_pipeline.producer_tail() + + # Update flag with release semantic with GPU scope + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = ( + barrier_flag_multicast.iterator + cta_tile_id_linear + ) + # Release flag MUST be `release` semantic + # Refer to CUDA PTX doc Memory Consistency Model for more detail + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + order="release", + scope="gpu", + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + # Post-epilogue: arrive flag for distributed communication + if cutlass.const_expr(self.all_reduce == "LDxSTMC"): + # 1D linear index of current output tile + cta_tile_id_linear = cute.Int32( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + + # Sync all thread in epilogue warp + # Ensure all thread in epilogue warp have issue store operation + epilogue_sync_barrier.arrive_and_wait() + + # Update flag with release semantic with GPU scope + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = ( + barrier_flag_multicast.iterator + cta_tile_id_linear + ) + # Release flag MUST be `release` semantic + # Refer to CUDA PTX doc Memory Consistency Model for more detail + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + order="release", + scope="gpu", + ) + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete (only if not using LDxSTMC which does per-tile tail) + if cutlass.const_expr(self.all_reduce != "LDxSTMC"): + c_pipeline.producer_tail() + else: + # Synchronize before TMEM dealloc + if cutlass.const_expr(self.all_reduce != "LDxSTMC"): + tmem_dealloc_barrier.arrive_and_wait() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # AllReduce warps + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.all_reduce == "LDxSTMC"): + if warp_idx >= self.comm_warp_id[0]: + # Per peer: (bM, bN, RestM, RestN, RestL) with that peer's local + # UC view (direct P2P address). Same tile bounds as the MC view — + # only the base pointer differs across the tuple. Used by LDxSTMC + # to issue per-peer ld.weak.global.v4.u32 loads. + gCommInUC_mnl_per_peer = tuple( + cute.local_tile( + t, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + for t in mCommInUC_mnl_per_peer + ) + # (bM, bN, RestM, RestN, RestL) with multicast memory as comm output + gCommOutMC_mnl = cute.local_tile( + mCommOutMC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + # Per-peer thread-MMA partitioning of the input UC views. + tCgCommInUC_per_peer = tuple( + thr_mma.partition_C(g) for g in gCommInUC_mnl_per_peer + ) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) with multicast memory + tCgCommOutMC = thr_mma.partition_C(gCommOutMC_mnl) + + m_cta_tiles_in_total = gCommOutMC_mnl.shape[2] * cute.size( + tiled_mma.thr_id.shape + ) + n_cta_tiles_in_total = gCommOutMC_mnl.shape[3] + + # + # Persistent tile scheduling loop for all reduce + # + comm_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.comm_sync_bar_id, + num_threads=32 * len(self.comm_warp_id), + ) + + # AllReduce communication pattern (w/ rank4 as example): + # + # [rank0 A] [rank1 A] [rank2 A] [rank3 A] -> RS -> [rank0+1+2+3 A] [rank0+1+2+3 A] [rank0+1+2+3 A] [rank0+1+2+3 A] + # [rank0 B] [rank1 B] [rank2 B] [rank3 B] -> RS -> [rank0+1+2+3 B] [rank0+1+2+3 B] [rank0+1+2+3 B] [rank0+1+2+3 B] + # [rank0 C] [rank1 C] [rank2 C] [rank3 C] -> RS -> [rank0+1+2+3 C] [rank0+1+2+3 C] [rank0+1+2+3 C] [rank0+1+2+3 C] + # [rank0 D] [rank1 D] [rank2 D] [rank3 D] -> RS -> [rank0+1+2+3 D] [rank0+1+2+3 D] [rank0+1+2+3 D] [rank0+1+2+3 D] + # rank0 rank1 rank2 rank3 -> RS -> rank0 rank1 rank2 rank3 + # + # Each rank's AR warp only handle 1/P of overall output tiles, and relies on other rank to broadcast data to it + # rank0: compute each tile's [0:1/4] and broadcast it to all rank + # rank1: compute each tile's [1/4:2:4] and broadcast it to all rank + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + + # (1) Per-tile flag wait. + cta_tile_id_linear = ( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + if warp_idx == self.comm_warp_id[0]: + with cute.arch.elect_one(): + utils.distributed.spin_lock_atom_cas_acquire_wait( + lock_ptr=barrier_flag_unicast.iterator + + cta_tile_id_linear, + expected_val=self.num_ranks, + reset_val=0, + scope="gpu", + ) + + comm_sync_barrier.arrive_and_wait() + + # (2) Per-tile data movement (per-peer ld.weak + accumulate + multimem.st). + LDxSTMC( + self, + cur_tile_coord=cur_tile_coord, + tiled_mma=tiled_mma, + comm_tidx=tidx, + tCgCommInUC_per_peer=tCgCommInUC_per_peer, + tCgCommOutMC=tCgCommOutMC, + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Ensure all threads in allreduce warp group have complete issuing all reduce operations. + comm_sync_barrier.arrive_and_wait() + + # + # Set Per SM Flag with Release + # + # This ensure + # 1. no rank early exit while other ranks are still issuing multimem.ld_reduce + # 2. each rank's prior multimem.st have become visible to all other ranks in the system (w/ .SYS scope) + if warp_idx == self.comm_warp_id[0]: + with cute.arch.elect_one(): + # Offset to last tile flag idx + total_num_cta_tile = m_cta_tiles_in_total * n_cta_tiles_in_total + # Linear id of current SM. + sm_id_linear = ( + cute.arch.block_idx()[0] + + cute.arch.block_idx()[1] * cute.arch.grid_dim()[0] + + cute.arch.block_idx()[2] + * cute.arch.grid_dim()[0] + * cute.arch.grid_dim()[1] + ) + # Release flag with sys scope + utils.distributed.multimem_red_add1( + lock_ptr=barrier_flag_multicast.iterator + + total_num_cta_tile + + sm_id_linear, + scope="sys", + order="release", + ) + # Acquire spin-lock wait flag with sys scope + utils.distributed.spin_lock_atom_cas_acquire_wait( + lock_ptr=barrier_flag_unicast.iterator + + total_num_cta_tile + + sm_id_linear, + expected_val=self.num_ranks, + reset_val=0, + scope="sys", + ) + + @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], + c_dtype: Type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + c_smem_layout: Union[cute.Layout, None], + ) -> Tuple[int, int, int]: + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 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 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + 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 C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_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 + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + swizzle_size: int, + raster_order: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_order == "m" + ) + 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, + arch: str, + ) -> 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, arch=arch) + + return num_tmem_alloc_cols + + def check_supported_dtypes( + self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] + ) -> bool: + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if ab_dtype not in valid_ab_dtypes: + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype: {ab_dtype}" + ) + + valid_acc_dtypes = {cutlass.Float32, cutlass.Float16, cutlass.Int32} + if self.acc_dtype not in valid_acc_dtypes: + raise testing.CantImplementError( + f"[dtype] Unsupported accumulator dtype: {self.acc_dtype}" + ) + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, # Float32 accumulator supports floating point AB types only + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and AB type + if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype {ab_dtype} for accumulator {self.acc_dtype}" + ) + + # Define compatibility mapping between accumulator type and C type + acc_c_compatibility = { + cutlass.Float32: { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + cutlass.Float16: { + cutlass.BFloat16, + cutlass.Float16, + }, + cutlass.Int32: { + cutlass.BFloat16, + cutlass.Float16, + cutlass.Float32, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + } + # Check compatibility between accumulator type and C type + if c_dtype not in acc_c_compatibility[self.acc_dtype]: + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for accumulator {self.acc_dtype}" + ) + + # check if c_dtype is supported by multimem all-reduce + if cutlass.const_expr( + c_dtype + not in { + cutlass.Float16, + cutlass.Float32, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + ): + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for multimem all-reduce" + ) + + return True + + def check_mma_tiler_and_cluster_shape(self) -> bool: + # Skip invalid mma tile shape + # 1SM: MMA tile {64, 128} x {64, 128, 192, 256} + # 2SM: MMA tile {128, 256} x {64, 128, 192, 256} + valid_mma_m = [64, 128] if not self.use_2cta_instrs else [128, 256] + valid_mma_n = [64, 128, 192, 256] + if self.mma_tiler_mn[0] not in valid_mma_m: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler M: {self.mma_tiler_mn[0]}, expected {valid_mma_m} for {'1SM' if not self.use_2cta_instrs else '2SM'}" + ) + if self.mma_tiler_mn[1] not in valid_mma_n: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler N: {self.mma_tiler_mn[1]}, expected {valid_mma_n}" + ) + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape M: {self.cluster_shape_mn[0]}, must be multiple of {2 if self.use_2cta_instrs else 1} for {'2SM' if self.use_2cta_instrs else '1SM'}" + ) + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16 + or self.cluster_shape_mn[0] <= 0 + or self.cluster_shape_mn[1] <= 0 + or not is_power_of_2(self.cluster_shape_mn[0]) + or not is_power_of_2(self.cluster_shape_mn[1]) + ): + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape: {self.cluster_shape_mn}, must be positive power of 2 with total <= 16" + ) + return True + + def check_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + ) -> bool: + def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: contiguous dimension must be 16B aligned for (M={m}, K={k}, N={n})" + ) + + if m % 128 != 0 and n % 128 != 0: + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: at least one of M={m} or N={n} must be divisible by 128" + ) + + if a_major not in ["m", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid a_major: {a_major}, expected 'm' or 'k'" + ) + if b_major not in ["n", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid b_major: {b_major}, expected 'n' or 'k'" + ) + if c_major not in ["m", "n"]: + raise testing.CantImplementError( + f"[alignment] Invalid c_major: {c_major}, expected 'm' or 'n'" + ) + + return True + + def check_epilogue_store_option(self, m: int, n: int) -> bool: + # Non TMA store version does not have predication, can not support OOB tiles + cta_tile_shape_mn = ( + self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_mn[1], + ) + if not self.use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + raise testing.CantImplementError( + f"[epilogue] Problem shape (M={m}, N={n}) must be divisible by cta tile shape {cta_tile_shape_mn} for non TMA store" + ) + # CTA swizzling improves the L2 cache utilization and reduces the number of cache misses. + # Make sure the swizzle size divides the cta/cga count since non TMA epilogue don't support OOB tiles. + # Swizzle only applies to the dimension orthogonal to the raster direction. + m_per_swizzle = (m // cta_tile_shape_mn[0]) // self.cluster_shape_mn[0] + n_per_swizzle = (n // cta_tile_shape_mn[1]) // self.cluster_shape_mn[1] + if self.raster_order == "m": + # Swizzle applies to N dimension only + if n_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape n={n} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if m_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape m={m} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_n" + ) + return True + + def check_valid_rank_config(self) -> bool: + if self.num_ranks not in [2, 4, 8] and self.all_reduce != "none": + raise testing.CantImplementError( + f"[rank_config] Invalid rank config: num_ranks={self.num_ranks}, must be 2, 4, or 8 for all-reduce" + ) + return True + + def check_valid_all_reduce_mode(self) -> bool: + if self.all_reduce not in ["LDxSTMC", "none"]: + raise testing.CantImplementError( + f"[all_reduce_mode] Invalid all-reduce mode: {self.all_reduce}, expected 'LDxSTMC' or 'none'" + ) + return True + + def check_even_divisiable_by_cluster_tilesize( + self, m: int, n: int, k: int, l: int + ) -> bool: + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + + cluster_shape_m = self.cluster_shape_mn[0] + cluster_shape_n = self.cluster_shape_mn[1] + + cluster_tile_m = cta_tile_m * cluster_shape_m + cluster_tile_n = cta_tile_n * cluster_shape_n + + if m % cluster_tile_m != 0 or n % cluster_tile_n != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_cluster_tile_size] Problem shape (M={m}, N={n}) not divisible by cluster tile size ({cluster_tile_m}, {cluster_tile_n})" + ) + return True + + def check_even_divisible_by_swizzle(self, m: int, n: int) -> bool: + # Only check for LDxSTMC mode which uses multimem operations + if self.all_reduce != "LDxSTMC": + return True + + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + + # Number of clusters in each dimension + num_clusters_m = m // (cta_tile_m * self.cluster_shape_mn[0]) + num_clusters_n = n // (cta_tile_n * self.cluster_shape_mn[1]) + + if self.raster_order == "m": + # Swizzle applies to N dimension only + if num_clusters_n % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in N dimension ({num_clusters_n}) " + f"must be divisible by swizzle_size={self.swizzle_size} for LDxSTMC with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if num_clusters_m % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in M dimension ({num_clusters_m}) " + f"must be divisible by swizzle_size={self.swizzle_size} for LDxSTMC with raster_order_n" + ) + + return True + + def can_implement( + self, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + ) -> bool: + # Skip unsupported types + self.check_supported_dtypes(ab_dtype, c_dtype) + + # Skip invalid mma tile shape and cluster shape + self.check_mma_tiler_and_cluster_shape() + + # Unpack mnkl for clarity in calling the epilog check + m, n, k, l = mnkl + + # Skip if GEMM not even divisiable by tilesize + # Right now, AR Fusion **ONLY** support even divisiable by tilesize + self.check_even_divisiable_by_cluster_tilesize(m, n, k, l) + + # Skip if swizzle doesn't evenly divide the problem (for LDxSTMC communication) + self.check_even_divisible_by_swizzle(m, n) + + # Skip illegal problem shape for load/store alignment + self.check_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ) + + # Skip invalid epilogue store option + self.check_epilogue_store_option(m, n) + + # Skip invalid all-reduce rank config + self.check_valid_rank_config() + + # Skip invalid all-reduce mode + self.check_valid_all_reduce_mode() + + return True + + +##################################################################### +#### KERNEL IMPLEMENTATION END HERE #### +##################################################################### + +_DISTRIBUTED_INITIALIZED = False + + +# Check for device p2p access +def check_cuda_p2p_access(): + import torch + + num_devices = torch.cuda.device_count() + cuda_p2p_access_available = True + for i in range(num_devices - 1): + for j in range(i + 1, num_devices): + if not torch.cuda.can_device_access_peer(i, j): + cuda_p2p_access_available = False + break + if not cuda_p2p_access_available: + break + return cuda_p2p_access_available + + +# Check for multicast support +def check_multicast_support(): + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + device_count = torch.cuda.device_count() + for device_id in range(device_count): + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device_id): + return False + return True + except ImportError: + return False + + +def init_distributed(): + global _DISTRIBUTED_INITIALIZED + + if not dist.is_initialized() and not _DISTRIBUTED_INITIALIZED: + # * Check before running * + if not torch.cuda.is_available(): + print("CUDA is not available") + exit(-1) + if torch.cuda.get_device_capability()[0] != 10: + print("Running on non-sm100 device") + exit(-1) + if torch.cuda.device_count() < 2: + print("Less than two devices") + exit(-1) + if torch.cuda.device_count() % 2 != 0: + print("Odd world size") + exit(-1) + for var in ("LOCAL_RANK", "RANK", "WORLD_SIZE"): + if var not in os.environ: + print(f"{var} not in environment variables") + exit(-1) + + # LOCAL_RANK = which of this host's GPUs we bind to. + # RANK = global identity across all hosts; used for NCCL init + # and for kernel rank_id (slab ownership in the AR). + local_rank = int(os.environ["LOCAL_RANK"]) + global_rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + world_size=world_size, + rank=global_rank, + device_id=device, + ) + + _DISTRIBUTED_INITIALIZED = True + + print( + f"[init] hostname={socket.gethostname()} global_rank={global_rank} " + f"local_rank={local_rank} world_size={world_size} device={device}" + ) + + global_rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", global_rank)) + device = torch.device("cuda", local_rank) + + # Return the GLOBAL rank — callers feed this into kernel rank_id. + return global_rank, world_size, device + + +def clean_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def test_simple_torch_allreduce(device): + def _simple_torch_allreduce(): + tensor = symm_mem.empty((16384), dtype=torch.float32, device=device) + symm_mem.rendezvous(tensor, group=dist.group.WORLD) + torch.ops.symm_mem.multimem_all_reduce_( + tensor, "sum", dist.group.WORLD.group_name + ) + + try: + _simple_torch_allreduce() + except Exception: + return False + return True + + +def allocate_tensors( + *, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + world_size: int, + num_workspace: int, + device: int, + slot_init_mode: Literal["test", "benchmark"], + global_rank: int = 0, +): + """Allocate every kernel-side tensor required by the LDxSTMC AllReduce + fused GEMM in one place. + + Symmetric/multicast tensors (``comm_in``, ``comm_out``, barrier flags) + are backed by ``torch.distributed._symmetric_memory`` so the kernel can + issue per-peer ``ld.weak.global`` loads (input side) and + ``multimem.st`` broadcasts (output side) on them. Returns a dict whose + keys exactly match the LDMCxSTMC sibling for symmetry. + """ + + def _create_multicast_tensor(*, torch_tensor_ref, dtype, leading_dim): + is_col_major = leading_dim == 0 + shape = torch_tensor_ref.shape + ref_dtype = torch_tensor_ref.dtype + + if is_col_major: + reversed_shape = tuple(reversed(shape)) + perm = tuple(range(len(shape) - 1, -1, -1)) + torch_symm_base = symm_mem.empty( + reversed_shape, device=device, dtype=ref_dtype + ) + torch_tensor_uc = torch_symm_base.permute(perm) + else: + torch_symm_base = symm_mem.empty(shape, device=device, dtype=ref_dtype) + torch_tensor_uc = torch_symm_base + + torch_tensor_uc.copy_(torch_tensor_ref) + + torch_tensor_hdl = symm_mem.rendezvous(torch_symm_base, group=dist.group.WORLD) + + cute_tensor_uc = from_dlpack(torch_tensor_uc, assumed_align=16) + cute_tensor_uc.element_type = dtype + cute_tensor_uc = cute_tensor_uc.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor_uc = cutlass_torch.convert_cute_tensor( + torch_tensor_uc, + cute_tensor_uc, + dtype, + is_dynamic_layout=True, + ) + + # Multicast view: multimem.st on cute_tensor_mc fans out to every peer. + if is_col_major: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, reversed_shape, ref_dtype + ).permute(perm) + else: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, shape, ref_dtype + ) + cute_tensor_mc = from_dlpack(mc_torch, assumed_align=16) + cute_tensor_mc.element_type = dtype + cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim) + + # Per-peer UC views — one per rank in the world team. Each peer view + # has identical shape/strides to the local UC; only the base pointer + # differs (via symm_mem buffer_ptrs[r]). LDxSTMC issues one + # ld.weak.global.v4.u32 per peer per atom against this tuple. + peer_torch_anchors = [] + cute_tensors_uc_per_peer_list = [] + for r in range(world_size): + peer_ptr = torch_tensor_hdl.buffer_ptrs[r] + if is_col_major: + peer_torch_base = cutlass_torch.as_tensor( + peer_ptr, reversed_shape, ref_dtype + ) + peer_torch = peer_torch_base.permute(perm) + else: + peer_torch_base = cutlass_torch.as_tensor(peer_ptr, shape, ref_dtype) + peer_torch = peer_torch_base + peer_cute = from_dlpack(peer_torch, assumed_align=16) + peer_cute.element_type = dtype + peer_cute = peer_cute.mark_layout_dynamic(leading_dim=leading_dim) + peer_cute = cutlass_torch.convert_cute_tensor( + peer_torch, peer_cute, dtype, is_dynamic_layout=True + ) + peer_torch_anchors.append(peer_torch_base) + cute_tensors_uc_per_peer_list.append(peer_cute) + cute_tensors_uc_per_peer = tuple(cute_tensors_uc_per_peer_list) + + return ( + cute_tensor_uc, + cute_tensor_mc, + cute_tensors_uc_per_peer, + torch_tensor_uc, + mc_torch, + tuple(peer_torch_anchors), + ) + + torch.manual_seed(42 + global_rank) + + m, n, k, l = mnkl + + _init_lo, _init_hi = (-1, 1) if slot_init_mode == "test" else (-20, 20) + + torch_tensor_a_int_cpu = cutlass_torch.matrix( + l, + m, + k, + a_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + torch_tensor_b_int_cpu = cutlass_torch.matrix( + l, + n, + k, + b_major == "n", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + torch_tensor_c_int_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + torch_tensor_comm_out_c_int_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), + ) + + torch_tensor_a_cpu = torch_tensor_a_int_cpu.to(dtype=torch_dtype(ab_dtype)) + torch_tensor_b_cpu = torch_tensor_b_int_cpu.to(dtype=torch_dtype(ab_dtype)) + torch_tensor_c_cpu = torch_tensor_c_int_cpu.to(dtype=torch_dtype(c_dtype)) + torch_tensor_comm_out_c_cpu = torch_tensor_comm_out_c_int_cpu.to( + dtype=torch_dtype(c_dtype) + ) + + cute_tensor_a_list = [] + cute_tensor_b_list = [] + torch_tensor_a_gpu_list = [] + torch_tensor_b_gpu_list = [] + for _ in range(num_workspace): + cute_tensor_a_i, torch_tensor_a_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_a_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_b_i, torch_tensor_b_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_b_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_a_list.append(cute_tensor_a_i) + cute_tensor_b_list.append(cute_tensor_b_i) + torch_tensor_a_gpu_list.append(torch_tensor_a_gpu_i) + torch_tensor_b_gpu_list.append(torch_tensor_b_gpu_i) + + assert not torch.isnan(torch_tensor_a_gpu_list[0]).any() + assert not torch.isnan(torch_tensor_b_gpu_list[0]).any() + + ( + cute_tensor_c, + cute_tensor_comm_in_mc, + cute_tensor_comm_in_uc_per_peer, + torch_tensor_c_gpu, + _torch_tensor_c_mc_anchor, + _torch_tensor_c_peer_anchors, + ) = _create_multicast_tensor( + torch_tensor_ref=torch_tensor_c_cpu, + dtype=c_dtype, + leading_dim=(1 if c_major == "n" else 0), + ) + ( + _cute_tensor_comm_out_uc, + cute_tensor_comm_out_mc, + _cute_tensor_comm_out_uc_per_peer, + torch_tensor_comm_out_gpu, + _torch_tensor_comm_out_mc_anchor, + _torch_tensor_comm_out_peer_anchors, + ) = _create_multicast_tensor( + torch_tensor_ref=torch_tensor_comm_out_c_cpu, + dtype=c_dtype, + leading_dim=(1 if c_major == "n" else 0), + ) + + # Flag tensor sized for the worst-case kernel config: smallest cta_tile is + # 64x64, so max #tiles = (m/64) * (n/64). Pad by 160 for per-SM slots + # (B200 has 148; 160 is a safe upper bound). The kernel resets these flags + # internally on each launch, so a single allocation can be reused across + # (mma_tiler, use_2cta) candidates in a benchmark sweep. + num_flags = (m // 64) * (n // 64) + 160 + + torch_tensor_flag_unicast = symm_mem.empty( + (num_flags,), dtype=torch.int32, device=device + ) + torch_tensor_flag_unicast.fill_(0) + torch_flag_hdl = symm_mem.rendezvous( + torch_tensor_flag_unicast, group=dist.group.WORLD + ) + torch_tensor_flag_multicast = cutlass_torch.as_tensor( + torch_flag_hdl.multicast_ptr, + torch_tensor_flag_unicast.shape, + torch_tensor_flag_unicast.dtype, + ) + cute_tensor_flag_unicast = from_dlpack(torch_tensor_flag_unicast, assumed_align=16) + cute_tensor_flag_unicast.element_type = cutlass.Int32 + cute_tensor_flag_unicast = cute_tensor_flag_unicast.mark_layout_dynamic() + cute_tensor_flag_multicast = from_dlpack( + torch_tensor_flag_multicast, assumed_align=16 + ) + cute_tensor_flag_multicast.element_type = cutlass.Int32 + cute_tensor_flag_multicast = cute_tensor_flag_multicast.mark_layout_dynamic() + + return { + "cute_tensor_a_list": cute_tensor_a_list, + "cute_tensor_b_list": cute_tensor_b_list, + "cute_tensor_c": cute_tensor_c, + "cute_tensor_comm_in_mc": cute_tensor_comm_in_mc, + "cute_tensor_comm_in_uc_per_peer": cute_tensor_comm_in_uc_per_peer, + "cute_tensor_comm_out_mc": cute_tensor_comm_out_mc, + "cute_tensor_flag_unicast": cute_tensor_flag_unicast, + "cute_tensor_flag_multicast": cute_tensor_flag_multicast, + "torch_tensor_a_cpu": torch_tensor_a_cpu, + "torch_tensor_b_cpu": torch_tensor_b_cpu, + "torch_tensor_c_cpu": torch_tensor_c_cpu, + "torch_tensor_c_gpu": torch_tensor_c_gpu, + "torch_tensor_comm_out_gpu": torch_tensor_comm_out_gpu, + "torch_tensor_flag_unicast": torch_tensor_flag_unicast, + "_anchors": ( + torch_tensor_a_gpu_list, + torch_tensor_b_gpu_list, + torch_tensor_c_gpu, + torch_tensor_comm_out_gpu, + torch_tensor_flag_unicast, + torch_tensor_flag_multicast, + torch_flag_hdl, + _torch_tensor_c_mc_anchor, + _torch_tensor_c_peer_anchors, + _torch_tensor_comm_out_mc_anchor, + _torch_tensor_comm_out_peer_anchors, + ), + } + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + warmup_iterations: int = 0, + iterations: int = 1, + all_reduce: str = "LDxSTMC", + benchmark_or_test: str = "test", + use_cuda_graph: bool = False, + csv_report: str = "./csv_report.csv", +): + # * Init Distributed * + global_rank, world_size, device = init_distributed() + + # * Check based on AR mode * + if all_reduce == "LDxSTMC": + if not check_cuda_p2p_access(): + print("CUDA P2P access not available") + exit(-1) + if not check_multicast_support(): + print("Multicast support not available") + exit(-1) + if not test_simple_torch_allreduce(device): + print("Simple torch allreduce test failed") + exit(-1) + + # * Verbose * + def _verbose(): + print("Running Blackwell SM100 Persistent Dense GEMM AR test with:") + print(f"mnkl: {mnkl}") + print(f"global_rank: {global_rank}, world_size: {world_size}, device: {device}") + print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}") + print(f"Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") + print(f"Swizzle Size: {swizzle_size}") + print(f"Raster Order: {raster_order}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Use CUDA Graph: {'True' if use_cuda_graph else 'False'}") + print(f"Mode: {benchmark_or_test}") + print(f"Fused AllReduce Op: {all_reduce}") + + if global_rank == 0: + _verbose() + + # * Run Distributed * + if benchmark_or_test in ("benchmark", "benchmark_all"): + run_distributed_benchmark( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + warmup_iterations=warmup_iterations, + iterations=iterations, + all_reduce=all_reduce, + use_cuda_graph=use_cuda_graph, + benchmark_all=(benchmark_or_test == "benchmark_all"), + csv_report=csv_report, + ) + else: + run_distributed_test( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + all_reduce=all_reduce, + ) + + clean_distributed() + + +def run_distributed_test( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + all_reduce: str = "LDxSTMC", +): + # * CUDA * + # Non-default stream as we'll use CUDA graphs + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # * Build GEMM AllReduce Object * + gemm_all_reduce = Sm100PersistentDenseGemmAllReduceLDxSTMCKernel( + acc_dtype, + c_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + rank_id=global_rank, + num_ranks=world_size, + all_reduce=all_reduce, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + + # * Set Filename & Funcname for different rank * + # Use hash to keep filename short and avoid "File name too long" error + try: + import hashlib + + param_str = f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}_{'2Sm' if use_2cta_instrs else '1Sm'}_{'EpiTma' if use_tma_store else 'EpiStg'}" + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_ar_{param_hash}_{all_reduce}_r{global_rank}_w{world_size}" + ) + except Exception: + unique_name = f"sm100_gemm_ar_{all_reduce}_r{global_rank}_w{world_size}" + gemm_all_reduce.__call__.set_name_prefix(unique_name) + + # * can_implement * + can_implement = gemm_all_reduce.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, all_reduce={all_reduce}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + # * Allocate tensors * + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + world_size=world_size, + num_workspace=1, + device=device, + slot_init_mode="test", + global_rank=global_rank, + ) + cute_tensor_a = tensors["cute_tensor_a_list"][0] + cute_tensor_b = tensors["cute_tensor_b_list"][0] + cute_tensor_c = tensors["cute_tensor_c"] + cute_tensor_comm_in_uc_per_peer = tensors["cute_tensor_comm_in_uc_per_peer"] + cute_tensor_comm_out_mc = tensors["cute_tensor_comm_out_mc"] + cute_tensor_flag_unicast = tensors["cute_tensor_flag_unicast"] + cute_tensor_flag_multicast = tensors["cute_tensor_flag_multicast"] + torch_tensor_a_cpu = tensors["torch_tensor_a_cpu"] + torch_tensor_b_cpu = tensors["torch_tensor_b_cpu"] + torch_tensor_c = tensors["torch_tensor_c_gpu"] + torch_tensor_comm_out = tensors["torch_tensor_comm_out_gpu"] + torch_tensor_flag_unicast = tensors["torch_tensor_flag_unicast"] + + # * Compile * + compiled_gemm = cute.compile( + gemm_all_reduce, + a=cute_tensor_a, + b=cute_tensor_b, + c=cute_tensor_c, + comm_in_unicast_tensors_per_peer=cute_tensor_comm_in_uc_per_peer, + comm_out_multicast_tensor=cute_tensor_comm_out_mc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + stream=stream, + max_active_clusters=utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ), + ) + + # * Run * + compiled_gemm( + a=cute_tensor_a, + b=cute_tensor_b, + c=cute_tensor_c, + comm_in_unicast_tensors_per_peer=cute_tensor_comm_in_uc_per_peer, + comm_out_multicast_tensor=cute_tensor_comm_out_mc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + stream=stream, + ) + + # * Compare * + # FP8 nan to saturate + def _fp8_nan_to_saturate(tensor_fp8: torch.Tensor) -> torch.Tensor: + if tensor_fp8.dtype == torch.float8_e4m3fn: + # Map to u8 raw tensor for easier byte handle + tensor_u8 = tensor_fp8.view(torch.uint8) + # pos & neg nan + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + # Map pos nan to max, neg nan to min + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 448 + tensor_f32[is_neg_nan] = -448 + # Cast back to f8 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + elif tensor_fp8.dtype == torch.float8_e5m2: + # Map to u8 raw tensor for easier byte handle + tensor_u8 = tensor_fp8.view(torch.uint8) + # pos & neg nan + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + # Map pos nan to max, neg nan to min + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 57344 + tensor_f32[is_neg_nan] = -57344 + # Cast back to f8 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + else: + # Not supported dtype + tensor_f8_out = tensor_fp8 + return tensor_f8_out + + def _compare_gemm( + torch_tensor_a_cpu: torch.Tensor, + torch_tensor_b_cpu: torch.Tensor, + torch_tensor_c: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + # GEMM Ref + gemm_ref_f32 = torch.einsum( + "mkl,nkl->mnl", + torch_tensor_a_cpu.to(dtype=torch.float32), + torch_tensor_b_cpu.to(dtype=torch.float32), + ) + + # For fp8_e4m3, clamp value to [-448, 448] to avoid nan and then cast to fp8_e4m3 + if c_dtype == cutlass.Float8E4M3FN: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-448, max=448) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + # For fp8_e5m2, clamp value to [-57344, 57344] to avoid nan and then cast to fp8_e5m2 + elif c_dtype == cutlass.Float8E5M2: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-57344, max=57344) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + else: + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + + # Cast ref to fp32 for comparison + gemm_ref_f32_c_f32 = gemm_ref_f32_c.to(dtype=torch.float32) + + # Cast kernel output to fp32 for comparison + gemm_out_f32 = torch_tensor_c.cpu().to(dtype=torch.float32) + + # Compare + torch.testing.assert_close( + gemm_out_f32, + gemm_ref_f32_c_f32, + atol=atol, + rtol=rtol, + ) + + def _compare_comm( + torch_tensor_comm_in: torch.Tensor, + torch_tensor_comm_out: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + comm_ref = torch_tensor_comm_in + + # Sync across all rank + dist.barrier() + + # Reference AR — NCCL dist.all_reduce for all supported dtypes. + if c_dtype in [cutlass.Float32, cutlass.BFloat16, cutlass.Float16]: + comm_ref = comm_ref.contiguous() + dist.all_reduce(comm_ref, op=torch.distributed.ReduceOp.SUM) + comm_ref_f32 = comm_ref.to(dtype=torch.float32).cpu() + # For fp8, cast to higher precision, do AR, then cast back to fp8 and handle nan + elif c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_ref_fp32_acc = comm_ref.to(dtype=torch.float32).contiguous() + dist.all_reduce(comm_ref_fp32_acc, op=torch.distributed.ReduceOp.SUM) + # Add small epsilon to mimic multimem.ld_reduce round up behavior where torch cast round down + epsilon = 1e-2 + comm_ref_f32_acc_epsilon = comm_ref_fp32_acc + epsilon + comm_ref_c = comm_ref_f32_acc_epsilon.to(dtype=torch_dtype(c_dtype)) + comm_ref_nan2sat = _fp8_nan_to_saturate(comm_ref_c) + comm_ref_f32 = comm_ref_nan2sat.to(dtype=torch.float32).cpu() + + # Fusion Kernel Result + # multimem.ld_reduce on fp8 will return nan if val is out of range + # Need to handle nan the same way as comm ref + comm_out = torch_tensor_comm_out.cpu() + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + # Handle nan + comm_out_nan2sat = _fp8_nan_to_saturate(comm_out) + # Cast to fp32 for comparison + comm_out_f32 = comm_out_nan2sat.to(dtype=torch.float32) + else: + # Cast to fp32 for comparison + comm_out_f32 = comm_out.to(dtype=torch.float32) + + # Compare + torch.testing.assert_close( + comm_ref_f32, + comm_out_f32, + atol=atol, + rtol=rtol, + ) + + # compare + # Impose a strick check on GEMM as input is uniform random int + gemm_atol = 1e-05 + gemm_rtol = 1e-05 + _compare_gemm( + torch_tensor_a_cpu=torch_tensor_a_cpu, + torch_tensor_b_cpu=torch_tensor_b_cpu, + torch_tensor_c=torch_tensor_c, + c_dtype=c_dtype, + atol=gemm_atol, + rtol=gemm_rtol, + ) + + if all_reduce == "LDxSTMC": + # Relaxed tolerance — reference uses NCCL dist.all_reduce, which + # reduces in a different order than the fused multimem kernel. + comm_atol = 5 + comm_rtol = 1e-1 + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_rtol = 1e-1 + comm_atol = 0.0 + _compare_comm( + torch_tensor_comm_in=torch_tensor_c, + torch_tensor_comm_out=torch_tensor_comm_out, + c_dtype=c_dtype, + atol=comm_atol, + rtol=comm_rtol, + ) + + # * Check flags are reset * + # The kernel uses the flag tensor for cross-CTA / cross-rank synchronization + # and is expected to leave every entry at 0 once it returns. If any entry is + # non-zero, a tile didn't drain its slot — re-running the kernel would block + # forever on that stale flag. + torch.cuda.synchronize() + flag_after = torch_tensor_flag_unicast.cpu() + nonzero_count = int(torch.count_nonzero(flag_after).item()) + assert nonzero_count == 0, ( + f"Kernel did not reset all barrier flags to zero: " + f"{nonzero_count}/{flag_after.numel()} entries non-zero " + f"(first 16 values: {flag_after[:16].tolist()})" + ) + + +def compile_and_benchmark( + *, + mnkl, + ab_dtype, + c_dtype, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + swizzle_size, + raster_order, + use_2cta_instrs, + use_tma_store, + all_reduce, + global_rank, + world_size, + num_workspace, + tensors, + torch_stream, + stream, + warmup_iterations, + iterations, + use_cuda_graph, +): + """Instantiate one LDxSTMC kernel candidate, compile it, and time the + warmup + benchmark loop. + + Reads cute views out of ``tensors`` (produced by :func:`allocate_tensors`) + and uses the caller-supplied ``torch_stream`` / ``stream`` pair so every + candidate in a sweep shares a single CUDA stream + event pool. + Returns the mean per-iteration time in microseconds over the benchmark + iterations (warmup timings are recorded but not included in the average). + """ + cute_tensor_a_list = tensors["cute_tensor_a_list"] + cute_tensor_b_list = tensors["cute_tensor_b_list"] + cute_tensor_c = tensors["cute_tensor_c"] + cute_tensor_comm_in_uc_per_peer = tensors["cute_tensor_comm_in_uc_per_peer"] + cute_tensor_comm_out_mc = tensors["cute_tensor_comm_out_mc"] + cute_tensor_flag_unicast = tensors["cute_tensor_flag_unicast"] + cute_tensor_flag_multicast = tensors["cute_tensor_flag_multicast"] + + # Build kernel + can_implement first — before any cross-rank op, so a + # deterministic skip keeps all ranks in lockstep for the next candidate. + gemm_all_reduce = Sm100PersistentDenseGemmAllReduceLDxSTMCKernel( + acc_dtype, + c_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + rank_id=global_rank, + num_ranks=world_size, + all_reduce=all_reduce, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + try: + import hashlib + + param_str = ( + f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}" + f"_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}" + f"_{'2Sm' if use_2cta_instrs else '1Sm'}" + f"_{'EpiTma' if use_tma_store else 'EpiStg'}" + ) + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_ar_{param_hash}_{all_reduce}_r{global_rank}_w{world_size}" + ) + except Exception: + unique_name = f"sm100_gemm_ar_{all_reduce}_r{global_rank}_w{world_size}" + gemm_all_reduce.__call__.set_name_prefix(unique_name) + + can_implement = gemm_all_reduce.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, all_reduce={all_reduce}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + def make_kernel_kwargs(ws): + return dict( + a=cute_tensor_a_list[ws], + b=cute_tensor_b_list[ws], + c=cute_tensor_c, + comm_in_unicast_tensors_per_peer=cute_tensor_comm_in_uc_per_peer, + comm_out_multicast_tensor=cute_tensor_comm_out_mc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + stream=stream, + ) + + kernel_kwargs_list = [make_kernel_kwargs(i) for i in range(num_workspace)] + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + compiled_gemm = cute.compile( + gemm_all_reduce, + **kernel_kwargs_list[0], + max_active_clusters=max_active_clusters, + ) + + # Per-iteration runnable. Under --use_cuda_graph we capture one graph per + # workspace (pointers are baked in at capture time, so N_workspace graphs + # are required to rotate A/B). Without graphs we dispatch the compiled + # kernel directly with the appropriate kernel_kwargs entry. + if use_cuda_graph: + graphs = [] + for ws in range(num_workspace): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch_stream): + compiled_gemm(**kernel_kwargs_list[ws]) + graphs.append(g) + + def run_one_iter(i): + graphs[i % num_workspace].replay() + else: + + def run_one_iter(i): + compiled_gemm(**kernel_kwargs_list[i % num_workspace]) + + # One loop for warmup + benchmark. Every iteration is bracketed by its own + # event pair; only the trailing `iterations` event pairs contribute to the + # reported average so warmup transients don't skew the result. + + # Cross-rank sync before timing. torch.cuda.synchronize() drains any + # pending GPU work (graph captures land on the default stream; timing + # runs on torch_stream) so dist.barrier()'s host-side rendezvous is + # not racing in-flight kernels from the previous candidate. Without + # this, the slowest rank's leftover work would skew the measured + # window of every other rank in a benchmark_all sweep. + torch.cuda.synchronize() + dist.barrier() + + total_iters = warmup_iterations + iterations + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + with torch.cuda.stream(torch_stream): + for i in range(total_iters): + start_events[i].record() + run_one_iter(i) + end_events[i].record() + end_events[-1].synchronize() + + per_iter_us = [s.elapsed_time(e) * 1000.0 for s, e in zip(start_events, end_events)] + bench_times = per_iter_us[warmup_iterations:] + avg_us = sum(bench_times) / iterations + print( + f"Rank {global_rank} Execution time: {avg_us} us " + f"(avg of {iterations} iters, after {warmup_iterations} warmup)\n", + flush=True, + ) + return avg_us + + +def run_distributed_benchmark( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + warmup_iterations: int = 0, + iterations: int = 1, + all_reduce: str = "LDxSTMC", + use_cuda_graph: bool = False, + benchmark_all: bool = False, + csv_report: str = "./csv_report.csv", +): + # Allocate every kernel-side tensor (A/B/C/comm_out/flags) up front — + # shared across all candidates in the --benchmark_all loop. + num_workspace = 10 + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + world_size=world_size, + num_workspace=num_workspace, + device=device, + slot_init_mode="benchmark", + global_rank=global_rank, + ) + + # One non-default stream shared across every candidate. `torch_stream` is + # kept alive because we capture it into a CUDA graph; creating a fresh + # stream per candidate would churn events + graph context for no reason. + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # Build the candidate list of args for compile_and_benchmark. For a single + # run the list has one entry; --benchmark_all sweeps: + # raster_order ∈ {m, n} + # swizzle_size ∈ {1, 2, 4, 8} + # use_2cta_instrs, cluster_shape_mn, mma_tiler_mn jointly, with the + # constraint that cluster_m ∈ {2,4,8} requires 2cta (smaller tiles and + # single-cta clusters are only valid without 2cta). + base_kwargs = dict( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + use_tma_store=use_tma_store, + all_reduce=all_reduce, + global_rank=global_rank, + world_size=world_size, + num_workspace=num_workspace, + tensors=tensors, + torch_stream=torch_stream, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cuda_graph=use_cuda_graph, + ) + if benchmark_all: + raster_candidates = ["m", "n"] + swizzle_candidates = [1, 2, 4, 8] + # All cluster shapes from the allowed set; each is filtered per 2cta + # option below by the "cluster_m ∈ {2,4,8} needs 2cta" rule. + cluster_all = [ + (1, 1), + (1, 2), + (2, 1), + (2, 2), + ] + two_cta_options = [ + { + "use_2cta_instrs": True, + "clusters": [c for c in cluster_all if c[0] in (2, 4, 8)], + "mma_tilers": [ + (256, 256), + (256, 192), + (256, 128), + (256, 64), + (128, 256), + (128, 192), + (128, 128), + (128, 64), + ], + }, + { + "use_2cta_instrs": False, + "clusters": cluster_all, + "mma_tilers": [ + (128, 256), + (128, 192), + (128, 128), + (128, 64), + (64, 256), + (64, 192), + (64, 128), + (64, 64), + ], + }, + ] + candidates = [] + for opt in two_cta_options: + for cluster in opt["clusters"]: + for tile in opt["mma_tilers"]: + for r in raster_candidates: + for s in swizzle_candidates: + candidates.append( + { + **base_kwargs, + "use_2cta_instrs": opt["use_2cta_instrs"], + "cluster_shape_mn": cluster, + "mma_tiler_mn": tile, + "raster_order": r, + "swizzle_size": s, + } + ) + else: + candidates = [ + { + **base_kwargs, + "use_2cta_instrs": use_2cta_instrs, + "cluster_shape_mn": cluster_shape_mn, + "mma_tiler_mn": mma_tiler_mn, + "raster_order": raster_order, + "swizzle_size": swizzle_size, + } + ] + + if global_rank == 0: + print(f"Total number of candidate configs: {len(candidates)}", flush=True) + + # Run each candidate. Tensors (A/B/C/comm_out/flags) are reused across + # candidates; the kernel resets flags internally. Skip any combo the kernel + # can't implement — can_implement is deterministic so all ranks skip the + # same ones, keeping rendezvous collective operations lockstep. + # ~2s sleep between candidates to let the GPU cool down. + results = [] + for i, cand in enumerate(candidates): + if i > 0: + time.sleep(2) + try: + avg_us = compile_and_benchmark(**cand) + except testing.CantImplementError: + continue + results.append((cand, avg_us)) + + # Each rank writes its own CSV — one row per candidate, one config field + # per column. The path comes from --csv_report; per-rank suffix is appended + # before the extension so concurrent ranks don't clobber the same file. + global_rank = dist.get_rank() + m, n, k, l = mnkl + fieldnames = [ + "gemm_m", + "gemm_n", + "gemm_k", + "gemm_l", + "ab_dtype", + "c_dtype", + "acc_dtype", + "a_major", + "b_major", + "c_major", + "mma_tiler_m", + "mma_tiler_n", + "cluster_m", + "cluster_n", + "use_2cta_instrs", + "use_tma_store", + "all_reduce", + "use_cuda_graph", + "raster_order", + "swizzle_size", + "num_workspace", + "warmup_iterations", + "iterations", + "world_size", + "global_rank", + "benchmark_time_us", + ] + csv_path = Path(csv_report) + csv_path = csv_path.with_name( + f"{csv_path.stem}_rank{global_rank}_world{world_size}{csv_path.suffix}" + ) + csv_path.parent.mkdir(parents=True, exist_ok=True) + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for cand, avg_us in results: + writer.writerow( + { + "gemm_m": m, + "gemm_n": n, + "gemm_k": k, + "gemm_l": l, + "ab_dtype": str(ab_dtype), + "c_dtype": str(c_dtype), + "acc_dtype": str(acc_dtype), + "a_major": a_major, + "b_major": b_major, + "c_major": c_major, + "mma_tiler_m": cand["mma_tiler_mn"][0], + "mma_tiler_n": cand["mma_tiler_mn"][1], + "cluster_m": cand["cluster_shape_mn"][0], + "cluster_n": cand["cluster_shape_mn"][1], + "use_2cta_instrs": cand["use_2cta_instrs"], + "use_tma_store": use_tma_store, + "all_reduce": all_reduce, + "use_cuda_graph": use_cuda_graph, + "raster_order": cand["raster_order"], + "swizzle_size": cand["swizzle_size"], + "num_workspace": num_workspace, + "warmup_iterations": warmup_iterations, + "iterations": iterations, + "world_size": world_size, + "global_rank": global_rank, + "benchmark_time_us": avg_us, + } + ) + + return results[0][1] if len(results) == 1 else results + + +def prepare_parser(): + parser = argparse.ArgumentParser( + description="Example of SM100 Persistent Dense GEMM AR on Blackwell." + ) + + def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser.add_argument( + "--mnkl", + type=_parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=_parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=_parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--use_tma_store", action="store_true", help="Use tma store or not" + ) + parser.add_argument( + "--benchmark_or_test", + type=str, + choices=["test", "benchmark", "benchmark_all"], + default="test", + help="Select run mode: 'test' for a single-launch correctness check, " + "'benchmark' for cudaEvent timing + per-rank CSV on the given config, " + "'benchmark_all' for the same with a sweep across kernel-config variants.", + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--use_cuda_graph", + action="store_true", + default=False, + help="Capture one iteration into a CUDA graph and replay per iteration", + ) + parser.add_argument( + "--all_reduce", + choices=["LDxSTMC", "none"], + type=str, + default="LDxSTMC", + help="Allreduce algorithm to fuse with gemm", + ) + parser.add_argument( + "--swizzle_size", + type=int, + default=1, + help="Swizzling size in the unit of cluster for improving L2 cache hit rate", + ) + parser.add_argument( + "--raster_order", + type=str, + choices=["m", "n"], + default="m", + help="Rasterization order of clusters", + ) + parser.add_argument( + "--csv_report", + type=str, + default="./csv_report.csv", + help="Path for the per-rank benchmark CSV. Each rank appends " + "_rank{R}_world{W} before the extension to avoid collisions.", + ) + + return parser + + +if __name__ == "__main__": + parser = prepare_parser() + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + run( + args.mnkl, + args.ab_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.swizzle_size, + args.raster_order, + args.use_2cta_instrs, + args.use_tma_store, + args.warmup_iterations, + args.iterations, + args.all_reduce, + args.benchmark_or_test, + args.use_cuda_graph, + args.csv_report, + ) + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + print( + f"PASS (hostname={socket.gethostname()}, global_rank={rank}, local_rank={local_rank})", + flush=True, + ) diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_blockscaled_all_reduce_ldmcxstmc_blackwell.py b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_blockscaled_all_reduce_ldmcxstmc_blackwell.py new file mode 100644 index 000000000..d69eea338 --- /dev/null +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_blockscaled_all_reduce_ldmcxstmc_blackwell.py @@ -0,0 +1,3941 @@ +# Copyright (c) 2025 - 2026 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. + +import os +import argparse +import csv +import socket +import time +from pathlib import Path +from math import gcd + +from typing import Literal, Tuple, Type, Union + +import cuda.bindings.driver as cuda + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import from_dlpack, make_ptr +from cutlass.cutlass_dsl import BaseDSL + +""" +SM100 BlockScale GEMM + AllReduce (LDMCxSTMC) Fusion Kernel + +* Test (correctness only — single launch) +python3 -m torch.distributed.run --nproc-per-node=auto \ + examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_blockscaled_all_reduce_ldmcxstmc_blackwell.py \ + --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN --sf_dtype Float8E4M3FN --sf_vec_size 16 \ + --c_dtype Float16 --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --all_reduce LDMCxSTMC \ + --benchmark_or_test test + +* Benchmark One +python3 -m torch.distributed.run --nproc-per-node=auto \ + examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_blockscaled_all_reduce_ldmcxstmc_blackwell.py \ + --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN --sf_dtype Float8E4M3FN --sf_vec_size 16 \ + --c_dtype Float16 --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --swizzle_size 1 --raster_order m \ + --all_reduce LDMCxSTMC \ + --benchmark_or_test benchmark --use_cuda_graph --csv_report ./csv_report_blockscaled_ar.csv + +* Benchmark All (sweeps mma_tiler_mn x cluster_shape_mn; sf_vec_size/dtypes fixed per invocation) +python3 -m torch.distributed.run --nproc-per-node=auto \ + examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_blockscaled_all_reduce_ldmcxstmc_blackwell.py \ + --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN --sf_dtype Float8E4M3FN --sf_vec_size 16 \ + --c_dtype Float16 \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --benchmark_or_test benchmark_all --use_cuda_graph --csv_report ./csv_report_blockscaled_ar.csv + +NOTE: +* benchmark_all sweeps (mma_tiler_mn, cluster_shape_mn, raster_order, swizzle_size). use_2cta_instrs is implied by + mma_tiler_mn[0]==256. dtype/sf_vec_size/majors are FIXED per invocation. +""" + +##################################################################### +#### KERNEL IMPLEMENTATION START HERE #### +##################################################################### + + +@cute.jit +def LDMCxSTMC( + kernel, + cur_tile_coord, + tiled_mma: cute.TiledMma, + comm_tidx: cutlass.Int32, + tCgCommInMC: cute.Tensor, + tCgCommOutMC: cute.Tensor, +) -> None: + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # partition and slice at tile level + tCgCommInMC_slice = tCgCommInMC[((None, None), 0, 0, *mma_tile_coord_mnl)] + tCgCommOutMC_slice = tCgCommOutMC[((None, None), 0, 0, *mma_tile_coord_mnl)] + + # + # Partition s.t. each rank handle 1/P of this tile + # For row-major (N-contiguous): split along M to keep memory contiguous per rank + # For col-major (M-contiguous): split along N to keep memory contiguous per rank + # + cta_mma_tile_m = kernel.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + cta_tile_n = kernel.mma_tiler[1] + total_comm_threads = len(kernel.comm_warp_id) * cute.arch.WARP_SIZE + + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + # Row-major (N-contiguous): split along M + m_local_rank = int(cta_mma_tile_m / kernel.num_ranks) + tCgCommInMC_slice_partitioned = cute.zipped_divide( + tCgCommInMC_slice, (m_local_rank, cta_tile_n) + ) + tCgCommOutMC_slice_partitioned = cute.zipped_divide( + tCgCommOutMC_slice, (m_local_rank, cta_tile_n) + ) + # Data processed by local rank: select M chunk + tCgCommInMC_local_rank = cute.slice_( + tCgCommInMC_slice_partitioned, ((None, None), (kernel.rank_id, 0)) + ) + tCgCommOutMC_local_rank = cute.slice_( + tCgCommOutMC_slice_partitioned, ((None, None), (kernel.rank_id, 0)) + ) + local_tile_m = m_local_rank + local_tile_n = cta_tile_n + else: + # Col-major (M-contiguous): split along N + n_local_rank = int(cta_tile_n / kernel.num_ranks) + tCgCommInMC_slice_partitioned = cute.zipped_divide( + tCgCommInMC_slice, (cta_mma_tile_m, n_local_rank) + ) + tCgCommOutMC_slice_partitioned = cute.zipped_divide( + tCgCommOutMC_slice, (cta_mma_tile_m, n_local_rank) + ) + # Data processed by local rank: select N chunk + tCgCommInMC_local_rank = cute.slice_( + tCgCommInMC_slice_partitioned, ((None, None), (0, kernel.rank_id)) + ) + tCgCommOutMC_local_rank = cute.slice_( + tCgCommOutMC_slice_partitioned, ((None, None), (0, kernel.rank_id)) + ) + local_tile_m = cta_mma_tile_m + local_tile_n = n_local_rank + + # + # Determine vectorization length based on data size per thread + # Always use 128-bit instructions (guaranteed by _pick_num_comm_warp_for_128b) + # + elements_per_thread = (local_tile_m * local_tile_n) // total_comm_threads + + # Always use 128-bit atom size + # Guaranteed by _pick_num_comm_warp_for_128b which selects T to ensure + # elements_per_thread is divisible by 128-bit atom (16B = 128 bits) + atom_num_elements_128 = 128 // kernel.c_dtype.width + atom_num_elements = atom_num_elements_128 + + # COMPILE-TIME ASSERTION: Verify 128-bit PTX instruction usage + # _pick_num_comm_warp_for_128b selects num_warps such that every thread + # in the AllReduce warp group handles exactly atom_num_elements_128 elements, + # which means exactly one 128-bit LDMC/STMC instruction per AR slab iteration. + assert elements_per_thread % atom_num_elements_128 == 0, ( + f"KERNEL ASSERTION: 128-bit PTX instruction invariant violated. " + f"elements_per_thread={elements_per_thread} must be divisible by " + f"atom_num_elements_128={atom_num_elements_128}. " + f"This ensures every thread executes exactly one 128-bit instruction. " + f"Configuration: local_tile=({local_tile_m},{local_tile_n}), " + f"total_comm_threads={total_comm_threads}, c_dtype={kernel.c_dtype}" + ) + + # + # Create thr_copy_fake with dynamically selected atom size + # For row-major: threads spread along N (contiguous), load along N + # For col-major: threads spread along M (contiguous), load along M + # + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + # Row-major: N is contiguous, spread threads along N. + # Use GCD to handle non-power-of-2 tile widths (e.g. cta_n=192) + # that don't naively divide the thread count. + max_thr_n = local_tile_n // atom_num_elements + atom_thr_n = gcd(max_thr_n, total_comm_threads) + atom_thr_m = total_comm_threads // atom_thr_n + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) + val_layout = cute.make_layout( + (1, atom_num_elements), stride=(atom_num_elements, 1) + ) + else: + # Col-major: M is contiguous, spread threads along M. + max_thr_m = local_tile_m // atom_num_elements + atom_thr_m = gcd(max_thr_m, total_comm_threads) + atom_thr_n = total_comm_threads // atom_thr_m + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(1, atom_thr_m)) + val_layout = cute.make_layout( + (atom_num_elements, 1), stride=(1, atom_num_elements) + ) + + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), kernel.c_dtype) + tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + thr_copy_fake = tiled_copy_fake.get_slice(comm_tidx - kernel.comm_warp_id[0] * 32) + + # + # Partition per thread + # + tTR_gCommInMC = thr_copy_fake.partition_S(tCgCommInMC_local_rank) + tTR_gCommOutMC = thr_copy_fake.partition_S(tCgCommOutMC_local_rank) + _atom, loop_m, loop_n = tTR_gCommInMC.shape + + # + # All Reduce LDMCxSTMC — adaptive variant based on ldmcxstmc_default_inflight_depth: + # - ldmcxstmc_default_inflight_depth=2: pair-issue (2-element chunks) + # - ldmcxstmc_default_inflight_depth=4: quad-issue (4-element chunks) + # The chunk loop is staged as scf.for with explicit `unroll=1` hint. + # + total_iters = loop_m * loop_n + + if cutlass.const_expr(kernel.ldmcxstmc_default_inflight_depth == 2): + # Pair-issue variant (inflight2) + num_pairs = total_iters // 2 + for k in cutlass.range(num_pairs, unroll=1): + i0 = (2 * k) // loop_n + j0 = (2 * k) % loop_n + i1 = (2 * k + 1) // loop_n + j1 = (2 * k + 1) % loop_n + + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + + # Trailing odd iteration (singleton) when total_iters is odd. + if cutlass.const_expr(total_iters % 2 == 1): + last = total_iters - 1 + i = last // loop_n + j = last % loop_n + in_ = tTR_gCommInMC[None, i, j].iterator + out_ = tTR_gCommOutMC[None, i, j].iterator + regs = utils.distributed.multimem_ld_reduce( + in_, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out_, *regs) + else: + # Quad-issue variant (inflight4, default) + num_quads = total_iters // 4 + for k in cutlass.range(num_quads, unroll=1): + i0 = (4 * k) // loop_n + j0 = (4 * k) % loop_n + i1 = (4 * k + 1) // loop_n + j1 = (4 * k + 1) % loop_n + i2 = (4 * k + 2) // loop_n + j2 = (4 * k + 2) % loop_n + i3 = (4 * k + 3) // loop_n + j3 = (4 * k + 3) % loop_n + + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + in2 = tTR_gCommInMC[None, i2, j2].iterator + out2 = tTR_gCommOutMC[None, i2, j2].iterator + in3 = tTR_gCommInMC[None, i3, j3].iterator + out3 = tTR_gCommOutMC[None, i3, j3].iterator + + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs2 = utils.distributed.multimem_ld_reduce( + in2, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs3 = utils.distributed.multimem_ld_reduce( + in3, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + utils.distributed.multimem_st(out2, *regs2) + utils.distributed.multimem_st(out3, *regs3) + + # Tail: total_iters % 4 == 1 → 1 trailing LDMC + 1 STMC. + if cutlass.const_expr(total_iters % 4 == 1): + last = total_iters - 1 + i = last // loop_n + j = last % loop_n + in_ = tTR_gCommInMC[None, i, j].iterator + out_ = tTR_gCommOutMC[None, i, j].iterator + regs = utils.distributed.multimem_ld_reduce( + in_, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out_, *regs) + # Tail: total_iters % 4 == 2 → 2 trailing LDMCs + 2 STMCs. + if cutlass.const_expr(total_iters % 4 == 2): + flat0 = total_iters - 2 + flat1 = total_iters - 1 + i0 = flat0 // loop_n + j0 = flat0 % loop_n + i1 = flat1 // loop_n + j1 = flat1 % loop_n + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + # Tail: total_iters % 4 == 3 → 3 trailing LDMCs + 3 STMCs. + if cutlass.const_expr(total_iters % 4 == 3): + flat0 = total_iters - 3 + flat1 = total_iters - 2 + flat2 = total_iters - 1 + i0 = flat0 // loop_n + j0 = flat0 % loop_n + i1 = flat1 // loop_n + j1 = flat1 % loop_n + i2 = flat2 // loop_n + j2 = flat2 % loop_n + in0 = tTR_gCommInMC[None, i0, j0].iterator + out0 = tTR_gCommOutMC[None, i0, j0].iterator + in1 = tTR_gCommInMC[None, i1, j1].iterator + out1 = tTR_gCommOutMC[None, i1, j1].iterator + in2 = tTR_gCommInMC[None, i2, j2].iterator + out2 = tTR_gCommOutMC[None, i2, j2].iterator + regs0 = utils.distributed.multimem_ld_reduce( + in0, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs1 = utils.distributed.multimem_ld_reduce( + in1, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + regs2 = utils.distributed.multimem_ld_reduce( + in2, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + utils.distributed.multimem_st(out0, *regs0) + utils.distributed.multimem_st(out1, *regs1) + utils.distributed.multimem_st(out2, *regs2) + + +class Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel: + """Persistent block-scaled GEMM + AllReduce (LDMCxSTMC) fusion kernel for Blackwell GPUs. + + **Example:** + gemm = Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel( + sf_vec_size=16, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + c_dtype=cutlass.Float16, + rank_id=rank, + num_ranks=world_size, + ) + gemm(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, layouts, problem_mnkl, + max_active_clusters, stream, comm_in_multicast_tensor=..., + comm_out_multicast_tensor=..., barrier_flag_unicast=..., + barrier_flag_multicast=...) + """ + + def __init__( + self, + sf_vec_size: int, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + c_dtype: Type[cutlass.Numeric] = cutlass.Float16, + rank_id: int = 0, + num_ranks: int = 8, + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + ): + self.acc_dtype = cutlass.Float32 + self.sf_vec_size = sf_vec_size + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn + self.swizzle_size = swizzle_size + self.raster_order = raster_order + arch_enum = BaseDSL._get_dsl().get_arch_enum() + self.arch = f"sm_{arch_enum.major}{arch_enum.minor}" + self.mma_tiler_mn = mma_tiler_mn + self.c_dtype = c_dtype + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.rank_id = rank_id + self.num_ranks = num_ranks + num_comm_warps = self._pick_num_comm_warp_for_128b(c_dtype) + self.comm_warp_id = tuple(range(6, 6 + num_comm_warps)) + cta_tile_m = mma_tiler_mn[0] // 2 if mma_tiler_mn[0] == 256 else mma_tiler_mn[0] + cta_tile_n = mma_tiler_mn[1] + self.ldmcxstmc_default_inflight_depth = 4 + self.comm_sync_bar_id = 4 + self.all_reduce = "LDMCxSTMC" + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilogue_warp_id, + *self.comm_warp_id, + ) + ) + # Set barrier id for epilogue sync and tmem ptr sync + self.epilogue_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_warp * len(self.epilogue_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_warp + * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols(self.arch) + + def _pick_num_comm_warp_for_128b(self, c_dtype): + """Pick the LARGEST comm-warp count W ∈ {4, 3, 2, 1} (T = W·32 ∈ + {128, 96, 64, 32}) such that the AR slab is evenly distributed + across T threads AND every thread gets a full 128 b LDMC payload + (= ``atom_128`` elements, where ``atom_128 = 16 / sizeof(c_dtype)``). + + Iterates W = 4 → 1 and returns the first match, so the largest + viable W is returned. Smaller W is only chosen when strictly + necessary, because larger W means more parallel AR threads and + proportionally fewer per-thread PTX instructions. + """ + # cta-tile dims (the per-CTA M·N output the AR walks). 2cta halves M. + cta_mma_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + slab_elements = (cta_mma_tile_m * cta_tile_n) // self.num_ranks + atom_128 = 16 // (c_dtype.width // 8) + for num_warps in (4, 3, 2, 1): + T = num_warps * 32 + if slab_elements % T: + continue + if (slab_elements // T) % atom_128 == 0: + return num_warps + raise ValueError( + f"no comm-warp count enables 128 b PTX for c_dtype={c_dtype}, " + f"cta=({cta_mma_tile_m},{cta_tile_n}), P={self.num_ranks}" + ) + + def _setup_attributes(self): + # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) + self.mma_inst_shape_mn = ( + self.mma_tiler[0], + self.mma_tiler[1], + ) + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = utils.sm100.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = utils.sm100.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + 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_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + 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.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + self.epi_tile_n = cute.size(self.epi_tile[1]) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.smem_alloc_b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + ) + + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.smem_alloc_b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case + self.overlapping_accum = self.num_acc_stage == 1 + + sf_atom_mn = 32 + self.num_sfa_tmem_cols = ( + self.cta_tile_shape_mnk[0] // sf_atom_mn + ) * mma_inst_tile_k + self.num_sfb_tmem_cols = ( + self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn + ) * mma_inst_tile_k + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage + if not self.overlapping_accum + else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols + ) + + self.iter_acc_early_release_in_epilogue = ( + self.num_sf_tmem_cols // self.epi_tile_n + ) + + @cute.jit + def __call__( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + layouts: cutlass.Constexpr[ + Tuple[OperandMajorMode, OperandMajorMode, utils.LayoutEnum] + ], + problem_mnkl: Tuple[int, int, int, int], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + comm_in_multicast_tensor: cute.Tensor = None, + comm_out_multicast_tensor: cute.Tensor = None, + barrier_flag_unicast: cute.Tensor = None, + barrier_flag_multicast: cute.Tensor = None, + ): + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a_ptr.value_type + self.b_dtype: Type[cutlass.Numeric] = b_ptr.value_type + self.sf_dtype: Type[cutlass.Numeric] = sfa_ptr.value_type + self.c_dtype: Type[cutlass.Numeric] = c_ptr.value_type + self.mxf8f6f4 = self.needs_unpack_tma(self.a_dtype, self.b_dtype) + self.smem_alloc_a_dtype = ( + cutlass.Int8 if (self.mxf8f6f4 and self.a_dtype.width < 8) else self.a_dtype + ) + self.smem_alloc_b_dtype = ( + cutlass.Int8 if (self.mxf8f6f4 and self.b_dtype.width < 8) else self.b_dtype + ) + m, n, k, l = problem_mnkl + self.a_major_mode, self.b_major_mode, self.c_layout = layouts + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + a_layout = cute.make_ordered_layout((m, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.a_major_mode == OperandMajorMode.K): + a_layout = cute.make_ordered_layout( + (cute.assume(m, 32), k, l), order=(1, 0, 2) + ) + b_layout = cute.make_ordered_layout((n, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.b_major_mode == OperandMajorMode.K): + b_layout = cute.make_ordered_layout( + (cute.assume(n, 32), k, l), order=(1, 0, 2) + ) + c_layout = cute.make_ordered_layout((cute.assume(m, 32), n, l), order=(0, 1, 2)) + if cutlass.const_expr(self.c_layout == utils.LayoutEnum.ROW_MAJOR): + c_layout = cute.make_ordered_layout( + (m, cute.assume(n, 32), l), order=(1, 0, 2) + ) + a_tensor = cute.make_tensor(a_ptr, a_layout) + b_tensor = cute.make_tensor(b_ptr, b_layout) + c_tensor = cute.make_tensor(c_ptr, c_layout) + + mCommInMC_mnl = comm_in_multicast_tensor + mCommOutMC_mnl = comm_out_multicast_tensor + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, self.sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, self.sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) + + tiled_mma = utils.sm100.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + # For 2CTA blockscaled kernels, SFB needs to be replicated across peer CTAs. + tiled_mma_sfb = utils.sm100.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.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_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_a_dtype + if (self.mxf8f6f4 and self.a_dtype.width < 8) + else None, + ) + + # Setup TMA load for B + b_op = utils.sm100.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_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_b_dtype + if (self.mxf8f6f4 and self.b_dtype.width < 8) + else None, + ) + + # Setup TMA load for SFA + sfa_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Setup TMA load for SFB + sfb_op = utils.sm100.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb_tensor, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + x = tma_tensor_sfb.stride[0][1] + y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) + + new_shape = ( + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), + tma_tensor_sfb.shape[1], + tma_tensor_sfb.shape[2], + ) + # Use right multiplication for ScaledBasis (3 * x instead of x * 3) + x_times_3 = 3 * x + new_stride = ( + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), + tma_tensor_sfb.stride[1], + tma_tensor_sfb.stride[2], + ) + tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) + tma_tensor_sfb = cute.make_tensor( + tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout + ) + + 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) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Setup TMA store for C + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c_tensor, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c_tensor, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + self.swizzle_size, + self.raster_order, + max_active_clusters, + ) + + self.buffer_align_bytes = 1024 + + # 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: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.smem_alloc_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.smem_alloc_b_dtype, + cute.cosize(self.b_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sSFA: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sSFB: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma=tiled_mma, + tiled_mma_sfb=tiled_mma_sfb, + tma_atom_a=tma_atom_a, + mA_mkl=tma_tensor_a, + tma_atom_b=tma_atom_b, + mB_nkl=tma_tensor_b, + tma_atom_sfa=tma_atom_sfa, + mSFA_mkl=tma_tensor_sfa, + tma_atom_sfb=tma_atom_sfb, + mSFB_nkl=tma_tensor_sfb, + tma_atom_c=tma_atom_c, + mC_mnl=tma_tensor_c, + mCommInMC_mnl=mCommInMC_mnl, + mCommOutMC_mnl=mCommOutMC_mnl, + barrier_flag_unicast=barrier_flag_unicast, + barrier_flag_multicast=barrier_flag_multicast, + cluster_layout_vmnk=self.cluster_layout_vmnk, + cluster_layout_sfb_vmnk=self.cluster_layout_sfb_vmnk, + a_smem_layout_staged=self.a_smem_layout_staged, + b_smem_layout_staged=self.b_smem_layout_staged, + sfa_smem_layout_staged=self.sfa_smem_layout_staged, + sfb_smem_layout_staged=self.sfb_smem_layout_staged, + c_smem_layout_staged=self.c_smem_layout_staged, + epi_tile=self.epi_tile, + tile_sched_params=self.tile_sched_params, + epilogue_op=epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + mCommInMC_mnl: cute.Tensor, + mCommOutMC_mnl: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """ + 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_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + 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 + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_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) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Warp) + 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, + enable_multicast_signaling=True, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = ( + self.threads_per_warp + * len(self.epilogue_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, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/SFA/SFB/C + # + # (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 + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_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 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_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) + ) + # (bM, bK, RestM, RestK, RestL) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_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) + thr_mma_sfb = tiled_mma_sfb.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_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # 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), RestN, 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), + ) + + # TMA load scaled factor A partition_S/D + sfa_cta_layout = a_cta_layout + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load scaled factor B partition_S/D + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # 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]) + if cutlass.const_expr(self.overlapping_accum): + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, num_acc_stage_overlapped) + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + # (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 + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Construct the scheduler (shared across all warp specializations) + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + 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]) + ] + + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + + # Apply SFB slicing hack when cta_tile_shape_n=64 + slice_n = mma_tile_coord_mnl[1] + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + slice_n = mma_tile_coord_mnl[1] // 2 + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_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/SFA/SFB + 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, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfb_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_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.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # Make accumulator tmem tensor + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Make SFA tmem tensor + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # Make SFB tmem tensor + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + # + # Partition for S2T copy of SFA/SFB + # + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + 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], + ) + + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_producer_state.phase ^ 1 + else: + acc_stage_index = acc_producer_state.index + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_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) + + # Apply TMEM pointer offset hack when cta_tile_shape_n=192 or cta_tile_shape_n=64 + tCtSFB_mma = tCtSFB + if cutlass.const_expr(self.cta_tile_shape_mnk[1] in {64, 192}): + # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB) + offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) + + # Copy SFA/SFB from smem to tmem + s2t_stage_coord = ( + None, + None, + None, + None, + ab_consumer_state.index, + ) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, ab_consumer_state.index) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[tile_crd], tCtSFA], + [tCrB[tile_crd], tCtSFB_mma], + tCtAcc, + ) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_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 + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Number of CTA tiles along M — for per-tile flag linearization + m_cta_tiles_in_total = gC_mnl.shape[2] * cute.size(tiled_mma.thr_id.shape) + + # + # Partition for epilogue + # + epi_tidx = tidx + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilogue_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilogue_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilogue_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + + 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, + self.threads_per_warp * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_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_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = True if acc_stage_index == 0 else False + else: + acc_stage_index = acc_consumer_state.index + + # 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_stage_index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + 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(subtile_cnt): + real_subtile_idx = subtile_idx + if cutlass.const_expr(self.overlapping_accum): + if reverse_subtile: + # Subtile always iterates on N dimension as we only have 4x1DP tmem load pattern for cta_tile_m = 128 cases. + real_subtile_idx = ( + self.cta_tile_shape_mnk[1] // self.epi_tile_n + - 1 + - subtile_idx + ) + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled + # + if cutlass.const_expr(self.overlapping_accum): + if subtile_idx == self.iter_acc_early_release_in_epilogue: + # Fence for TMEM load + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + 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( + "async.shared", + space="cta", + ) + self.epilogue_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, real_subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilogue_sync_barrier.arrive_and_wait() + + # AllReduce: per-tile — drain TMA then signal flag + # (c_pipeline.producer_tail drains ALL pending TMA stores up to this point) + c_pipeline.producer_tail() + cta_tile_id_linear = cute.Int32( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + utils.distributed.multimem_red_add1( + lock_ptr=barrier_flag_multicast.iterator + + cta_tile_id_linear, + order="release", + scope="gpu", + ) + + # + # Async arrive accumulator buffer empty + # + if cutlass.const_expr(not self.overlapping_accum): + 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 + # + tmem.relinquish_alloc_permit() + self.epilogue_sync_barrier.arrive_and_wait() + tmem.free(acc_tmem_ptr) + # Note: producer_tail() is now called per-tile in the AR path + + # /////////////////////////////////////////////////////////////////////////////// + # AllReduce warps + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.all_reduce == "LDMCxSTMC"): + if warp_idx >= self.comm_warp_id[0]: + # (bM, bN, RestM, RestN, RestL) with multicast memory as comm input + gCommInMC_mnl = cute.local_tile( + mCommInMC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + # (bM, bN, RestM, RestN, RestL) with multicast memory as comm output + gCommOutMC_mnl = cute.local_tile( + mCommOutMC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) with multicast memory + tCgCommInMC = thr_mma.partition_C(gCommInMC_mnl) + tCgCommOutMC = thr_mma.partition_C(gCommOutMC_mnl) + + m_cta_tiles_in_total = gC_mnl.shape[2] * cute.size( + tiled_mma.thr_id.shape + ) + n_cta_tiles_in_total = gC_mnl.shape[3] + + # + # Persistent tile scheduling loop for all reduce + # + comm_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.comm_sync_bar_id, + num_threads=32 * len(self.comm_warp_id), + ) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + + # (1) Per-tile flag wait. + cta_tile_id_linear = ( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + if warp_idx == self.comm_warp_id[0]: + with cute.arch.elect_one(): + utils.distributed.spin_lock_atom_cas_acquire_wait( + lock_ptr=barrier_flag_unicast.iterator + + cta_tile_id_linear, + expected_val=self.num_ranks, + reset_val=0, + scope="gpu", + ) + + comm_sync_barrier.arrive_and_wait() + + # (2) Per-tile data movement (multimem.ld_reduce + multimem.st). + LDMCxSTMC( + self, + cur_tile_coord=cur_tile_coord, + tiled_mma=tiled_mma, + comm_tidx=tidx, + tCgCommInMC=tCgCommInMC, + tCgCommOutMC=tCgCommOutMC, + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Ensure all threads in allreduce warp group have complete issuing all reduce operations. + comm_sync_barrier.arrive_and_wait() + + # + # Set Per SM Flag with Release + # + if warp_idx == self.comm_warp_id[0]: + with cute.arch.elect_one(): + # Offset to last tile flag idx + total_num_cta_tile = m_cta_tiles_in_total * n_cta_tiles_in_total + # Linear id of current SM. + sm_id_linear = ( + cute.arch.block_idx()[0] + + cute.arch.block_idx()[1] * cute.arch.grid_dim()[0] + + cute.arch.block_idx()[2] + * cute.arch.grid_dim()[0] + * cute.arch.grid_dim()[1] + ) + # Release flag with sys scope + utils.distributed.multimem_red_add1( + lock_ptr=barrier_flag_multicast.iterator + + total_num_cta_tile + + sm_id_linear, + scope="sys", + order="release", + ) + # Acquire spin-lock wait flag with sys scope + utils.distributed.spin_lock_atom_cas_acquire_wait( + lock_ptr=barrier_flag_unicast.iterator + + total_num_cta_tile + + sm_id_linear, + expected_val=self.num_ranks, + reset_val=0, + scope="sys", + ) + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for smem-to-tmem scale factor load; partition smem source and tmem destination.""" + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def epilogue_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for tmem-to-register load; partition tmem source and register destination.""" + # Make tiledCopy for tensor memory load + copy_atom_t2r = utils.sm100.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_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) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilogue_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for register-to-smem store; partition register source and smem destination.""" + copy_atom_r2s = utils.sm100.get_smem_store_op( + self.c_layout, self.c_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_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilogue_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for TMA global memory store; partition smem source and gmem destination.""" + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 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_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_c, bSG_sC, 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, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + ) -> Tuple[int, int, int]: + """Compute (num_acc_stage, num_ab_stage, num_c_stage) based on smem capacity heuristics.""" + # ACC stages + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + + # Default C stages + num_c_stage = 2 + + # Calculate smem layout and size for one stage of A, B, SFA, SFB and C + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + + c_smem_layout_staged_one = utils.sm100.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 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) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + 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/SFA/SFB stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B/SFA/SFB stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + swizzle_size: int, + raster_order: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Compute (tile_sched_params, grid) using the persistent tile scheduler.""" + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_order == "m" + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def needs_unpack_tma( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + ) -> bool: + """Return True if TMA must use UNPACK_U8 (U4_UNPACK_U8 / U6_UNPACK_U8). + + Required when operand widths differ (mxf8f6f4 mixed-precision) or either + operand is 6-bit (no packed U6 TMA format exists). + """ + if a_dtype.width != b_dtype.width: + return True + if a_dtype.width == 6 or b_dtype.width == 6: + return True + return False + + @staticmethod + def is_valid_dtypes_and_scale_factor_vec_size( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """Return True if (a_dtype, b_dtype, sf_dtype, sf_vec_size, c_dtype) is a supported combination.""" + supported_ab_dtypes = { + cutlass.Float4E2M1FN, + cutlass.Float6E2M3FN, + cutlass.Float6E3M2FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + } + + # Check A/B element types + if a_dtype not in supported_ab_dtypes or b_dtype not in supported_ab_dtypes: + return False + + # Check SF element type + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + return False + + # sf_vec_size rules: + # * 16 is only supported for Float4E2M1FN x Float4E2M1FN (NVF4 / MXF4 fp4-pair) + # * 32 is required for every other A/B combination (MXF8, mxf8f6f4 mixed, MXF4-pair with MX scaling) + # SF dtype pairing with sf_vec_size: + # * sf_vec_size == 16 requires sf_dtype in {Float8E4M3FN (NVF4), Float8E8M0FNU (MXF4)} + # * sf_vec_size == 32 requires sf_dtype == Float8E8M0FNU (MX scaling) + both_fp4 = a_dtype is cutlass.Float4E2M1FN and b_dtype is cutlass.Float4E2M1FN + if sf_vec_size == 16: + if not both_fp4: + return False + elif sf_vec_size == 32: + if sf_dtype is not cutlass.Float8E8M0FNU: + return False + else: + return False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + return False + + return True + + @staticmethod + def is_valid_layouts( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + ) -> bool: + """Return True if the layout/dtype combination is valid.""" + is_valid = True + + # FP4 operands can only be k-major (checked per operand) + if a_dtype is cutlass.Float4E2M1FN and a_major != "k": + is_valid = False + if b_dtype is cutlass.Float4E2M1FN and b_major != "k": + is_valid = False + + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """Return True if mma_tiler_mn and cluster_shape_mn satisfy hardware constraints.""" + is_valid = True + # Skip invalid mma tile shape + if mma_tiler_mn[0] not in [128, 256]: + is_valid = False + if mma_tiler_mn[1] not in [64, 128, 192, 256]: + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + # Special cluster shape check for scale factor multicasts. + # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: Tuple[int, int], + ) -> bool: + """Return True if tensor alignment satisfies TMA 16B contiguous-dim and UNPACK constraints.""" + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + # TMA requires the contiguous inner dimension to be a multiple of + # 16 B (= 128 bits). Work in bits so non-byte-aligned widths + # (e.g. 6-bit) are handled correctly: 16 * 8 // dtype.width is + # wrong when dtype.width does not divide 128 (it returns 21 for + # 6-bit instead of the real requirement K*6 % 128 == 0). + return (num_major_elements * dtype.width) % (16 * 8) == 0 + + def check_contigous_128_alignment(dtype, is_mode0_major, tensor_shape): + # we only need to check alignment for subbyte dtype + if dtype.width >= 8: + return True + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 128 + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + # When an operand is loaded via the UNPACK TMA variant + # (U4_UNPACK_U8 or U6_UNPACK_U8), its inner tensor dimension in bytes + # must be a multiple of 64B (4-bit) or 96B (6-bit); both work out to + # a multiple of 128 elements along the contiguous dim. The check only + # applies to sub-byte operands and only when the pair triggers UNPACK. + if Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.needs_unpack_tma( + a_dtype, b_dtype + ) and ( + not check_contigous_128_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contigous_128_alignment(b_dtype, b_major == "n", (n, k, l)) + ): + is_valid = False + # Additional UNPACK constraint for any sub-byte operand on its contig + # axis: when a sub-byte A is m-major (contig=M) or a sub-byte B is + # n-major (contig=N), the MMA tile's contig dim (after 2CTA M-split, + # which splits M for both A and B on the non-multicast atom path) + # must be a multiple of 128 elements to satisfy the 64B (fp4) / + # 96B (fp6) inner-dim requirement of U4_/U6_UNPACK_U8. Observed + # failures: (128,192)/(1,1)/m-n-m (1CTA) and (256,128)/(2,2)/m-n-m + # (2CTA) trigger CUDA illegal instruction for fp6 when mma_tiler_N + # (or N/2 after 2CTA split) is not a 128-multiple; the same rule + # applies to any other sub-byte operand on its non-K contig axis. + use_2cta_instrs = mma_tiler_mn[0] == 256 + cta_div = 2 if use_2cta_instrs else 1 + if ( + Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.needs_unpack_tma( + a_dtype, b_dtype + ) + and a_major == "m" + and a_dtype.width < 8 + and (mma_tiler_mn[0] // cta_div) % 128 != 0 + ): + is_valid = False + if ( + Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.needs_unpack_tma( + a_dtype, b_dtype + ) + and b_major == "n" + and b_dtype.width < 8 + and (mma_tiler_mn[1] // cta_div) % 128 != 0 + ): + is_valid = False + return is_valid + + @staticmethod + def can_implement( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + sf_vec_size: int, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + num_ranks: int = 8, + ) -> bool: + """Return True if all dtype, layout, tiler, cluster-shape, alignment, and AR-warp checks pass.""" + # Unpack parameters + m, n, k, l = mnkl + can_implement = True + # Skip unsupported types + if not Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.is_valid_dtypes_and_scale_factor_vec_size( + a_dtype, b_dtype, sf_dtype, sf_vec_size, c_dtype + ): + can_implement = False + # Skip unsupported layouts + if not Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.is_valid_layouts( + a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn, cluster_shape_mn + ): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.is_valid_tensor_alignment( + m, + n, + k, + l, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + ): + can_implement = False + + # AR: verify comm warp count is achievable + class _Tmp: + pass + + tmp = _Tmp() + tmp.mma_tiler_mn = mma_tiler_mn + tmp.use_2cta_instrs = mma_tiler_mn[0] == 256 + tmp.num_ranks = num_ranks + try: + Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel._pick_num_comm_warp_for_128b( + tmp, c_dtype + ) + except (ValueError, AssertionError) as e: + raise testing.CantImplementError(str(e)) + return can_implement + + +# Helper function to convert scale factor tensor from MKL layout to (32, 4, restM, 4, restK, l) format +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_ptr: cute.Pointer, + sf_mma_ptr: cute.Pointer, + mn: int, + sf_k: int, + l: int, + mma_shape: tuple, +): + mma_permute_order = (3, 4, 1, 5, 2, 0) + permuted_shape = tuple(mma_shape[i] for i in mma_permute_order) + cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5)) + + sf_ref_tensor = cute.make_tensor( + sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k)) + ) + sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout) + + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + pass + + +# Helper function for ceil division +def ceil_div(a, b): + return (a + b - 1) // b + + +# Convert scale factor tensors from (m, k, l) to (32, 4, restM, 4, restK, l) format +def create_and_reorder_scale_factor_tensor( + l, mn, k, sf_vec_size, sf_dtype, torch_tensor +): + """ + Create the CUTE-format scale factor tensor on CUDA based on the reference tensor. + """ + sf_k = ceil_div(k, sf_vec_size) + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, # batch size + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + # Generate a random int8 tensor, then convert to float8_e4m3fn + cute_tensor = torch.ones(mma_shape, dtype=cutlass_torch.dtype(sf_dtype)).permute( + 3, 4, 1, 5, 2, 0 + ) + + # Call the helper function to do layout conversion + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + make_ptr( + sf_dtype, + torch_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + sf_dtype, + cute_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + mn, + sf_k, + l, + mma_shape, + ) + return cute_tensor.cuda() + + +# Compile the persistent dense blockscaled GEMM operation +def scaled_mm( + gemm_obj: Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + options: str = "", +): + # Construct CuTe Pointers + a_ptr = make_ptr(a_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(b_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + + a_major_mode = OperandMajorMode.K if a_major == "k" else OperandMajorMode.MN + b_major_mode = OperandMajorMode.K if b_major == "k" else OperandMajorMode.MN + c_layout = ( + utils.LayoutEnum.ROW_MAJOR if c_major == "n" else utils.LayoutEnum.COL_MAJOR + ) + return cute.compile( + gemm_obj, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + c_ptr, + (a_major_mode, b_major_mode, c_layout), + (cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + max_active_clusters, + stream, + epilogue_op, + options=options, + ) + + +def scaled_mm_ar( + gemm_obj: Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + comm_in_multicast_tensor: cute.Tensor, + comm_out_multicast_tensor: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, + epilogue_op: cutlass.Constexpr = lambda x: x, + options: str = "", +): + # Comm tensors (dynamic layout cute.Tensor) passed directly — cannot use make_ptr for these. + a_ptr = make_ptr(a_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(b_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + + a_major_mode = OperandMajorMode.K if a_major == "k" else OperandMajorMode.MN + b_major_mode = OperandMajorMode.K if b_major == "k" else OperandMajorMode.MN + c_layout = ( + utils.LayoutEnum.ROW_MAJOR if c_major == "n" else utils.LayoutEnum.COL_MAJOR + ) + return cute.compile( + gemm_obj, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + c_ptr, + (a_major_mode, b_major_mode, c_layout), + (cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + max_active_clusters, + stream, + epilogue_op, + comm_in_multicast_tensor, + comm_out_multicast_tensor, + barrier_flag_unicast, + barrier_flag_multicast, + options=options, + ) + + +def is_emulated_dtype( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +) -> bool: + # torch natively represents A/B only when both operands share the same + # dtype and the (dtype, sf_dtype) pair matches a supported non-emulated case. + if a_dtype != b_dtype: + return True + if c_dtype in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + }: + if a_dtype == cutlass.Float4E2M1FN and sf_dtype == cutlass.Float8E4M3FN: + return False + if a_dtype == cutlass.Float8E4M3FN and sf_dtype == cutlass.Float8E8M0FNU: + return False + + return True + + +# Convert scale factor tensor from MKL layout to blocked layout +def to_blocked(input_matrix): + rows, cols = input_matrix.shape + # Please ensure rows and cols are multiples of 128 and 4 respectively + n_row_blocks = ceil_div(rows, 128) + n_col_blocks = ceil_div(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + # Pad the input matrix if necessary + if padded_rows != rows or padded_cols != cols: + # For FP8 types, convert to float32 for padding, then convert back + original_dtype = input_matrix.dtype + input_float32 = input_matrix.to(torch.float32) + padded = torch.nn.functional.pad( + input_float32, + (0, padded_cols - cols, 0, padded_rows - rows), + mode="constant", + value=0, + ) + # Convert back to original dtype if needed + if original_dtype != input_float32.dtype: + padded = padded.to(original_dtype) + else: + padded = input_matrix + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + return rearranged.flatten() + + +# Reference implementation of the persistent dense blockscaled GEMM operation (emulated version) +def reference_scaled_mm_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: Tuple[int, int, int, int], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], +): + m, n, k, l = mnkl + sfa_expanded = ( + torch.repeat_interleave(sfa, sf_vec_size, dim=1)[:, :k, :] + .to(dtype=torch.float32) + .cuda() + ) + sfb_expanded = ( + torch.repeat_interleave(sfb, sf_vec_size, dim=1)[:, :k, :] + .to(dtype=torch.float32) + .cuda() + ) + res_a = torch.einsum("mkl,mkl->mkl", a, sfa_expanded) + res_b = torch.einsum("nkl,nkl->nkl", b, sfb_expanded) + # Cast res_a and res_b to float32 for einsum to avoid NotImplementedError on 'Byte' + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + c_ref = ref.to(dtype=cutlass_torch.dtype(c_dtype)) + return c_ref + + +# Reference implementation of the persistent dense blockscaled GEMM operation (non-emulated version) +def reference_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: Tuple[int, int, int, int], + c_dtype: Type[cutlass.Numeric], +): + m, n, k, l = mnkl + c_ref = torch.clone(c) + for l_idx in range(l): + # Convert the scale factor tensor to blocked format + scale_a = to_blocked(sfa[:, :, l_idx]) + scale_b = to_blocked(sfb[:, :, l_idx]) + # Ensure a_slice is row-major (M, K) with stride (K, 1) + a_slice = a[:, :, l_idx].contiguous() + # Ensure b_slice is row-major (N, K) so that transpose gives column-major (K, N) + b_slice = b[:, :, l_idx].contiguous() + # (m, k) @ (n, k).T -> (m, n) + res = torch._scaled_mm( + a_slice, + b_slice.transpose(0, 1), + scale_a.cuda(), + scale_b.cuda(), + bias=None, + out_dtype=c_ref.dtype, + ) + c_ref[:, :, l_idx] = res + return c_ref + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (emulated version) +def construct_cute_pointers_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +): + a_cute, _ = cutlass_torch.cute_tensor_like( + a.cpu(), + a_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + a_cute = cutlass_torch.convert_cute_tensor( + a, + a_cute, + a_dtype, + is_dynamic_layout=True, + ) + b_cute, _ = cutlass_torch.cute_tensor_like( + b.cpu(), + b_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + b_cute = cutlass_torch.convert_cute_tensor( + b, + b_cute, + b_dtype, + is_dynamic_layout=True, + ) + a_ptr = a_cute.iterator + b_ptr = b_cute.iterator + + sfa_ptr = make_ptr( + sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (non-emulated version) +def construct_cute_pointers( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +): + a_ptr = make_ptr(a_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(b_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr( + sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr + + +# Use uint8 and uint32 to emulate unsupported +# dtype in torch +def prepare_tensors_emulated( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + ab_init_lo: int = -2, + ab_init_hi: int = 2, + init_normal: bool = False, + normal_mean: float = 0.0, + normal_std: float = 1.0, +): + m, n, k, l = mnkl + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB with values in [1, 3) + sfa = ( + torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + sfb = ( + torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + + # Create tensor A/B + if a_major == "k": + a = torch.empty((l, m, k), dtype=torch.float32, device="cuda").permute(1, 2, 0) + else: + a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(2, 1, 0) + if b_major == "k": + b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(1, 2, 0) + else: + b = torch.empty((l, k, n), dtype=torch.float32, device="cuda").permute(2, 1, 0) + + # Initialize A/B tensors with either normal distribution or random integers + for tensor in [a, b]: + if init_normal: + tensor.normal_(mean=normal_mean, std=normal_std) + else: + tensor.copy_( + torch.randint( + ab_init_lo, + ab_init_hi, + tensor.shape, + dtype=torch.float32, + device="cuda", + ) + ) + + if c_major == "n": + c = torch.empty( + (l, m, n), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(1, 2, 0) + else: + c = torch.empty( + (l, n, m), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(2, 1, 0) + return a, b, c, sfa, sfb + + +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + ab_init_lo: int = -2, + ab_init_hi: int = 2, + init_normal: bool = False, + normal_mean: float = 0.0, + normal_std: float = 1.0, +): + m, n, k, l = mnkl + + if a_dtype == cutlass.Float4E2M1FN and b_dtype == cutlass.Float4E2M1FN: + # Using int8 for torch.float4_e2m1fn_x2 tensor allocation + # Thus the size of k needs to be halved in this case. + k_fct = 2 + else: + k_fct = 1 + + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB + sfa = ( + torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + sfb = ( + torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + + # Create tensor A/B + if a_major == "k": + a = torch.empty((l, m, k // k_fct), dtype=torch.int8, device="cuda").permute( + 1, 2, 0 + ) + else: + a = torch.empty((l, k, m), dtype=torch.int8, device="cuda").permute(2, 1, 0) + if b_major == "k": + b = torch.empty((l, n, k // k_fct), dtype=torch.int8, device="cuda").permute( + 1, 2, 0 + ) + else: + b = torch.empty((l, k, n), dtype=torch.int8, device="cuda").permute(2, 1, 0) + + # Initialize A/B tensors with random integers + # Note: int8 types always use random init (normal distribution not supported), + # consistent with use_normal_init pattern in generate_tensors + for tensor in [a, b]: + tensor.copy_( + torch.randint( + ab_init_lo, ab_init_hi, tensor.shape, dtype=torch.int8, device="cuda" + ) + ) + + # Create and initialize tensor C + if c_major == "n": + c_shape, c_perm = (l, m, n), (1, 2, 0) + else: + c_shape, c_perm = (l, n, m), (2, 1, 0) + + c = torch.randint( + -2, 2, c_shape, dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(c_perm) + + if a_dtype == cutlass.Float4E2M1FN and b_dtype == cutlass.Float4E2M1FN: + a = a.view(dtype=torch.float4_e2m1fn_x2) + b = b.view(dtype=torch.float4_e2m1fn_x2) + else: + a = a.to(dtype=cutlass_torch.dtype(a_dtype)) + b = b.to(dtype=cutlass_torch.dtype(b_dtype)) + + c = c.to(dtype=cutlass_torch.dtype(c_dtype)) + return a, b, c, sfa, sfb + + +# This will show how to convert torch tensor +# and pass to CuTe kernel + +##################################################################### +#### KERNEL IMPLEMENTATION END HERE #### +##################################################################### + +_DISTRIBUTED_INITIALIZED = False + + +def check_cuda_p2p_access(): + num_devices = torch.cuda.device_count() + for i in range(num_devices - 1): + for j in range(i + 1, num_devices): + if not torch.cuda.can_device_access_peer(i, j): + return False + return True + + +def check_multicast_support(): + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + for device_id in range(torch.cuda.device_count()): + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device_id): + return False + return True + except ImportError: + return False + + +def init_distributed(): + global _DISTRIBUTED_INITIALIZED + + if not dist.is_initialized() and not _DISTRIBUTED_INITIALIZED: + if not torch.cuda.is_available(): + print("CUDA is not available") + exit(-1) + if torch.cuda.get_device_capability()[0] != 10: + print("Running on non-sm100 device") + exit(-1) + if torch.cuda.device_count() < 2: + print("Less than two devices") + exit(-1) + if torch.cuda.device_count() % 2 != 0: + print("Odd world size") + exit(-1) + for var in ("LOCAL_RANK", "RANK", "WORLD_SIZE"): + if var not in os.environ: + print(f"{var} not in environment variables") + exit(-1) + + local_rank = int(os.environ["LOCAL_RANK"]) + global_rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + world_size=world_size, + rank=global_rank, + device_id=device, + ) + _DISTRIBUTED_INITIALIZED = True + print( + f"[init] hostname={socket.gethostname()} global_rank={global_rank} " + f"local_rank={local_rank} world_size={world_size} device={device}" + ) + + global_rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", global_rank)) + device = torch.device("cuda", local_rank) + return global_rank, world_size, device + + +def clean_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def test_simple_torch_allreduce(device): + try: + tensor = symm_mem.empty((16384,), dtype=torch.float32, device=device) + symm_mem.rendezvous(tensor, group=dist.group.WORLD) + torch.ops.symm_mem.multimem_all_reduce_( + tensor, "sum", dist.group.WORLD.group_name + ) + except Exception: + return False + return True + + +def allocate_tensors( + *, + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + num_workspace: int, + slot_init_mode: Literal["test", "benchmark"], + device=None, + global_rank: int = 0, + world_size: int = 1, +): + """Allocate num_workspace slabs of A/B/SFA/SFB for L2-cold rotation, + plus symm_mem C (comm_in) and comm_out and flag tensors for AR. + + Returns a dict with lists of cute pointers ready to pass to the compiled kernel, + plus workspace-0 raw torch tensors for reference checking. + """ + + def _create_multicast_tensor(*, torch_tensor_ref, dtype, leading_dim): + is_col_major = leading_dim == 0 + shape = torch_tensor_ref.shape + ref_dtype = torch_tensor_ref.dtype + + if is_col_major: + reversed_shape = tuple(reversed(shape)) + perm = tuple(range(len(shape) - 1, -1, -1)) + torch_symm_base = symm_mem.empty( + reversed_shape, device=device, dtype=ref_dtype + ) + torch_tensor_uc = torch_symm_base.permute(perm) + else: + torch_symm_base = symm_mem.empty(shape, device=device, dtype=ref_dtype) + torch_tensor_uc = torch_symm_base + + torch_tensor_uc.copy_(torch_tensor_ref) + + torch_tensor_hdl = symm_mem.rendezvous(torch_symm_base, group=dist.group.WORLD) + + cute_tensor_uc = from_dlpack(torch_tensor_uc, assumed_align=16) + cute_tensor_uc.element_type = dtype + cute_tensor_uc = cute_tensor_uc.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor_uc = cutlass_torch.convert_cute_tensor( + torch_tensor_uc, + cute_tensor_uc, + dtype, + is_dynamic_layout=True, + ) + + if is_col_major: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, reversed_shape, ref_dtype + ).permute(perm) + else: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, shape, ref_dtype + ) + cute_tensor_mc = from_dlpack(mc_torch, assumed_align=16) + cute_tensor_mc.element_type = dtype + cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim) + + return cute_tensor_uc, cute_tensor_mc, torch_tensor_uc + + torch.manual_seed(42 + global_rank) + ab_init_lo, ab_init_hi = (-1, 1) if slot_init_mode == "test" else (-5, 5) + m, n, k, l = mnkl + emulated = is_emulated_dtype(a_dtype, b_dtype, sf_dtype, c_dtype) + _prepare = prepare_tensors_emulated if emulated else prepare_tensors + + a_ptr_list, b_ptr_list, sfa_ptr_list, sfb_ptr_list = [], [], [], [] + # Keep every torch tensor alive — cute pointers hold raw addresses, not refs. + anchors = [] + + for _ in range(num_workspace): + a, b, c, sfa_raw, sfb_raw = _prepare( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + ab_init_lo=ab_init_lo, + ab_init_hi=ab_init_hi, + ) + sfa_r = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa_raw + ) + sfb_r = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb_raw + ) + + if emulated: + a_ptr, b_ptr, _c_ptr_unused, sfa_ptr, sfb_ptr, a_cute, b_cute = ( + construct_cute_pointers_emulated( + a, b, sfa_r, sfb_r, c, a_dtype, b_dtype, sf_dtype, c_dtype + ) + ) + anchors.append((a, b, c, sfa_raw, sfb_raw, sfa_r, sfb_r, a_cute, b_cute)) + else: + a_ptr, b_ptr, _c_ptr_unused, sfa_ptr, sfb_ptr = construct_cute_pointers( + a, b, sfa_r, sfb_r, c, a_dtype, b_dtype, sf_dtype, c_dtype + ) + anchors.append((a, b, c, sfa_raw, sfb_raw, sfa_r, sfb_r)) + + a_ptr_list.append(a_ptr) + b_ptr_list.append(b_ptr) + sfa_ptr_list.append(sfa_ptr) + sfb_ptr_list.append(sfb_ptr) + + # Use workspace-0 C as the symm_mem comm_in tensor (C is written by GEMM, then AR'd) + c_ref_cpu = anchors[0][2].cpu() + comm_out_cpu = torch.zeros_like(c_ref_cpu) + + leading_dim = 1 if c_major == "n" else 0 + ( + cute_tensor_c, + cute_tensor_comm_in_mc, + torch_tensor_c_gpu, + ) = _create_multicast_tensor( + torch_tensor_ref=c_ref_cpu, + dtype=c_dtype, + leading_dim=leading_dim, + ) + ( + _cute_tensor_comm_out_uc, + cute_tensor_comm_out_mc, + torch_tensor_comm_out_gpu, + ) = _create_multicast_tensor( + torch_tensor_ref=comm_out_cpu, + dtype=c_dtype, + leading_dim=leading_dim, + ) + + # Flag tensor: num_tiles + per-SM slots (160 is safe upper bound for B200) + num_flags = (m // 64) * (n // 64) + 160 + + torch_tensor_flag_unicast = symm_mem.empty( + (num_flags,), dtype=torch.int32, device=device + ) + torch_tensor_flag_unicast.fill_(0) + torch_flag_hdl = symm_mem.rendezvous( + torch_tensor_flag_unicast, group=dist.group.WORLD + ) + torch_tensor_flag_multicast = cutlass_torch.as_tensor( + torch_flag_hdl.multicast_ptr, + torch_tensor_flag_unicast.shape, + torch_tensor_flag_unicast.dtype, + ) + cute_tensor_flag_unicast = from_dlpack(torch_tensor_flag_unicast, assumed_align=16) + cute_tensor_flag_unicast.element_type = cutlass.Int32 + cute_tensor_flag_unicast = cute_tensor_flag_unicast.mark_layout_dynamic() + cute_tensor_flag_multicast = from_dlpack( + torch_tensor_flag_multicast, assumed_align=16 + ) + cute_tensor_flag_multicast.element_type = cutlass.Int32 + cute_tensor_flag_multicast = cute_tensor_flag_multicast.mark_layout_dynamic() + + return { + "a_ptr_list": a_ptr_list, + "b_ptr_list": b_ptr_list, + "sfa_ptr_list": sfa_ptr_list, + "sfb_ptr_list": sfb_ptr_list, + "cute_tensor_c": cute_tensor_c, + "cute_tensor_comm_in_mc": cute_tensor_comm_in_mc, + "cute_tensor_comm_out_mc": cute_tensor_comm_out_mc, + "cute_tensor_flag_unicast": cute_tensor_flag_unicast, + "cute_tensor_flag_multicast": cute_tensor_flag_multicast, + # Workspace-0 raw tensors for reference check in run_distributed_test. + "a_raw": anchors[0][0], + "b_raw": anchors[0][1], + "torch_tensor_c_gpu": torch_tensor_c_gpu, + "torch_tensor_comm_out_gpu": torch_tensor_comm_out_gpu, + "torch_tensor_flag_unicast": torch_tensor_flag_unicast, + "sfa_raw": anchors[0][3], + "sfb_raw": anchors[0][4], + "_anchors": anchors, + "emulated": emulated, + "_flag_handle": torch_flag_hdl, + } + + +def run( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + warmup_iterations: int = 0, + iterations: int = 1, + benchmark_or_test: str = "test", + use_cuda_graph: bool = False, + csv_report: str = "./csv_report.csv", +): + global_rank, world_size, device = init_distributed() + + if not check_cuda_p2p_access(): + print("CUDA P2P access not available") + exit(-1) + if not check_multicast_support(): + print("Multicast support not available") + exit(-1) + if not test_simple_torch_allreduce(device): + print("Simple torch allreduce test failed") + exit(-1) + + def _verbose(): + print( + "Running SM100 BlockScale GEMM + AllReduce (LDMCxSTMC) Fusion Kernel with:" + ) + print(f"mnkl: {mnkl}") + print(f"global_rank: {global_rank}, world_size: {world_size}, device: {device}") + print( + f"A dtype: {a_dtype}, B dtype: {b_dtype}, SF dtype: {sf_dtype}, sf_vec_size: {sf_vec_size}" + ) + print(f"C dtype: {c_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}") + print(f"Cluster Shape (M, N): {cluster_shape_mn}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Use CUDA Graph: {use_cuda_graph}") + print(f"Mode: {benchmark_or_test}") + print("AllReduce: LDMCxSTMC") + + if global_rank == 0: + _verbose() + + if benchmark_or_test in ("benchmark", "benchmark_all"): + run_distributed_benchmark( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + a_dtype=a_dtype, + b_dtype=b_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cuda_graph=use_cuda_graph, + benchmark_all=(benchmark_or_test == "benchmark_all"), + csv_report=csv_report, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + else: + run_distributed_test( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + a_dtype=a_dtype, + b_dtype=b_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + + clean_distributed() + + +def run_distributed_test( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (128, 128), + cluster_shape_mn: Tuple[int, int] = (1, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", +): + m, n, k, l = mnkl + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + gemm = Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + c_dtype=c_dtype, + rank_id=global_rank, + num_ranks=world_size, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + try: + import hashlib + + param_str = ( + f"a{a_dtype}_b{b_dtype}_sf{sf_dtype}_sfv{sf_vec_size}_c{c_dtype}" + f"_{a_major}{b_major}{c_major}" + f"_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}" + ) + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = f"sm100_blockscale_dense_gemm_ar_ar_{param_hash}_r{global_rank}_w{world_size}" + except Exception: + unique_name = f"sm100_blockscale_dense_gemm_ar_ar_r{global_rank}_w{world_size}" + gemm.__call__.set_name_prefix(unique_name) + + if not Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + c_dtype, + a_major, + b_major, + c_major, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + num_ranks=world_size, + ): + raise testing.CantImplementError( + f"Unsupported config: a={a_dtype} b={b_dtype} sf={sf_dtype} sfv={sf_vec_size} " + f"c={c_dtype} tile={mma_tiler_mn} cluster={cluster_shape_mn}" + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + tensors = allocate_tensors( + mnkl=mnkl, + a_dtype=a_dtype, + b_dtype=b_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + num_workspace=1, + slot_init_mode="test", + device=device, + global_rank=global_rank, + world_size=world_size, + ) + + compiled_gemm = scaled_mm_ar( + gemm, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, + max_active_clusters, + stream, + tensors["cute_tensor_comm_in_mc"], + tensors["cute_tensor_comm_out_mc"], + tensors["cute_tensor_flag_unicast"], + tensors["cute_tensor_flag_multicast"], + ) + + compiled_gemm( + tensors["a_ptr_list"][0], + tensors["b_ptr_list"][0], + tensors["sfa_ptr_list"][0], + tensors["sfb_ptr_list"][0], + tensors["cute_tensor_c"], + (m, n, k, l), + stream, + lambda x: x, + tensors["cute_tensor_comm_in_mc"], + tensors["cute_tensor_comm_out_mc"], + tensors["cute_tensor_flag_unicast"], + tensors["cute_tensor_flag_multicast"], + ) + torch.cuda.synchronize() + + emulated = tensors["emulated"] + a_raw = tensors["a_raw"] + b_raw = tensors["b_raw"] + sfa_raw = tensors["sfa_raw"] + sfb_raw = tensors["sfb_raw"] + torch_tensor_c = tensors["torch_tensor_c_gpu"] + torch_tensor_comm_out = tensors["torch_tensor_comm_out_gpu"] + torch_tensor_flag_unicast = tensors["torch_tensor_flag_unicast"] + + # FP8 nan to saturate helper + def _fp8_nan_to_saturate(tensor_fp8: torch.Tensor) -> torch.Tensor: + if tensor_fp8.dtype == torch.float8_e4m3fn: + tensor_u8 = tensor_fp8.view(torch.uint8) + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 448 + tensor_f32[is_neg_nan] = -448 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + elif tensor_fp8.dtype == torch.float8_e5m2: + tensor_u8 = tensor_fp8.view(torch.uint8) + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 57344 + tensor_f32[is_neg_nan] = -57344 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + else: + tensor_f8_out = tensor_fp8 + return tensor_f8_out + + # Compare GEMM output (blockscaled reference) + def _compare_gemm_bs( + a_raw, b_raw, sfa_raw, sfb_raw, torch_tensor_c, c_dtype, atol, rtol + ): + if emulated: + c_ref = reference_scaled_mm_emulated( + a_raw, + b_raw, + sfa_raw, + sfb_raw, + torch_tensor_c, + mnkl, + sf_vec_size, + c_dtype, + ) + else: + c_ref = reference_scaled_mm( + a_raw, b_raw, sfa_raw, sfb_raw, torch_tensor_c, mnkl, c_dtype + ) + c_out_f32 = torch_tensor_c.cpu().to(dtype=torch.float32) + c_ref_f32 = c_ref.cpu().to(dtype=torch.float32) + torch.testing.assert_close(c_out_f32, c_ref_f32, atol=atol, rtol=rtol) + + def _compare_comm(torch_tensor_comm_in, torch_tensor_comm_out, c_dtype, atol, rtol): + dist.barrier() + if c_dtype in [cutlass.Float32, cutlass.BFloat16, cutlass.Float16]: + # Use Float32 accumulation for an accurate reference (avoids Float16 overflow + # discrepancy between symm_mem contiguous path vs regular tensor path in NCCL). + comm_ref_f32_acc = torch_tensor_comm_in.to(dtype=torch.float32).contiguous() + dist.all_reduce(comm_ref_f32_acc, op=torch.distributed.ReduceOp.SUM) + comm_ref_f32 = comm_ref_f32_acc.cpu() + comm_out_f32 = torch_tensor_comm_out.cpu().to(dtype=torch.float32) + # multimem.ld_reduce reduces in Float16 and may overflow to ±inf for large values. + # Clamp both reference and output to the dtype's representable range so that + # hardware overflow (→ ±inf) compares equal to the clamped reference. + if c_dtype == cutlass.Float16: + float16_max = 65504.0 + comm_ref_f32 = comm_ref_f32.clamp(-float16_max, float16_max) + comm_out_f32 = comm_out_f32.nan_to_num( + nan=0.0, posinf=float16_max, neginf=-float16_max + ) + elif c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_ref_fp32_acc = torch_tensor_comm_in.to( + dtype=torch.float32 + ).contiguous() + dist.all_reduce(comm_ref_fp32_acc, op=torch.distributed.ReduceOp.SUM) + epsilon = 1e-2 + comm_ref_f32_acc_epsilon = comm_ref_fp32_acc + epsilon + comm_ref_c = comm_ref_f32_acc_epsilon.to(dtype=cutlass_torch.dtype(c_dtype)) + comm_ref_nan2sat = _fp8_nan_to_saturate(comm_ref_c) + comm_ref_f32 = comm_ref_nan2sat.to(dtype=torch.float32).cpu() + comm_out = torch_tensor_comm_out.cpu() + comm_out_nan2sat = _fp8_nan_to_saturate(comm_out) + comm_out_f32 = comm_out_nan2sat.to(dtype=torch.float32) + torch.testing.assert_close(comm_ref_f32, comm_out_f32, atol=atol, rtol=rtol) + + tol = 0.0 if c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN) else 1e-1 + _compare_gemm_bs( + a_raw, b_raw, sfa_raw, sfb_raw, torch_tensor_c, c_dtype, atol=tol, rtol=tol + ) + + # Compare AllReduce output + comm_atol = 5 + comm_rtol = 1e-1 + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_rtol = 1e-1 + comm_atol = 0.0 + _compare_comm( + torch_tensor_comm_in=torch_tensor_c, + torch_tensor_comm_out=torch_tensor_comm_out, + c_dtype=c_dtype, + atol=comm_atol, + rtol=comm_rtol, + ) + + # Check that flags were reset to zero + torch.cuda.synchronize() + flag_after = torch_tensor_flag_unicast.cpu() + nonzero_count = int(torch.count_nonzero(flag_after).item()) + assert nonzero_count == 0, ( + f"Kernel did not reset all barrier flags to zero: " + f"{nonzero_count}/{flag_after.numel()} entries non-zero " + f"(first 16 values: {flag_after[:16].tolist()})" + ) + + +def compile_and_benchmark( + *, + mnkl, + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + global_rank, + world_size, + device, + num_workspace, + tensors, + torch_stream, + stream, + warmup_iterations, + iterations, + use_cuda_graph, +): + """Compile one blockscaled-GEMM + AR candidate and time it. Returns avg_us.""" + m, n, k, l = mnkl + + print(f"Swizzle Size: {swizzle_size}") + print(f"Raster Order: {raster_order}") + gemm = Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + c_dtype=c_dtype, + rank_id=global_rank, + num_ranks=world_size, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + try: + import hashlib + + param_str = ( + f"a{a_dtype}_b{b_dtype}_sf{sf_dtype}_sfv{sf_vec_size}_c{c_dtype}" + f"_{a_major}{b_major}{c_major}" + f"_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}" + ) + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = f"sm100_blockscale_dense_gemm_ar_ar_{param_hash}_r{global_rank}_w{world_size}" + except Exception: + unique_name = f"sm100_blockscale_dense_gemm_ar_ar_r{global_rank}_w{world_size}" + gemm.__call__.set_name_prefix(unique_name) + + if not Sm100PersistentBlockScaleDenseGemmAllReduceLDMCxSTMCKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + sf_dtype, + c_dtype, + a_major, + b_major, + c_major, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + num_ranks=world_size, + ): + raise testing.CantImplementError( + f"Unsupported config: a={a_dtype} b={b_dtype} sf={sf_dtype} sfv={sf_vec_size} " + f"c={c_dtype} tile={mma_tiler_mn} cluster={cluster_shape_mn}" + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + compiled_gemm = scaled_mm_ar( + gemm, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, + max_active_clusters, + stream, + tensors["cute_tensor_comm_in_mc"], + tensors["cute_tensor_comm_out_mc"], + tensors["cute_tensor_flag_unicast"], + tensors["cute_tensor_flag_multicast"], + ) + + a_ptr_list = tensors["a_ptr_list"] + b_ptr_list = tensors["b_ptr_list"] + sfa_ptr_list = tensors["sfa_ptr_list"] + sfb_ptr_list = tensors["sfb_ptr_list"] + cute_c = tensors["cute_tensor_c"] + comm_in_mc = tensors["cute_tensor_comm_in_mc"] + comm_out_mc = tensors["cute_tensor_comm_out_mc"] + flag_uc = tensors["cute_tensor_flag_unicast"] + flag_mc = tensors["cute_tensor_flag_multicast"] + + def make_kernel_args(ws): + return ( + a_ptr_list[ws], + b_ptr_list[ws], + sfa_ptr_list[ws], + sfb_ptr_list[ws], + cute_c, + (m, n, k, l), + stream, + lambda x: x, + comm_in_mc, + comm_out_mc, + flag_uc, + flag_mc, + ) + + kernel_args_list = [make_kernel_args(i) for i in range(num_workspace)] + + if use_cuda_graph: + graphs = [] + for ws in range(num_workspace): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch_stream): + compiled_gemm(*kernel_args_list[ws]) + graphs.append(g) + + def run_one_iter(i): + graphs[i % num_workspace].replay() + else: + + def run_one_iter(i): + compiled_gemm(*kernel_args_list[i % num_workspace]) + + torch.cuda.synchronize() + dist.barrier() + + total_iters = warmup_iterations + iterations + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + with torch.cuda.stream(torch_stream): + for i in range(total_iters): + start_events[i].record() + run_one_iter(i) + end_events[i].record() + end_events[-1].synchronize() + + per_iter_us = [s.elapsed_time(e) * 1000.0 for s, e in zip(start_events, end_events)] + avg_us = sum(per_iter_us[warmup_iterations:]) / iterations + print( + f"Rank {global_rank} Execution time: {avg_us} us " + f"(avg of {iterations} iters, after {warmup_iterations} warmup)\n", + flush=True, + ) + return avg_us + + +def run_distributed_benchmark( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + warmup_iterations: int = 0, + iterations: int = 1, + use_cuda_graph: bool = False, + benchmark_all: bool = False, + csv_report: str = "./csv_report.csv", +): + num_workspace = 10 + tensors = allocate_tensors( + mnkl=mnkl, + a_dtype=a_dtype, + b_dtype=b_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + num_workspace=num_workspace, + slot_init_mode="benchmark", + device=device, + global_rank=global_rank, + world_size=world_size, + ) + + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + base_kwargs = dict( + mnkl=mnkl, + a_dtype=a_dtype, + b_dtype=b_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + global_rank=global_rank, + world_size=world_size, + device=device, + num_workspace=num_workspace, + tensors=tensors, + torch_stream=torch_stream, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cuda_graph=use_cuda_graph, + ) + + if benchmark_all: + raster_candidates = ["m", "n"] + swizzle_candidates = [1, 2, 4, 8] + # use_2cta_instrs is implied by mma_tiler_mn[0] == 256. + # 2CTA requires cluster_m in {2,4,8}; 1CTA accepts any cluster. + two_cta_options = [ + { + "clusters": [(2, 1), (2, 2)], + "mma_tilers": [(256, 64), (256, 128), (256, 192), (256, 256)], + }, + { + "clusters": [(1, 1), (1, 2), (2, 1), (2, 2)], + "mma_tilers": [(128, 64), (128, 128), (128, 192), (128, 256)], + }, + ] + candidates = [] + for opt in two_cta_options: + for cluster in opt["clusters"]: + for tile in opt["mma_tilers"]: + for r in raster_candidates: + for s in swizzle_candidates: + candidates.append( + { + **base_kwargs, + "mma_tiler_mn": tile, + "cluster_shape_mn": cluster, + "raster_order": r, + "swizzle_size": s, + } + ) + else: + candidates = [ + { + **base_kwargs, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "raster_order": raster_order, + "swizzle_size": swizzle_size, + } + ] + + if global_rank == 0: + print(f"Total number of candidate configs: {len(candidates)}", flush=True) + + results = [] + for i, cand in enumerate(candidates): + if i > 0: + time.sleep(2) + try: + avg_us = compile_and_benchmark(**cand) + except testing.CantImplementError: + continue + results.append((cand, avg_us)) + + m, n, k, l = mnkl + fieldnames = [ + "gemm_m", + "gemm_n", + "gemm_k", + "gemm_l", + "a_dtype", + "b_dtype", + "sf_dtype", + "sf_vec_size", + "c_dtype", + "a_major", + "b_major", + "c_major", + "mma_tiler_m", + "mma_tiler_n", + "cluster_m", + "cluster_n", + "use_cuda_graph", + "swizzle_size", + "raster_order", + "num_workspace", + "warmup_iterations", + "iterations", + "world_size", + "global_rank", + "benchmark_time_us", + ] + csv_path = Path(csv_report) + csv_path = csv_path.with_name( + f"{csv_path.stem}_rank{global_rank}_world{world_size}{csv_path.suffix}" + ) + csv_path.parent.mkdir(parents=True, exist_ok=True) + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for cand, avg_us in results: + writer.writerow( + { + "gemm_m": m, + "gemm_n": n, + "gemm_k": k, + "gemm_l": l, + "a_dtype": str(a_dtype), + "b_dtype": str(b_dtype), + "sf_dtype": str(sf_dtype), + "sf_vec_size": sf_vec_size, + "c_dtype": str(c_dtype), + "a_major": a_major, + "b_major": b_major, + "c_major": c_major, + "mma_tiler_m": cand["mma_tiler_mn"][0], + "mma_tiler_n": cand["mma_tiler_mn"][1], + "cluster_m": cand["cluster_shape_mn"][0], + "cluster_n": cand["cluster_shape_mn"][1], + "use_cuda_graph": use_cuda_graph, + "swizzle_size": cand["swizzle_size"], + "raster_order": cand["raster_order"], + "num_workspace": num_workspace, + "warmup_iterations": warmup_iterations, + "iterations": iterations, + "world_size": world_size, + "global_rank": global_rank, + "benchmark_time_us": avg_us, + } + ) + + return results[0][1] if len(results) == 1 else results + + +def prepare_parser(): + parser = argparse.ArgumentParser( + description="SM100 BlockScale GEMM + AllReduce (LDMCxSTMC) Fusion Kernel." + ) + + def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser.add_argument( + "--mnkl", type=_parse_comma_separated_ints, default=(512, 512, 256, 1) + ) + parser.add_argument( + "--mma_tiler_mn", type=_parse_comma_separated_ints, default=(128, 128) + ) + parser.add_argument( + "--cluster_shape_mn", type=_parse_comma_separated_ints, default=(1, 1) + ) + parser.add_argument("--a_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--b_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E4M3FN) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) + parser.add_argument("--a_major", choices=["k", "m"], default="k") + parser.add_argument("--b_major", choices=["k", "n"], default="k") + parser.add_argument("--c_major", choices=["n", "m"], default="n") + parser.add_argument( + "--benchmark_or_test", + type=str, + choices=["test", "benchmark", "benchmark_all"], + default="test", + help="Select run mode: 'test' for correctness, 'benchmark' for timing, 'benchmark_all' for a sweep.", + ) + parser.add_argument("--warmup_iterations", type=int, default=0) + parser.add_argument("--iterations", type=int, default=1) + parser.add_argument( + "--use_cuda_graph", + action="store_true", + default=False, + help="Capture one iteration into a CUDA graph and replay per iteration", + ) + parser.add_argument( + "--swizzle_size", + type=int, + default=1, + help="Swizzling size in the unit of cluster for improving L2 cache hit rate", + ) + parser.add_argument( + "--raster_order", + type=str, + choices=["m", "n"], + default="m", + help="Rasterization order of clusters", + ) + parser.add_argument( + "--csv_report", + type=str, + default="./csv_report.csv", + help="Path for the per-rank benchmark CSV. Each rank appends _rank{R}_world{W} before the extension.", + ) + parser.add_argument( + "--all_reduce", + type=str, + default="LDMCxSTMC", + help="AllReduce algorithm (currently only LDMCxSTMC is supported).", + ) + return parser + + +if __name__ == "__main__": + parser = prepare_parser() + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + run( + mnkl=args.mnkl, + a_dtype=args.a_dtype, + b_dtype=args.b_dtype, + sf_dtype=args.sf_dtype, + sf_vec_size=args.sf_vec_size, + c_dtype=args.c_dtype, + a_major=args.a_major, + b_major=args.b_major, + c_major=args.c_major, + mma_tiler_mn=args.mma_tiler_mn, + cluster_shape_mn=args.cluster_shape_mn, + swizzle_size=args.swizzle_size, + raster_order=args.raster_order, + warmup_iterations=args.warmup_iterations, + iterations=args.iterations, + benchmark_or_test=args.benchmark_or_test, + use_cuda_graph=args.use_cuda_graph, + csv_report=args.csv_report, + ) + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + print( + f"PASS (hostname={socket.gethostname()}, global_rank={rank}, local_rank={local_rank})", + flush=True, + ) diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_reduce_scatter_blackwell.py b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_reduce_scatter_blackwell.py index 27254b2fd..083dbd2e0 100644 --- a/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_reduce_scatter_blackwell.py +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_reduce_scatter_blackwell.py @@ -26,13 +26,16 @@ # 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. -import argparse import os -from typing import Optional, Tuple, Type, Union +import argparse +import csv +import socket +import time +from pathlib import Path + +from math import gcd +from typing import Optional, Tuple, Type, Union, Literal -import numpy as np -import torch -import torch.distributed as dist import cuda.bindings.driver as cuda try: from cuda.core import Device @@ -40,268 +43,408 @@ except ImportError: from cuda.core.experimental import Device from cuda.pathfinder import load_nvidia_dynamic_lib +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.torch as cutlass_torch +from cutlass.torch import dtype as torch_dtype +from cutlass import testing import cutlass.utils as utils import cutlass.pipeline as pipeline -import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cutlass_dsl import BaseDSL from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack -from cutlass.cute.typing import Int32, Float16, BFloat16, Float32, Float8E4M3FN, Float8E5M2 -from cutlass.cutlass_dsl import T -from cutlass._mlir.dialects import llvm - -try: - import nvshmem.core -except ImportError as exc: - raise ImportError( - "nvshmem4py is required but not installed. Please install it using:\n" - " For CUDA 12: pip install nvshmem4py-cu12\n" - " For CUDA 13: pip install nvshmem4py-cu13\n" - "Note: nvshmem4py version >= 0.1.3 is recommended." - ) from None - -try: - load_nvidia_dynamic_lib("nvshmem_host") -except RuntimeError as exc: - raise ImportError( - "nvshmem lib is required but not installed. Please install it using:\n" - " For CUDA 12: pip install nvidia-nvshmem-cu12\n" - " For CUDA 13: pip install nvidia-nvshmem-cu13\n" - ) from None +from cutlass._mlir import ir +from cutlass._mlir.dialects import vector """ -A high-performance distributed 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 C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") -- Matrix C_mc is a multicast C matrix that changes can be broadcasted to all GPUs by multimem instructions. +SM100 GEMM + ReduceScatter (LDMC) Fusion Kernel -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 - - Support all-reduce epilogue with multimem instructions to distribute the workload to all GPUs +* Test (correctness only — single launch, no warmup/iteration loop) +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_reduce_scatter_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs --reduce_scatter LDMC \ + --benchmark_or_test test -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)) -4. Reduce scatter epilogue: - - Load and reduce the 128bit data from all ranks by multimem instructions. - - Broadcast the reduced data to all ranks by multimem instructions. - - current implementation only supports two_shot all-reduce which means each rank only computes a portion of - the output tensor and broadcast the result to all ranks. - - the all-reduce epilogue is only supported when use_tma_store is True. - - the all-reduce epilogue is only supported when c_dtype is Float16, Float32, BFloat16, Float8E4M3FN, Float8E5M2. +* Benchmark One +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_reduce_scatter_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --use_2cta_instrs --reduce_scatter LDMC \ + --benchmark_or_test benchmark --use_cuda_graph --csv_report ./csv_report.csv -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. +* Benchmark All +python3 -m torch.distributed.run --nproc-per-node=auto examples/python/CuTeDSL/cute/blackwell/kernel/distributed/distributed_gemm_reduce_scatter_blackwell.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --a_major k --b_major k --c_major n \ + --mnkl 8192,8192,8192,1 --warmup_iterations 10 --iterations 20 \ + --use_tma_store --reduce_scatter LDMC \ + --benchmark_or_test benchmark_all --use_cuda_graph --csv_report ./csv_report.csv -Input arguments to this example is same as dense_gemm.py. +NOTE: +* multimem.ld_reduce does NOT handle NaN issue for f8. -.. code-block:: bash - - torchrun --nproc-per-node 8 examples/distributed/distributed_gemm_reduce_scatter_blackwell.py \ - --ab_dtype Float8E4M3FN --c_dtype Float16 --acc_dtype Float32 \ - --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ - --mnkl 16384,4080,4096,1 --warmup_iterations 3 --iterations 10 \ - --use_tma_store --use_2cta_instrs --reduce_scatter two_shot - -To collect performance with NSYS profiler: - -.. code-block:: bash - - nsys profile --gpu-metrics-devices=cuda-visible \ - torchrun --nproc-per-node 8 examples/distributed/distributed_gemm_reduce_scatter_blackwell.py \ - --ab_dtype Float8E5M2 --c_dtype Float16 --acc_dtype Float32 \ - --mma_tiler_mn 256,256 --cluster_shape_mn 2,1 \ - --mnkl 16384,4096,4096,1 \ - --use_tma_store --use_2cta_instrs --warmup_iterations 3 --iterations 10 \ - --reduce_scatter two_shot - - -Constraints are same as dense_gemm_persistent.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 -* when reduce_scatter, M must be multiple of 128, world_size must be 2, 4, 8 +ALGO: +-> Producer +for each output tile T + * store gemm output [data] to local rank global memory with .weak semantics (tma is default .weak semantic) + * [flag] update with .release semantics and .gpu scope, UNICAST +1 to owner(T)'s flag slot only + (each tile has exactly ONE owner rank in RS, unlike AR where every rank consumes) +-> Consumer +for each output tile T owned by this rank (only owner(T) runs this; P CTAs on owner cooperate, splitting T into P m-stripes) + * spin lock wait [flag] for >= N with .relaxed / .acquire and .gpu scope (all N producers have written to local rank gmem) + * multimem.ld_reduce on this CTA's m-stripe of T: send reduce request to switch, switch send request to all rank's gmem, reduce on switch, return to requester rank + * st [data] (UC) to local rank's RS output (each rank keeps only its 1/N shard — NO multimem.st because RS doesn't broadcast back) + * consumer-arrival count: atomic_add +1 with .relaxed / .gpu; last CTA (result == 2N-1) resets slot to 0 (placed after data movement so warp[0]'s atomic_add overlaps with other warps' data path) +before kernel exit + * [flag] update with .release semantics and .gpu scope (prior store is to LOCAL gmem only; .gpu suffices — AR needs .sys because its prior multimem.st targets all peers) + * spin lock wait [flag] with .relaxed / .acquire and .gpu scope """ +##################################################################### +#### KERNEL IMPLEMENTATION START HERE #### +##################################################################### -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] - :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results - :type use_tma_store: bool - :param reduce_scatter: reduce scatter mode, can be "two_shot" - :type reduce_scatter: str +@cute.jit +def LDMC( + kernel, + cur_tile_coord, + tiled_mma: cute.TiledMma, + comm_tidx: cutlass.Int32, + tCgCommInMC: cute.Tensor, + tCgCommOutUC: cute.Tensor, + m_cta_tiles_per_rank, +) -> None: + # Multiple CTAs per rank cooperatively process its RS output tiles + # + # Naive RS Implementation (inefficient): + # - On rank0: Only CTAs computing partition A tiles would execute load multiple CTA + store operations + # - On rank0: CTAs computing partitions B/C/D tiles would remain idle (since rank0 only stores partition A) for its RS part + # - Result: Limited in-flight load multiple CTA instruction as there's only 1/P CTA tile per rank issuing load multiple CTA instruction + # + # Our Implementation (efficient): + # - On rank0: All CTAs computing partition A/B/C/D tile would execute load multiple CTA + store operations + # - On rank0: partition A's output tile 0 is computed by + # partition A's output tile 0 + # partition B's output tile 0 + # partition C's output tile 0 + # partition D's output tile 0 + # where each RS output tile is computed by P number of CTAs per rank. + # - For row-major: split of work is done horizontally (along M) across each partition. + # - For col-major: split of work is done vertically (along N) across each partition. + # Code Explain + # % m_tiles_per_rank: get output tile's within partition tileid + # for output tile in rank0 partition B, get it RS tile in partition A + # + self.rank_id * m_tiles_per_rank: offset to current rank's RS output tile + # for rank0, this offset to start of partition A + # for rank1, this offset to start of partition B + # for rank2, this offset to start of partition C + # for rank3, this offset to start of partition D - :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 + # RS tile coordinate that current CTA handle (likely different output tile) + input_tile_coord = ( + cur_tile_coord[0] % m_cta_tiles_per_rank + + kernel.rank_id * m_cta_tiles_per_rank, + cur_tile_coord[1], + cur_tile_coord[2], + ) + mma_tile_coord_mnl_input = ( + ( + (cur_tile_coord[0] % m_cta_tiles_per_rank) + // cute.size(tiled_mma.thr_id.shape) + ) + + kernel.rank_id * (m_cta_tiles_per_rank // cute.size(tiled_mma.thr_id.shape)), + cur_tile_coord[1], + cur_tile_coord[2], + ) + mma_tile_coord_mnl_output = ( + ( + (cur_tile_coord[0] % m_cta_tiles_per_rank) + // cute.size(tiled_mma.thr_id.shape) + ), + cur_tile_coord[1], + cur_tile_coord[2], + ) - :note: Supported A/B data types: - - TFloat32 - - Float16/BFloat16 - - Int8/Uint8 - - Float8E4M3FN/Float8E5M2 + # partition and slice at tile level + tCgCommInMC_slice = tCgCommInMC[((None, None), 0, 0, *mma_tile_coord_mnl_input)] + tCgCommOutUC_slice = tCgCommOutUC[((None, None), 0, 0, *mma_tile_coord_mnl_output)] - :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) + # + # Partition s.t. each tile handle 1/P of processed tile + # For row-major (N-contiguous): each CTA handles a horizontal stripe (M/P, N) + # For col-major (M-contiguous): each CTA handles a vertical stripe (M, N/P) + # + cta_mma_tile_m = kernel.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) + cta_tile_n = kernel.mma_tiler[1] + total_comm_threads = len(kernel.comm_warp_id) * cute.arch.WARP_SIZE - :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) + # chunk_id determines which portion of the tile this CTA processes + chunk_id = cur_tile_coord[0] // m_cta_tiles_per_rank - :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 + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + # Row-major (N-contiguous): split along M + m_local_rank = int(cta_mma_tile_m / kernel.num_ranks) + tCgCommInMC_slice_partitioned = cute.zipped_divide( + tCgCommInMC_slice, (m_local_rank, cta_tile_n) + ) + tCgCommOutUC_slice_partitioned = cute.zipped_divide( + tCgCommOutUC_slice, (m_local_rank, cta_tile_n) + ) + # Data processed by local CTA: select M chunk + tCgCommInMC_local_rank = cute.slice_( + tCgCommInMC_slice_partitioned, ((None, None), (chunk_id, 0)) + ) + tCgCommOutUC_local_rank = cute.slice_( + tCgCommOutUC_slice_partitioned, ((None, None), (chunk_id, 0)) + ) + local_tile_m = m_local_rank + local_tile_n = cta_tile_n + else: + # Col-major (M-contiguous): split along N + n_local_rank = int(cta_tile_n / kernel.num_ranks) + tCgCommInMC_slice_partitioned = cute.zipped_divide( + tCgCommInMC_slice, (cta_mma_tile_m, n_local_rank) + ) + tCgCommOutUC_slice_partitioned = cute.zipped_divide( + tCgCommOutUC_slice, (cta_mma_tile_m, n_local_rank) + ) + # Data processed by local CTA: select N chunk + tCgCommInMC_local_rank = cute.slice_( + tCgCommInMC_slice_partitioned, ((None, None), (0, chunk_id)) + ) + tCgCommOutUC_local_rank = cute.slice_( + tCgCommOutUC_slice_partitioned, ((None, None), (0, chunk_id)) + ) + local_tile_m = cta_mma_tile_m + local_tile_n = n_local_rank - 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) + # + # Determine vectorization length based on data size per thread + # Always use 128-bit instructions (guaranteed by _pick_num_comm_warp_for_128b) + # + elements_per_thread = (local_tile_m * local_tile_n) // total_comm_threads + + # Always use 128-bit atom size + # Guaranteed by _pick_num_comm_warp_for_128b which selects T to ensure + # elements_per_thread is divisible by 128-bit atom (16B = 128 bits) + atom_num_elements_128 = 128 // kernel.c_dtype.width + atom_num_elements = atom_num_elements_128 + + # COMPILE-TIME ASSERTION: Verify 128-bit PTX instruction usage + # _pick_num_comm_warp_for_128b selects num_warps such that every thread + # in the RS warp group handles exactly atom_num_elements_128 elements, + # which means exactly one 128-bit LDMC instruction per RS slab iteration. + assert elements_per_thread % atom_num_elements_128 == 0, ( + f"KERNEL ASSERTION: 128-bit PTX instruction invariant violated. " + f"elements_per_thread={elements_per_thread} must be divisible by " + f"atom_num_elements_128={atom_num_elements_128}. " + f"This ensures every thread executes exactly one 128-bit instruction. " + f"Configuration: local_tile=({local_tile_m},{local_tile_n}), " + f"total_comm_threads={total_comm_threads}, c_dtype={kernel.c_dtype}" + ) + + # + # Create thr_copy_fake with dynamically selected atom size + # For row-major: threads spread along N (contiguous), load along N + # For col-major: threads spread along M (contiguous), load along M + # + if cutlass.const_expr(kernel.c_layout.is_n_major_c()): + # Row-major: N is contiguous, spread threads along N. + # Use GCD to handle non-power-of-2 tile widths (e.g. cta_n=192) + # that don't naively divide the thread count. + max_thr_n = local_tile_n // atom_num_elements + atom_thr_n = gcd(max_thr_n, total_comm_threads) + atom_thr_m = total_comm_threads // atom_thr_n + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) + val_layout = cute.make_layout( + (1, atom_num_elements), stride=(atom_num_elements, 1) + ) + else: + # Col-major: M is contiguous, spread threads along M. + max_thr_m = local_tile_m // atom_num_elements + atom_thr_m = gcd(max_thr_m, total_comm_threads) + atom_thr_n = total_comm_threads // atom_thr_m + thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(1, atom_thr_m)) + val_layout = cute.make_layout( + (atom_num_elements, 1), stride=(1, atom_num_elements) + ) + + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), kernel.c_dtype) + tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + thr_copy_fake = tiled_copy_fake.get_slice(comm_tidx - kernel.comm_warp_id[0] * 32) + + # + # Partition per thread + # + tTR_gCommInMC = thr_copy_fake.partition_S(tCgCommInMC_local_rank) + tTR_gCommOutUC = thr_copy_fake.partition_S(tCgCommOutUC_local_rank) + _atom, loop_m, loop_n = tTR_gCommInMC.shape + + # + # Reduce Scatter LDMC — always 128-bit (guaranteed by _pick_num_comm_warp_for_128b) + # + tmp_results = [] + for i in cutlass.range_constexpr(loop_m): + for j in cutlass.range_constexpr(loop_n): + comm_in_mc_ptr = tTR_gCommInMC[None, i, j].iterator + regs = utils.distributed.multimem_ld_reduce( + comm_in_mc_ptr, dtype=kernel.c_dtype, num_elements=atom_num_elements + ) + tmp_results.append(regs) + + # + # Store to unicast memory with 128-bit vector + # + for i in cutlass.range_constexpr(loop_m): + for j in cutlass.range_constexpr(loop_n): + comm_out_uc_ptr = tTR_gCommOutUC[None, i, j].iterator.llvm_ptr + x, y, z, w = tmp_results[i * loop_n + j] + vec_type = ir.VectorType.get([4], cutlass.Int32.mlir_type) + vec = vector.from_elements(vec_type, [x, y, z, w]) + cute.arch.store(comm_out_uc_ptr, vec) + + +class Sm100PersistentDenseGemmReduceScatterLDMCKernel: + """ + **Example:** + gemm = Sm100PersistentDenseGemmReduceScatterLDMCKernel( + acc_dtype=cutlass.Float32, + c_dtype=cutlass.BFloat16, + use_2cta_instrs=True, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + use_tma_store=True, + rank_id=rank_id, + num_ranks=num_ranks, + reduce_scatter="LDMC", + swizzle_size=1, + raster_order="m", + ) + compiled_gemm = cute.compile( + gemm, + a=a, b=b, c=c, + comm_in_multicast_tensor=comm_in_mc, + comm_out_unicast_tensor=comm_out_uc, + barrier_flag_unicast=flag_unicast, + barrier_flag_multicast=flag_multicast, + barrier_flag_unicast_per_peer=flag_unicast_per_peer, + stream=stream, + max_active_clusters=max_active_clusters, + ) + compiled_gemm( + a=a, b=b, c=c, + comm_in_multicast_tensor=comm_in_mc, + comm_out_unicast_tensor=comm_out_uc, + barrier_flag_unicast=flag_unicast, + barrier_flag_multicast=flag_multicast, + barrier_flag_unicast_per_peer=flag_unicast_per_peer, + stream=stream, + ) """ def __init__( self, acc_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], use_2cta_instrs: bool, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], use_tma_store: bool, - reduce_scatter="two_shot", + rank_id: int, + num_ranks: int, + reduce_scatter: str = "LDMC", + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", ): - """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: - - use_tma_store: Boolean indicating whether to use Tensor Memory Access (TMA) for storing results. - - :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] - :param use_tma_store: Use Tensor Memory Access (TMA) or normal store for output C tensor. - :type use_tma_store: bool - :param reduce_scatter: reduce scatter mode, can be "two_shot" - :type reduce_scatter: str - """ - self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.c_dtype: Type[cutlass.Numeric] = c_dtype self.use_2cta_instrs = use_2cta_instrs self.cluster_shape_mn = cluster_shape_mn + self.swizzle_size = swizzle_size + self.raster_order = raster_order # K dimension is deferred in _setup_attributes self.mma_tiler_mn = mma_tiler_mn self.mma_tiler = (*mma_tiler_mn, 1) self.use_tma_store = use_tma_store + # Capture the JIT target arch so SMEM/TMEM sizing tracks the actual + # hardware (sm_100 / sm_103 / sm_107 / ...). Matches the FMHA pattern. + arch_enum = BaseDSL._get_dsl().get_arch_enum() + self.arch = f"sm_{arch_enum.major}{arch_enum.minor}" self.cta_group = ( tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE ) - self.reduce_scatter = reduce_scatter - self.occupancy = 1 # Set specialized warp ids - self.epilog_warp_id = ( - 0, - 1, - 2, - 3, - ) + self.epilogue_warp_id = (0, 1, 2, 3) self.mma_warp_id = 4 self.tma_warp_id = 5 self.reduce_scatter = reduce_scatter - self.reduce_scatter_warp_id = (6, 7, 8, 9) + if reduce_scatter == "LDMC": + self.rank_id = rank_id + self.num_ranks = num_ranks + num_comm_warps = self._pick_num_comm_warp_for_128b(c_dtype) + self.comm_warp_id = tuple(range(6, 6 + num_comm_warps)) + else: + # Other RS mode not support by now. + self.comm_warp_id = () + self.rank_id = 0 + self.num_ranks = 1 + self.threads_per_cta = 32 * len( ( self.mma_warp_id, self.tma_warp_id, - *self.epilog_warp_id, - *self.reduce_scatter_warp_id, + *self.epilogue_warp_id, + *self.comm_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.reduce_scatter_sync_bar_id = 3 - self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.comm_sync_bar_id = 4 - self.num_ranks = 1 - self.rank_id = 0 - self.num_ranks = torch.distributed.get_world_size() - self.rank_id = torch.distributed.get_rank() + def _pick_num_comm_warp_for_128b(self, c_dtype): + """Pick the LARGEST comm-warp count W ∈ {4, 3, 2, 1} (T = W·32 ∈ + {128, 96, 64, 32}) such that the RS slab is evenly distributed + across T threads AND every thread gets a full 128 b LDMC payload + (= ``atom_128`` elements, where ``atom_128 = 16 / sizeof(c_dtype)``). - def is_valid(self): - mma_m, mma_n = self.mma_tile_shape_mn - if (mma_m // (2 if self.use_2cta_instrs else 1)) not in [64, 128]: - return False - if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: - return False - if self.cluster_shape_mn[0] == 4 and self.cluster_shape_mn[1] == 4: - return False - return True - - self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + Iterates W = 4 → 1 and returns the first match, so the largest + viable W is returned. Smaller W is only chosen when strictly + necessary, because larger W means more parallel RS threads and + proportionally fewer per-thread PTX instructions. + """ + # cta-tile dims (the per-CTA M·N output the RS walks). 2cta halves M. + cta_mma_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + slab_elements = (cta_mma_tile_m * cta_tile_n) // self.num_ranks + atom_128 = 16 // (c_dtype.width // 8) + for num_warps in (4, 3, 2, 1): + T = num_warps * 32 + if slab_elements % T: + continue + if (slab_elements // T) % atom_128 == 0: + return num_warps + raise ValueError( + f"no comm-warp count enables 128 b PTX for c_dtype={c_dtype}, " + f"cta=({cta_mma_tile_m},{cta_tile_n}), P={self.num_ranks}" + ) def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -317,9 +460,9 @@ class PersistentDenseGemmKernel: - Computing A/B/C shared memory layout - Computing tensor memory allocation columns """ - # Configure tiled mma - tiled_mma = sm100_utils.make_trivial_tiled_mma( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, + self.b_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, @@ -327,7 +470,6 @@ class PersistentDenseGemmKernel: 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 = ( @@ -341,21 +483,18 @@ class PersistentDenseGemmKernel: self.mma_tiler[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 if cutlass.const_expr(self.use_tma_store): - self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, @@ -364,47 +503,41 @@ class PersistentDenseGemmKernel: else: self.epi_tile = self.cta_tile_shape_mnk[:2] - # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + c_smem_layout = None + if cutlass.const_expr(self.use_tma_store): + c_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( tiled_mma, self.mma_tiler, self.a_dtype, self.b_dtype, - self.epi_tile, self.c_dtype, - self.c_layout, self.smem_capacity, self.occupancy, self.use_tma_store, + c_smem_layout, ) - # Compute A/B/C 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.a_smem_layout_staged = utils.sm100.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.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage ) - self.c_smem_layout_staged = ( - sm100_utils.make_smem_layout_epi( - self.c_dtype, - self.c_layout, - self.epi_tile, - self.num_c_stage, + + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage ) - if self.use_tma_store - else None - ) - # 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 + tiled_mma, self.mma_tiler, self.num_acc_stage, self.arch ) @cute.jit @@ -413,48 +546,25 @@ class PersistentDenseGemmKernel: a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, - max_active_clusters: cutlass.Constexpr, + comm_in_multicast_tensor: cute.Tensor, + comm_out_unicast_tensor: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, + barrier_flag_unicast_per_peer: Tuple[cute.Tensor, ...], stream: cuda.CUstream, + max_active_clusters: cutlass.Constexpr, epilogue_op: cutlass.Constexpr = lambda x: x, - c_mc: cute.Tensor = None, - c_peer_tensors: list = None, - barrier_flag: cute.Tensor = None, - barrier_flag_mc: cute.Tensor = None, ): - """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 synchronously - - :param a: Input tensor A - :type a: cute.Tensor - :param b: Input tensor B - :type b: cute.Tensor - :param c: Output tensor C - :type c: cute.Tensor - :param c_mc: Output symmetric tensor C_mc, any write or read to a multicast tensor will be broadcasted to all GPUs - :type c_mc: cute.Tensor - :param c_peer_tensors: List of peer tensors for all-reduce - :type c_peer_tensors: List[cute.Tensor] - :param barrier_flag: Barrier flag for all-reduce - :type barrier_flag: cute.Tensor - :param barrier_flag_mc: Barrier flag for multicast tensor - :type barrier_flag_mc: 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 + # Setup static attributes before smem/grid/tma computation. + # NOTE: self.c_dtype was set in __init__ from the explicit c_dtype + # arg. We only sanity-check here that the C tensor matches that hint. self.a_dtype: Type[cutlass.Numeric] = a.element_type self.b_dtype: Type[cutlass.Numeric] = b.element_type - self.c_dtype: Type[cutlass.Numeric] = c.element_type + if cutlass.const_expr(c.element_type != self.c_dtype): + raise TypeError( + f"C tensor element_type ({c.element_type}) does not match " + f"the c_dtype declared at __init__ ({self.c_dtype})" + ) 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.c_layout = utils.LayoutEnum.from_tensor(c) @@ -463,21 +573,23 @@ class PersistentDenseGemmKernel: 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( + tiled_mma = utils.sm100.make_trivial_tiled_mma( self.a_dtype, + self.b_dtype, self.a_major_mode, self.b_major_mode, self.acc_dtype, self.cta_group, self.mma_tiler[:2], ) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + 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( + a_op = utils.sm100.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)) @@ -494,7 +606,7 @@ class PersistentDenseGemmKernel: ) # Setup TMA load for B - b_op = sm100_utils.cluster_shape_to_tma_atom_B( + b_op = utils.sm100.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)) @@ -518,84 +630,58 @@ class PersistentDenseGemmKernel: tma_atom_c = None tma_tensor_c = None if cutlass.const_expr(self.use_tma_store): - epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), - c, - epi_smem_layout, - self.epi_tile, + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile ) # Compute grid size self.tile_sched_params, grid = self._compute_grid( - c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters - ) - - self.buffer_align_bytes = 1024 - - c_smem_size = ( - cute.cosize(self.c_smem_layout_staged.outer) if self.use_tma_store else 0 + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + self.swizzle_size, + self.raster_order, + max_active_clusters, ) # 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] + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] tmem_dealloc_mbar: cutlass.Int64 tmem_holding_buf: cutlass.Int32 - # (EPI_TILE_M, EPI_TILE_N, STAGE) - sC: cute.struct.Align[ - cute.struct.MemRange[ - self.c_dtype, - c_smem_size, - ], - self.buffer_align_bytes, - ] - # (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 synchronously self.kernel( - tiled_mma, - tma_atom_a, - tma_tensor_a, - tma_atom_b, - tma_tensor_b, - tma_atom_c, - tma_tensor_c if self.use_tma_store else c, - self.cluster_layout_vmnk, - self.a_smem_layout_staged, - self.b_smem_layout_staged, - self.c_smem_layout_staged, - self.epi_tile, - self.tile_sched_params, - epilogue_op, - c_mc, - c_peer_tensors, - barrier_flag, - barrier_flag_mc, + tiled_mma=tiled_mma, + tma_atom_a=tma_atom_a, + mA_mkl=tma_tensor_a, + tma_atom_b=tma_atom_b, + mB_nkl=tma_tensor_b, + tma_atom_c=tma_atom_c, + mC_mnl=(tma_tensor_c if self.use_tma_store else c), + mCommInMC_mnl=comm_in_multicast_tensor, + mCommOutUC_mnl=comm_out_unicast_tensor, + barrier_flag_unicast=barrier_flag_unicast, + barrier_flag_multicast=barrier_flag_multicast, + barrier_flag_unicast_per_peer=barrier_flag_unicast_per_peer, + cluster_layout_vmnk=self.cluster_layout_vmnk, + a_smem_layout_staged=self.a_smem_layout_staged, + b_smem_layout_staged=self.b_smem_layout_staged, + c_smem_layout_staged=self.c_smem_layout_staged, + epi_tile=self.epi_tile, + tile_sched_params=self.tile_sched_params, + epilogue_op=epilogue_op, ).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 @@ -611,6 +697,11 @@ class PersistentDenseGemmKernel: mB_nkl: cute.Tensor, tma_atom_c: Optional[cute.CopyAtom], mC_mnl: cute.Tensor, + mCommInMC_mnl: cute.Tensor, + mCommOutUC_mnl: cute.Tensor, + barrier_flag_unicast: cute.Tensor, + barrier_flag_multicast: cute.Tensor, + barrier_flag_unicast_per_peer: Tuple[cute.Tensor, ...], cluster_layout_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, b_smem_layout_staged: cute.ComposedLayout, @@ -618,14 +709,7 @@ class PersistentDenseGemmKernel: epi_tile: cute.Tile, tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, - c_mc: cute.Tensor, - c_peer_tensors: list, - barrier_flag: cute.Tensor, - barrier_flag_mc: cute.Tensor, ): - """ - GPU device kernel performing the Persistent batched GEMM computation. - """ warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) @@ -645,7 +729,6 @@ class PersistentDenseGemmKernel: # # 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( @@ -663,27 +746,23 @@ class PersistentDenseGemmKernel: 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.ptr - # 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( + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Warp) + ab_producer, ab_consumer = 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, - ) + enable_multicast_signaling=True, + defer_sync=True, + ).make_participants() # Initialize acc_pipeline (barrier) and states acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - num_acc_consumer_threads = len(self.epilog_warp_id) * ( + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( 2 if use_2cta_instrs else 1 ) acc_pipeline_consumer_group = pipeline.CooperativeGroup( @@ -695,40 +774,47 @@ class PersistentDenseGemmKernel: producer_group=acc_pipeline_producer_group, consumer_group=acc_pipeline_consumer_group, cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, ) + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) # 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() + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) # Cluster arrive after barrier init - if cute.size(self.cluster_shape_mn) > 1: - cute.arch.cluster_arrive_relaxed() + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # # Setup smem tensor A/B/C # - # (EPI_TILE_M, EPI_TILE_N, STAGE) - sC = ( - storage.sC.get_tensor( - c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner - ) - if self.use_tma_store - else None - ) # (MMA, MMA_M, MMA_K, STAGE) - sA = storage.sA.get_tensor( - a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + 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 + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, ) # @@ -819,12 +905,17 @@ class PersistentDenseGemmKernel: # # 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 - ) + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() # # Specialized TMA load warp @@ -834,15 +925,6 @@ class PersistentDenseGemmKernel: # # 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 @@ -865,44 +947,36 @@ class PersistentDenseGemmKernel: ] # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt - ab_producer_state.reset_count() - peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire( - ab_producer_state - ) + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + # # Tma load loop # for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): # Conditionally wait for AB buffer empty - ab_pipeline.producer_acquire( - ab_producer_state, peek_ab_empty_status - ) + handle = ab_producer.acquire_and_advance(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), + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, 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), + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, mcast_mask=b_full_mcast_mask, ) # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 - ab_producer_state.advance() peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire( - ab_producer_state - ) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() # # Advance to next tile @@ -913,43 +987,23 @@ class PersistentDenseGemmKernel: # # Wait A/B buffer empty # - ab_pipeline.producer_tail(ab_producer_state) + ab_producer.tail() # # 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, - ) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) # (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 ) @@ -968,12 +1022,10 @@ class PersistentDenseGemmKernel: tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] # Peek (try_wait) AB buffer full for k_tile = 0 - ab_consumer_state.reset_count() + ab_consumer.reset() peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt and is_leader_cta: - peek_ab_full_status = ab_pipeline.consumer_try_wait( - ab_consumer_state - ) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() # # Wait for accumulator buffer empty @@ -981,52 +1033,28 @@ class PersistentDenseGemmKernel: 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_tile in range(k_tile_cnt): if is_leader_cta: # Conditionally wait for AB buffer full - ab_pipeline.consumer_wait( - ab_consumer_state, peek_ab_full_status - ) + handle = ab_consumer.wait_and_advance(peek_ab_full_status) # tCtAcc += tCrA * tCrB - num_kblocks = cute.size(tCrA, mode=[2]) - for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): - kblock_coord = ( - None, - None, - kblock_idx, - ab_consumer_state.index, - ) - - cute.gemm( - tiled_mma, - tCtAcc, - tCrA[kblock_coord], - tCrB[kblock_coord], - tCtAcc, - ) - # Enable accumulate on tCtAcc after first kblock - tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm( + tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc + ) # Async arrive AB buffer empty - ab_pipeline.consumer_release(ab_consumer_state) + handle.release() - # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 - ab_consumer_state.advance() - peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt: - if is_leader_cta: - peek_ab_full_status = ab_pipeline.consumer_try_wait( - ab_consumer_state - ) + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() # # Async arrive accumulator buffer full @@ -1045,6 +1073,17 @@ class PersistentDenseGemmKernel: # Wait for accumulator buffer empty # acc_pipeline.producer_tail(acc_producer_state) + + sC = None + if cutlass.const_expr(self.use_tma_store): + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + # # Specialized epilogue warps # @@ -1052,98 +1091,41 @@ class PersistentDenseGemmKernel: # # 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, - ) + tmem.allocate(self.num_tmem_alloc_cols) # # 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, - ) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) # (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, - ) = self.epilog_tmem_copy_and_partition( - epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs - ) - - tTR_rC = None - tiled_copy_r2s = None - simt_atom = None - tRS_rC = None - tRS_sC = None - bSG_sC = None - bSG_gC_partitioned = None - tTR_gC_partitioned = None - if cutlass.const_expr(self.use_tma_store): - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( - tiled_copy_t2r, tTR_rC, epi_tidx, sC - ) - ( - tma_atom_c, - bSG_sC, - bSG_gC_partitioned, - ) = self.epilog_gmem_copy_and_partition( - epi_tidx, tma_atom_c, tCgC, epi_tile, sC - ) - else: - ( - simt_atom, - tTR_rC, - tTR_gC_partitioned, - ) = self.epilog_gmem_copy_and_partition( - epi_tidx, tiled_copy_t2r, tCgC, epi_tile, sC - ) + m_cta_tiles_in_total = gC_mnl.shape[2] * cute.size(tiled_mma.thr_id.shape) + n_cta_tiles_in_total = gC_mnl.shape[3] + m_cta_tiles_per_rank = m_cta_tiles_in_total // self.num_ranks # - # Persistent tile scheduling loop + # Persistent tile scheduling loop for epilogue # - 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 ) - - c_pipeline = None if cutlass.const_expr(self.use_tma_store): - # Threads/warps participating in tma store pipeline + assert tma_atom_c is not None and sC is not None c_producer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, - 32 * len(self.epilog_warp_id) + 32 * len(self.epilogue_warp_id), ) c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, - producer_group=c_producer_group, + num_stages=self.num_c_stage, producer_group=c_producer_group ) + epilogue_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + while work_tile.is_valid_tile: # Get tile coord from tile scheduler cur_tile_coord = work_tile.tile_idx @@ -1152,345 +1134,242 @@ class PersistentDenseGemmKernel: cur_tile_coord[1], cur_tile_coord[2], ) - # - # Slice to per mma tile index - # - bSG_gC = None - tTR_gC = None - if cutlass.const_expr(self.use_tma_store): - # ((ATOM_V, REST_V), EPI_M, EPI_N) - bSG_gC = bSG_gC_partitioned[ - ( - None, - None, - None, - *mma_tile_coord_mnl, - ) - ] - else: - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_gC = tTR_gC_partitioned[ - ( - None, - None, - 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) - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - if cutlass.const_expr(self.use_tma_store): - bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) - else: - tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_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(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - if cutlass.const_expr(self.use_tma_store): - # - # Convert to C type - # - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tRS_rC.store(acc_vec) - # - # Store C to shared memory - # - c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage - 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("async.shared", space="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 C to global memory - # - if warp_idx == self.epilog_warp_id[0]: - cute.copy( - tma_atom_c, - bSG_sC[(None, c_buffer)], - bSG_gC[(None, subtile_idx)], - ) - # 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, - ) - else: - # - # Convert to C type - # - acc_vec = tTR_rAcc.load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tTR_rC.store(acc_vec) - - # - # Store C to global memory - # - cute.copy( - simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)] - ) - - # - # Async arrive accumulator buffer empty - # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - - # Reduce Scatter - if cutlass.const_expr(self.reduce_scatter == "two_shot"): - tile_id = Int32( - tile_sched._current_work_linear_idx - * cute.size(self.cluster_shape_mn) + cute.arch.block_idx_in_cluster() - ) - if warp_idx == self.epilog_warp_id[0]: - cute.arch.cp_async_bulk_wait_group(0, read=False) - # System barrier to make sure that data from each GPU is in memory before reduce scatter - with cute.arch.elect_one(): - flag = barrier_flag_mc.iterator + tile_id - utils.distributed.multimem_red_add1( - lock_ptr=flag, - scope="gpu", - order="release", - ) - - # - # Advance to next tile + # Pre-advance to next tile # tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + + # Post-epilogue: arrive flag for distributed communication + if cutlass.const_expr(self.reduce_scatter == "LDMC"): + # 1D linear index of current output tile + cta_tile_id_linear = cute.Int32( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + + # Wait for C store complete + # Unlike regular epilogue where we only wait C store complete once at end of each kernel. + # Here we need to wait for C store complete for each output tile before we set the release flag. + c_pipeline.producer_tail() + + # Update flag with release semantic, unicast to the + # tile's single owner. Each output tile in RS has exactly + # one owner (P CTAs from that one rank consume it), so + # broadcasting +1 to all peers would leave dead state on + # non-owners that nothing ever resets. The constexpr loop + # materializes an N-way dispatch because the DSL can't + # index a tuple of cute Tensors by a runtime Int32. + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + owner_rank = cur_tile_coord[0] // m_cta_tiles_per_rank + for r in cutlass.range_constexpr(self.num_ranks): + if owner_rank == r: + utils.distributed.red_add1( + lock_ptr=barrier_flag_unicast_per_peer[ + r + ].iterator + + cta_tile_id_linear, + order="release", + scope="gpu", + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + # Post-epilogue: arrive flag for distributed communication + if cutlass.const_expr(self.reduce_scatter == "LDMC"): + # 1D linear index of current output tile + cta_tile_id_linear = cute.Int32( + cur_tile_coord[0] + cur_tile_coord[1] * m_cta_tiles_in_total + ) + + # Sync all thread in epilogue warp + # Ensure all thread in epilogue warp have issue store operation + epilogue_sync_barrier.arrive_and_wait() + + # Producer +1 is routed unicast to the tile's owner — + # see note in the TMA-store branch above. + if warp_idx == self.epilogue_warp_id[0]: + with cute.arch.elect_one(): + owner_rank = cur_tile_coord[0] // m_cta_tiles_per_rank + for r in cutlass.range_constexpr(self.num_ranks): + if owner_rank == r: + utils.distributed.red_add1( + lock_ptr=barrier_flag_unicast_per_peer[ + r + ].iterator + + cta_tile_id_linear, + order="release", + scope="gpu", + ) + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete (only if not using LDMC which does per-tile tail) + if cutlass.const_expr(self.reduce_scatter != "LDMC"): + c_pipeline.producer_tail() + else: + # Synchronize before TMEM dealloc + if cutlass.const_expr(self.reduce_scatter != "LDMC"): + tmem_dealloc_barrier.arrive_and_wait() + # # 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 - # - if cutlass.const_expr(self.use_tma_store): - c_pipeline.producer_tail() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) # /////////////////////////////////////////////////////////////////////////////// - # Reduce Scatter warps + # ReduceScatter warps # /////////////////////////////////////////////////////////////////////////////// - if cutlass.const_expr(self.reduce_scatter == "two_shot"): - if warp_idx >= self.reduce_scatter_warp_id[0]: - # /////////////////////////////////////////////////////////////////////////////// - # Add persistent tile loop - # /////////////////////////////////////////////////////////////////////////////// - - rank_id = self.rank_id - num_ranks = Int32(self.num_ranks) - lane_id = cute.arch.lane_idx() - - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + if cutlass.const_expr(self.reduce_scatter == "LDMC"): + if warp_idx >= self.comm_warp_id[0]: + # (bM, bN, RestM, RestN, RestL) with multicast memory as comm input + gCommInMC_mnl = cute.local_tile( + mCommInMC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), ) - work_tile = tile_sched.initial_work_tile_info() - - # we want 128bit ld/st for better performance - atom_val = 128 // c_mc.element_type.width - atom_thr_n = self.mma_tiler[1] // atom_val - atom_thr_m = len(self.reduce_scatter_warp_id) * cute.arch.WARP_SIZE // atom_thr_n - thr_layout = cute.make_layout((atom_thr_m, atom_thr_n), stride=(atom_thr_n, 1)) - val_layout = cute.make_layout((1, atom_val), stride=(atom_val, 1)) - - copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_mc.element_type) - tiled_copy_fake = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) - thr_copy_fake = tiled_copy_fake.get_slice(tidx-self.reduce_scatter_warp_id[0]*32) - - # predicate tensor - idC = cute.make_identity_tensor(c_mc.shape) - - # partition and slice at tile level - gC_mc = cute.local_tile( - c_mc, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + # (bM, bN, RestM, RestN, RestL) with local memory as comm output + gCommOutUC_mnl = cute.local_tile( + mCommOutUC_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), ) - cC = cute.local_tile( - idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) with multicast memory + tCgCommInMC = thr_mma.partition_C(gCommInMC_mnl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) with peer memory + tCgCommOutUC = thr_mma.partition_C(gCommOutUC_mnl) + + m_cta_tiles_in_total = gCommInMC_mnl.shape[2] * cute.size( + tiled_mma.thr_id.shape + ) + n_cta_tiles_in_total = gCommInMC_mnl.shape[3] + + # + # Persistent tile scheduling loop for reduce scatter + # + comm_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.comm_sync_bar_id, + num_threads=32 * len(self.comm_warp_id), ) - m_tiles_in_total = gC_mc.shape[2] - n_tiles_in_total = gC_mc.shape[3] - m_tiles_per_rank = m_tiles_in_total // self.num_ranks + m_cta_tiles_per_rank = m_cta_tiles_in_total // self.num_ranks + + # ReduceScatter communication pattern (w/ rank4 as example): + # + # [rank0 A] [rank1 A] [rank2 A] [rank3 A] -> RS -> [rank0+1+2+3 A] + # [rank0 B] [rank1 B] [rank2 B] [rank3 B] -> RS -> [rank0+1+2+3 B] + # [rank0 C] [rank1 C] [rank2 C] [rank3 C] -> RS -> [rank0+1+2+3 C] + # [rank0 D] [rank1 D] [rank2 D] [rank3 D] -> RS -> [rank0+1+2+3 D] + # rank0 rank1 rank2 rank3 -> RS -> rank0 rank1 rank2 rank3 + # + # Each rank's RS output is 1/P (where is number of rank) of RS input (RS input = GEMM output) + # Each rank's RS warp only issue load multiple CTA instruction on its corresponding RS output part while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - mma_tile_coord_mnl = ( - ((cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)) % m_tiles_per_rank) + self.rank_id * m_tiles_per_rank, - cur_tile_coord[1], - cur_tile_coord[2], + + # Flag index uses the m-chunk-remap so this rank's CTAs + # walk through their owned slice of input tiles. Matches + # the input_tile_coord computation inside LDMC. + input_tile_m = ( + cur_tile_coord[0] % m_cta_tiles_per_rank + + self.rank_id * m_cta_tiles_per_rank + ) + flag_offset = ( + input_tile_m + cur_tile_coord[1] * m_cta_tiles_in_total ) - chunk_id = ((cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)) // m_tiles_per_rank) - tile_id = mma_tile_coord_mnl[0] + mma_tile_coord_mnl[1] * m_tiles_in_total - tile_id = tile_id * cute.size(tiled_mma.thr_id.shape) - - if not is_leader_cta: - tile_id = tile_id + 1 - - # System barrier to make sure that data from each GPU is in memory before Reduce Scatter - flag = barrier_flag.iterator + tile_id - - if warp_idx == self.reduce_scatter_warp_id[0]: - if lane_id == 0: - res = 0 - while res < self.num_ranks: - res = cute.arch.load(flag.llvm_ptr, cutlass.Int32, sem="relaxed", scope="gpu") - cute.arch.barrier( - barrier_id=self.reduce_scatter_sync_bar_id, - number_of_threads=32 * len(self.reduce_scatter_warp_id), - ) - if warp_idx == self.reduce_scatter_warp_id[0]: - if lane_id == 0: - res = cute.arch.atomic_add( - flag.llvm_ptr, - Int32(1), - sem="relaxed", - scope="sys", + # (1) Per-tile producer-arrival wait. + if warp_idx == self.comm_warp_id[0]: + with cute.arch.elect_one(): + utils.distributed.spin_lock_ld_lt_relaxed_wait( + lock_ptr=barrier_flag_unicast.iterator + flag_offset, + expected_val=self.num_ranks, + scope="gpu", ) - res = cute.arch.load(flag.llvm_ptr, cutlass.Int32, sem="relaxed", scope="sys") - if res == self.num_ranks*2: - cute.arch.store(flag.llvm_ptr, Int32(0), sem="relaxed", scope="sys") - tCgC_mc = thr_mma.partition_C(gC_mc) - tCpC = thr_mma.partition_C(cC) - tCgC_mc_slice = tCgC_mc[((None, None), 0, 0, *mma_tile_coord_mnl)] - tCpC_slice = tCpC[((None, None), 0, 0, *mma_tile_coord_mnl)] + comm_sync_barrier.arrive_and_wait() - cta_mma_tile_m = self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) - m_local_rank = int(cta_mma_tile_m / self.num_ranks) - tCgC_mc_slice_partitioned = cute.zipped_divide(tCgC_mc_slice, (m_local_rank, self.mma_tiler[1])) - tCpC_slice_partitioned = cute.zipped_divide(tCpC_slice, (m_local_rank, self.mma_tiler[1])) - - tCgC_mc_local_rank = cute.slice_(tCgC_mc_slice_partitioned, ((None, None), (chunk_id, 0))) - tCpC_local_rank = cute.slice_(tCpC_slice_partitioned, ((None, None), (chunk_id, 0))) - frgC_mc = thr_copy_fake.partition_S(tCgC_mc_local_rank) - frpC = thr_copy_fake.partition_S(tCpC_local_rank) - - m_idx = gC_mc.shape[0] * mma_tile_coord_mnl[0] - m_per_rank = c_mc.shape[0] // self.num_ranks - dst_rank_this_tile = m_idx // m_per_rank - - peer = c_peer_tensors[rank_id] - - gC_peer = cute.local_tile( - peer, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + # (2) Per-tile data movement. + LDMC( + self, + cur_tile_coord=cur_tile_coord, + tiled_mma=tiled_mma, + comm_tidx=tidx, + tCgCommInMC=tCgCommInMC, + tCgCommOutUC=tCgCommOutUC, + m_cta_tiles_per_rank=m_cta_tiles_per_rank, ) - tCgC_peer = thr_mma.partition_C(gC_peer) - tCgC_peer_slice = tCgC_peer[((None, None), 0, 0, *mma_tile_coord_mnl)] - cta_mma_tile_m = self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape) - m_local_rank = int(cta_mma_tile_m / self.num_ranks) - tCgC_peer_slice_partitioned = cute.zipped_divide(tCgC_peer_slice, (m_local_rank, self.mma_tiler[1])) - tCgC_peer_local_rank = cute.slice_(tCgC_peer_slice_partitioned, ((None, None), (chunk_id, 0))) - frgC_peer = thr_copy_fake.partition_S(tCgC_peer_local_rank) - atom, loop_m, loop_n = frgC_mc.shape - tmp_results = cute.make_rmem_tensor((4, loop_m, loop_n), cutlass.Int32) - local_chunk_lower_bound = (rank_id * m_per_rank, c_mc.shape[1], c_mc.shape[2]) - local_chunk_upper_bound = ((rank_id + 1) * m_per_rank, c_mc.shape[1], c_mc.shape[2]) - - for i in cutlass.range_constexpr(loop_m): - for j in cutlass.range_constexpr(loop_n): - if cute.elem_less(frpC[0, i, j], local_chunk_upper_bound) and not cute.elem_less(frpC[0, i, j], local_chunk_lower_bound): - mc_ptr = frgC_mc[None, i, j].iterator - if cutlass.const_expr(self.c_dtype == Float16): - x, y, z, w = utils.distributed.multimem_ld_reduce_8xf16(mc_ptr) - elif cutlass.const_expr(self.c_dtype == Float32): - x, y, z, w = utils.distributed.multimem_ld_reduce_4xf32(mc_ptr) - elif cutlass.const_expr(self.c_dtype == BFloat16): - x, y, z, w = utils.distributed.multimem_ld_reduce_8xbf16(mc_ptr) - elif cutlass.const_expr(self.c_dtype == Float8E4M3FN): - x, y, z, w = utils.distributed.multimem_ld_reduce_16xe4m3(mc_ptr) - elif cutlass.const_expr(self.c_dtype == Float8E5M2): - x, y, z, w = utils.distributed.multimem_ld_reduce_16xe5m2(mc_ptr) - tmp_results[0, i, j] = x - tmp_results[1, i, j] = y - tmp_results[2, i, j] = z - tmp_results[3, i, j] = w - - for i in cutlass.range_constexpr(loop_m): - for j in cutlass.range_constexpr(loop_n): - if cute.elem_less(frpC[0, i, j], local_chunk_upper_bound) and not cute.elem_less(frpC[0, i, j], local_chunk_lower_bound): - ptr_int = frgC_peer[None, i, j].iterator.toint().ir_value() - x, y, z, w = tmp_results[0, i, j].ir_value(), tmp_results[1, i, j].ir_value(), tmp_results[2, i, j].ir_value(), tmp_results[3, i, j].ir_value() - llvm.inline_asm( - T.i32(), - [ptr_int, x, y, z, w], - "st.global.sys.relaxed.v4.f32 [$1], {$2, $3, $4, $5};", - "=r,l,r,r,r,r", - has_side_effects=True, - asm_dialect=0 + # (3) Per-tile consumer-arrival counting. P CTAs from this + # rank cooperatively process the tile; each does +1 and the + # last (result == 2N-1) resets the slot to 0. + if warp_idx == self.comm_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile_ptr = ( + barrier_flag_unicast.iterator + flag_offset + ) + result = cute.arch.atomic_add( + ptr=flag_curr_tile_ptr.llvm_ptr, + val=cutlass.Int32(1), + sem="relaxed", + scope="gpu", + ) + if result == self.num_ranks * 2 - 1: + cute.arch.store( + flag_curr_tile_ptr.llvm_ptr, + cutlass.Int32(0), + sem="relaxed", + scope="gpu", ) + # # Advance to next tile + # tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() - cute.arch.barrier( - barrier_id=self.reduce_scatter_sync_bar_id, - number_of_threads=32 * len(self.reduce_scatter_warp_id), - ) - # System barrier to make sure all the peer memory transfers are completed. - last_flag_idx = cute.size( - tile_sched.params.problem_layout_ncluster_mnl - ) * cute.size(self.cluster_shape_mn) - if warp_idx == self.reduce_scatter_warp_id[0]: + # Ensure all threads in reduce scatter warp group have complete issuing reduce scatter operations. + comm_sync_barrier.arrive_and_wait() + + # + # Set Per SM Flag with Release + # + # This ensure + # 1. no rank early exit while other ranks are still issuing multimem.ld_reduce + if warp_idx == self.comm_warp_id[0]: with cute.arch.elect_one(): # Offset to last tile flag idx - last_tile_id_linear = cute.size( - tile_sched.params.problem_layout_ncluster_mnl - ) * cute.size(self.cluster_shape_mn) + total_num_cta_tile = m_cta_tiles_in_total * n_cta_tiles_in_total # Linear id of current SM. sm_id_linear = ( cute.arch.block_idx()[0] @@ -1499,267 +1378,56 @@ class PersistentDenseGemmKernel: * cute.arch.grid_dim()[0] * cute.arch.grid_dim()[1] ) - # Release flag with sys scope + # Release flag with gpu scope utils.distributed.multimem_red_add1( - lock_ptr=barrier_flag_mc.iterator - + last_tile_id_linear + lock_ptr=barrier_flag_multicast.iterator + + total_num_cta_tile + sm_id_linear, - scope="sys", + scope="gpu", order="release", ) - # Relaxed spin-lock wait flag with sys scope - utils.distributed.spin_lock_atom_cas_relaxed_wait( - lock_ptr=barrier_flag.iterator - + last_tile_id_linear + # Acquire spin-lock wait flag with gpu scope + utils.distributed.spin_lock_atom_cas_acquire_wait( + lock_ptr=barrier_flag_unicast.iterator + + total_num_cta_tile + sm_id_linear, expected_val=self.num_ranks, reset_val=0, - scope="sys", + scope="gpu", ) - - def epilog_tmem_copy_and_partition( - self, - tidx: cutlass.Int32, - tAcc: cute.Tensor, - gC_mnl: cute.Tensor, - epi_tile: cute.Tile, - use_2cta_instrs: Union[cutlass.Boolean, bool], - ) -> Tuple[cute.TiledCopy, 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 gC_mnl: The global tensor C - :type gC_mnl: cute.Tensor - :param epi_tile: The epilogue tiler - :type epi_tile: 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.c_layout, - self.c_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) - gC_mnl_epi = cute.flat_divide( - gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile - ) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) - tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) - # (T2R, T2R_M, T2R_N) - tTR_rAcc = cute.make_rmem_tensor( - tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype - ) - return tiled_copy_t2r, tTR_tAcc, tTR_rAcc - - def epilog_smem_copy_and_partition( - self, - tiled_copy_t2r: cute.TiledCopy, - tTR_rC: cute.Tensor, - tidx: cutlass.Int32, - sC: cute.Tensor, - ) -> Tuple[cute.TiledCopy, 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_rC: The partitioned accumulator tensor - :type tTR_rC: cute.Tensor - :param tidx: The thread index in epilogue warp groups - :type tidx: cutlass.Int32 - :param sC: The shared memory tensor to be copied and partitioned - :type sC: cute.Tensor - :type sepi: cute.Tensor - - :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: - - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) - - tRS_rC: The partitioned tensor C (register source) - - tRS_sC: The partitioned tensor C (smem destination) - :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] - """ - copy_atom_r2s = sm100_utils.get_smem_store_op( - self.c_layout, self.c_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_sC = thr_copy_r2s.partition_D(sC) - # (R2S, R2S_M, R2S_N) - tRS_rC = tiled_copy_r2s.retile(tTR_rC) - return tiled_copy_r2s, tRS_rC, tRS_sC - - def epilog_gmem_copy_and_partition( - self, - tidx: cutlass.Int32, - atom: Union[cute.CopyAtom, cute.TiledCopy], - gC_mnl: cute.Tensor, - epi_tile: cute.Tile, - sC: cute.Tensor, - ) -> Tuple[cute.CopyAtom, 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 atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version - :type atom: cute.CopyAtom or cute.TiledCopy - :param gC_mnl: The global tensor C - :type gC_mnl: cute.Tensor - :param epi_tile: The epilogue tiler - :type epi_tile: cute.Tile - :param sC: The shared memory tensor to be copied and partitioned - :type sC: cute.Tensor - - :return: A tuple containing either: - - For TMA store: (tma_atom_c, bSG_sC, bSG_gC) where: - - tma_atom_c: The TMA copy atom - - bSG_sC: The partitioned shared memory tensor C - - bSG_gC: The partitioned global tensor C - - For non-TMA store: (simt_atom, tTR_rC, tTR_gC) where: - - simt_atom: The SIMT copy atom - - tTR_rC: The register tensor C - - tTR_gC: The partitioned global tensor C - :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] - """ - # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) - gC_epi = cute.flat_divide( - gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile - ) - if cutlass.const_expr(self.use_tma_store): - tma_atom_c = atom - sC_for_tma_partition = cute.group_modes(sC, 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_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_c, bSG_sC, bSG_gC - else: - tiled_copy_t2r = atom - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) - thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) - tTR_gC = thr_copy_t2r.partition_D(gC_epi) - # (T2R, T2R_M, T2R_N) - tTR_rC = cute.make_rmem_tensor( - tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype - ) - simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) - return simt_atom, tTR_rC, tTR_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, c_dtype: Type[cutlass.Numeric], - c_layout: utils.LayoutEnum, smem_capacity: int, occupancy: int, use_tma_store: bool, + c_smem_layout: Union[cute.Layout, None], ) -> Tuple[int, int, int]: - """Computes the number of stages for A/B/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. - :type epi_tile: cute.Tile - :param c_dtype: Data type of operand C (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 - :param use_tma_store: Whether TMA store is enabled. - :type use_tma_store: bool - - :return: A tuple containing the computed number of stages for: - (ACC stages, A/B operand stages, C stages) - :rtype: tuple[int, int, int] - """ # Default ACC stages num_acc_stage = 2 # Default C stages num_c_stage = 2 if use_tma_store else 0 - # Calculate smem layout and size for one stage of A, B, 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 + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 ) - 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 - ) - c_smem_layout_staged_one = ( - sm100_utils.make_smem_layout_epi( - c_dtype, - c_layout, - epi_tile, - 1, - ) - if use_tma_store - else None + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 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 - c_bytes_per_stage = ( - cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) - if use_tma_store - else 0 - ) + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) c_bytes = c_bytes_per_stage * num_c_stage # Calculate A/B stages: @@ -1786,31 +1454,17 @@ class PersistentDenseGemmKernel: c: cute.Tensor, cta_tile_shape_mnk: Tuple[int, int, int], cluster_shape_mn: Tuple[int, int], + swizzle_size: int, + raster_order: Literal["m", "n"], 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 C. - - :param c: The output tensor C - :type c: 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]] - """ c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) gc = cute.zipped_divide(c, tiler=c_shape) num_ctas_mnl = gc[(0, (None, None, None))].shape cluster_shape_mnl = (*cluster_shape_mn, 1) tile_sched_params = utils.PersistentTileSchedulerParams( - num_ctas_mnl, cluster_shape_mnl + num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_order == "m" ) grid = utils.StaticPersistentTileScheduler.get_grid_shape( tile_sched_params, max_active_clusters @@ -1823,44 +1477,17 @@ class PersistentDenseGemmKernel: tiled_mma: cute.TiledMma, mma_tiler: Tuple[int, int, int], num_acc_stage: int, + arch: str, ) -> 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) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) return num_tmem_alloc_cols - def is_valid_dtypes( - self, - ab_dtype: Type[cutlass.Numeric], - c_dtype: Type[cutlass.Numeric], + def check_supported_dtypes( + self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] ) -> bool: - """ - Check if the dtypes are valid - - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] - :param acc_dtype: The data type of the accumulator - :type acc_dtype: Type[cutlass.Numeric] - :param c_dtype: The data type of the output tensor - :type c_dtype: Type[cutlass.Numeric] - - :return: True if the dtypes are valid, False otherwise - :rtype: bool - """ valid_ab_dtypes = { cutlass.Float16, cutlass.BFloat16, @@ -1871,10 +1498,15 @@ class PersistentDenseGemmKernel: cutlass.Float8E5M2, } if ab_dtype not in valid_ab_dtypes: - return False + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype: {ab_dtype}" + ) - if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: - return False + valid_acc_dtypes = {cutlass.Float32, cutlass.Float16, cutlass.Int32} + if self.acc_dtype not in valid_acc_dtypes: + raise testing.CantImplementError( + f"[dtype] Unsupported accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and AB type acc_ab_compatibility = { @@ -1894,7 +1526,9 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and AB type if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"[dtype] Unsupported AB dtype {ab_dtype} for accumulator {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and C type acc_c_compatibility = { @@ -1923,32 +1557,46 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and C type if c_dtype not in acc_c_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for accumulator {self.acc_dtype}" + ) - # check if c_dtype is supported by multimem all-reduce - if cutlass.const_expr(c_dtype not in {cutlass.Float16, cutlass.Float32, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2}): - return False + # check if c_dtype is supported by multimem reduce-scatter + if cutlass.const_expr( + c_dtype + not in { + cutlass.Float16, + cutlass.Float32, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + ): + raise testing.CantImplementError( + f"[dtype] Unsupported C dtype {c_dtype} for multimem reduce-scatter" + ) return True - def is_valid_mma_tiler_and_cluster_shape(self) -> bool: - """Check if the mma tiler and cluster shape are valid. - - :return: True if the mma tiler and cluster shape are valid, False otherwise - :rtype: bool - """ - is_valid = True + def check_mma_tiler_and_cluster_shape(self) -> bool: # Skip invalid mma tile shape - if not ( - (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) - or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) - ): - is_valid = False - if self.mma_tiler_mn[1] not in range(32, 257, 32): - is_valid = False + # 1SM: MMA tile {64, 128} x {64, 128, 192, 256} + # 2SM: MMA tile {128, 256} x {64, 128, 192, 256} + valid_mma_m = [64, 128] if not self.use_2cta_instrs else [128, 256] + valid_mma_n = [64, 128, 192, 256] + if self.mma_tiler_mn[0] not in valid_mma_m: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler M: {self.mma_tiler_mn[0]}, expected {valid_mma_m} for {'1SM' if not self.use_2cta_instrs else '2SM'}" + ) + if self.mma_tiler_mn[1] not in valid_mma_n: + raise testing.CantImplementError( + f"[mma_tile] Invalid mma tiler N: {self.mma_tiler_mn[1]}, expected {valid_mma_n}" + ) # Skip illegal cluster shape if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: - is_valid = False + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape M: {self.cluster_shape_mn[0]}, must be multiple of {2 if self.use_2cta_instrs else 1} for {'2SM' if self.use_2cta_instrs else '1SM'}" + ) # Skip invalid cluster shape is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 if ( @@ -1958,10 +1606,12 @@ class PersistentDenseGemmKernel: or not is_power_of_2(self.cluster_shape_mn[0]) or not is_power_of_2(self.cluster_shape_mn[1]) ): - is_valid = False - return is_valid + raise testing.CantImplementError( + f"[cluster] Invalid cluster shape: {self.cluster_shape_mn}, must be positive power of 2 with total <= 16" + ) + return True - def is_valid_tensor_alignment( + def check_tensor_alignment( self, m: int, n: int, @@ -1969,236 +1619,579 @@ class PersistentDenseGemmKernel: l: int, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], ) -> bool: - """ - Check if the tensor alignment is valid - - :param m: The number of rows in the A tensor - :type m: int - :param n: The number of columns in the B tensor - :type n: int - :param k: The number of columns in the A tensor - :type k: int - :param l: The number of columns in the C tensor - :type l: int - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] - :param c_dtype: The data type of the output tensor - :type c_dtype: Type[cutlass.Numeric] - :param a_major: The major axis of the A tensor - :type a_major: str - :param b_major: The major axis of the B tensor - :type b_major: str - :param c_major: The major axis of the C tensor - :type c_major: str - - :return: True if the problem shape is valid, False otherwise - :rtype: bool - """ - is_valid = True - - def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): major_mode_idx = 0 if is_mode0_major else 1 num_major_elements = tensor_shape[major_mode_idx] num_contiguous_elements = 16 * 8 // dtype.width return num_major_elements % num_contiguous_elements == 0 if ( - not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) - or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) - or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) ): - is_valid = False + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: contiguous dimension must be 16B aligned for (M={m}, K={k}, N={n})" + ) - return is_valid + if m % 128 != 0 and n % 128 != 0: + raise testing.CantImplementError( + f"[alignment] Invalid tensor alignment: at least one of M={m} or N={n} must be divisible by 128" + ) - def is_valid_epilog_store_option( - self, - m: int, - n: int, - ) -> bool: - """ - Check if the epilogue store option is valid + if a_major not in ["m", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid a_major: {a_major}, expected 'm' or 'k'" + ) + if b_major not in ["n", "k"]: + raise testing.CantImplementError( + f"[alignment] Invalid b_major: {b_major}, expected 'n' or 'k'" + ) + if c_major not in ["m", "n"]: + raise testing.CantImplementError( + f"[alignment] Invalid c_major: {c_major}, expected 'm' or 'n'" + ) - :param m: The number of rows in the A tensor - :type m: int - :param n: The number of columns in the B tensor - :type n: int + return True - :return: True if the epilogue store option is valid, False otherwise - :rtype: bool - """ - - is_valid = True - # None TMA store version does not have predication, can not support OOB tiles + def check_epilogue_store_option(self, m: int, n: int) -> bool: + # Non TMA store version does not have predication, can not support OOB tiles cta_tile_shape_mn = ( self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), self.mma_tiler_mn[1], ) if not self.use_tma_store: if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): - is_valid = False - return is_valid + raise testing.CantImplementError( + f"[epilogue] Problem shape (M={m}, N={n}) must be divisible by cta tile shape {cta_tile_shape_mn} for non TMA store" + ) + # CTA swizzling improves the L2 cache utilization and reduces the number of cache misses. + # Make sure the swizzle size divides the cta/cluster count since non TMA epilogue don't support OOB tiles. + # Swizzle only applies to the dimension orthogonal to the raster direction. + m_per_swizzle = (m // cta_tile_shape_mn[0]) // self.cluster_shape_mn[0] + n_per_swizzle = (n // cta_tile_shape_mn[1]) // self.cluster_shape_mn[1] + if self.raster_order == "m": + # Swizzle applies to N dimension only + if n_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape n={n} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if m_per_swizzle % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[epilogue] Problem shape m={m} must be divisible by swizzle size {self.swizzle_size} for non TMA store with raster_order_n" + ) + return True - def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool: - """Check if the given tensors can be implemented by this kernel. + def check_valid_rank_config(self) -> bool: + if self.num_ranks not in [2, 4, 8] and self.reduce_scatter != "none": + raise testing.CantImplementError( + f"[rank_config] Invalid rank config: num_ranks={self.num_ranks}, must be 2, 4, or 8 for reduce-scatter" + ) + return True - :param a: Input tensor A - :type a: cute.Tensor - :param b: Input tensor B - :type b: cute.Tensor - :param c: Output tensor C - :type c: cute.Tensor + def check_valid_reduce_scatter_mode(self) -> bool: + if self.reduce_scatter not in ["LDMC", "none"]: + raise testing.CantImplementError( + f"[reduce_scatter_mode] Invalid reduce-scatter mode: {self.reduce_scatter}, expected 'LDMC' or 'none'" + ) + return True - :return: True if the gemm supports the given config, False otherwise + def check_reduce_scatter_tile_divisibility(self, m: int) -> bool: + """ + Check if the number of M CTA tiles is evenly divisible by the number of ranks + and that each rank has enough tiles to support the MMA tile coordination logic. + + For reduce-scatter with LDMC mode, the M dimension tiles are distributed across + ranks. Each rank must receive an equal number of tiles, and there must be at + least one tile per rank. + + Additionally, for 2CTA mode, each rank must have at least 2 CTA tiles (one per CTA + in the cluster) because the MMA tile coordinate calculation divides by the number + of CTAs per MMA tile (cute.size(tiled_mma.thr_id.shape)). + + The constraints are: + 1. (M / cta_tile_m) % num_ranks == 0 + 2. (M / cta_tile_m) >= num_ranks + 3. (M / cta_tile_m) / num_ranks >= num_ctas_per_mma_tile (2 for 2CTA, 1 otherwise) + + where cta_tile_m = mma_tiler_mn[0] / (2 if use_2cta_instrs else 1) + + :param m: The M dimension of the problem + :type m: int + :return: True if valid, raises CantImplementError otherwise :rtype: bool """ - m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] + if self.reduce_scatter == "none": + return True - # infer a_major, b_major, c_major - is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() - is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() - is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c() - a_major = "m" if is_m_major_a else "k" - b_major = "n" if is_n_major_b else "k" - c_major = "m" if is_m_major_c else "n" + # Calculate CTA tile size in M dimension + # For 2CTA mode, mma_tiler covers 2 CTAs, so cta_tile_m = mma_tiler_mn[0] / 2 + num_ctas_per_mma_tile = 2 if self.use_2cta_instrs else 1 + cta_tile_m = self.mma_tiler_mn[0] // num_ctas_per_mma_tile - can_implement = True + # Number of CTA tiles in M dimension + num_m_cta_tiles = m // cta_tile_m + + # Check divisibility and minimum tile requirement + if num_m_cta_tiles % self.num_ranks != 0: + raise testing.CantImplementError( + f"[rs_tile_divisibility] Number of M CTA tiles ({num_m_cta_tiles}) must be " + f"evenly divisible by num_ranks ({self.num_ranks}) for reduce-scatter. " + f"Got M={m}, cta_tile_m={cta_tile_m}. " + f"Consider using M >= {cta_tile_m * self.num_ranks} (multiple of {cta_tile_m * self.num_ranks})" + ) + + if num_m_cta_tiles < self.num_ranks: + raise testing.CantImplementError( + f"[rs_tile_divisibility] Number of M CTA tiles ({num_m_cta_tiles}) must be " + f">= num_ranks ({self.num_ranks}) for reduce-scatter. " + f"Got M={m}, cta_tile_m={cta_tile_m}. " + f"Minimum M required: {cta_tile_m * self.num_ranks}" + ) + + # For 2CTA mode, each rank needs at least num_ctas_per_mma_tile CTA tiles + # because the MMA tile coordinate calculation divides m_cta_tiles_per_rank + # by num_ctas_per_mma_tile (cute.size(tiled_mma.thr_id.shape)) + m_cta_tiles_per_rank = num_m_cta_tiles // self.num_ranks + if m_cta_tiles_per_rank < num_ctas_per_mma_tile: + min_m_required = cta_tile_m * self.num_ranks * num_ctas_per_mma_tile + raise testing.CantImplementError( + f"[rs_tile_divisibility] For {'2CTA' if self.use_2cta_instrs else '1CTA'} mode, " + f"m_cta_tiles_per_rank ({m_cta_tiles_per_rank}) must be >= {num_ctas_per_mma_tile}. " + f"Got M={m}, cta_tile_m={cta_tile_m}, num_ranks={self.num_ranks}. " + f"Minimum M required: {min_m_required}" + ) + return True + + def check_even_divisiable_by_cluster_tilesize( + self, m: int, n: int, k: int, l: int + ) -> bool: + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + + cluster_shape_m = self.cluster_shape_mn[0] + cluster_shape_n = self.cluster_shape_mn[1] + + cluster_tile_m = cta_tile_m * cluster_shape_m + cluster_tile_n = cta_tile_n * cluster_shape_n + + if m % cluster_tile_m != 0 or n % cluster_tile_n != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_cluster_tile_size] Problem shape (M={m}, N={n}) not divisible by cluster tile size ({cluster_tile_m}, {cluster_tile_n})" + ) + return True + + def check_even_divisible_by_swizzle(self, m: int, n: int) -> bool: + # Only check for LDMC mode which uses multimem operations + if self.reduce_scatter != "LDMC": + return True + + cta_tile_m = self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1) + cta_tile_n = self.mma_tiler_mn[1] + + # Number of clusters in each dimension + num_clusters_m = m // (cta_tile_m * self.cluster_shape_mn[0]) + num_clusters_n = n // (cta_tile_n * self.cluster_shape_mn[1]) + + if self.raster_order == "m": + # Swizzle applies to N dimension only + if num_clusters_n % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in N dimension ({num_clusters_n}) " + f"must be divisible by swizzle_size={self.swizzle_size} for LDMC with raster_order_m" + ) + else: + # Swizzle applies to M dimension only + if num_clusters_m % self.swizzle_size != 0: + raise testing.CantImplementError( + f"[problem_size_divisible_swizzle_size] Number of clusters in M dimension ({num_clusters_m}) " + f"must be divisible by swizzle_size={self.swizzle_size} for LDMC with raster_order_n" + ) + + return True + + def can_implement( + self, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + ) -> bool: # Skip unsupported types - if not self.is_valid_dtypes(a.element_type, c.element_type): - can_implement = False + self.check_supported_dtypes(ab_dtype, c_dtype) + # Skip invalid mma tile shape and cluster shape - if not self.is_valid_mma_tiler_and_cluster_shape(): - can_implement = False + self.check_mma_tiler_and_cluster_shape() + + # Unpack mnkl for clarity in calling the epilog check + m, n, k, l = mnkl + + # Skip if GEMM not even divisiable by tilesize + # Right now, RS Fusion **ONLY** support even divisiable by tilesize + self.check_even_divisiable_by_cluster_tilesize(m, n, k, l) + + # Skip if swizzle doesn't evenly divide the problem (for LDMC communication) + self.check_even_divisible_by_swizzle(m, n) + + # Skip if M tiles not evenly divisible by num_ranks for reduce-scatter + self.check_reduce_scatter_tile_divisibility(m) + # Skip illegal problem shape for load/store alignment - if not self.is_valid_tensor_alignment( - m, n, k, l, a.element_type, c.element_type, a_major, b_major, c_major - ): - can_implement = False + self.check_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ) + # Skip invalid epilogue store option - if not self.is_valid_epilog_store_option(m, n): - can_implement = False - if dist.get_world_size() not in [2, 4, 8]: - can_implement = False - if m % 128 != 0: - # cannot support OOB tiles when m is not divisible by 128 - can_implement = False + self.check_epilogue_store_option(m, n) - return can_implement + # Skip invalid reduce-scatter rank config + self.check_valid_rank_config() + + # Skip invalid reduce-scatter mode + self.check_valid_reduce_scatter_mode() + + return True -def create_mc_tensor(torch_tensor_cpu, dtype, leading_dim, is_dynamic_layout=True): - torch_tensor_gpu_local = nvshmem.core.tensor(torch_tensor_cpu.shape, dtype=torch_tensor_cpu.dtype) - torch_tensor_gpu_local.copy_(torch_tensor_cpu) - torch_tensor_gpu_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, torch_tensor_gpu_local) - cute_tensor_mc = from_dlpack( - torch_tensor_gpu_mc, - assumed_align=16, - ) +##################################################################### +#### KERNEL IMPLEMENTATION END HERE #### +##################################################################### - cute_tensor_c_peer_torch_tensors = [nvshmem.core.get_peer_tensor(torch_tensor_gpu_local, rank) for rank in range(dist.get_world_size())] - cute_tensor_c_peer_tensors = [from_dlpack(t) for t in cute_tensor_c_peer_torch_tensors] +_DISTRIBUTED_INITIALIZED = False - if is_dynamic_layout: + +# Check for device p2p access +def check_cuda_p2p_access(): + import torch + + num_devices = torch.cuda.device_count() + cuda_p2p_access_available = True + for i in range(num_devices - 1): + for j in range(i + 1, num_devices): + if not torch.cuda.can_device_access_peer(i, j): + cuda_p2p_access_available = False + break + if not cuda_p2p_access_available: + break + return cuda_p2p_access_available + + +# Check for multicast support +def check_multicast_support(): + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + device_count = torch.cuda.device_count() + for device_id in range(device_count): + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device_id): + return False + return True + except ImportError: + return False + + +def init_distributed(): + global _DISTRIBUTED_INITIALIZED + + if not dist.is_initialized() and not _DISTRIBUTED_INITIALIZED: + # * Check before running * + if not torch.cuda.is_available(): + print("CUDA is not available") + exit(-1) + if torch.cuda.get_device_capability()[0] != 10: + print("Running on non-sm100 device") + exit(-1) + if torch.cuda.device_count() < 2: + print("Less than two devices") + exit(-1) + if torch.cuda.device_count() % 2 != 0: + print("Odd world size") + exit(-1) + for var in ("LOCAL_RANK", "RANK", "WORLD_SIZE"): + if var not in os.environ: + print(f"{var} not in environment variables") + exit(-1) + + # LOCAL_RANK = which of this host's GPUs we bind to. + # RANK = global identity across all hosts; used for NCCL init + # and for kernel rank_id (slab ownership in the RS). + local_rank = int(os.environ["LOCAL_RANK"]) + global_rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + + dist.init_process_group( + backend="nccl", + world_size=world_size, + rank=global_rank, + device_id=device, + ) + + _DISTRIBUTED_INITIALIZED = True + + print( + f"[init] hostname={socket.gethostname()} global_rank={global_rank} " + f"local_rank={local_rank} world_size={world_size} device={device}" + ) + + global_rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", global_rank)) + device = torch.device("cuda", local_rank) + + # Return the GLOBAL rank — callers feed this into kernel rank_id. + return global_rank, world_size, device + + +def clean_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def test_simple_torch_allreduce(device): + def _simple_torch_allreduce(): + tensor = symm_mem.empty((16384), dtype=torch.float32, device=device) + symm_mem.rendezvous(tensor, group=dist.group.WORLD) + torch.ops.symm_mem.multimem_all_reduce_( + tensor, "sum", dist.group.WORLD.group_name + ) + + try: + _simple_torch_allreduce() + except Exception: + return False + return True + + +def allocate_tensors( + *, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + world_size: int, + num_workspace: int, + device: int, + slot_init_mode: Literal["test", "benchmark"], + global_rank: int = 0, +): + def _create_multicast_tensor(*, torch_tensor_ref, dtype, leading_dim): + is_col_major = leading_dim == 0 + shape = torch_tensor_ref.shape + ref_dtype = torch_tensor_ref.dtype + + if is_col_major: + reversed_shape = tuple(reversed(shape)) + perm = tuple(range(len(shape) - 1, -1, -1)) + torch_symm_base = symm_mem.empty( + reversed_shape, device=device, dtype=ref_dtype + ) + torch_tensor_uc = torch_symm_base.permute(perm) + else: + torch_symm_base = symm_mem.empty(shape, device=device, dtype=ref_dtype) + torch_tensor_uc = torch_symm_base + + torch_tensor_uc.copy_(torch_tensor_ref) + + torch_tensor_hdl = symm_mem.rendezvous(torch_symm_base, group=dist.group.WORLD) + + cute_tensor_uc = from_dlpack(torch_tensor_uc, assumed_align=16) + cute_tensor_uc.element_type = dtype + cute_tensor_uc = cute_tensor_uc.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor_uc = cutlass_torch.convert_cute_tensor( + torch_tensor_uc, + cute_tensor_uc, + dtype, + is_dynamic_layout=True, + ) + + if is_col_major: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, reversed_shape, ref_dtype + ).permute(perm) + else: + mc_torch = cutlass_torch.as_tensor( + torch_tensor_hdl.multicast_ptr, shape, ref_dtype + ) + cute_tensor_mc = from_dlpack(mc_torch, assumed_align=16) + cute_tensor_mc.element_type = dtype cute_tensor_mc = cute_tensor_mc.mark_layout_dynamic(leading_dim=leading_dim) - cute_tensor = from_dlpack(torch_tensor_gpu_local, assumed_align=16) - cute_tensor.element_type = dtype - if is_dynamic_layout: - cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) - cute_tensor = cutlass_torch.convert_cute_tensor( - torch_tensor_gpu_local, - cute_tensor, - dtype, - is_dynamic_layout=is_dynamic_layout, + return cute_tensor_uc, cute_tensor_mc, torch_tensor_uc + + torch.manual_seed(42 + global_rank) + + m, n, k, l = mnkl + + _init_lo, _init_hi = (-1, 1) if slot_init_mode == "test" else (-20, 20) + + torch_tensor_a_int_cpu = cutlass_torch.matrix( + l, + m, + k, + a_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), ) - return cute_tensor, cute_tensor_mc, torch_tensor_gpu_local, torch_tensor_gpu_mc, cute_tensor_c_peer_torch_tensors, cute_tensor_c_peer_tensors - - -def create_tensors( - l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype -): - torch.manual_seed(1111) - - a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) - b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype).to(torch.float32).normal_().round_().to(dtype=cutlass_torch.dtype(ab_dtype)) - c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype) - - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + torch_tensor_b_int_cpu = cutlass_torch.matrix( + l, + n, + k, + b_major == "n", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + torch_tensor_c_int_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), ) - c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( - c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + # comm_out is sharded along M: each rank holds m/world_size rows. + torch_tensor_comm_out_c_int_cpu = cutlass_torch.matrix( + l, + m // world_size, + n, + c_major == "m", + cutlass.Int32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=_init_lo, max_val=_init_hi), ) - c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc, c_peer_torch_tensors, c_peer_tensors = create_mc_tensor( - c_torch_cpu, c_dtype, (1 if c_major == "n" else 0), is_dynamic_layout=True + torch_tensor_a_cpu = torch_tensor_a_int_cpu.to(dtype=torch_dtype(ab_dtype)) + torch_tensor_b_cpu = torch_tensor_b_int_cpu.to(dtype=torch_dtype(ab_dtype)) + torch_tensor_c_cpu = torch_tensor_c_int_cpu.to(dtype=torch_dtype(c_dtype)) + torch_tensor_comm_out_c_cpu = torch_tensor_comm_out_c_int_cpu.to( + dtype=torch_dtype(c_dtype) ) - return ( - a_tensor, - b_tensor, - c_tensor, - c_tensor_mc, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - c_torch_gpu, - c_torch_gpu_mc, - c_peer_torch_tensors, - c_peer_tensors, + cute_tensor_a_list = [] + cute_tensor_b_list = [] + torch_tensor_a_gpu_list = [] + torch_tensor_b_gpu_list = [] + for _ in range(num_workspace): + cute_tensor_a_i, torch_tensor_a_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_a_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_b_i, torch_tensor_b_gpu_i = cutlass_torch.cute_tensor_like( + torch_tensor_b_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + cute_tensor_a_list.append(cute_tensor_a_i) + cute_tensor_b_list.append(cute_tensor_b_i) + torch_tensor_a_gpu_list.append(torch_tensor_a_gpu_i) + torch_tensor_b_gpu_list.append(torch_tensor_b_gpu_i) + + assert not torch.isnan(torch_tensor_a_gpu_list[0]).any() + assert not torch.isnan(torch_tensor_b_gpu_list[0]).any() + + ( + cute_tensor_c, + cute_tensor_comm_in_mc, + torch_tensor_c_gpu, + ) = _create_multicast_tensor( + torch_tensor_ref=torch_tensor_c_cpu, + dtype=c_dtype, + leading_dim=(1 if c_major == "n" else 0), ) - -def compare( - a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance -): - # Copy gpu result back - kernel_result = c_torch_gpu.cpu() - - # Compute reference result - ref = torch.einsum( - "mkl,nkl->mnl", - a_torch_cpu.to(dtype=torch.float32), - b_torch_cpu.to(dtype=torch.float32), + # comm_out is per-rank only — LDMC reduce-scatter writes through unicast. + # Allocated as a regular cuda tensor (no symm_mem / multicast required). + cute_tensor_comm_out_uc, torch_tensor_comm_out_gpu = cutlass_torch.cute_tensor_like( + torch_tensor_comm_out_c_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 ) - acc_type = cutlass.Float32 - if c_dtype.width == 16: - acc_dtype = cutlass.Float32 - elif c_dtype.width == 8: - acc_dtype = cutlass.Float16 - - torch_acc_dtype = cutlass_torch.dtype(acc_dtype) - torch_c_dtype = cutlass_torch.dtype(c_dtype) + # Flag tensor sized for the worst-case kernel config: smallest cta_tile is + # 64x64, so max #tiles = (m/64) * (n/64). Pad by 160 for per-SM slots + # (B200 has 148; 160 is a safe upper bound). The kernel resets these flags + # internally on each launch, so a single allocation can be reused across + # (mma_tiler, use_2cta) candidates in a benchmark sweep. + num_flags = (m // 64) * (n // 64) + 160 - # note that in this example, f8 reduce sum is done by f16 acc precision and f16 by f32 - ref = ref.to(torch_c_dtype).to(torch_acc_dtype).cuda() - torch.distributed.all_reduce(ref, op=torch.distributed.ReduceOp.SUM) - ref = ref.to(torch_c_dtype) + torch_tensor_flag_unicast = symm_mem.empty( + (num_flags,), dtype=torch.int32, device=device + ) + torch_tensor_flag_unicast.fill_(0) + torch_flag_hdl = symm_mem.rendezvous( + torch_tensor_flag_unicast, group=dist.group.WORLD + ) + torch_tensor_flag_multicast = cutlass_torch.as_tensor( + torch_flag_hdl.multicast_ptr, + torch_tensor_flag_unicast.shape, + torch_tensor_flag_unicast.dtype, + ) + cute_tensor_flag_unicast = from_dlpack(torch_tensor_flag_unicast, assumed_align=16) + cute_tensor_flag_unicast.element_type = cutlass.Int32 + cute_tensor_flag_unicast = cute_tensor_flag_unicast.mark_layout_dynamic() + cute_tensor_flag_multicast = from_dlpack( + torch_tensor_flag_multicast, assumed_align=16 + ) + cute_tensor_flag_multicast.element_type = cutlass.Int32 + cute_tensor_flag_multicast = cute_tensor_flag_multicast.mark_layout_dynamic() - # convert to higher precision to satisfy assert_close function - ref_result = ref.cpu().to(torch_acc_dtype) - kernel_result = c_torch_gpu.view(torch_c_dtype).cpu().to(torch_acc_dtype) + # Per-peer UC views of the flag tensor. Each rank r's local buffer is + # exposed via flag_hdl.buffer_ptrs[r]; from any rank these are direct + # NVLink P2P device pointers. Used by the producer epilogue to unicast + # +1 to the owning rank's flag slot only (Option C — see + # reducescatter_flag_fix.md). + peer_flag_torch_anchors = [] + cute_tensor_flag_uc_per_peer_list = [] + for r in range(world_size): + peer_ptr = torch_flag_hdl.buffer_ptrs[r] + peer_flag_torch = cutlass_torch.as_tensor( + peer_ptr, + torch_tensor_flag_unicast.shape, + torch_tensor_flag_unicast.dtype, + ) + peer_flag_cute = from_dlpack(peer_flag_torch, assumed_align=16) + peer_flag_cute.element_type = cutlass.Int32 + peer_flag_cute = peer_flag_cute.mark_layout_dynamic() + peer_flag_torch_anchors.append(peer_flag_torch) + cute_tensor_flag_uc_per_peer_list.append(peer_flag_cute) + cute_tensor_flag_unicast_per_peer = tuple(cute_tensor_flag_uc_per_peer_list) - max_val = torch.finfo(kernel_result.dtype).max - min_val = torch.finfo(kernel_result.dtype).min - kernel_result = torch.nan_to_num(kernel_result, nan=max_val, posinf=max_val, neginf=min_val) - ref_result = torch.nan_to_num(ref_result, nan=max_val, posinf=max_val, neginf=min_val) - - rank_id = dist.get_rank() - chunk_per_rank = kernel_result.shape[0] // dist.get_world_size() - start_idx = rank_id * chunk_per_rank - end_idx = start_idx + chunk_per_rank - - # Assert close results - torch.testing.assert_close(kernel_result[start_idx:end_idx,:], ref_result[start_idx:end_idx,:], atol=tolerance, rtol=1e-05) + return { + "cute_tensor_a_list": cute_tensor_a_list, + "cute_tensor_b_list": cute_tensor_b_list, + "cute_tensor_c": cute_tensor_c, + "cute_tensor_comm_in_mc": cute_tensor_comm_in_mc, + "cute_tensor_comm_out_uc": cute_tensor_comm_out_uc, + "cute_tensor_flag_unicast": cute_tensor_flag_unicast, + "cute_tensor_flag_multicast": cute_tensor_flag_multicast, + "cute_tensor_flag_unicast_per_peer": cute_tensor_flag_unicast_per_peer, + "torch_tensor_a_cpu": torch_tensor_a_cpu, + "torch_tensor_b_cpu": torch_tensor_b_cpu, + "torch_tensor_c_cpu": torch_tensor_c_cpu, + "torch_tensor_c_gpu": torch_tensor_c_gpu, + "torch_tensor_comm_out_gpu": torch_tensor_comm_out_gpu, + "torch_tensor_flag_unicast": torch_tensor_flag_unicast, + "_anchors": ( + torch_tensor_a_gpu_list, + torch_tensor_b_gpu_list, + torch_tensor_c_gpu, + torch_tensor_comm_out_gpu, + torch_tensor_flag_unicast, + torch_tensor_flag_multicast, + torch_flag_hdl, + tuple(peer_flag_torch_anchors), + ), + } def run( @@ -2206,298 +2199,812 @@ def run( ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], acc_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", use_2cta_instrs: bool = True, use_tma_store: bool = True, - tolerance: float = 1e-01, warmup_iterations: int = 0, iterations: int = 1, - skip_ref_check: bool = False, - use_cold_l2: bool = False, - reduce_scatter: str = "two_shot", - **kwargs, + reduce_scatter: str = "LDMC", + benchmark_or_test: str = "test", + use_cuda_graph: bool = False, + csv_report: str = "./csv_report.csv", ): - """Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + # * Init Distributed * + global_rank, world_size, device = init_distributed() - This function prepares input tensors, configures and launches the persistent GEMM kernel, - optionally performs reference validation, and benchmarks the execution performance. + # * Check based on RS mode * + if reduce_scatter == "LDMC": + if not check_cuda_p2p_access(): + print("CUDA P2P access not available") + exit(-1) + if not check_multicast_support(): + print("Multicast support not available") + exit(-1) + if not test_simple_torch_allreduce(device): + print("Simple torch allreduce test failed") + exit(-1) - :param mnkl: Problem size (M, N, K, L) - :type mnkl: Tuple[int, int, int, int] - :param ab_dtype: Data type for input tensors A and B - :type ab_dtype: Type[cutlass.Numeric] - :param c_dtype: Data type for output tensor C - :type c_dtype: Type[cutlass.Numeric] - :param acc_dtype: Data type for accumulation during matrix multiplication - :type acc_dtype: Type[cutlass.Numeric] - :param a_major/b_major/c_major: Memory layout of tensor A/B/C - :type a_major/b_major/c_major: str - :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the - default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters. - :type mma_tiler_mn: Tuple[int, int], optional - :param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the - default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters. - :type cluster_shape_mn: Tuple[int, int], optional - :param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner - will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. - :type use_2cta_instrs: bool, optional - :param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use - the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters. - :type use_tma_store: bool, optional - :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 - :type tolerance: float, optional - :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 - :type warmup_iterations: int, optional - :param iterations: Number of benchmark iterations to run, defaults to 1 - :type iterations: int, optional - :param skip_ref_check: Whether to skip reference result validation, defaults to False - :type skip_ref_check: bool, optional - :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False - :type use_cold_l2: bool, optional - :param reduce_scatter: reduce scatter mode, can be "two_shot" - :type reduce_scatter: str, optional - :raises RuntimeError: If CUDA GPU is not available - :raises ValueError: If the configuration is invalid or unsupported by the kernel - :return: Execution time of the GEMM kernel - :rtype: float - """ - print(f"Running Blackwell Persistent Dense GEMM test with:") - print(f"mnkl: {mnkl}") - print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") - print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") - print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") - print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") - print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") - print(f"Tolerance: {tolerance}") - print(f"Warmup iterations: {warmup_iterations}") - print(f"Iterations: {iterations}") - print(f"Skip reference checking: {skip_ref_check}") - print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") - print(f"Fused Reduce Scatter Op: {reduce_scatter}") + # * Verbose * + def _verbose(): + print("Running Blackwell SM100 Persistent Dense GEMM RS test with:") + print(f"mnkl: {mnkl}") + print(f"global_rank: {global_rank}, world_size: {world_size}, device: {device}") + print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}") + print(f"Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") + print(f"Swizzle Size: {swizzle_size}") + print(f"Raster Order: {raster_order}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Use CUDA Graph: {'True' if use_cuda_graph else 'False'}") + print(f"Mode: {benchmark_or_test}") + print(f"Fused ReduceScatter Op: {reduce_scatter}") - # Unpack parameters - m, n, k, l = mnkl + if global_rank == 0: + _verbose() - if not torch.cuda.is_available(): - raise RuntimeError("GPU is required to run this example!") + # * Run Distributed * + if benchmark_or_test in ("benchmark", "benchmark_all"): + run_distributed_benchmark( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + warmup_iterations=warmup_iterations, + iterations=iterations, + reduce_scatter=reduce_scatter, + use_cuda_graph=use_cuda_graph, + benchmark_all=(benchmark_or_test == "benchmark_all"), + csv_report=csv_report, + ) + else: + run_distributed_test( + global_rank=global_rank, + world_size=world_size, + device=device, + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_order=raster_order, + use_2cta_instrs=use_2cta_instrs, + use_tma_store=use_tma_store, + reduce_scatter=reduce_scatter, + ) - # Get current CUDA stream from PyTorch - torch_stream = torch.cuda.current_stream() - # Get the raw stream pointer as a CUstream - current_stream = cuda.CUstream(torch_stream.cuda_stream) - ( - a_tensor, - b_tensor, - c_tensor, - c_tensor_mc, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - c_torch_gpu, - c_torch_gpu_mc, - c_peer_torch_tensors, - c_peer_tensors, - ) = create_tensors( - l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype - ) + clean_distributed() - # Build GEMM object - gemm = PersistentDenseGemmKernel( + +def run_distributed_test( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + reduce_scatter: str = "LDMC", +): + # * CUDA * + # Non-default stream as we'll use CUDA graphs + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # * Build GEMM ReduceScatter Object * + gemm_reduce_scatter = Sm100PersistentDenseGemmReduceScatterLDMCKernel( acc_dtype, + c_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store, + rank_id=global_rank, + num_ranks=world_size, reduce_scatter=reduce_scatter, + swizzle_size=swizzle_size, + raster_order=raster_order, ) - # Check if configuration can be implemented - can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor) - if not can_implement: - raise ValueError( - f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " - f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " - f"use_tma_store = {use_tma_store}" + # * Set Filename & Funcname for different rank * + # Use hash to keep filename short and avoid "File name too long" error + try: + import hashlib + + param_str = f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}_{'2Sm' if use_2cta_instrs else '1Sm'}_{'EpiTma' if use_tma_store else 'EpiStg'}" + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_rs_{param_hash}_{reduce_scatter}_r{global_rank}_w{world_size}" ) + except Exception: + unique_name = f"sm100_gemm_rs_{reduce_scatter}_r{global_rank}_w{world_size}" + gemm_reduce_scatter.__call__.set_name_prefix(unique_name) + + # * can_implement * + can_implement = gemm_reduce_scatter.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, reduce_scatter={reduce_scatter}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + # * Allocate tensors * + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + world_size=world_size, + num_workspace=1, + device=device, + slot_init_mode="test", + global_rank=global_rank, + ) + cute_tensor_a = tensors["cute_tensor_a_list"][0] + cute_tensor_b = tensors["cute_tensor_b_list"][0] + cute_tensor_c = tensors["cute_tensor_c"] + cute_tensor_comm_in_mc = tensors["cute_tensor_comm_in_mc"] + cute_tensor_comm_out_uc = tensors["cute_tensor_comm_out_uc"] + cute_tensor_flag_unicast = tensors["cute_tensor_flag_unicast"] + cute_tensor_flag_multicast = tensors["cute_tensor_flag_multicast"] + cute_tensor_flag_unicast_per_peer = tensors["cute_tensor_flag_unicast_per_peer"] + torch_tensor_a_cpu = tensors["torch_tensor_a_cpu"] + torch_tensor_b_cpu = tensors["torch_tensor_b_cpu"] + torch_tensor_c = tensors["torch_tensor_c_gpu"] + torch_tensor_comm_out = tensors["torch_tensor_comm_out_gpu"] + torch_tensor_flag_unicast = tensors["torch_tensor_flag_unicast"] + + # * Compile * + compiled_gemm = cute.compile( + gemm_reduce_scatter, + a=cute_tensor_a, + b=cute_tensor_b, + c=cute_tensor_c, + comm_in_multicast_tensor=cute_tensor_comm_in_mc, + comm_out_unicast_tensor=cute_tensor_comm_out_uc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + barrier_flag_unicast_per_peer=cute_tensor_flag_unicast_per_peer, + stream=stream, + max_active_clusters=utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ), + ) + + # * Run * + compiled_gemm( + a=cute_tensor_a, + b=cute_tensor_b, + c=cute_tensor_c, + comm_in_multicast_tensor=cute_tensor_comm_in_mc, + comm_out_unicast_tensor=cute_tensor_comm_out_uc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + barrier_flag_unicast_per_peer=cute_tensor_flag_unicast_per_peer, + stream=stream, + ) + + # * Compare * + # FP8 nan to saturate + def _fp8_nan_to_saturate(tensor_fp8: torch.Tensor) -> torch.Tensor: + if tensor_fp8.dtype == torch.float8_e4m3fn: + # Map to u8 raw tensor for easier byte handle + tensor_u8 = tensor_fp8.view(torch.uint8) + # pos & neg nan + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + # Map pos nan to max, neg nan to min + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 448 + tensor_f32[is_neg_nan] = -448 + # Cast back to f8 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + elif tensor_fp8.dtype == torch.float8_e5m2: + # Map to u8 raw tensor for easier byte handle + tensor_u8 = tensor_fp8.view(torch.uint8) + # pos & neg nan + is_pos_nan = tensor_u8 == 0x7F + is_neg_nan = tensor_u8 == 0xFF + # Map pos nan to max, neg nan to min + tensor_f32 = tensor_fp8.to(dtype=torch.float32) + tensor_f32[is_pos_nan] = 57344 + tensor_f32[is_neg_nan] = -57344 + # Cast back to f8 + tensor_f8_out = tensor_f32.to(dtype=tensor_fp8.dtype) + else: + # Not supported dtype + tensor_f8_out = tensor_fp8 + return tensor_f8_out + + def _compare_gemm( + torch_tensor_a_cpu: torch.Tensor, + torch_tensor_b_cpu: torch.Tensor, + torch_tensor_c: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + # GEMM Ref + gemm_ref_f32 = torch.einsum( + "mkl,nkl->mnl", + torch_tensor_a_cpu.to(dtype=torch.float32), + torch_tensor_b_cpu.to(dtype=torch.float32), + ) + + # For fp8_e4m3, clamp value to [-448, 448] to avoid nan and then cast to fp8_e4m3 + if c_dtype == cutlass.Float8E4M3FN: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-448, max=448) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + # For fp8_e5m2, clamp value to [-57344, 57344] to avoid nan and then cast to fp8_e5m2 + elif c_dtype == cutlass.Float8E5M2: + gemm_ref_f32 = torch.clamp(gemm_ref_f32, min=-57344, max=57344) + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + else: + gemm_ref_f32_c = gemm_ref_f32.to(dtype=torch_dtype(c_dtype)) + + # Cast ref to fp32 for comparison + gemm_ref_f32_c_f32 = gemm_ref_f32_c.to(dtype=torch.float32) + + # Cast kernel output to fp32 for comparison + gemm_out_f32 = torch_tensor_c.cpu().to(dtype=torch.float32) + + # Compare + torch.testing.assert_close( + gemm_out_f32, + gemm_ref_f32_c_f32, + atol=atol, + rtol=rtol, + ) + + def _compare_comm( + torch_tensor_comm_in: torch.Tensor, + torch_tensor_comm_out: torch.Tensor, + c_dtype: Type[cutlass.Numeric], + atol: float, + rtol: float, + ): + comm_ref = torch_tensor_comm_in + + # Sync across all rank + dist.barrier() + + # Reference reduce-scatter — perform an all-reduce on the full comm_in + # (C), then take this rank's slice. Matches what the LDMC fused kernel + # writes into comm_out for this rank. + if c_dtype in [cutlass.Float32, cutlass.BFloat16, cutlass.Float16]: + comm_ref = comm_ref.contiguous() + dist.all_reduce(comm_ref, op=torch.distributed.ReduceOp.SUM) + comm_ref_f32 = comm_ref.to(dtype=torch.float32).cpu() + # For fp8, cast to higher precision, do AR, then cast back to fp8 and handle nan + elif c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_ref_fp32_acc = comm_ref.to(dtype=torch.float32).contiguous() + dist.all_reduce(comm_ref_fp32_acc, op=torch.distributed.ReduceOp.SUM) + # Add small epsilon to mimic multimem.ld_reduce round up behavior where torch cast round down + epsilon = 1e-2 + comm_ref_f32_acc_epsilon = comm_ref_fp32_acc + epsilon + comm_ref_c = comm_ref_f32_acc_epsilon.to(dtype=torch_dtype(c_dtype)) + comm_ref_nan2sat = _fp8_nan_to_saturate(comm_ref_c) + comm_ref_f32 = comm_ref_nan2sat.to(dtype=torch.float32).cpu() + + # Fusion Kernel Result + # multimem.ld_reduce on fp8 will return nan if val is out of range + # Need to handle nan the same way as comm ref + comm_out = torch_tensor_comm_out.cpu() + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + # Handle nan + comm_out_nan2sat = _fp8_nan_to_saturate(comm_out) + # Cast to fp32 for comparison + comm_out_f32 = comm_out_nan2sat.to(dtype=torch.float32) + else: + # Cast to fp32 for comparison + comm_out_f32 = comm_out.to(dtype=torch.float32) + + # Compare + + # Slice the all-reduced reference to this rank's chunk. + chunk_per_rank = comm_ref_f32.shape[0] // world_size + start_idx = global_rank * chunk_per_rank + end_idx = start_idx + chunk_per_rank + + torch.testing.assert_close( + comm_ref_f32[start_idx:end_idx, :], + comm_out_f32, + atol=atol, + rtol=rtol, + ) + + # compare + # Impose a strick check on GEMM as input is uniform random int + gemm_atol = 1e-05 + gemm_rtol = 1e-05 + _compare_gemm( + torch_tensor_a_cpu=torch_tensor_a_cpu, + torch_tensor_b_cpu=torch_tensor_b_cpu, + torch_tensor_c=torch_tensor_c, + c_dtype=c_dtype, + atol=gemm_atol, + rtol=gemm_rtol, + ) + + if reduce_scatter == "LDMC": + # Relaxed tolerance — reference uses NCCL dist.all_reduce, which + # reduces in a different order than the fused multimem kernel. + comm_atol = 5 + comm_rtol = 1e-1 + if c_dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + comm_rtol = 1e-1 + comm_atol = 0.0 + _compare_comm( + torch_tensor_comm_in=torch_tensor_c, + torch_tensor_comm_out=torch_tensor_comm_out, + c_dtype=c_dtype, + atol=comm_atol, + rtol=comm_rtol, + ) + + # * Check flags are reset * + # The kernel uses the flag tensor for cross-CTA / cross-rank synchronization + # and is expected to leave every entry at 0 once it returns. If any entry is + # non-zero, a tile didn't drain its slot — re-running the kernel would block + # forever on that stale flag. + torch.cuda.synchronize() + flag_after = torch_tensor_flag_unicast.cpu() + nonzero_count = int(torch.count_nonzero(flag_after).item()) + assert nonzero_count == 0, ( + f"Kernel did not reset all barrier flags to zero: " + f"{nonzero_count}/{flag_after.numel()} entries non-zero" + ) + + +def compile_and_benchmark( + *, + mnkl, + ab_dtype, + c_dtype, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + swizzle_size, + raster_order, + use_2cta_instrs, + use_tma_store, + reduce_scatter, + global_rank, + world_size, + num_workspace, + tensors, + torch_stream, + stream, + warmup_iterations, + iterations, + use_cuda_graph, +): + """Instantiate one LDMC kernel candidate, compile it, and time the + warmup + benchmark loop. + + Reads cute views out of ``tensors`` (produced by :func:`allocate_tensors`) + and uses the caller-supplied ``torch_stream`` / ``stream`` pair so every + candidate in a sweep shares a single CUDA stream + event pool. + Returns the mean per-iteration time in microseconds over the benchmark + iterations (warmup timings are recorded but not included in the average). + """ + cute_tensor_a_list = tensors["cute_tensor_a_list"] + cute_tensor_b_list = tensors["cute_tensor_b_list"] + cute_tensor_c = tensors["cute_tensor_c"] + cute_tensor_comm_in_mc = tensors["cute_tensor_comm_in_mc"] + cute_tensor_comm_out_uc = tensors["cute_tensor_comm_out_uc"] + cute_tensor_flag_unicast = tensors["cute_tensor_flag_unicast"] + cute_tensor_flag_multicast = tensors["cute_tensor_flag_multicast"] + cute_tensor_flag_unicast_per_peer = tensors["cute_tensor_flag_unicast_per_peer"] + + # Build kernel + can_implement first — before any cross-rank op, so a + # deterministic skip keeps all ranks in lockstep for the next candidate. + gemm_reduce_scatter = Sm100PersistentDenseGemmReduceScatterLDMCKernel( + acc_dtype, + c_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + rank_id=global_rank, + num_ranks=world_size, + reduce_scatter=reduce_scatter, + swizzle_size=swizzle_size, + raster_order=raster_order, + ) + try: + import hashlib + + param_str = ( + f"ab{ab_dtype}_acc{acc_dtype}_c{c_dtype}_{a_major}{b_major}{c_major}" + f"_tile{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_cluster{cluster_shape_mn[0]}x{cluster_shape_mn[1]}" + f"_{'2Sm' if use_2cta_instrs else '1Sm'}" + f"_{'EpiTma' if use_tma_store else 'EpiStg'}" + ) + param_hash = hashlib.md5(param_str.encode()).hexdigest()[:8] + unique_name = ( + f"sm100_gemm_rs_{param_hash}_{reduce_scatter}_r{global_rank}_w{world_size}" + ) + except Exception: + unique_name = f"sm100_gemm_rs_{reduce_scatter}_r{global_rank}_w{world_size}" + gemm_reduce_scatter.__call__.set_name_prefix(unique_name) + + can_implement = gemm_reduce_scatter.can_implement( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"[general] The current config is invalid/unsupported: use_2cta_instrs={use_2cta_instrs}, " + f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, " + f"use_tma_store={use_tma_store}, reduce_scatter={reduce_scatter}, world_size={world_size}, " + f"ab_dtype={ab_dtype}, c_dtype={c_dtype}, acc_dtype={acc_dtype}" + ) + + def make_kernel_kwargs(ws): + return dict( + a=cute_tensor_a_list[ws], + b=cute_tensor_b_list[ws], + c=cute_tensor_c, + comm_in_multicast_tensor=cute_tensor_comm_in_mc, + comm_out_unicast_tensor=cute_tensor_comm_out_uc, + barrier_flag_unicast=cute_tensor_flag_unicast, + barrier_flag_multicast=cute_tensor_flag_multicast, + barrier_flag_unicast_per_peer=cute_tensor_flag_unicast_per_peer, + stream=stream, + ) + + kernel_kwargs_list = [make_kernel_kwargs(i) for i in range(num_workspace)] + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) - - def create_barrier_flags(): - cta_tile_shape_mn = ( - mma_tiler_mn[0] // (2 if use_2cta_instrs else 1), - mma_tiler_mn[1], - ) - problem_shape_ntile_mn = (m // cta_tile_shape_mn[0], n // cta_tile_shape_mn[1]) - num_tiles = problem_shape_ntile_mn[0] * problem_shape_ntile_mn[1] - num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count - - barrier_flag_torch = nvshmem.core.tensor( - (num_tiles + num_sms,), dtype=torch.int32 - ) - barrier_flag_torch.fill_(0) - barrier_flag_torch_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, barrier_flag_torch) - - barrier_flag = from_dlpack(barrier_flag_torch) - barrier_flag = barrier_flag.mark_layout_dynamic() - barrier_flag_mc = from_dlpack(barrier_flag_torch_mc) - barrier_flag_mc = barrier_flag_mc.mark_layout_dynamic() - - return barrier_flag_torch, barrier_flag_torch_mc, barrier_flag, barrier_flag_mc - - barrier_flag_torch, barrier_flag_torch_mc, barrier_flag, barrier_flag_mc = create_barrier_flags() compiled_gemm = cute.compile( - gemm, - a_tensor, - b_tensor, - c_tensor, - max_active_clusters, - current_stream, - c_mc=c_tensor_mc, - c_peer_tensors=c_peer_tensors, - barrier_flag=barrier_flag, - barrier_flag_mc=barrier_flag_mc, + gemm_reduce_scatter, + **kernel_kwargs_list[0], + max_active_clusters=max_active_clusters, ) - if not skip_ref_check: - compiled_gemm( - a_tensor, - b_tensor, - c_tensor, - current_stream, - c_mc=c_tensor_mc, - c_peer_tensors=c_peer_tensors, - barrier_flag=barrier_flag, - barrier_flag_mc=barrier_flag_mc, - ) - compare( - a_torch_cpu, - b_torch_cpu, - c_torch_gpu, - c_dtype, - tolerance, - ) - - nvshmem.core.free_tensor(c_torch_gpu_mc) - nvshmem.core.free_tensor(c_torch_gpu) - for i in range(len(c_peer_tensors)): - if i != dist.get_rank(): - nvshmem.core.free_tensor(c_peer_torch_tensors[i]) - nvshmem.core.free_tensor(barrier_flag_torch_mc) - nvshmem.core.free_tensor(barrier_flag_torch) + # Per-iteration runnable. Under --use_cuda_graph we capture one graph per + # workspace (pointers are baked in at capture time, so N_workspace graphs + # are required to rotate A/B). Without graphs we dispatch the compiled + # kernel directly with the appropriate kernel_kwargs entry. + if use_cuda_graph: + graphs = [] + for ws in range(num_workspace): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=torch_stream): + compiled_gemm(**kernel_kwargs_list[ws]) + graphs.append(g) - free_func_and_tensor_pairs = [] - def add_free_func_and_tensor(free_func, tensor): - free_func_and_tensor_pairs.append((free_func, tensor)) + def run_one_iter(i): + graphs[i % num_workspace].replay() + else: - def generate_tensors(): - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, _ = cutlass_torch.cute_tensor_like( - c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, c_tensor_mc, c_torch_gpu, c_torch_gpu_mc, c_peer_torch_tensors, c_peer_tensors = create_mc_tensor( - c_torch_cpu, - c_dtype, - (1 if c_major == "n" else 0), - is_dynamic_layout=True, - ) - barrier_flag_torch, barrier_flag_torch_mc, barrier_flag, barrier_flag_mc = create_barrier_flags() - add_free_func_and_tensor(nvshmem.core.free_tensor, c_torch_gpu_mc) - add_free_func_and_tensor(nvshmem.core.free_tensor, c_torch_gpu) - for i in range(len(c_peer_torch_tensors)): - if i != dist.get_rank(): - add_free_func_and_tensor(nvshmem.core.free_tensor, c_peer_torch_tensors[i]) - add_free_func_and_tensor(nvshmem.core.free_tensor, barrier_flag_torch_mc) - add_free_func_and_tensor(nvshmem.core.free_tensor, barrier_flag_torch) - return testing.JitArguments( - a_tensor, - b_tensor, - c_tensor, - current_stream, - c_mc=c_tensor_mc, - c_peer_tensors=c_peer_tensors, - barrier_flag=barrier_flag, - barrier_flag_mc=barrier_flag_mc, - ) + def run_one_iter(i): + compiled_gemm(**kernel_kwargs_list[i % num_workspace]) - workspace_count = 1 - if use_cold_l2: - one_workspace_bytes = ( - a_torch_cpu.numel() * a_torch_cpu.element_size() - + b_torch_cpu.numel() * b_torch_cpu.element_size() - + c_torch_cpu.numel() * c_torch_cpu.element_size() - ) - workspace_count = testing.get_workspace_count( - one_workspace_bytes, warmup_iterations, iterations - ) + # One loop for warmup + benchmark. Every iteration is bracketed by its own + # event pair; only the trailing `iterations` event pairs contribute to the + # reported average so warmup transients don't skew the result. - exec_time = testing.benchmark( - compiled_gemm, - workspace_generator=generate_tensors, - workspace_count=workspace_count, - stream=current_stream, + # Cross-rank sync before timing. torch.cuda.synchronize() drains any + # pending GPU work (graph captures land on the default stream; timing + # runs on torch_stream) so dist.barrier()'s host-side rendezvous is + # not racing in-flight kernels from the previous candidate. Without + # this, the slowest rank's leftover work would skew the measured + # window of every other rank in a benchmark_all sweep. + torch.cuda.synchronize() + dist.barrier() + + total_iters = warmup_iterations + iterations + start_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(total_iters)] + with torch.cuda.stream(torch_stream): + for i in range(total_iters): + start_events[i].record() + run_one_iter(i) + end_events[i].record() + end_events[-1].synchronize() + + per_iter_us = [s.elapsed_time(e) * 1000.0 for s, e in zip(start_events, end_events)] + bench_times = per_iter_us[warmup_iterations:] + avg_us = sum(bench_times) / iterations + print( + f"Rank {global_rank} Execution time: {avg_us} us " + f"(avg of {iterations} iters, after {warmup_iterations} warmup)\n", + flush=True, + ) + return avg_us + + +def run_distributed_benchmark( + global_rank: int, + world_size: int, + device: int, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: Literal["k", "m"], + b_major: Literal["k", "n"], + c_major: Literal["n", "m"], + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + swizzle_size: int = 1, + raster_order: Literal["m", "n"] = "m", + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + warmup_iterations: int = 0, + iterations: int = 1, + reduce_scatter: str = "LDMC", + use_cuda_graph: bool = False, + benchmark_all: bool = False, + csv_report: str = "./csv_report.csv", +): + # Allocate every kernel-side tensor (A/B/C/comm_out/flags) up front — + # shared across all candidates in the --benchmark_all loop. + num_workspace = 10 + tensors = allocate_tensors( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + world_size=world_size, + num_workspace=num_workspace, + device=device, + slot_init_mode="benchmark", + global_rank=global_rank, + ) + + # One non-default stream shared across every candidate. `torch_stream` is + # kept alive because we capture it into a CUDA graph; creating a fresh + # stream per candidate would churn events + graph context for no reason. + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # Build the candidate list of args for compile_and_benchmark. For a single + # run the list has one entry; --benchmark_all sweeps: + # raster_order ∈ {m, n} + # swizzle_size ∈ {1, 2, 4, 8} + # use_2cta_instrs, cluster_shape_mn, mma_tiler_mn jointly, with the + # constraint that cluster_m ∈ {2,4,8} requires 2cta (smaller tiles and + # single-cta clusters are only valid without 2cta). + base_kwargs = dict( + mnkl=mnkl, + ab_dtype=ab_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + a_major=a_major, + b_major=b_major, + c_major=c_major, + use_tma_store=use_tma_store, + reduce_scatter=reduce_scatter, + global_rank=global_rank, + world_size=world_size, + num_workspace=num_workspace, + tensors=tensors, + torch_stream=torch_stream, + stream=stream, warmup_iterations=warmup_iterations, iterations=iterations, + use_cuda_graph=use_cuda_graph, ) - for free_func, tensor in free_func_and_tensor_pairs: - free_func(tensor) - print(f"exec_time: {exec_time}\n") + if benchmark_all: + raster_candidates = ["m", "n"] + swizzle_candidates = [1, 2, 4, 8] + # All cluster shapes from the allowed set; each is filtered per 2cta + # option below by the "cluster_m ∈ {2,4,8} needs 2cta" rule. + cluster_all = [ + (1, 1), + (1, 2), + (2, 1), + (2, 2), + ] + two_cta_options = [ + { + "use_2cta_instrs": True, + "clusters": [c for c in cluster_all if c[0] in (2, 4, 8)], + "mma_tilers": [ + (256, 256), + (256, 192), + (256, 128), + (256, 64), + (128, 256), + (128, 192), + (128, 128), + (128, 64), + ], + }, + { + "use_2cta_instrs": False, + "clusters": cluster_all, + "mma_tilers": [ + (128, 256), + (128, 192), + (128, 128), + (128, 64), + (64, 256), + (64, 192), + (64, 128), + (64, 64), + ], + }, + ] + candidates = [] + for opt in two_cta_options: + for cluster in opt["clusters"]: + for tile in opt["mma_tilers"]: + for r in raster_candidates: + for s in swizzle_candidates: + candidates.append( + { + **base_kwargs, + "use_2cta_instrs": opt["use_2cta_instrs"], + "cluster_shape_mn": cluster, + "mma_tiler_mn": tile, + "raster_order": r, + "swizzle_size": s, + } + ) + else: + candidates = [ + { + **base_kwargs, + "use_2cta_instrs": use_2cta_instrs, + "cluster_shape_mn": cluster_shape_mn, + "mma_tiler_mn": mma_tiler_mn, + "raster_order": raster_order, + "swizzle_size": swizzle_size, + } + ] - return exec_time # Return execution time in microseconds + if global_rank == 0: + print(f"Total number of candidate configs: {len(candidates)}", flush=True) + + # Run each candidate. Tensors (A/B/C/comm_out/flags) are reused across + # candidates; the kernel resets flags internally. Skip any combo the kernel + # can't implement — can_implement is deterministic so all ranks skip the + # same ones, keeping rendezvous collective operations lockstep. + # ~2s sleep between candidates to let the GPU cool down. + results = [] + for i, cand in enumerate(candidates): + if i > 0: + time.sleep(2) + try: + avg_us = compile_and_benchmark(**cand) + except testing.CantImplementError: + continue + results.append((cand, avg_us)) + + # Each rank writes its own CSV — one row per candidate, one config field + # per column. The path comes from --csv_report; per-rank suffix is appended + # before the extension so concurrent ranks don't clobber the same file. + global_rank = dist.get_rank() + m, n, k, l = mnkl + fieldnames = [ + "gemm_m", + "gemm_n", + "gemm_k", + "gemm_l", + "ab_dtype", + "c_dtype", + "acc_dtype", + "a_major", + "b_major", + "c_major", + "mma_tiler_m", + "mma_tiler_n", + "cluster_m", + "cluster_n", + "use_2cta_instrs", + "use_tma_store", + "reduce_scatter", + "use_cuda_graph", + "raster_order", + "swizzle_size", + "num_workspace", + "warmup_iterations", + "iterations", + "world_size", + "global_rank", + "benchmark_time_us", + ] + csv_path = Path(csv_report) + csv_path = csv_path.with_name( + f"{csv_path.stem}_rank{global_rank}_world{world_size}{csv_path.suffix}" + ) + csv_path.parent.mkdir(parents=True, exist_ok=True) + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for cand, avg_us in results: + writer.writerow( + { + "gemm_m": m, + "gemm_n": n, + "gemm_k": k, + "gemm_l": l, + "ab_dtype": str(ab_dtype), + "c_dtype": str(c_dtype), + "acc_dtype": str(acc_dtype), + "a_major": a_major, + "b_major": b_major, + "c_major": c_major, + "mma_tiler_m": cand["mma_tiler_mn"][0], + "mma_tiler_n": cand["mma_tiler_mn"][1], + "cluster_m": cand["cluster_shape_mn"][0], + "cluster_n": cand["cluster_shape_mn"][1], + "use_2cta_instrs": cand["use_2cta_instrs"], + "use_tma_store": use_tma_store, + "reduce_scatter": reduce_scatter, + "use_cuda_graph": use_cuda_graph, + "raster_order": cand["raster_order"], + "swizzle_size": cand["swizzle_size"], + "num_workspace": num_workspace, + "warmup_iterations": warmup_iterations, + "iterations": iterations, + "world_size": world_size, + "global_rank": global_rank, + "benchmark_time_us": avg_us, + } + ) + + return results[0][1] if len(results) == 1 else results -def torchrun_uid_init_bcast(): - """ - Initialize NVSHMEM using UniqueID with `torchrun` as the launcher - - It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting - """ - # Set Torch device - local_rank = int(os.environ['LOCAL_RANK']) - torch.cuda.set_device(local_rank) - - # nvshmem4py requires a cuda.core Device at init time - dev = Device(local_rank) - dev.set_current() - global stream - stream = dev.create_stream() - - # Initialize torch.distributed process group - dist.init_process_group( - backend="cpu:gloo,cuda:nccl", +def prepare_parser(): + parser = argparse.ArgumentParser( + description="Example of SM100 Persistent Dense GEMM AR on Blackwell." ) - # Extract rank, nranks from process group - num_ranks = dist.get_world_size() - - # Create an empty uniqueid for all ranks - uid = nvshmem.core.get_unique_id(empty=(local_rank != 0)) - uid_bytes = uid._data.view(np.uint8).copy() - uid_tensor = torch.from_numpy(uid_bytes).cuda() - dist.broadcast(uid_tensor, src=0) - dist.barrier() - uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype) - - nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid") - - -def torchrun_finalize(): - nvshmem.core.finalize() - dist.destroy_process_group() - - -if __name__ == "__main__": - - def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: try: return tuple(int(x.strip()) for x in s.split(",")) except ValueError: @@ -2505,25 +3012,21 @@ if __name__ == "__main__": "Invalid format. Expected comma-separated integers." ) - parser = argparse.ArgumentParser( - description="Example of Dense Persistent GEMM on Blackwell." - ) - parser.add_argument( "--mnkl", - type=parse_comma_separated_ints, + type=_parse_comma_separated_ints, default=(256, 256, 512, 1), help="mnkl dimensions (comma-separated)", ) parser.add_argument( "--mma_tiler_mn", - type=parse_comma_separated_ints, + type=_parse_comma_separated_ints, default=(128, 128), help="Mma tile shape (comma-separated)", ) parser.add_argument( "--cluster_shape_mn", - type=parse_comma_separated_ints, + type=_parse_comma_separated_ints, default=(1, 1), help="Cluster shape (comma-separated)", ) @@ -2542,7 +3045,13 @@ if __name__ == "__main__": "--use_tma_store", action="store_true", help="Use tma store or not" ) parser.add_argument( - "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + "--benchmark_or_test", + type=str, + choices=["test", "benchmark", "benchmark_all"], + default="test", + help="Select run mode: 'test' for a single-launch correctness check, " + "'benchmark' for cudaEvent timing + per-rank CSV on the given config, " + "'benchmark_all' for the same with a sweep across kernel-config variants.", ) parser.add_argument( "--warmup_iterations", type=int, default=0, help="Warmup iterations" @@ -2554,21 +3063,44 @@ if __name__ == "__main__": help="Number of iterations to run the kernel", ) parser.add_argument( - "--skip_ref_check", action="store_true", help="Skip reference checking" - ) - parser.add_argument( - "--use_cold_l2", + "--use_cuda_graph", action="store_true", default=False, - help="Use circular buffer tensor sets to ensure L2 cold cache", + help="Capture one iteration into a CUDA graph and replay per iteration", ) parser.add_argument( "--reduce_scatter", - choices=["two_shot"], + choices=["LDMC", "none"], type=str, - default="two_shot", - help="Reduce Scatter algorithm to fuse with gemm", + default="LDMC", + help="ReduceScatter algorithm to fuse with gemm", ) + parser.add_argument( + "--swizzle_size", + type=int, + default=1, + help="Swizzling size in the unit of cluster for improving L2 cache hit rate", + ) + parser.add_argument( + "--raster_order", + type=str, + choices=["m", "n"], + default="m", + help="Rasterization order of clusters", + ) + parser.add_argument( + "--csv_report", + type=str, + default="./csv_report.csv", + help="Path for the per-rank benchmark CSV. Each rank appends " + "_rank{R}_world{W} before the extension to avoid collisions.", + ) + + return parser + + +if __name__ == "__main__": + parser = prepare_parser() args = parser.parse_args() @@ -2580,14 +3112,6 @@ if __name__ == "__main__": if len(args.cluster_shape_mn) != 2: parser.error("--cluster_shape_mn must contain exactly 2 values") - - if args.c_dtype.width == 8: - import warnings - warnings.warn("f8 output is easy to overflow and can be nan here") - - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - torchrun_uid_init_bcast() run( args.mnkl, @@ -2599,16 +3123,20 @@ if __name__ == "__main__": args.c_major, args.mma_tiler_mn, args.cluster_shape_mn, + args.swizzle_size, + args.raster_order, args.use_2cta_instrs, args.use_tma_store, - args.tolerance, args.warmup_iterations, args.iterations, - args.skip_ref_check, - args.use_cold_l2, - reduce_scatter=args.reduce_scatter, + args.reduce_scatter, + args.benchmark_or_test, + args.use_cuda_graph, + args.csv_report, + ) + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + print( + f"PASS (hostname={socket.gethostname()}, global_rank={rank}, local_rank={local_rank})", + flush=True, ) - - torchrun_finalize() - -