# Copyright (c) 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 argparse from typing import Type, Tuple from dataclasses import dataclass import torch import cutlass from cutlass import ( cute as cute, utils as utils, ) from cutlass.cute import experimental as cute_ext from cutlass.cute.runtime import from_dlpack from cutlass.cute.nvgpu import cpasync, tcgen05 import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils """ This is an implementation of dense block scaled GEMM. """ class BlockScaledDenseGemmKernel: def __init__( self, mma_inst_mn: tuple[int, int], mma_dtype: tuple[Type[cutlass.Numeric], Type[cutlass.Numeric]], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, epilogue_op=lambda x: x, ): self.ab_dtype, self.acc_dtype = mma_dtype self.sf_dtype = sf_dtype self.sf_vec_size = sf_vec_size self.mma_inst_shape_mn = mma_inst_mn self.use_2cta_instrs = False self.cta_group = ( tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE ) self.epilogue_op = epilogue_op # TODO: instead of using max shared memory, we should define a SharedStorage and then # query its size. self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") # Stages self.num_acc_stages = 1 if self.mma_inst_shape_mn[1] == 256 else 2 # TODO: provide a computation for this so that it is not fixed; # fitting as many stages as there is available shared memory self.num_main_stages = 4 self.tma_store_stages = 4 @cute.experimental.jit def __call__( self, mA: cute.Tensor, mSFA: cute.Tensor, mB: cute.Tensor, mSFB: cute.Tensor, mC: cute.Tensor, ): tile_mn = (*self.mma_inst_shape_mn, 1) div = cute.tiled_divide(mC, tile_mn) grid = (div.shape[1], div.shape[2], div.shape[3]) self.kernel(mA, mSFA, mB, mSFB, mC).launch( grid=grid, # Using a total of 6 warps (1x load + 1x mma + 4x epilogue) block=(192, 1, 1), cluster=(1, 1, 1), smem=self.smem_capacity, ) @cute.experimental.kernel def kernel( self, mA: cute.Tensor, mSFA: cute.Tensor, mB: cute.Tensor, mSFB: cute.Tensor, mC: cute.Tensor, ): # Prologue tidx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) cta_m, cta_n, cta_l = cute.arch.block_idx() a_dtype: Type[cutlass.Numeric] = mA.element_type sf_dtype: Type[cutlass.Numeric] = mSFA.element_type c_dtype: Type[cutlass.Numeric] = mC.element_type a_major_mode = utils.LayoutEnum.from_tensor(mA).mma_major_mode() b_major_mode = utils.LayoutEnum.from_tensor(mB).mma_major_mode() d_layout = utils.LayoutEnum.from_tensor(mC) tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( a_dtype, a_major_mode, b_major_mode, sf_dtype, self.sf_vec_size, self.cta_group, self.mma_inst_shape_mn, ) mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) mma_inst_tile_k = 4 mma_tiler_mnk = ( self.mma_inst_shape_mn[0], self.mma_inst_shape_mn[1], mma_inst_shape_k * mma_inst_tile_k, ) tiler_mk = (mma_tiler_mnk[0], mma_tiler_mnk[2]) tiler_nk = (mma_tiler_mnk[1], mma_tiler_mnk[2]) tiler_mn = (mma_tiler_mnk[0], mma_tiler_mnk[1]) # ((Atom_M, Rest_M),(Atom_K, Rest_K), RestL) sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(mA.shape, self.sf_vec_size) sfa_tensor = cute.make_tensor(mSFA.iterator, sfa_layout) # ((Atom_N, Rest_N),(Atom_K, Rest_K), RestL) sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(mB.shape, self.sf_vec_size) sfb_tensor = cute.make_tensor(mSFB.iterator, sfb_layout) gA = cute.zipped_divide(mA, tiler_mk) gB = cute.zipped_divide(mB, tiler_nk) gSFA = cute.zipped_divide(sfa_tensor, tiler_mk) gSFB = cute.zipped_divide(sfb_tensor, tiler_nk) gC = cute.zipped_divide(mC, tiler_mn) gA_tile = gA[(None, None), (cta_m, None, cta_l)] gB_tile = gB[(None, None), (cta_n, None, cta_l)] gSFA_tile = gSFA[(None, None), (cta_m, None, cta_l)] gSFB_tile = gSFB[(None, None), (cta_n, None, cta_l)] gC_tile = gC[(None, None), (cta_m, cta_n, cta_l)] # Shared memory layouts for A/B/SFA/SFB/D # (MMA, MMA_M, MMA_K, PIPE) a_smem_layout_staged = sm100_utils.make_smem_layout_a( tiled_mma, mma_tiler_mnk, self.ab_dtype, self.num_main_stages, ) # (MMA, MMA_N, MMA_K, PIPE) b_smem_layout_staged = sm100_utils.make_smem_layout_b( tiled_mma, mma_tiler_mnk, self.ab_dtype, self.num_main_stages, ) # (MMA, MMA_M, MMA_K, PIPE) sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( tiled_mma, mma_tiler_mnk, self.sf_vec_size, self.num_main_stages, ) # (MMA, MMA_N, MMA_K, PIPE) sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( tiled_mma, mma_tiler_mnk, self.sf_vec_size, self.num_main_stages, ) cta_tile_shape_mnk = cute.shape_div( mma_tiler_mnk, (cute.size(tiled_mma.thr_id.shape), 1, 1) ) epi_tile = sm100_utils.compute_epilogue_tile_shape( cta_tile_shape_mnk, self.use_2cta_instrs, d_layout, c_dtype, ) smem_epi_staged_layout = sm100_utils.make_smem_layout_epi( c_dtype, d_layout, epi_tile, self.tma_store_stages, ) # UMMA ACC TMEM Layout # ((MMA_M, MMA_N), REST_MMA_M, REST_MMA_N, ACC_STAGES) tmem_accs_layout = cute_ext.make_tmem_layout_acc( tiled_mma, mma_tiler_mnk, self.num_acc_stages ) sfa_tmem_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, mma_tiler_mnk, self.sf_vec_size, cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), ) sfb_tmem_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, mma_tiler_mnk, self.sf_vec_size, cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), ) # Allocate UMMA Buffers buffer_smem_a = cute_ext.allocate( self.ab_dtype, cute.AddressSpace.smem, a_smem_layout_staged, alignment=1024, ) buffer_smem_b = cute_ext.allocate( self.ab_dtype, cute.AddressSpace.smem, b_smem_layout_staged, alignment=1024, ) buffer_smem_sfa = cute_ext.allocate( self.sf_dtype, cute.AddressSpace.smem, sfa_smem_layout_staged, alignment=1024, ) buffer_smem_sfb = cute_ext.allocate( self.sf_dtype, cute.AddressSpace.smem, sfb_smem_layout_staged, alignment=1024, ) buffer_tmem_accs = cute_ext.allocate( self.acc_dtype, cute.AddressSpace.tmem, tmem_accs_layout, alignment=16, ) buffer_tmem_sfa = cute_ext.allocate( self.sf_dtype, cute.AddressSpace.tmem, sfa_tmem_layout, alignment=16, ) buffer_tmem_sfb = cute_ext.allocate( self.sf_dtype, cute.AddressSpace.tmem, sfb_tmem_layout, alignment=16, ) buffer_tmem_sfa_compact = cute.filter_zeros(buffer_tmem_sfa) buffer_tmem_sfb_compact = cute.filter_zeros(buffer_tmem_sfb) # Make S2T CopyAtom and tiledCopy copy_atom_s2t = cute.make_copy_atom( tcgen05.Cp4x32x128bOp(self.cta_group), self.sf_dtype, ) tiled_copy_s2t_sfa = cute.nvgpu.tcgen05.make_s2t_copy( copy_atom_s2t, buffer_tmem_sfa_compact ) tiled_copy_s2t_sfb = cute.nvgpu.tcgen05.make_s2t_copy( copy_atom_s2t, buffer_tmem_sfb_compact ) # Allocate SMEM buffer for C buffer_smem_d = cute_ext.allocate( c_dtype, cute.AddressSpace.smem, smem_epi_staged_layout, alignment=1024, ) # Create the TMEM load atom copy_atom_t2r = sm100_utils.get_tmem_load_op( cta_tile_shape_mnk, d_layout, c_dtype, self.acc_dtype, epi_tile, self.use_2cta_instrs, ) # Derive tiled_copy_t2r from the allocated TMEM buffer accumulators = cute.zipped_divide(buffer_tmem_accs, ((epi_tile), 1)) acc_epi_div = accumulators[((None, None), 0), 0] tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) # Derive per-thread RMEM layout for the T2R epilogue copy gC_tile_epi = cute.flat_divide(gC_tile, epi_tile) acc_epi_rmem_layout = cute_ext.make_t2r_rmem_layout( tiled_copy_t2r, gC_tile_epi, tidx ) # Allocate RMEM buffers buffer_rmem_t2r = cute_ext.allocate( self.acc_dtype, cute.AddressSpace.rmem, acc_epi_rmem_layout, alignment=32, ) buffer_rmem_r2s = cute_ext.allocate( c_dtype, cute.AddressSpace.rmem, acc_epi_rmem_layout, alignment=32, ) # TMA -> UMMA mainloop_pipe = cute_ext.TMAToUMMAPipeline.create( num_stages=self.num_main_stages, mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, ) # UMMA -> TMEM acc_pipe = cute_ext.UMMAtoAsyncPipeline.create( num_stages=self.num_acc_stages, mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS, consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R, consumer_arv_count=128, ) # warp assignment: [0]-tma_store, [0-3]-epi, [4]-mma, [5]-tma_load tma_store_warp_id = 0 mma_warp_id = 4 tma_load_warp_id = 5 is_tma_load_warp = warp_idx == tma_load_warp_id is_mma_warp = warp_idx == mma_warp_id is_epi_warp = warp_idx < 4 # SMEM -> GMEM tma_store_pipe = cute_ext.TMAStorePipeline( stages=self.tma_store_stages, arv_count=128, barrier_id=1, tma_warp_id=tma_store_warp_id, ) k_tile_size = cute.size(gA, mode=[1, 1]) if is_tma_load_warp: for k_tile_idx in cutlass.range(0, k_tile_size, 1, unroll=1): gA_k = gA_tile[None, None, k_tile_idx] gB_k = gB_tile[None, None, k_tile_idx] gSFA_k = gSFA_tile[None, None, k_tile_idx] gSFB_k = gSFB_tile[None, None, k_tile_idx] # Scoped state management - pipeline object manages state internally ( producer_stage_token, stage_idx, ) = mainloop_pipe.producer_acquire_and_get_stage() mbar = cute_ext.get_mbarrier(producer_stage_token) ## producer_body begin ## buffer_smem_a_sliced = buffer_smem_a[None, None, None, stage_idx] buffer_smem_b_sliced = buffer_smem_b[None, None, None, stage_idx] buffer_smem_sfa_sliced = buffer_smem_sfa[None, None, None, stage_idx] buffer_smem_sfb_sliced = buffer_smem_sfb[None, None, None, stage_idx] a_cta_v_map = cute_ext.get_cta_v_map_ab( mA, mma_tiler_mnk, tiled_mma, "A" ) b_cta_v_map = cute_ext.get_cta_v_map_ab( mB, mma_tiler_mnk, tiled_mma, "B" ) sfa_cta_v_map = cute_ext.get_cta_v_map_ab( sfa_tensor, mma_tiler_mnk, tiled_mma, "SFA" ) sfb_cta_v_map = cute_ext.get_cta_v_map_ab( sfb_tensor, mma_tiler_mnk, tiled_mma, "SFB" ) cute_ext.tma_load( gA_k, buffer_smem_a_sliced, mbar, cta_v_map=a_cta_v_map, ) cute_ext.tma_load( gB_k, buffer_smem_b_sliced, mbar, cta_v_map=b_cta_v_map, ) cute_ext.tma_load( gSFA_k, buffer_smem_sfa_sliced, mbar, cta_v_map=sfa_cta_v_map, ) cute_ext.tma_load( gSFB_k, buffer_smem_sfb_sliced, mbar, cta_v_map=sfb_cta_v_map, ) ## producer_body end ## mainloop_pipe.producer_commit_and_advance() if is_mma_warp: producer_stage_token, acc_stage_idx = ( acc_pipe.producer_acquire_and_get_stage() ) ## acc_producer_body begin ## accumulators_sliced = buffer_tmem_accs[None, None, None, acc_stage_idx] mma_atom = cute.make_mma_atom(tiled_mma.op) mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False) filtered_buffer_smem_sfa = cute.filter_zeros(buffer_smem_sfa) filtered_buffer_smem_sfb = cute.filter_zeros(buffer_smem_sfb) for k_tile_idx in cutlass.range(0, k_tile_size, 1, unroll=1): # Scoped state management - pipeline object manages consumer state internally ( _, mainloop_idx, ) = mainloop_pipe.consumer_wait_and_get_stage() ## tma_consumer_body begin ## buffer_smem_a_sliced_stage = buffer_smem_a[ (None, None, None, mainloop_idx) ] buffer_smem_b_sliced_stage = buffer_smem_b[ (None, None, None, mainloop_idx) ] filtered_buffer_smem_sfa_sliced_stage = filtered_buffer_smem_sfa[ (None, None, None, mainloop_idx) ] filtered_buffer_smem_sfb_sliced_stage = filtered_buffer_smem_sfb[ (None, None, None, mainloop_idx) ] # Copy SFA/SFB from SMEM to TMEM (UTCCP) src_partitioned_SFA = cute_ext.partition( filtered_buffer_smem_sfa_sliced_stage, cute.Int32(0), layout_tv=tiled_copy_s2t_sfa.layout_src_tv_tiled, tiler=cute.core._pack_tile(tiled_copy_s2t_sfa.tiler_mn), ) dst_partitioned_SFA = cute_ext.partition( buffer_tmem_sfa_compact, cute.Int32(0), layout_tv=tiled_copy_s2t_sfa.layout_dst_tv_tiled, tiler=cute.core._pack_tile(tiled_copy_s2t_sfa.tiler_mn), ) cute_ext.copy( src_partitioned_SFA, dst_partitioned_SFA, copy_atom=copy_atom_s2t ) src_partitioned_SFB = cute_ext.partition( filtered_buffer_smem_sfb_sliced_stage, cute.Int32(0), layout_tv=tiled_copy_s2t_sfb.layout_src_tv_tiled, tiler=cute.core._pack_tile(tiled_copy_s2t_sfb.tiler_mn), ) dst_partitioned_SFB = cute_ext.partition( buffer_tmem_sfb_compact, cute.Int32(0), layout_tv=tiled_copy_s2t_sfb.layout_dst_tv_tiled, tiler=cute.core._pack_tile(tiled_copy_s2t_sfb.tiler_mn), ) cute_ext.copy( src_partitioned_SFB, dst_partitioned_SFB, copy_atom=copy_atom_s2t ) for k_block_idx in cutlass.range(mma_inst_tile_k, unroll_full=True): buffer_smem_a_sliced = buffer_smem_a_sliced_stage[ None, None, k_block_idx ] buffer_smem_b_sliced = buffer_smem_b_sliced_stage[ None, None, k_block_idx ] cute_ext.dot_block_scaled( mma_atom, cute.append_ones(buffer_smem_a_sliced, up_to_rank=3), buffer_tmem_sfa[None, None, k_block_idx], cute.append_ones(buffer_smem_b_sliced, up_to_rank=3), buffer_tmem_sfb[None, None, k_block_idx], accumulators_sliced, ) mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True) ## tma_consumer_body end ## mainloop_pipe.consumer_release_and_advance() ## acc_producer_body end ## acc_pipe.producer_commit_and_advance() if is_epi_warp: _, acc_stage_idx = acc_pipe.consumer_wait_and_get_stage() ## acc_consume_body begin ## tmem_acc_stage = buffer_tmem_accs[ (None, None), 0, 0, acc_stage_idx ] # (MMA_M, MMA_N) # (EPI_TILE_M, EPI_TILE_N, EPI_REST_M, EPI_REST_N) # we have an implicit assumption that EPI_REST_M == 1 tmem_acc_epi_stage = cute.flat_divide(tmem_acc_stage, epi_tile) subtile_cnt = cute.size(tmem_acc_epi_stage.shape, mode=[3]) # EPI_REST_N for subtile_idx in range(subtile_cnt): # TMEM -> RMEM cute_ext.partition_and_copy( thr_copy_t2r, tmem_acc_epi_stage[(None, None, 0, subtile_idx)], buffer_rmem_t2r, ) # RMEM -> RMEM buffer_rmem_r2s.store( self.epilogue_op(buffer_rmem_t2r.load().to(c_dtype)) ) # Acquire pipeline stage and synchronize before RMEM->SMEM copy tma_store_pipe.acquire_sync() tma_store_idx = tma_store_pipe.get_index() # RMEM -> SMEM tiled_copy_r2s = cute.make_tiled_copy_D( cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype), tiled_copy_t2r, ) cute_ext.partition_and_copy( tiled_copy_r2s.get_slice(tidx), buffer_rmem_r2s, buffer_smem_d[None, None, tma_store_idx], ) # Fence SMEM writes and synchronize before TMA store tma_store_pipe.commit_sync() # SMEM -> GMEM (only designated TMA store warp performs TMA store) if warp_idx == tma_store_warp_id: c_cta_v_map = cute_ext.get_cta_v_map_c(mC, epi_tile) cute_ext.tma_store( buffer_smem_d[None, None, tma_store_idx], gC_tile_epi[(None, None, 0, subtile_idx)], cta_v_map=c_cta_v_map, ) # Release pipeline stage and advance tma_store_pipe.release_advance() tma_store_pipe.tail() acc_pipe.consumer_release_and_advance() @cute.experimental.jit def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( sf_ref_tensor: cute.Tensor, sf_mma_tensor: cute.Tensor, ): """ Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout """ # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) # group to ((32, 4, rest_m), (4, rest_k), l) 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] # TODO: add residual support (C) @dataclass class BlockScaledGemmTestbed: """ Testbed for block-scaled GEMM operations on Blackwell (SM100) architecture. This class manages test data and tensors for block-scaled matrix multiplication: D = (A * scale_factor_A) @ (B * scale_factor_B) The testbed maintains three representations of each tensor: 1. Reference tensors (f32 on CPU) - used for reference computation and validation 2. CUTE tensors - device tensors passed directly to CUDA kernels 3. PyTorch tensors - mirrors of CUTE tensors for host-side operations Attributes: a_ref, b_ref: Reference input matrices (f32 format) sfa_ref, sfb_ref: Reference scale factors for A and B matrices (f32 format) d_ref: Reference output matrice (f32 format) a_tensor, b_tensor: CUTE tensors for input matrices (device) sfa_tensor, sfb_tensor: CUTE tensors for scale factors (device) d_tensor: CUTE tensors for output (device) a_torch, b_torch: PyTorch mirrors of A and B CUTE tensors sfa_torch, sfb_torch: PyTorch mirrors of scale factor CUTE tensors d_torch: PyTorch mirrors of D CUTE tensors The class provides: - Automatic tensor creation with proper layouts and alignment - Scale factor tensor generation with block-scaled MMA layout - Reference checking via einsum-based computation Example: testbed = BlockScaledGemmTestbed( MNKL=(128, 128, 64, 1), mma_dtypes=(cutlass.Float16, cutlass.Float16, cutlass.Float32), c_dtypes=(cutlass.Float16), sf_dtype=cutlass.Float16, sf_vec_size=32, a_major='m', b_major='n', d_major='m' ) # ... run kernel with testbed.a_tensor, testbed.b_tensor, etc. testbed.reference_check() # Validate results """ import torch # Reference tensors (all are in f32 format for simplicity of # reference checks) a_ref: torch.Tensor b_ref: torch.Tensor sfa_ref: torch.Tensor sfb_ref: torch.Tensor # CUTE tensors (to be passed to the device kernel) a_tensor: cute.Tensor b_tensor: cute.Tensor sfa_tensor: cute.Tensor sfb_tensor: cute.Tensor d_tensor: cute.Tensor # PyTorch tensors (mirrors the CUTE tensors above); these tensors # can be used on the host, for example if certain trivial epilogue # needs to be performed. a_torch: torch.Tensor b_torch: torch.Tensor sfa_torch: torch.Tensor sfb_torch: torch.Tensor d_torch: torch.Tensor def __init__( self, MNKL: Tuple[int, int, int, int], mma_dtypes: tuple[ Type[cutlass.Numeric], Type[cutlass.Numeric], Type[cutlass.Numeric] ], c_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, a_major: str, b_major: str, d_major: str, ): import cutlass.torch as cutlass_torch self.d_major = d_major # Problem size (M, N, K, L) = MNKL a_dtype, b_dtype, _ = mma_dtypes assert a_major in ("m", "k"), f"a_major must be 'm' or 'k', got {a_major}" assert b_major in ("n", "k"), f"b_major must be 'n' or 'k', got {b_major}" assert d_major in ("m", "n"), f"d_major must be 'm' or 'n', got {d_major}" self.a_ref = cutlass_torch.matrix(L, M, K, a_major == "m", cutlass.Float32) self.b_ref = cutlass_torch.matrix(L, N, K, b_major == "n", cutlass.Float32) self.d_temp = cutlass_torch.matrix(L, M, N, d_major == "m", cutlass.Float32) self.a_tensor, self.a_torch = cutlass_torch.cute_tensor_like( self.a_ref, a_dtype, is_dynamic_layout=True, assumed_align=16 ) self.b_tensor, self.b_torch = cutlass_torch.cute_tensor_like( self.b_ref, b_dtype, is_dynamic_layout=True, assumed_align=16 ) self.d_tensor, self.d_torch = cutlass_torch.cute_tensor_like( self.d_temp, c_dtype, is_dynamic_layout=True, assumed_align=16 ) # Mark tensor with element divisibility for 16B alignment self.a_tensor.mark_compact_shape_dynamic( mode=1 if a_major == "k" else 0, stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), divisibility=32 if a_dtype == cutlass.Float4E2M1FN else 16, ) self.b_tensor.mark_compact_shape_dynamic( mode=1 if b_major == "k" else 0, stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), divisibility=32 if b_dtype == cutlass.Float4E2M1FN else 16, ) self.d_tensor.mark_compact_shape_dynamic( mode=1 if d_major == "k" else 0, stride_order=(2, 0, 1) if d_major == "n" else (2, 1, 0), divisibility=32 if c_dtype == cutlass.Float4E2M1FN else 16, ) self.sfa_ref, self.sfa_tensor, self.sfa_torch = self.create_scale_factor_tensor( L, M, K, sf_vec_size, sf_dtype ) self.sfb_ref, self.sfb_tensor, self.sfb_torch = self.create_scale_factor_tensor( L, N, K, sf_vec_size, sf_dtype ) # Create scale factor tensor @staticmethod def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): import torch import cutlass.torch as cutlass_torch def ceil_div(a, b): return (a + b - 1) // b sf_k = ceil_div(k, sf_vec_size) ref_shape = (l, mn, sf_k) ref_permute_order = (1, 2, 0) # MKL atom_m = (32, 4) atom_k = 4 mma_shape = ( l, ceil_div(mn, atom_m[0] * atom_m[1]), ceil_div(sf_k, atom_k), atom_m[0], atom_m[1], atom_k, ) mma_permute_order = (3, 4, 1, 5, 2, 0) # M(32x4xrest_m)xK(4xrest_k)xL # Create f32 ref torch tensor (cpu) ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( ref_shape, torch.float32, permute_order=ref_permute_order, init_type=cutlass_torch.TensorInitType.RANDOM, init_config=cutlass_torch.RandomInitConfig( min_val=1, max_val=3, ), ) # Create f32 cute torch tensor (cpu) cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( mma_shape, torch.float32, permute_order=mma_permute_order, init_type=cutlass_torch.TensorInitType.RANDOM, init_config=cutlass_torch.RandomInitConfig( min_val=0, max_val=1, ), ) # convert ref f32 tensor to cute f32 tensor cvt_sf_MKL_to_M32x4xrm_K4xrk_L( from_dlpack(ref_f32_torch_tensor_cpu), from_dlpack(cute_f32_torch_tensor_cpu), ) cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() # reshape makes memory contiguous ref_f32_torch_tensor_cpu = ( ref_f32_torch_tensor_cpu.permute(2, 0, 1) .unsqueeze(-1) .expand(l, mn, sf_k, sf_vec_size) .reshape(l, mn, sf_k * sf_vec_size) .permute(*ref_permute_order) ) # prune to mkl for reference check. ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] # Create dtype cute torch tensor (cpu) cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( cute_f32_torch_tensor_cpu, dtype, is_dynamic_layout=True, assumed_align=16, ) # Convert f32 cute tensor to dtype cute tensor cute_tensor = cutlass_torch.convert_cute_tensor( cute_f32_torch_tensor, cute_tensor, dtype, is_dynamic_layout=True, ) return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor # Transfers results back to CPU and uses PyTorch's methods to do # reference checks def reference_check(self): import torch # Compute reference result, simulate block-scaled GEMV via 2 FFMA # based elementwise multiplication and 1 FFMA based matmul computations res_a = torch.einsum("mkl,mkl->mkl", self.a_ref, self.sfa_ref) res_b = torch.einsum("nkl,nkl->nkl", self.b_ref, self.sfb_ref) ref_output = torch.einsum("mkl,nkl->mnl", res_a, res_b) # Convert d back to f32 for comparison. d_epi_device = self.d_temp.cuda() cute.testing.convert( self.d_tensor, from_dlpack(d_epi_device, assumed_align=16).mark_layout_dynamic( leading_dim=(1 if self.d_major == "n" else 0) ), ) # abs(actual - expected) <= atol + rtol * abs(expected) torch.testing.assert_close( d_epi_device.cpu(), ref_output, atol=1e-01, rtol=1e-02 ) print("Reference check finished.") def run( mnkl: Tuple[int, int, int, int], mma_inst_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, d_dtype: Type[cutlass.Numeric], acc_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, d_major: str, ): """Execute a batched block scaled dense GEMM operation on Blackwell architecture. This function prepares input tensors, configures and launches the GEMM kernel, and performs reference validation. :param mnkl: Problem size (M, N, K, L) :type mnkl: Tuple[int, int, int, int] :param mma_inst_mn: MMA instruction shape. :type mma_inst_mn: Tuple[int, int] :param cluster_shape_mn: Cluster shape. :type cluster_shape_mn: Tuple[int, int] :param ab_dtype: Data type for input tensors A and B :type ab_dtype: Type[Numeric] :param sf_dtype: Data type for scale factors (SFA/SFB) :type sf_dtype: Type[Numeric] :param sf_vec_size: Vector size for the scale factor :type sf_vec_size: int :param c_dtype: Data type for output tensor D :type c_dtype: Type[Numeric] :param acc_dtype: Accumulator data type (precision) :type acc_dtype: Type[Numeric] :param a_major: Major-ness of A tensor (m or k) :type a_major: str :param b_major: Major-ness of B tensor (n or k) :type b_major: str :param d_major: Major-ness of D tensor (m or n) :type d_major: str """ print("Running Blackwell Dense Block Scaled GEMM test with:") print(f"mnkl: {mnkl}") print(f"A: {ab_dtype}, B: {ab_dtype}, D: {d_dtype}, Acc dtype: {acc_dtype}") print(f"Block scaled MMA with SF: {sf_dtype}, vector size: {sf_vec_size}") print(f"Matrix majors - A: {a_major}-major, B: {b_major}-major, D: {d_major}-major") print( f"Mma Tiler (M, N): {mma_inst_mn}, Cluster Shape: {cluster_shape_mn[0]}x{cluster_shape_mn[1]}x1" ) import torch # TODO: add can_implement to exclude unsupported/un-implemented test cases if cluster_shape_mn != (1, 1): raise RuntimeError("Only 1x1x1 cluster shapes are supported right now.") if mma_inst_mn != (128, 128): raise RuntimeError("MMA instruction shape not supported yet.") if ab_dtype not in (cutlass.Float8E4M3FN, cutlass.Float8E5M2): raise RuntimeError("Input data type not supported.") if sf_dtype not in (cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN): raise RuntimeError("Scale factor data type not supported.") if not torch.cuda.is_available(): raise RuntimeError("A GPU is required to run this example!") # Manual seed torch.manual_seed(111) # Create tensors tb = BlockScaledGemmTestbed( mnkl, (ab_dtype, ab_dtype, acc_dtype), d_dtype, sf_dtype, sf_vec_size, a_major, b_major, d_major, ) # JIT-Compile the device kernel block_scaled_gemm = BlockScaledDenseGemmKernel( mma_inst_mn=mma_inst_mn, mma_dtype=(ab_dtype, acc_dtype), sf_dtype=sf_dtype, sf_vec_size=sf_vec_size, ) compiled_kernel = cute.experimental.compile( block_scaled_gemm, tb.a_tensor, tb.sfa_tensor, tb.b_tensor, tb.sfb_tensor, tb.d_tensor, ) # Launch the device kernel compiled_kernel( tb.a_tensor, tb.sfa_tensor, tb.b_tensor, tb.sfb_tensor, tb.d_tensor, ) tb.reference_check() if __name__ == "__main__": 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 = argparse.ArgumentParser( description="Example of Sm100 Dense BlockScaled GEMM." ) parser.add_argument( "--mnkl", type=parse_comma_separated_ints, default=(512, 256, 256, 1), help="mnkl dimensions (comma-separated)", ) parser.add_argument( "--mma_inst_mn", type=parse_comma_separated_ints, default=(128, 128), help="Mma instruction 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.Float8E4M3FN) parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) parser.add_argument("--sf_vec_size", type=int, default=32) parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) parser.add_argument("--d_dtype", type=cutlass.dtype, default=cutlass.Float16) 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("--d_major", choices=["n", "m"], type=str, default="n") args = parser.parse_args() if len(args.mnkl) != 4: parser.error("--mnkl must contain exactly 4 values") if len(args.mma_inst_mn) != 2: parser.error("--mma_inst_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.mma_inst_mn, args.cluster_shape_mn, args.ab_dtype, args.sf_dtype, args.sf_vec_size, args.d_dtype, args.acc_dtype, args.a_major, args.b_major, args.d_major, )