From 3336173eab7b70480a7902f034b5e60e1014f8bd Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Mon, 24 Nov 2025 21:15:18 -0800 Subject: [PATCH 01/11] Update cuteDSL to 4.3.0 and several perf opt for cuteDSL nvf4 gemm Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 7 +- .../dense_blockscaled_gemm_persistent.py | 239 ++++++++++++++++-- 2 files changed, 224 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 01468e0fa1f7..6b7913123e8a 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -81,12 +81,14 @@ def get_valid_tactics( # full shamoo mma_tiler_mn_candidates = [ + (256, 64), + (128, 64), + (256, 192), + (128, 192), (256, 128), (128, 128), (128, 256), (256, 256), - (256, 64), - (128, 64), ] cluster_shape_mn_candidates = [ (1, 1), @@ -258,6 +260,7 @@ def forward( max_active_clusters, stream, swap_ab, + options=f"--opt-level 2", ) self.__class__.kernel_cache[cache_key] = compiled_gemm diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py index 407e239ed1be..10295f2906b9 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -61,6 +61,7 @@ from .custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync from .utils import is_power_of_2 +# from testing import benchmark class Sm100BlockScaledPersistentDenseGemmKernel: """Implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types @@ -81,7 +82,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: * Float8E4M3FN/Float8E5M2 - Constraints: * MMA tiler M must be 128 or 256 (use_2cta_instrs). - * MMA tiler N must be 128 or 256. + * MMA tiler N must be 64/128/192/256 * Cluster shape M must be a multiple of 2 if MMA tiler M is 256. * Cluster shape M/N must be positive and a power of 2, with total cluster size <= 16. * Cluster shape M/N must be <= 4 for scale factor multicasts due to limited scale factor size. @@ -100,6 +101,7 @@ def __init__( sf_vec_size: int, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], + use_prefetch: bool = False, ): """Initializes the configuration for a Blackwell dense GEMM kernel. @@ -125,7 +127,7 @@ def __init__( self.cluster_shape_mn = cluster_shape_mn # K dimension is deferred in _setup_attributes self.mma_tiler = (*mma_tiler_mn, 1) - + self.use_prefetch = use_prefetch self.cta_group = (tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE) @@ -215,6 +217,11 @@ def _setup_attributes(self): 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], + ) # Compute cluster layout self.cluster_layout_vmnk = cute.tiled_divide( @@ -243,6 +250,8 @@ def _setup_attributes(self): self.c_dtype, ) + self.epi_tile_n = cute.size(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 = self._compute_stages( tiled_mma, @@ -292,6 +301,19 @@ def _setup_attributes(self): self.num_c_stage, ) + 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 + + # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue + self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n + + # TODO: [alel] Currently set prefetch dist to num_ab_stage, we may have more option for prefetch dist auto tuning + self.prefetch_dist = self.num_ab_stage + @cute.jit def __call__( self, @@ -433,6 +455,30 @@ def __call__( 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) @@ -541,6 +587,7 @@ class SharedStorage: block=[self.threads_per_cta, 1, 1], cluster=(*self.cluster_shape_mn, 1), smem=self.shared_storage.size_in_bytes(), + min_blocks_per_mp=1, stream=stream, ) return @@ -804,8 +851,29 @@ def kernel( # (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)) + 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 @@ -860,6 +928,26 @@ def kernel( tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] + if cutlass.const_expr(self.use_prefetch): + # Prefetch both A and B (default behavior) + for pf_k_block in cutlass.range(0, min(self.prefetch_dist, k_block_cnt), unroll=1): + cute.prefetch( + tma_atom_a, + tAgA_slice[(None, pf_k_block)], + ) + cute.prefetch( + tma_atom_b, + tBgB_slice[(None, pf_k_block)], + ) + cute.prefetch( + tma_atom_sfa, + tAgSFA_slice[(None, pf_k_block)], + ) + cute.prefetch( + tma_atom_sfb, + tBgSFB_slice[(None, pf_k_block)], + ) + # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt ab_producer_state.reset_count() peek_ab_empty_status = cutlass.Boolean(1) @@ -908,6 +996,28 @@ def kernel( mcast_mask=sfb_full_mcast_mask, ) + # Prefetch: Rolling prefetch for next tiles + if cutlass.const_expr(self.use_prefetch): + if k_block < k_block_cnt - self.prefetch_dist: + future_k_block = ab_producer_state.count + self.prefetch_dist + # Prefetch both A and B (default behavior) + cute.prefetch( + tma_atom_a, + tAgA_slice[(None, future_k_block)], + ) + cute.prefetch( + tma_atom_b, + tBgB_slice[(None, future_k_block)], + ) + cute.prefetch( + tma_atom_sfa, + tAgSFA_slice[(None, future_k_block)], + ) + cute.prefetch( + tma_atom_sfb, + tBgSFB_slice[(None, future_k_block)], + ) + # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1 ab_producer_state.advance() peek_ab_empty_status = cutlass.Boolean(1) @@ -954,7 +1064,7 @@ def kernel( # Make SFA tmem tensor sfa_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + acc_tmem_ptr + self.num_accumulator_tmem_cols, dtype=self.sf_dtype, ) # (MMA, MMA_M, MMA_K) @@ -968,9 +1078,7 @@ def kernel( # Make SFB tmem tensor sfb_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + - tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + - tcgen05.find_tmem_tensor_col_offset(tCtSFA), + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype, ) # (MMA, MMA_N, MMA_K) @@ -1010,10 +1118,15 @@ def kernel( cur_tile_coord[2], ) + 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_producer_state.index)] + acc_stage_index)] # Peek (try_wait) AB buffer full for k_block = 0 ab_consumer_state.reset_count() @@ -1029,13 +1142,21 @@ def kernel( acc_pipeline.producer_acquire(acc_producer_state) tCtSFB_mma = tCtSFB - if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 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(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) + 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) + elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): # Move in increments of 64 columns of SFB offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) shifted_ptr = cute.recast_ptr( - acc_tmem_ptr + - tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + - tcgen05.find_tmem_tensor_col_offset(tCtSFA) + offset, + 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) @@ -1230,10 +1351,16 @@ def kernel( *mma_tile_coord_mnl, )] + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(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_consumer_state.index)] + acc_stage_index)] # # Wait for accumulator buffer full @@ -1249,12 +1376,26 @@ def kernel( 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: + 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, subtile_idx)] + 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() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Convert to C type # @@ -1291,7 +1432,7 @@ def kernel( cute.copy( tma_atom_c, bSG_sC[(None, c_buffer)], - bSG_gC[(None, subtile_idx)], + bSG_gC[(None, real_subtile_idx)], ) # Fence and barrier to make sure shared memory store is visible to TMA store c_pipeline.producer_commit() @@ -1304,9 +1445,10 @@ def kernel( # # Async arrive accumulator buffer empty # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() + if cutlass.const_expr(not self.overlapping_accum): + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Advance to next tile @@ -1755,7 +1897,7 @@ def is_valid_mma_tiler_and_cluster_shape( # Skip invalid mma tile shape if not mma_tiler_mn[0] in [128, 256]: is_valid = False - if not mma_tiler_mn[1] in [64, 128, 256]: + if not mma_tiler_mn[1] 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: @@ -1995,11 +2137,13 @@ def run( c_major: str, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], + use_prefetch: bool = False, tolerance: float = 1e-01, warmup_iterations: int = 0, iterations: int = 1, skip_ref_check: bool = False, use_cold_l2: bool = False, + # use_cupti: bool = False, **kwargs, ): """Runs and benchmarks the persistent batched dense block-scaled GEMM. @@ -2048,11 +2192,13 @@ def run( print( f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}" ) + print(f"Use prefetch: {'True' if use_prefetch 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"Use CUPTI: {'True' if use_cupti else 'False'}") # Unpack parameters m, n, k, l = mnkl @@ -2205,6 +2351,7 @@ def ceil_div(a, b): sf_vec_size, mma_tiler_mn, cluster_shape_mn, + use_prefetch, ) # Compute max active clusters on current device @@ -2226,6 +2373,7 @@ def ceil_div(a, b): 1.0, # alpha max_active_clusters, current_stream, + options=f"--opt-level 2", ) # Compute reference result @@ -2322,6 +2470,15 @@ def generate_tensors(): warmup_iterations=warmup_iterations, iterations=iterations, ) + # exec_time = benchmark( + # compiled_gemm, + # workspace_generator=generate_tensors, + # workspace_count=workspace_count, + # stream=current_stream, + # warmup_iterations=warmup_iterations, + # iterations=iterations, + # use_cupti=use_cupti, + # ) return exec_time # Return execution time in microseconds @@ -2369,6 +2526,12 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: 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_prefetch", + action="store_true", + default=False, + help="Enable TMA prefetch for both A and B matrices (default: False)", + ) parser.add_argument("--tolerance", type=float, default=1e-01, @@ -2392,6 +2555,18 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: default=False, help="Use circular buffer tensor sets to ensure L2 cold cache", ) + # parser.add_argument( + # "--use_cupti", + # action="store_true", + # default=False, + # help="Use CUPTI to measure execution time", + # ) + # parser.add_argument( + # "--print_duration", + # action="store_true", + # default=False, + # help="Print execution time", + # ) args = parser.parse_args() @@ -2415,10 +2590,34 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: args.c_major, args.mma_tiler_mn, args.cluster_shape_mn, + args.use_prefetch, args.tolerance, args.warmup_iterations, args.iterations, args.skip_ref_check, args.use_cold_l2, ) + + # TODO: [alel] Would like to involve cupti testing here, but need to double check how to add cupti wheel in framework + # exec_time = run( + # args.mnkl, + # args.ab_dtype, + # args.sf_dtype, + # args.sf_vec_size, + # args.c_dtype, + # args.a_major, + # args.b_major, + # args.c_major, + # args.mma_tiler_mn, + # args.cluster_shape_mn, + # args.use_prefetch, + # args.tolerance, + # args.warmup_iterations, + # args.iterations, + # args.skip_ref_check, + # args.use_cold_l2, + # #args.use_cupti, + # ) + # if args.print_duration: + # print(f"Execution time: {exec_time} us") print("PASS") From 744af5d58177950f68d87e84765a615aeeaacb6b Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Wed, 26 Nov 2025 22:41:33 -0800 Subject: [PATCH 02/11] minor Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 6b7913123e8a..eceed9b779b9 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -81,12 +81,12 @@ def get_valid_tactics( # full shamoo mma_tiler_mn_candidates = [ - (256, 64), (128, 64), - (256, 192), - (128, 192), - (256, 128), + (256, 64), (128, 128), + (256, 128), + (128, 192), + (256, 192), (128, 256), (256, 256), ] From db24ad64dc8b3d57d1f521831e7e685e0a560c25 Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Thu, 27 Nov 2025 00:21:40 -0800 Subject: [PATCH 03/11] add unit test for cute dsl kernels Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .../dense_blockscaled_gemm_persistent.py | 50 +- .../test_dense_blockscaled_gemm_persistent.py | 543 ++++++++++++++++++ tests/scripts/cute_dsl_kernels/testing.py | 446 ++++++++++++++ 3 files changed, 990 insertions(+), 49 deletions(-) create mode 100644 tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py create mode 100644 tests/scripts/cute_dsl_kernels/testing.py diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py index 10295f2906b9..8e926a62735b 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -61,8 +61,6 @@ from .custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync from .utils import is_power_of_2 -# from testing import benchmark - class Sm100BlockScaledPersistentDenseGemmKernel: """Implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types and Blackwell GPU architectural features, including persistent tile scheduling and warp specialization. @@ -311,7 +309,7 @@ def _setup_attributes(self): # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n - # TODO: [alel] Currently set prefetch dist to num_ab_stage, we may have more option for prefetch dist auto tuning + # TODO: [alel] Currently set prefetch dist to num_ab_stage, we may have more options for prefetch dist auto tuning self.prefetch_dist = self.num_ab_stage @cute.jit @@ -2143,7 +2141,6 @@ def run( iterations: int = 1, skip_ref_check: bool = False, use_cold_l2: bool = False, - # use_cupti: bool = False, **kwargs, ): """Runs and benchmarks the persistent batched dense block-scaled GEMM. @@ -2198,7 +2195,6 @@ def run( 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"Use CUPTI: {'True' if use_cupti else 'False'}") # Unpack parameters m, n, k, l = mnkl @@ -2470,15 +2466,6 @@ def generate_tensors(): warmup_iterations=warmup_iterations, iterations=iterations, ) - # exec_time = benchmark( - # compiled_gemm, - # workspace_generator=generate_tensors, - # workspace_count=workspace_count, - # stream=current_stream, - # warmup_iterations=warmup_iterations, - # iterations=iterations, - # use_cupti=use_cupti, - # ) return exec_time # Return execution time in microseconds @@ -2555,18 +2542,6 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: default=False, help="Use circular buffer tensor sets to ensure L2 cold cache", ) - # parser.add_argument( - # "--use_cupti", - # action="store_true", - # default=False, - # help="Use CUPTI to measure execution time", - # ) - # parser.add_argument( - # "--print_duration", - # action="store_true", - # default=False, - # help="Print execution time", - # ) args = parser.parse_args() @@ -2597,27 +2572,4 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: args.skip_ref_check, args.use_cold_l2, ) - - # TODO: [alel] Would like to involve cupti testing here, but need to double check how to add cupti wheel in framework - # exec_time = run( - # args.mnkl, - # args.ab_dtype, - # args.sf_dtype, - # args.sf_vec_size, - # args.c_dtype, - # args.a_major, - # args.b_major, - # args.c_major, - # args.mma_tiler_mn, - # args.cluster_shape_mn, - # args.use_prefetch, - # args.tolerance, - # args.warmup_iterations, - # args.iterations, - # args.skip_ref_check, - # args.use_cold_l2, - # #args.use_cupti, - # ) - # if args.print_duration: - # print(f"Execution time: {exec_time} us") print("PASS") diff --git a/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py b/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py new file mode 100644 index 000000000000..4a8e6bf207da --- /dev/null +++ b/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py @@ -0,0 +1,543 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# This file is copied and modified from cutlass example https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py + + +import argparse +import sys +from pathlib import Path +from typing import Tuple, Type + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +import torch +from cutlass.cute.runtime import from_dlpack + +try: + from tensorrt_llm._torch.cute_dsl_kernels.blackwell import dense_blockscaled_gemm_persistent as kernel_module +except (ModuleNotFoundError, ImportError): + sys.path.insert(0, str(Path(__file__).parents[3] / "tensorrt_llm/_torch/cute_dsl_kernels")) + from blackwell import dense_blockscaled_gemm_persistent as kernel_module + +Sm100BlockScaledPersistentDenseGemmKernel = kernel_module.Sm100BlockScaledPersistentDenseGemmKernel +cvt_sf_MKL_to_M32x4xrm_K4xrk_L = kernel_module.cvt_sf_MKL_to_M32x4xrm_K4xrk_L + +try: + from .testing import benchmark +except ImportError: + from testing import benchmark + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_prefetch: bool = False, + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + use_cupti: bool = False, + **kwargs, +): + """Runs and benchmarks the persistent batched dense block-scaled GEMM. + + This function prepares input tensors, launches the kernel, optionally + validates the results against a reference implementation, and measures + performance. + + Args: + mnkl (Tuple[int, int, int, int]): The dimensions (M, N, K, L) of the + GEMM problem. + ab_dtype (Type[cutlass.Numeric]): Data type for input tensors A and B. + sf_dtype (Type[cutlass.Numeric]): Data type for the scale factor tensors. + sf_vec_size (int): Vector size for the scale factor tensors. + c_dtype (Type[cutlass.Numeric]): Data type for the output tensor C. + a_major (str): The major layout of tensor A ('k' or 'm'). + b_major (str): The major layout of tensor B ('k' or 'n'). + c_major (str): The major layout of tensor C ('n' or 'm'). + mma_tiler_mn (Tuple[int, int]): The shape of the MMA tile. + cluster_shape_mn (Tuple[int, int]): The shape of the CTA cluster. + tolerance (float, optional): Tolerance for result validation. + Defaults to 1e-01. + warmup_iterations (int, optional): Number of warmup runs. Defaults to 0. + iterations (int, optional): Number of benchmark iterations. + Defaults to 1. + skip_ref_check (bool, optional): If True, skips result validation. + Defaults to False. + use_cold_l2 (bool, optional): If True, uses a circular buffer to + ensure a cold L2 cache. Defaults to False. + use_cupti (bool, optional): If True, uses CUPTI to measure execution time. + Defaults to False. + **kwargs: Additional keyword arguments. + + Returns: + float: The execution time of the kernel in microseconds. + + Raises: + RuntimeError: If no CUDA-capable GPU is available. + TypeError: If the configuration is not supported by the kernel. + """ + print(f"Running Sm100 Persistent Dense BlockScaled GEMM test with:") + print(f"mnkl: {mnkl}") + print( + f"AB dtype: {ab_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}, Cluster Shape (M, N): {cluster_shape_mn}" + ) + print(f"Use prefetch: {'True' if use_prefetch 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"Use CUPTI: {'True' if use_cupti else 'False'}") + + # Unpack parameters + m, n, k, l = mnkl + + # Skip unsupported testcase + if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement( + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + c_major, + ): + raise TypeError( + f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" + ) + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Create tensor A/B/C + a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) + b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) + c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + + a_tensor, a_torch = cutlass_torch.cute_tensor_like(a_ref, + ab_dtype, + is_dynamic_layout=True, + assumed_align=16) + b_tensor, b_torch = cutlass_torch.cute_tensor_like(b_ref, + ab_dtype, + is_dynamic_layout=True, + assumed_align=16) + c_tensor, c_torch = cutlass_torch.cute_tensor_like(c_ref, + c_dtype, + is_dynamic_layout=True, + assumed_align=16) + + # Mark tensor to be byte aligned + 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, + ) + + # Create scale factor tensor SFA/SFB + def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): + + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (l, mn, sf_k) + + 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, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # 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 + + sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype) + sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype) + + # Configure gemm kernel + gemm = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + use_prefetch, + ) + + # Compute max active clusters on current device + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + # Compile gemm kernel + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + 1.0, # alpha + max_active_clusters, + current_stream, + options=f"--opt-level 2", + ) + + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, 1.0, + current_stream) + print("Verifying results...") + res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) + res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + + # Convert c back to f32 for comparison. + c_ref_device = c_ref.cuda() + cute.testing.convert( + c_tensor, + from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=(1 if c_major == "n" else 0)), + ) + c_ref = c_ref_device.cpu() + + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Convert ref : f32 -> f8 -> f32 + ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, + device="cuda").permute(1, 2, 0) + ref_f8 = from_dlpack( + ref_f8_, assumed_align=16).mark_layout_dynamic(leading_dim=1) + ref_f8.element_type = c_dtype + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, + 0).cuda() + ref_tensor = from_dlpack( + ref_device, assumed_align=16).mark_layout_dynamic(leading_dim=1) + cute.testing.convert(ref_tensor, ref_f8) + cute.testing.convert(ref_f8, ref_tensor) + ref = ref_device.cpu() + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + + def generate_tensors(): + a_tensor, _ = cutlass_torch.cute_tensor_like(a_ref, + ab_dtype, + is_dynamic_layout=True, + assumed_align=16) + b_tensor, _ = cutlass_torch.cute_tensor_like(b_ref, + ab_dtype, + is_dynamic_layout=True, + assumed_align=16) + c_tensor, _ = cutlass_torch.cute_tensor_like(c_ref, + c_dtype, + is_dynamic_layout=True, + assumed_align=16) + + # Mark tensor to be byte aligned + 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, + ) + + _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, + sf_dtype) + _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, + sf_dtype) + return cute.testing.JitArguments(a_tensor, b_tensor, sfa_tensor, + sfb_tensor, c_tensor, 1.0, + current_stream) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = (a_torch.numel() * a_torch.element_size() + + b_torch.numel() * b_torch.element_size() + + sfa_torch.numel() * sfa_torch.element_size() + + sfb_torch.numel() * sfb_torch.element_size() + + c_torch.numel() * c_torch.element_size()) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations) + + exec_time = benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cupti=use_cupti, + ) + + return exec_time # Return execution time in microseconds + + +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="Performance test for Sm100 Dense Persistent 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_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, 4), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", + type=cutlass.dtype, + default=cutlass.Float4E2M1FN) + parser.add_argument("--sf_dtype", + type=cutlass.dtype, + default=cutlass.Float8E8M0FNU) + 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"], 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_prefetch", + action="store_true", + default=False, + help="Enable TMA prefetch for both A and B matrices (default: False)", + ) + parser.add_argument("--tolerance", + type=float, + default=1e-01, + help="Tolerance for validation") + 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("--skip_ref_check", + action="store_true", + help="Skip reference checking") + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + parser.add_argument( + "--use_cupti", + action="store_true", + default=False, + help="Use CUPTI to measure execution time", + ) + parser.add_argument( + "--print_duration", + action="store_true", + default=False, + help="Print execution time", + ) + + 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") + + exec_time = run( + args.mnkl, + args.ab_dtype, + args.sf_dtype, + args.sf_vec_size, + args.c_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_prefetch, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + args.use_cupti, + ) + if args.print_duration: + print(f"Execution time: {exec_time} us") + print("PASS") + diff --git a/tests/scripts/cute_dsl_kernels/testing.py b/tests/scripts/cute_dsl_kernels/testing.py new file mode 100644 index 000000000000..ba72b5b61531 --- /dev/null +++ b/tests/scripts/cute_dsl_kernels/testing.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +import functools +import inspect +import logging +import os +from itertools import product +from time import time +from typing import Type, Union, Callable, Optional, Dict, List, Any + +import cuda.bindings.driver as cuda_driver +import cuda.bindings.runtime as cuda_runtime + +import cutlass.base_dsl.jit_executor +from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op + +from cutlass.cute.typing import Numeric, Int8, Boolean + +import cutlass.cute as cute +from cutlass.cute import nvgpu + +from cutlass._mlir.dialects import builtin, cf, nvvm, vector +from functools import partial + +from cutlass.cute.testing import JitArguments +from cupti import cupti + +class CuptiProfiler: + """A class for managing CUPTI profiling measurements with start, stop, and duration methods. + + This class provides a clean interface for measuring CUDA kernel execution times + using CUPTI (CUDA Profiling Tools Interface). It encapsulates the complexity + of buffer management, callback registration, and activity tracking. + + Example usage: + profiler = CuptiProfiler() + profiler.start() + # ... run your CUDA kernels ... + profiler.stop() + duration = profiler.get_duration() # Returns total duration in milliseconds + """ + + def __init__(self, buffer_size: int = 8 * 1024 * 1024): + """Initialize the CUPTI profiler. + + Args: + buffer_size: Size of the CUPTI buffer in bytes (default: 8MB) + """ + self.buffer_size = buffer_size + self.timings = [] + self._is_active = False + self._buffer_requested_callback = None + self._buffer_completed_callback = None + + def _buffer_requested(self): + """Internal callback for CUPTI buffer requests.""" + max_num_records = 0 + return self.buffer_size, max_num_records + + def _buffer_completed(self, activities: list): + """Internal callback for processing completed CUPTI activities.""" + for activity in activities: + start = activity.start if hasattr(activity, "start") else "nil" + end = activity.end if hasattr(activity, "end") else "nil" + duration = end - start if start != "nil" and end != "nil" else "nil" + name = activity.name[:100] if hasattr(activity, "name") else "unknown" + # Convert to milliseconds + if duration != "nil": + self.timings.append((name, duration / 1e6)) + # print(f"Activity: {name}, Duration: {duration / 1e6} ms") + + def start(self): + """Start CUPTI profiling. + + Enables CUPTI activity tracking for concurrent kernels and registers + the necessary callbacks for buffer management. + + Raises: + ValueError: If CUPTI activity cannot be enabled + """ + if self._is_active: + raise RuntimeError("CUPTI profiler is already active") + + # Clear previous timings + self.timings = [] + + try: + cupti.activity_enable(cupti.ActivityKind.CONCURRENT_KERNEL) + except cupti.cuptiError as e: + raise ValueError( + f"Error while enabling Activity Kind {cupti.ActivityKind.CONCURRENT_KERNEL.name}: {e}" + ) + + # Register callbacks + self._buffer_requested_callback = self._buffer_requested + self._buffer_completed_callback = partial(self._buffer_completed) + + cupti.activity_register_callbacks( + self._buffer_requested_callback, self._buffer_completed_callback + ) + + self._is_active = True + + def stop(self): + """Stop CUPTI profiling. + + Flushes all activities, disables CUPTI tracking, and finalizes the profiler. + This method should be called after the kernels you want to measure have completed. + """ + if not self._is_active: + raise RuntimeError("CUPTI profiler is not active") + + # Flush all activities and cleanup + cupti.activity_flush_all(0) + cupti.activity_disable(cupti.ActivityKind.CONCURRENT_KERNEL) + cupti.finalize() + + self._is_active = False + + def get_duration(self) -> float: + """Get the total duration of all measured activities in milliseconds. + + Returns: + Total duration in milliseconds. Returns 0.0 if no activities were recorded. + """ + return sum(timing[1] for timing in self.timings) + +def _cuda_success( + err: Union[tuple, cuda_runtime.cudaError_t, cuda_driver.CUresult], message: str +): + """ + Helper function to check CUDA API errors. + """ + if isinstance(err, tuple): + _cuda_success(err[0], message) + elif isinstance(err, cuda_runtime.cudaError_t): + error_message = cuda_runtime.cudaGetErrorString(err)[1].decode("utf-8") + if err != cuda_runtime.cudaError_t.cudaSuccess: + raise RuntimeError(f"{message} : {error_message}") + elif isinstance(err, cuda_driver.CUresult): + if err != cuda_driver.CUresult.CUDA_SUCCESS: + error_message = cuda_driver.cuGetErrorString(err)[1].decode("utf-8") + raise RuntimeError(f"{message} : {error_message}") + else: + raise TypeError( + f"{err} is an unexpected type : it should be a cudaError_t or CUresult" + ) + +def _does_kernel_use_stream( + kernel: Callable, stream: cuda_driver.CUstream, *args, **kwargs +): + """ + This function checks if the kernel uses the provided non-default stream. + It does this by capturing the stream and then checking if any kernels were launched. + :param kernel: The kernel to check + :type kernel: Callable + :param stream: The stream to check + :type stream: cuda_driver.CUstream + :return: True if the kernel uses the stream, False otherwise + :rtype: bool + """ + + assert int(stream) != int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT), ( + "Stream must be a non-default stream" + ) + + err = cuda_runtime.cudaStreamBeginCapture( + stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal + ) + _cuda_success(err, "Error on stream capture") + + kernel(*args, **kwargs) + + err, graph = cuda_runtime.cudaStreamEndCapture(stream) + _cuda_success(err, "Error on stream capture") + + # Get number of nodes in warmup graph to check it matches what is expected + err, _, num_nodes = cuda_runtime.cudaGraphGetNodes(graph) + _cuda_success(err, "Error on querying graph") + return num_nodes > 0 + +def benchmark( + callable: Callable, + *, + warmup_iterations: int = 10, + iterations: int = 100, + stream: Optional[cuda_driver.CUstream] = None, + kernel_arguments: Optional[JitArguments] = None, + workspace_generator: Optional[Callable[[], JitArguments]] = None, + workspace_count: int = 1, + use_cuda_graphs: bool = False, + use_cupti: bool = False, +) -> float: + """Benchmarks a callable function with the specified parameters. + + For example, + .. code-block:: python + + from cutlass.cute.testing import benchmark + + @cute.jit + def user_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda_driver.CUstream): + # contents of the function + pass + + time_us = benchmark(user_function, kernel_arguments=JitArguments(a, b, c, stream) + warmup_iterations=10, iterations=100 + stream=stream) + + To prevent skewing results by repeately accessing the L2 cache, use the workspace_count and workspace_generator + parameters to cycle through a number of different workspaces. + + .. code-block:: python + + from cutlass.cute.testing import benchmark + + @cute.jit + def user_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor): + # contents of the function + pass + + def workspace_generator(): + # create a, b, and c + return JitArguments(a, b, c) + + time_us = benchmark(user_function, + workspace_generator=workspace_generator, + workspace_count=10, + warmup_iterations=10000, + iterations=1000) + + To benchmark you may always configure the function being profiled (callable), the warmup iterations, and + the number of profiling iterations. + + Whenever the kernel being benchmarked runs in a non-default stream, the stream must be provided through the stream parameter. + + To use CUDA graphs, the callable must be a compiled @cute.jit annotated function. + When using CUDA graphs, the kernel must be launched in a non-default stream. + + :param callable: The function to benchmark + :type callable: Callable + :param warmup_iterations: Number of warmup iterations, defaults to 10 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations, defaults to 100 + :type iterations: int, optional + :param stream: Stream kernel is launched in, defaults to CUDA stream default + :type stream: CUstream, None + :param kernel_arguments: Kernel arguments to launch callable with, defaults to None + :type kernel_arguments: JitArguments, None + :param workspace_generator: Function that returns kernel arguments, defaults to None + :type workspace_generator: Callable + :param workspace_count: Number of workspaces (arguments) to loop through, looping through enough workspaces will keep the L2 cache cold + :type workspace_count: int, optional + :param use_cuda_graphs: Whether to use cuda graphs, defaults to False + :type use_cuda_graphs: bool, optional + + :return: The benchmark time in microseconds + :rtype: float + """ + + if stream is None: + stream = cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) + + if workspace_count < 1: + raise ValueError("workspace_count must be at least 1") + + time_us = float("nan") + if workspace_generator == None: + # If no workspace generator is provided, we need a single workspace + if workspace_count != 1: + raise ValueError("Need a single workspace if not providing a generator") + + # If no workspace generator is provided, we need a kernel_argument + if kernel_arguments == None: + raise ValueError( + "Please pass a kernel argument if not providing a generator" + ) + workspace_generator = lambda: kernel_arguments + + workspaces = [workspace_generator() for _ in range(workspace_count)] + + for workspace in workspaces: + if type(workspace) != JitArguments: + raise TypeError( + "workspace_generator and/or kernel_arguments should use JitArguments type" + ) + + def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): + for _ in range(iterations): + current_workspace = workspaces[workspace_index] + callable(*current_workspace.args, **current_workspace.kwargs) + workspace_index = (workspace_index + 1) % workspace_count + return workspace_index + + # Create CUDA events for timing + err, start_event = cuda_driver.cuEventCreate( + cuda_driver.CUevent_flags.CU_EVENT_DEFAULT + ) + _cuda_success(err, "Error on creating event") + err, end_event = cuda_driver.cuEventCreate( + cuda_driver.CUevent_flags.CU_EVENT_DEFAULT + ) + _cuda_success(err, "Error on creating event") + + elapsed_time = float("nan") + + if use_cuda_graphs: + # Check if the callable is a JitCompiledFunction or JitExecutor + # These are functions that can be called to launch kernels + compiled_types = ( + cutlass.base_dsl.jit_executor.JitCompiledFunction, + cutlass.base_dsl.jit_executor.JitExecutor, + ) + if not isinstance(callable, compiled_types): + raise TypeError("Function must be precompiled to be used with CUDA Graphs") + + # Check if the stream is a non-default stream + if int(stream) == int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT): + raise ValueError( + "Measuring with CUDA Graphs requires executing in a non-default stream" + ) + + workspace_index = 0 + + # Capture warmup graph + err = cuda_runtime.cudaStreamBeginCapture( + stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal + ) + _cuda_success(err, "Error on stream capture") + + workspace_index = _loop_and_call_kernel(warmup_iterations) + err, gwarm = cuda_runtime.cudaStreamEndCapture(stream) + _cuda_success(err, "Error on stream capture") + + # Get number of nodes in warmup graph to check it matches what is expected + err, _, num_nodes = cuda_runtime.cudaGraphGetNodes(gwarm) + _cuda_success(err, "Error on querying graph") + # Assertion is >= since we may launch multiple kernels in one host function + if num_nodes < warmup_iterations: + raise ValueError( + "CUDA stream passed to benchmark does not match the stream the kernel was launched in" + ) + + # Capture profiling graph + err = cuda_runtime.cudaStreamBeginCapture( + stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal + ) + _cuda_success(err, "Error on stream capture") + _loop_and_call_kernel(iterations, workspace_index) + err, gprofile = cuda_runtime.cudaStreamEndCapture(stream) + _cuda_success(err, "Error on stream capture") + + # Instantiate graphs + err, gwarm = cuda_runtime.cudaGraphInstantiate(gwarm, 0) + _cuda_success(err, "Error on graph instantiation") + err, gprofile = cuda_runtime.cudaGraphInstantiate(gprofile, 0) + _cuda_success(err, "Error on graph instantiation") + + # Launch warmup graph + err = cuda_runtime.cudaGraphLaunch(gwarm, stream) + _cuda_success(err, "Error on graph launch") + + # Record start time + err = cuda_driver.cuEventRecord(start_event, stream) + _cuda_success(err, "Error on recording event") + + # Launch profiling graph + err = cuda_runtime.cudaGraphLaunch(gprofile, stream) + _cuda_success(err, "Error on graph launch") + + # Record end time + err = cuda_driver.cuEventRecord(end_event, stream) + _cuda_success(err, "Error on recording event") + err = cuda_driver.cuEventSynchronize(end_event) + _cuda_success(err, "Error on synchronizing event") + + # Get elapsed time + err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event) + _cuda_success(err, "Error on querying event") + + # Destroy graphs + err = cuda_runtime.cudaGraphExecDestroy(gwarm) + _cuda_success(err, "Error on destroying graph") + err = cuda_runtime.cudaGraphExecDestroy(gprofile) + _cuda_success(err, "Error on destroying graph") + + + elif use_cupti: + # Use the new CuptiProfiler class + profiler = CuptiProfiler() + + # Warmup + workspace_index = _loop_and_call_kernel(warmup_iterations) + + profiler.start() + + _loop_and_call_kernel(iterations, workspace_index) + # Synchronize device + err = cuda_runtime.cudaDeviceSynchronize() + _cuda_success(err, "Error on synchronizing device") + + profiler.stop() + elapsed_time = profiler.get_duration() + + else: + if int(stream) != int( + cuda_driver.CUstream_flags.CU_STREAM_DEFAULT + ) and not _does_kernel_use_stream( + callable, stream, *workspaces[0].args, **workspaces[0].kwargs + ): + raise ValueError( + "CUDA stream passed to benchmark does not match the stream the kernel was launched in" + ) + + # Not using graphs + # Warmup + workspace_index = _loop_and_call_kernel(warmup_iterations) + # Record start event + err = cuda_driver.cuEventRecord(start_event, stream) + _cuda_success(err, "Error on recording event") + _loop_and_call_kernel(iterations, workspace_index) + # Record end event + err = cuda_driver.cuEventRecord(end_event, stream) + _cuda_success(err, "Error on recording event") + # Synchronize end event + err = cuda_driver.cuEventSynchronize(end_event) + _cuda_success(err, "Error on synchronizing event") + err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event) + _cuda_success(err, "Error on querying event") + + # Destroy events + err = cuda_driver.cuEventDestroy(start_event) + _cuda_success(err, "Error on destroying event") + err = cuda_driver.cuEventDestroy(end_event) + _cuda_success(err, "Error on destroying event") + + return elapsed_time / iterations * 1e3 From 41bcb5c931c46024ac60e3c7163ea7a4ca157f49 Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Thu, 27 Nov 2025 01:18:32 -0800 Subject: [PATCH 04/11] add prefetch auto tune knob Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 69 ++++++++++--------- 1 file changed, 36 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index eceed9b779b9..c07dc89ee5a9 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -102,37 +102,38 @@ def get_valid_tactics( (4, 4), ] swap_ab_candidates = [True, False] + use_prefetch_candidates = [True, False] valid_tactics = [] - for swap_ab in swap_ab_candidates: - for mma_tiler_mn in mma_tiler_mn_candidates: - for cluster_shape_mn in cluster_shape_mn_candidates: - if swap_ab: - c_major = "m" - kernel_m = n - kernel_n = m - else: - c_major = "n" - kernel_m = m - kernel_n = n - - if self.__class__.kernel_class.can_implement( - cutlass.Float4E2M1FN, # ab_dtype, - cutlass.Float8E4M3FN, # sf_dtype - sf_vec_size, # sf_vec_size, - cutlass.BFloat16, # c_dtype, - mma_tiler_mn, - cluster_shape_mn, - kernel_m, - kernel_n, - real_k, - batch_size, - a_major, - b_major, - c_major, - ): - valid_tactics.append( - (mma_tiler_mn, cluster_shape_mn, swap_ab)) + for mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch in itertools.product( + mma_tiler_mn_candidates, cluster_shape_mn_candidates, + swap_ab_candidates, use_prefetch_candidates): + if swap_ab: + c_major = "m" + kernel_m = n + kernel_n = m + else: + c_major = "n" + kernel_m = m + kernel_n = n + + if self.__class__.kernel_class.can_implement( + cutlass.Float4E2M1FN, # ab_dtype, + cutlass.Float8E4M3FN, # sf_dtype + sf_vec_size, # sf_vec_size, + cutlass.BFloat16, # c_dtype, + mma_tiler_mn, + cluster_shape_mn, + kernel_m, + kernel_n, + real_k, + batch_size, + a_major, + b_major, + c_major, + ): + valid_tactics.append( + (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch)) return valid_tactics @@ -161,7 +162,7 @@ def forward( inputs[3]: Weight scale tensor of shape (n, k//16), dtype: fp8. inputs[4]: Alpha scaling factor. dtype: float32. inputs[5]: Output dtype, expected to be torch.bfloat16. - tactic: Tiling and cluster strategy, typically a tuple (mma_tiler_mn, cluster_shape_mn). + tactic: Tiling and cluster strategy, typically a tuple (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch). Returns: torch.Tensor: Output tensor of shape (m, n), dtype: bf16. @@ -169,13 +170,14 @@ def forward( sf_vec_size = 16 if isinstance(tactic, tuple): - mma_tiler_mn, cluster_shape_mn, swap_ab = tactic + mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch = tactic else: # fallback to default tactic - mma_tiler_mn, cluster_shape_mn, swap_ab = [ + mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch = [ (128, 128), (1, 1), False, + False, ] a_tensor, b_tensor, a_sf_tensor, b_sf_tensor = inputs @@ -211,7 +213,7 @@ def forward( torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) - cache_key = (sf_vec_size, mma_tiler_mn, cluster_shape_mn, swap_ab) + cache_key = (sf_vec_size, mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch) if swap_ab: kernel_a_ptr = b_ptr kernel_a_sf_ptr = b_sf_ptr @@ -236,6 +238,7 @@ def forward( sf_vec_size, mma_tiler_mn, cluster_shape_mn, + use_prefetch, ) # Compute max active clusters on current device hardware_info = cutlass.utils.HardwareInfo() From 837b58ac8b34b8f6c6503c00de69d920160a81ed Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Thu, 27 Nov 2025 23:46:23 -0800 Subject: [PATCH 05/11] refine Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- pyproject.toml | 2 + .../dense_blockscaled_gemm_persistent.py | 549 ++---------------- 2 files changed, 48 insertions(+), 503 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b453d975e36a..0606d871ffd3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -983,6 +983,8 @@ exclude = [ "tests/scripts/iteration_log_parser.py", "tests/scripts/perf-sanity/parse_benchmark_results.py", "tests/scripts/perf-sanity/run_benchmark_serve.py", + "tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py", + "tests/scripts/cute_dsl_kernels/testing.py", "tests/unittest/_torch/attention/sparse/test_dsa_indexer.py", "tests/unittest/_torch/attention/sparse/test_flash_mla.py", "tests/unittest/_torch/attention/sparse/test_rocketkv.py", diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py index 8e926a62735b..11e8e74d7a17 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -43,24 +43,21 @@ # This file is copied and modified from cutlass example https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py -import argparse from typing import Optional, Tuple, Type, Union import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils -import torch from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.runtime import from_dlpack from .custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync from .utils import is_power_of_2 + class Sm100BlockScaledPersistentDenseGemmKernel: """Implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types and Blackwell GPU architectural features, including persistent tile scheduling and warp specialization. @@ -301,10 +298,14 @@ def _setup_attributes(self): 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_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.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 # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n @@ -457,26 +458,16 @@ def __call__( 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] - ) + 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) + 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) @@ -852,26 +843,20 @@ def kernel( 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) - ) + 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] - ) - ) - ) + 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) - ) + cute.append(acc_shape, self.num_acc_stage)) # # Cluster wait before tensor memory alloc @@ -928,7 +913,10 @@ def kernel( if cutlass.const_expr(self.use_prefetch): # Prefetch both A and B (default behavior) - for pf_k_block in cutlass.range(0, min(self.prefetch_dist, k_block_cnt), unroll=1): + for pf_k_block in cutlass.range(0, + min(self.prefetch_dist, + k_block_cnt), + unroll=1): cute.prefetch( tma_atom_a, tAgA_slice[(None, pf_k_block)], @@ -1076,7 +1064,8 @@ def kernel( # Make SFB tmem tensor sfb_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + acc_tmem_ptr + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols, dtype=self.sf_dtype, ) # (MMA, MMA_N, MMA_K) @@ -1123,8 +1112,7 @@ def kernel( # Set tensor memory buffer for current tile # (MMA, MMA_M, MMA_N) - tCtAcc = tCtAcc_base[(None, None, None, - acc_stage_index)] + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] # Peek (try_wait) AB buffer full for k_block = 0 ab_consumer_state.reset_count() @@ -1142,10 +1130,12 @@ def kernel( tCtSFB_mma = tCtSFB if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 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(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) + offset = cutlass.Int32( + 2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32( + 0) shifted_ptr = cute.recast_ptr( - acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols - + offset, + 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) @@ -1153,8 +1143,8 @@ def kernel( # Move in increments of 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, + 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) @@ -1351,7 +1341,9 @@ def kernel( if cutlass.const_expr(self.overlapping_accum): acc_stage_index = acc_consumer_state.phase - reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False) + reverse_subtile = cutlass.Boolean( + True) if acc_stage_index == 0 else cutlass.Boolean( + False) else: acc_stage_index = acc_consumer_state.index @@ -1377,7 +1369,8 @@ def kernel( real_subtile_idx = subtile_idx if cutlass.const_expr(self.overlapping_accum): if reverse_subtile: - real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx + real_subtile_idx = self.cta_tile_shape_mnk[ + 1] // self.epi_tile_n - 1 - subtile_idx # # Load accumulator from tensor memory buffer to register # @@ -1385,14 +1378,15 @@ def kernel( cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) # - # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled + # Async arrive accumulator buffer empty earlier 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() with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) + acc_pipeline.consumer_release( + acc_consumer_state) acc_consumer_state.advance() # # Convert to C type @@ -2122,454 +2116,3 @@ def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( 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] - - -def run( - mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], - sf_dtype: Type[cutlass.Numeric], - sf_vec_size: int, - c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, - mma_tiler_mn: Tuple[int, int], - cluster_shape_mn: Tuple[int, int], - use_prefetch: bool = False, - tolerance: float = 1e-01, - warmup_iterations: int = 0, - iterations: int = 1, - skip_ref_check: bool = False, - use_cold_l2: bool = False, - **kwargs, -): - """Runs and benchmarks the persistent batched dense block-scaled GEMM. - - This function prepares input tensors, launches the kernel, optionally - validates the results against a reference implementation, and measures - performance. - - Args: - mnkl (Tuple[int, int, int, int]): The dimensions (M, N, K, L) of the - GEMM problem. - ab_dtype (Type[cutlass.Numeric]): Data type for input tensors A and B. - sf_dtype (Type[cutlass.Numeric]): Data type for the scale factor tensors. - sf_vec_size (int): Vector size for the scale factor tensors. - c_dtype (Type[cutlass.Numeric]): Data type for the output tensor C. - a_major (str): The major layout of tensor A ('k' or 'm'). - b_major (str): The major layout of tensor B ('k' or 'n'). - c_major (str): The major layout of tensor C ('n' or 'm'). - mma_tiler_mn (Tuple[int, int]): The shape of the MMA tile. - cluster_shape_mn (Tuple[int, int]): The shape of the CTA cluster. - tolerance (float, optional): Tolerance for result validation. - Defaults to 1e-01. - warmup_iterations (int, optional): Number of warmup runs. Defaults to 0. - iterations (int, optional): Number of benchmark iterations. - Defaults to 1. - skip_ref_check (bool, optional): If True, skips result validation. - Defaults to False. - use_cold_l2 (bool, optional): If True, uses a circular buffer to - ensure a cold L2 cache. Defaults to False. - **kwargs: Additional keyword arguments. - - Returns: - float: The execution time of the kernel in microseconds. - - Raises: - RuntimeError: If no CUDA-capable GPU is available. - TypeError: If the configuration is not supported by the kernel. - """ - print(f"Running Sm100 Persistent Dense BlockScaled GEMM test with:") - print(f"mnkl: {mnkl}") - print( - f"AB dtype: {ab_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}, Cluster Shape (M, N): {cluster_shape_mn}" - ) - print(f"Use prefetch: {'True' if use_prefetch 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'}") - - # Unpack parameters - m, n, k, l = mnkl - - # Skip unsupported testcase - if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement( - ab_dtype, - sf_dtype, - sf_vec_size, - c_dtype, - mma_tiler_mn, - cluster_shape_mn, - m, - n, - k, - l, - a_major, - b_major, - c_major, - ): - raise TypeError( - f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" - ) - - if not torch.cuda.is_available(): - raise RuntimeError("GPU is required to run this example!") - - torch.manual_seed(1111) - - # Create tensor A/B/C - a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) - b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) - c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) - - a_tensor, a_torch = cutlass_torch.cute_tensor_like(a_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - b_tensor, b_torch = cutlass_torch.cute_tensor_like(b_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - c_tensor, c_torch = cutlass_torch.cute_tensor_like(c_ref, - c_dtype, - is_dynamic_layout=True, - assumed_align=16) - - # Mark tensor to be byte aligned - 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - c_tensor.mark_compact_shape_dynamic( - mode=1 if c_major == "n" else 0, - stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), - divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, - ) - - # Create scale factor tensor SFA/SFB - def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): - - def ceil_div(a, b): - return (a + b - 1) // b - - sf_k = ceil_div(k, sf_vec_size) - ref_shape = (l, mn, sf_k) - - 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, - ) - - ref_permute_order = (1, 2, 0) - mma_permute_order = (3, 4, 1, 5, 2, 0) - - # 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 - - sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( - l, m, k, sf_vec_size, sf_dtype) - sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( - l, n, k, sf_vec_size, sf_dtype) - - # Configure gemm kernel - gemm = Sm100BlockScaledPersistentDenseGemmKernel( - sf_vec_size, - mma_tiler_mn, - cluster_shape_mn, - use_prefetch, - ) - - # Compute max active clusters on current device - hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters( - cluster_shape_mn[0] * cluster_shape_mn[1]) - - # Initialize Stream - current_stream = cutlass_torch.default_stream() - - # Compile gemm kernel - compiled_gemm = cute.compile( - gemm, - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - c_tensor, - 1.0, # alpha - max_active_clusters, - current_stream, - options=f"--opt-level 2", - ) - - # Compute reference result - if not skip_ref_check: - # Execute kernel once for reference checking - compiled_gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, 1.0, - current_stream) - print("Verifying results...") - res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) - res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) - ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) - - # Convert c back to f32 for comparison. - c_ref_device = c_ref.cuda() - cute.testing.convert( - c_tensor, - from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( - leading_dim=(1 if c_major == "n" else 0)), - ) - c_ref = c_ref_device.cpu() - - if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): - torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) - elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): - # Convert ref : f32 -> f8 -> f32 - ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, - device="cuda").permute(1, 2, 0) - ref_f8 = from_dlpack( - ref_f8_, assumed_align=16).mark_layout_dynamic(leading_dim=1) - ref_f8.element_type = c_dtype - ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, - 0).cuda() - ref_tensor = from_dlpack( - ref_device, assumed_align=16).mark_layout_dynamic(leading_dim=1) - cute.testing.convert(ref_tensor, ref_f8) - cute.testing.convert(ref_f8, ref_tensor) - ref = ref_device.cpu() - torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) - - def generate_tensors(): - a_tensor, _ = cutlass_torch.cute_tensor_like(a_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - b_tensor, _ = cutlass_torch.cute_tensor_like(b_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - c_tensor, _ = cutlass_torch.cute_tensor_like(c_ref, - c_dtype, - is_dynamic_layout=True, - assumed_align=16) - - # Mark tensor to be byte aligned - 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - 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=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - c_tensor.mark_compact_shape_dynamic( - mode=1 if c_major == "n" else 0, - stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), - divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, - ) - - _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, - sf_dtype) - _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, - sf_dtype) - return cute.testing.JitArguments(a_tensor, b_tensor, sfa_tensor, - sfb_tensor, c_tensor, 1.0, - current_stream) - - workspace_count = 1 - if use_cold_l2: - one_workspace_bytes = (a_torch.numel() * a_torch.element_size() + - b_torch.numel() * b_torch.element_size() + - sfa_torch.numel() * sfa_torch.element_size() + - sfb_torch.numel() * sfb_torch.element_size() + - c_torch.numel() * c_torch.element_size()) - workspace_count = cute.testing.get_workspace_count( - one_workspace_bytes, warmup_iterations, iterations) - - exec_time = cute.testing.benchmark( - compiled_gemm, - workspace_generator=generate_tensors, - workspace_count=workspace_count, - stream=current_stream, - warmup_iterations=warmup_iterations, - iterations=iterations, - ) - - return exec_time # Return execution time in microseconds - - -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 Persistent 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_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, 4), - help="Cluster shape (comma-separated)", - ) - parser.add_argument("--ab_dtype", - type=cutlass.dtype, - default=cutlass.Float4E2M1FN) - parser.add_argument("--sf_dtype", - type=cutlass.dtype, - default=cutlass.Float8E8M0FNU) - 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"], 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_prefetch", - action="store_true", - default=False, - help="Enable TMA prefetch for both A and B matrices (default: False)", - ) - parser.add_argument("--tolerance", - type=float, - default=1e-01, - help="Tolerance for validation") - 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("--skip_ref_check", - action="store_true", - help="Skip reference checking") - parser.add_argument( - "--use_cold_l2", - action="store_true", - default=False, - help="Use circular buffer tensor sets to ensure L2 cold cache", - ) - - 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.sf_dtype, - args.sf_vec_size, - args.c_dtype, - args.a_major, - args.b_major, - args.c_major, - args.mma_tiler_mn, - args.cluster_shape_mn, - args.use_prefetch, - args.tolerance, - args.warmup_iterations, - args.iterations, - args.skip_ref_check, - args.use_cold_l2, - ) - print("PASS") From 873b4cb08be784de8d895c4b0a8ceec53870cbe5 Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Thu, 27 Nov 2025 23:54:26 -0800 Subject: [PATCH 06/11] refine 1 Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .pre-commit-config.yaml | 2 ++ tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 311db068b505..25793e007c63 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -940,6 +940,8 @@ common-files: &common_files | tests/microbenchmarks/build_time_dashboard.py | tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py | tests/scripts/allreduce_perf/allreduce_perf_viz.py | + tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py | + tests/scripts/cute_dsl_kernels/testing.py | tests/scripts/iteration_log_parser.py | tests/scripts/perf-sanity/parse_benchmark_results.py | tests/scripts/perf-sanity/run_benchmark_serve.py | diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index c07dc89ee5a9..708c69a26041 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -213,7 +213,8 @@ def forward( torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) - cache_key = (sf_vec_size, mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch) + cache_key = (sf_vec_size, mma_tiler_mn, cluster_shape_mn, swap_ab, + use_prefetch) if swap_ab: kernel_a_ptr = b_ptr kernel_a_sf_ptr = b_sf_ptr From 060366862a183bfe444142905d859e9958eeb96a Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Thu, 27 Nov 2025 23:55:57 -0800 Subject: [PATCH 07/11] refine 2 Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .pre-commit-config.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 25793e007c63..a4745efb50d6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -940,11 +940,11 @@ common-files: &common_files | tests/microbenchmarks/build_time_dashboard.py | tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py | tests/scripts/allreduce_perf/allreduce_perf_viz.py | - tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py | - tests/scripts/cute_dsl_kernels/testing.py | tests/scripts/iteration_log_parser.py | tests/scripts/perf-sanity/parse_benchmark_results.py | tests/scripts/perf-sanity/run_benchmark_serve.py | + tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py | + tests/scripts/cute_dsl_kernels/testing.py | tests/unittest/_torch/attention/sparse/test_dsa_indexer.py | tests/unittest/_torch/attention/sparse/test_flash_mla.py | tests/unittest/_torch/attention/sparse/test_rocketkv.py | From 935e1de6ebe2cd59097f6dd9f7923a893eebe5af Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Fri, 28 Nov 2025 00:13:18 -0800 Subject: [PATCH 08/11] refine 3 Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- tests/scripts/cute_dsl_kernels/testing.py | 93 ++++++++++------------- 1 file changed, 39 insertions(+), 54 deletions(-) diff --git a/tests/scripts/cute_dsl_kernels/testing.py b/tests/scripts/cute_dsl_kernels/testing.py index ba72b5b61531..14e13df7ef28 100644 --- a/tests/scripts/cute_dsl_kernels/testing.py +++ b/tests/scripts/cute_dsl_kernels/testing.py @@ -9,30 +9,15 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -import functools -import inspect -import logging -import os -from itertools import product -from time import time -from typing import Type, Union, Callable, Optional, Dict, List, Any +from functools import partial +from typing import Callable, Optional, Union import cuda.bindings.driver as cuda_driver import cuda.bindings.runtime as cuda_runtime - import cutlass.base_dsl.jit_executor -from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op - -from cutlass.cute.typing import Numeric, Int8, Boolean - -import cutlass.cute as cute -from cutlass.cute import nvgpu - -from cutlass._mlir.dialects import builtin, cf, nvvm, vector -from functools import partial - -from cutlass.cute.testing import JitArguments from cupti import cupti +from cutlass.cute.testing import JitArguments + class CuptiProfiler: """A class for managing CUPTI profiling measurements with start, stop, and duration methods. @@ -72,7 +57,8 @@ def _buffer_completed(self, activities: list): start = activity.start if hasattr(activity, "start") else "nil" end = activity.end if hasattr(activity, "end") else "nil" duration = end - start if start != "nil" and end != "nil" else "nil" - name = activity.name[:100] if hasattr(activity, "name") else "unknown" + name = activity.name[:100] if hasattr(activity, + "name") else "unknown" # Convert to milliseconds if duration != "nil": self.timings.append((name, duration / 1e6)) @@ -104,9 +90,8 @@ def start(self): self._buffer_requested_callback = self._buffer_requested self._buffer_completed_callback = partial(self._buffer_completed) - cupti.activity_register_callbacks( - self._buffer_requested_callback, self._buffer_completed_callback - ) + cupti.activity_register_callbacks(self._buffer_requested_callback, + self._buffer_completed_callback) self._is_active = True @@ -134,9 +119,9 @@ def get_duration(self) -> float: """ return sum(timing[1] for timing in self.timings) -def _cuda_success( - err: Union[tuple, cuda_runtime.cudaError_t, cuda_driver.CUresult], message: str -): + +def _cuda_success(err: Union[tuple, cuda_runtime.cudaError_t, + cuda_driver.CUresult], message: str): """ Helper function to check CUDA API errors. """ @@ -155,9 +140,9 @@ def _cuda_success( f"{err} is an unexpected type : it should be a cudaError_t or CUresult" ) -def _does_kernel_use_stream( - kernel: Callable, stream: cuda_driver.CUstream, *args, **kwargs -): + +def _does_kernel_use_stream(kernel: Callable, stream: cuda_driver.CUstream, + *args, **kwargs): """ This function checks if the kernel uses the provided non-default stream. It does this by capturing the stream and then checking if any kernels were launched. @@ -170,12 +155,11 @@ def _does_kernel_use_stream( """ assert int(stream) != int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT), ( - "Stream must be a non-default stream" - ) + "Stream must be a non-default stream") err = cuda_runtime.cudaStreamBeginCapture( - stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal - ) + stream, + cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) _cuda_success(err, "Error on stream capture") kernel(*args, **kwargs) @@ -188,6 +172,7 @@ def _does_kernel_use_stream( _cuda_success(err, "Error on querying graph") return num_nodes > 0 + def benchmark( callable: Callable, *, @@ -268,22 +253,23 @@ def workspace_generator(): """ if stream is None: - stream = cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) + stream = cuda_driver.CUstream( + cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) if workspace_count < 1: raise ValueError("workspace_count must be at least 1") - time_us = float("nan") + float("nan") if workspace_generator == None: # If no workspace generator is provided, we need a single workspace if workspace_count != 1: - raise ValueError("Need a single workspace if not providing a generator") + raise ValueError( + "Need a single workspace if not providing a generator") # If no workspace generator is provided, we need a kernel_argument if kernel_arguments == None: raise ValueError( - "Please pass a kernel argument if not providing a generator" - ) + "Please pass a kernel argument if not providing a generator") workspace_generator = lambda: kernel_arguments workspaces = [workspace_generator() for _ in range(workspace_count)] @@ -303,12 +289,10 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): # Create CUDA events for timing err, start_event = cuda_driver.cuEventCreate( - cuda_driver.CUevent_flags.CU_EVENT_DEFAULT - ) + cuda_driver.CUevent_flags.CU_EVENT_DEFAULT) _cuda_success(err, "Error on creating event") err, end_event = cuda_driver.cuEventCreate( - cuda_driver.CUevent_flags.CU_EVENT_DEFAULT - ) + cuda_driver.CUevent_flags.CU_EVENT_DEFAULT) _cuda_success(err, "Error on creating event") elapsed_time = float("nan") @@ -321,7 +305,8 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): cutlass.base_dsl.jit_executor.JitExecutor, ) if not isinstance(callable, compiled_types): - raise TypeError("Function must be precompiled to be used with CUDA Graphs") + raise TypeError( + "Function must be precompiled to be used with CUDA Graphs") # Check if the stream is a non-default stream if int(stream) == int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT): @@ -333,8 +318,8 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): # Capture warmup graph err = cuda_runtime.cudaStreamBeginCapture( - stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal - ) + stream, + cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) _cuda_success(err, "Error on stream capture") workspace_index = _loop_and_call_kernel(warmup_iterations) @@ -352,8 +337,8 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): # Capture profiling graph err = cuda_runtime.cudaStreamBeginCapture( - stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal - ) + stream, + cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) _cuda_success(err, "Error on stream capture") _loop_and_call_kernel(iterations, workspace_index) err, gprofile = cuda_runtime.cudaStreamEndCapture(stream) @@ -384,7 +369,8 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): _cuda_success(err, "Error on synchronizing event") # Get elapsed time - err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event) + err, elapsed_time = cuda_driver.cuEventElapsedTime( + start_event, end_event) _cuda_success(err, "Error on querying event") # Destroy graphs @@ -393,7 +379,6 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): err = cuda_runtime.cudaGraphExecDestroy(gprofile) _cuda_success(err, "Error on destroying graph") - elif use_cupti: # Use the new CuptiProfiler class profiler = CuptiProfiler() @@ -413,10 +398,9 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): else: if int(stream) != int( - cuda_driver.CUstream_flags.CU_STREAM_DEFAULT - ) and not _does_kernel_use_stream( - callable, stream, *workspaces[0].args, **workspaces[0].kwargs - ): + cuda_driver.CUstream_flags.CU_STREAM_DEFAULT + ) and not _does_kernel_use_stream(callable, stream, *workspaces[0].args, + **workspaces[0].kwargs): raise ValueError( "CUDA stream passed to benchmark does not match the stream the kernel was launched in" ) @@ -434,7 +418,8 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): # Synchronize end event err = cuda_driver.cuEventSynchronize(end_event) _cuda_success(err, "Error on synchronizing event") - err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event) + err, elapsed_time = cuda_driver.cuEventElapsedTime( + start_event, end_event) _cuda_success(err, "Error on querying event") # Destroy events From 697f9a50c47f316aaa11775d5fb2de0d07851533 Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Fri, 28 Nov 2025 00:29:04 -0800 Subject: [PATCH 09/11] Refine 4 Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .../test_dense_blockscaled_gemm_persistent.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py b/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py index 4a8e6bf207da..3419694896e5 100644 --- a/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py +++ b/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py @@ -43,13 +43,11 @@ # This file is copied and modified from cutlass example https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py - import argparse import sys from pathlib import Path from typing import Tuple, Type -import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.torch as cutlass_torch @@ -57,9 +55,12 @@ from cutlass.cute.runtime import from_dlpack try: - from tensorrt_llm._torch.cute_dsl_kernels.blackwell import dense_blockscaled_gemm_persistent as kernel_module + from tensorrt_llm._torch.cute_dsl_kernels.blackwell import \ + dense_blockscaled_gemm_persistent as kernel_module except (ModuleNotFoundError, ImportError): - sys.path.insert(0, str(Path(__file__).parents[3] / "tensorrt_llm/_torch/cute_dsl_kernels")) + sys.path.insert( + 0, + str(Path(__file__).parents[3] / "tensorrt_llm/_torch/cute_dsl_kernels")) from blackwell import dense_blockscaled_gemm_persistent as kernel_module Sm100BlockScaledPersistentDenseGemmKernel = kernel_module.Sm100BlockScaledPersistentDenseGemmKernel @@ -432,7 +433,8 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: "Invalid format. Expected comma-separated integers.") parser = argparse.ArgumentParser( - description="Performance test for Sm100 Dense Persistent BlockScaled GEMM.") + description= + "Performance test for Sm100 Dense Persistent BlockScaled GEMM.") parser.add_argument( "--mnkl", @@ -540,4 +542,3 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: if args.print_duration: print(f"Execution time: {exec_time} us") print("PASS") - From d07c0db15cfded95450ee74556e342320f1d92a0 Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Fri, 28 Nov 2025 00:32:34 -0800 Subject: [PATCH 10/11] Refine 5 Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .../cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py b/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py index 3419694896e5..1298cc26f0f4 100644 --- a/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py +++ b/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py @@ -434,7 +434,8 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: parser = argparse.ArgumentParser( description= - "Performance test for Sm100 Dense Persistent BlockScaled GEMM.") + "Functionality and Performance Test for Sm100 Dense Persistent BlockScaled GEMM." + ) parser.add_argument( "--mnkl", From e78d1f5d5b7e1cc1f6820fda97cdaffcf6504c6d Mon Sep 17 00:00:00 2001 From: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> Date: Sun, 30 Nov 2025 18:23:55 -0800 Subject: [PATCH 11/11] clean up ruff Signed-off-by: Yuhan Li <51736452+liyuhannnnn@users.noreply.github.com> --- .pre-commit-config.yaml | 2 - pyproject.toml | 2 - ... run_dense_blockscaled_gemm_persistent.py} | 208 ++++++++---------- tests/scripts/cute_dsl_kernels/testing.py | 106 +++++---- 4 files changed, 143 insertions(+), 175 deletions(-) rename tests/scripts/cute_dsl_kernels/{test_dense_blockscaled_gemm_persistent.py => run_dense_blockscaled_gemm_persistent.py} (72%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a4745efb50d6..311db068b505 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -943,8 +943,6 @@ common-files: &common_files | tests/scripts/iteration_log_parser.py | tests/scripts/perf-sanity/parse_benchmark_results.py | tests/scripts/perf-sanity/run_benchmark_serve.py | - tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py | - tests/scripts/cute_dsl_kernels/testing.py | tests/unittest/_torch/attention/sparse/test_dsa_indexer.py | tests/unittest/_torch/attention/sparse/test_flash_mla.py | tests/unittest/_torch/attention/sparse/test_rocketkv.py | diff --git a/pyproject.toml b/pyproject.toml index 0606d871ffd3..b453d975e36a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -983,8 +983,6 @@ exclude = [ "tests/scripts/iteration_log_parser.py", "tests/scripts/perf-sanity/parse_benchmark_results.py", "tests/scripts/perf-sanity/run_benchmark_serve.py", - "tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py", - "tests/scripts/cute_dsl_kernels/testing.py", "tests/unittest/_torch/attention/sparse/test_dsa_indexer.py", "tests/unittest/_torch/attention/sparse/test_flash_mla.py", "tests/unittest/_torch/attention/sparse/test_rocketkv.py", diff --git a/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py b/tests/scripts/cute_dsl_kernels/run_dense_blockscaled_gemm_persistent.py similarity index 72% rename from tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py rename to tests/scripts/cute_dsl_kernels/run_dense_blockscaled_gemm_persistent.py index 1298cc26f0f4..81228de69dbd 100644 --- a/tests/scripts/cute_dsl_kernels/test_dense_blockscaled_gemm_persistent.py +++ b/tests/scripts/cute_dsl_kernels/run_dense_blockscaled_gemm_persistent.py @@ -55,12 +55,11 @@ from cutlass.cute.runtime import from_dlpack try: - from tensorrt_llm._torch.cute_dsl_kernels.blackwell import \ - dense_blockscaled_gemm_persistent as kernel_module + from tensorrt_llm._torch.cute_dsl_kernels.blackwell import ( + dense_blockscaled_gemm_persistent as kernel_module, + ) except (ModuleNotFoundError, ImportError): - sys.path.insert( - 0, - str(Path(__file__).parents[3] / "tensorrt_llm/_torch/cute_dsl_kernels")) + sys.path.insert(0, str(Path(__file__).parents[3] / "tensorrt_llm/_torch/cute_dsl_kernels")) from blackwell import dense_blockscaled_gemm_persistent as kernel_module Sm100BlockScaledPersistentDenseGemmKernel = kernel_module.Sm100BlockScaledPersistentDenseGemmKernel @@ -130,16 +129,12 @@ def run( RuntimeError: If no CUDA-capable GPU is available. TypeError: If the configuration is not supported by the kernel. """ - print(f"Running Sm100 Persistent Dense BlockScaled GEMM test with:") + print("Running Sm100 Persistent Dense BlockScaled GEMM test with:") print(f"mnkl: {mnkl}") - print( - f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}" - ) + print(f"AB dtype: {ab_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}, Cluster Shape (M, N): {cluster_shape_mn}" - ) + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") print(f"Use prefetch: {'True' if use_prefetch else 'False'}") print(f"Tolerance: {tolerance}") print(f"Warmup iterations: {warmup_iterations}") @@ -149,26 +144,28 @@ def run( print(f"Use CUPTI: {'True' if use_cupti else 'False'}") # Unpack parameters - m, n, k, l = mnkl + m, n, k, batch = mnkl # Skip unsupported testcase if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement( - ab_dtype, - sf_dtype, - sf_vec_size, - c_dtype, - mma_tiler_mn, - cluster_shape_mn, - m, - n, - k, - l, - a_major, - b_major, - c_major, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + batch, + a_major, + b_major, + c_major, ): raise TypeError( - f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" + f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, " + f"{mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {batch}, " + f"{a_major}, {b_major}, {c_major}" ) if not torch.cuda.is_available(): @@ -177,22 +174,19 @@ def run( torch.manual_seed(1111) # Create tensor A/B/C - a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) - b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) - c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) - - a_tensor, a_torch = cutlass_torch.cute_tensor_like(a_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - b_tensor, b_torch = cutlass_torch.cute_tensor_like(b_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - c_tensor, c_torch = cutlass_torch.cute_tensor_like(c_ref, - c_dtype, - is_dynamic_layout=True, - assumed_align=16) + a_ref = cutlass_torch.matrix(batch, m, k, a_major == "m", cutlass.Float32) + b_ref = cutlass_torch.matrix(batch, n, k, b_major == "n", cutlass.Float32) + c_ref = cutlass_torch.matrix(batch, m, n, c_major == "m", cutlass.Float32) + + a_tensor, a_torch = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) # Mark tensor to be byte aligned a_tensor.mark_compact_shape_dynamic( @@ -212,18 +206,17 @@ def run( ) # Create scale factor tensor SFA/SFB - def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): - + def create_scale_factor_tensor(batch, mn, k, sf_vec_size, dtype): 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_shape = (batch, mn, sf_k) atom_m = (32, 4) atom_k = 4 mma_shape = ( - l, + batch, ceil_div(mn, atom_m[0] * atom_m[1]), ceil_div(sf_k, atom_k), atom_m[0], @@ -266,9 +259,13 @@ def ceil_div(a, b): 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)) + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(batch, mn, sf_k, sf_vec_size) + .reshape(batch, 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, :] @@ -289,10 +286,8 @@ def ceil_div(a, b): ) return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor - sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( - l, m, k, sf_vec_size, sf_dtype) - sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( - l, n, k, sf_vec_size, sf_dtype) + sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor(batch, m, k, sf_vec_size, sf_dtype) + sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor(batch, n, k, sf_vec_size, sf_dtype) # Configure gemm kernel gemm = Sm100BlockScaledPersistentDenseGemmKernel( @@ -305,7 +300,8 @@ def ceil_div(a, b): # Compute max active clusters on current device hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( - cluster_shape_mn[0] * cluster_shape_mn[1]) + cluster_shape_mn[0] * cluster_shape_mn[1] + ) # Initialize Stream current_stream = cutlass_torch.default_stream() @@ -321,14 +317,13 @@ def ceil_div(a, b): 1.0, # alpha max_active_clusters, current_stream, - options=f"--opt-level 2", + options="--opt-level 2", ) # Compute reference result if not skip_ref_check: # Execute kernel once for reference checking - compiled_gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, 1.0, - current_stream) + compiled_gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, 1.0, current_stream) print("Verifying results...") res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) @@ -339,7 +334,8 @@ def ceil_div(a, b): cute.testing.convert( c_tensor, from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( - leading_dim=(1 if c_major == "n" else 0)), + leading_dim=(1 if c_major == "n" else 0) + ), ) c_ref = c_ref_device.cpu() @@ -347,33 +343,28 @@ def ceil_div(a, b): torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): # Convert ref : f32 -> f8 -> f32 - ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, - device="cuda").permute(1, 2, 0) - ref_f8 = from_dlpack( - ref_f8_, assumed_align=16).mark_layout_dynamic(leading_dim=1) + ref_f8_ = torch.empty(*(batch, m, n), dtype=torch.uint8, device="cuda").permute(1, 2, 0) + ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic(leading_dim=1) ref_f8.element_type = c_dtype - ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, - 0).cuda() - ref_tensor = from_dlpack( - ref_device, assumed_align=16).mark_layout_dynamic(leading_dim=1) + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) cute.testing.convert(ref_tensor, ref_f8) cute.testing.convert(ref_f8, ref_tensor) ref = ref_device.cpu() torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) def generate_tensors(): - a_tensor, _ = cutlass_torch.cute_tensor_like(a_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - b_tensor, _ = cutlass_torch.cute_tensor_like(b_ref, - ab_dtype, - is_dynamic_layout=True, - assumed_align=16) - c_tensor, _ = cutlass_torch.cute_tensor_like(c_ref, - c_dtype, - is_dynamic_layout=True, - assumed_align=16) + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, _ = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) # Mark tensor to be byte aligned a_tensor.mark_compact_shape_dynamic( @@ -392,23 +383,24 @@ def generate_tensors(): divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, ) - _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, - sf_dtype) - _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, - sf_dtype) - return cute.testing.JitArguments(a_tensor, b_tensor, sfa_tensor, - sfb_tensor, c_tensor, 1.0, - current_stream) + _, sfa_tensor, _ = create_scale_factor_tensor(batch, m, k, sf_vec_size, sf_dtype) + _, sfb_tensor, _ = create_scale_factor_tensor(batch, n, k, sf_vec_size, sf_dtype) + return cute.testing.JitArguments( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, 1.0, current_stream + ) workspace_count = 1 if use_cold_l2: - one_workspace_bytes = (a_torch.numel() * a_torch.element_size() + - b_torch.numel() * b_torch.element_size() + - sfa_torch.numel() * sfa_torch.element_size() + - sfb_torch.numel() * sfb_torch.element_size() + - c_torch.numel() * c_torch.element_size()) + one_workspace_bytes = ( + a_torch.numel() * a_torch.element_size() + + b_torch.numel() * b_torch.element_size() + + sfa_torch.numel() * sfa_torch.element_size() + + sfb_torch.numel() * sfb_torch.element_size() + + c_torch.numel() * c_torch.element_size() + ) workspace_count = cute.testing.get_workspace_count( - one_workspace_bytes, warmup_iterations, iterations) + one_workspace_bytes, warmup_iterations, iterations + ) exec_time = benchmark( compiled_gemm, @@ -429,12 +421,10 @@ 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.") + raise argparse.ArgumentTypeError("Invalid format. Expected comma-separated integers.") parser = argparse.ArgumentParser( - description= - "Functionality and Performance Test for Sm100 Dense Persistent BlockScaled GEMM." + description="Functionality and Performance Test for Sm100 Dense Persistent BlockScaled GEMM." ) parser.add_argument( @@ -455,16 +445,10 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: default=(1, 4), help="Cluster shape (comma-separated)", ) - parser.add_argument("--ab_dtype", - type=cutlass.dtype, - default=cutlass.Float4E2M1FN) - parser.add_argument("--sf_dtype", - type=cutlass.dtype, - default=cutlass.Float8E8M0FNU) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) parser.add_argument("--sf_vec_size", type=int, default=16) - parser.add_argument("--c_dtype", - type=cutlass.dtype, - default=cutlass.Float16) + parser.add_argument("--c_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("--c_major", choices=["n", "m"], type=str, default="n") @@ -474,23 +458,15 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: default=False, help="Enable TMA prefetch for both A and B matrices (default: False)", ) - parser.add_argument("--tolerance", - type=float, - default=1e-01, - help="Tolerance for validation") - parser.add_argument("--warmup_iterations", - type=int, - default=0, - help="Warmup iterations") + parser.add_argument("--tolerance", type=float, default=1e-01, help="Tolerance for validation") + 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("--skip_ref_check", - action="store_true", - help="Skip reference checking") + parser.add_argument("--skip_ref_check", action="store_true", help="Skip reference checking") parser.add_argument( "--use_cold_l2", action="store_true", diff --git a/tests/scripts/cute_dsl_kernels/testing.py b/tests/scripts/cute_dsl_kernels/testing.py index 14e13df7ef28..55fd37cc36cf 100644 --- a/tests/scripts/cute_dsl_kernels/testing.py +++ b/tests/scripts/cute_dsl_kernels/testing.py @@ -57,8 +57,7 @@ def _buffer_completed(self, activities: list): start = activity.start if hasattr(activity, "start") else "nil" end = activity.end if hasattr(activity, "end") else "nil" duration = end - start if start != "nil" and end != "nil" else "nil" - name = activity.name[:100] if hasattr(activity, - "name") else "unknown" + name = activity.name[:100] if hasattr(activity, "name") else "unknown" # Convert to milliseconds if duration != "nil": self.timings.append((name, duration / 1e6)) @@ -90,8 +89,9 @@ def start(self): self._buffer_requested_callback = self._buffer_requested self._buffer_completed_callback = partial(self._buffer_completed) - cupti.activity_register_callbacks(self._buffer_requested_callback, - self._buffer_completed_callback) + cupti.activity_register_callbacks( + self._buffer_requested_callback, self._buffer_completed_callback + ) self._is_active = True @@ -120,11 +120,8 @@ def get_duration(self) -> float: return sum(timing[1] for timing in self.timings) -def _cuda_success(err: Union[tuple, cuda_runtime.cudaError_t, - cuda_driver.CUresult], message: str): - """ - Helper function to check CUDA API errors. - """ +def _cuda_success(err: Union[tuple, cuda_runtime.cudaError_t, cuda_driver.CUresult], message: str): + """Helper function to check CUDA API errors.""" if isinstance(err, tuple): _cuda_success(err[0], message) elif isinstance(err, cuda_runtime.cudaError_t): @@ -136,16 +133,14 @@ def _cuda_success(err: Union[tuple, cuda_runtime.cudaError_t, error_message = cuda_driver.cuGetErrorString(err)[1].decode("utf-8") raise RuntimeError(f"{message} : {error_message}") else: - raise TypeError( - f"{err} is an unexpected type : it should be a cudaError_t or CUresult" - ) + raise TypeError(f"{err} is an unexpected type : it should be a cudaError_t or CUresult") -def _does_kernel_use_stream(kernel: Callable, stream: cuda_driver.CUstream, - *args, **kwargs): - """ - This function checks if the kernel uses the provided non-default stream. +def _does_kernel_use_stream(kernel: Callable, stream: cuda_driver.CUstream, *args, **kwargs): + """Check if the kernel uses the provided non-default stream. + It does this by capturing the stream and then checking if any kernels were launched. + :param kernel: The kernel to check :type kernel: Callable :param stream: The stream to check @@ -153,13 +148,13 @@ def _does_kernel_use_stream(kernel: Callable, stream: cuda_driver.CUstream, :return: True if the kernel uses the stream, False otherwise :rtype: bool """ - assert int(stream) != int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT), ( - "Stream must be a non-default stream") + "Stream must be a non-default stream" + ) err = cuda_runtime.cudaStreamBeginCapture( - stream, - cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) + stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal + ) _cuda_success(err, "Error on stream capture") kernel(*args, **kwargs) @@ -208,25 +203,31 @@ def user_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda_d from cutlass.cute.testing import benchmark + @cute.jit def user_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor): # contents of the function pass + def workspace_generator(): # create a, b, and c return JitArguments(a, b, c) - time_us = benchmark(user_function, - workspace_generator=workspace_generator, - workspace_count=10, - warmup_iterations=10000, - iterations=1000) + + time_us = benchmark( + user_function, + workspace_generator=workspace_generator, + workspace_count=10, + warmup_iterations=10000, + iterations=1000, + ) To benchmark you may always configure the function being profiled (callable), the warmup iterations, and the number of profiling iterations. - Whenever the kernel being benchmarked runs in a non-default stream, the stream must be provided through the stream parameter. + Whenever the kernel being benchmarked runs in a non-default stream, the stream must be provided + through the stream parameter. To use CUDA graphs, the callable must be a compiled @cute.jit annotated function. When using CUDA graphs, the kernel must be launched in a non-default stream. @@ -243,7 +244,8 @@ def workspace_generator(): :type kernel_arguments: JitArguments, None :param workspace_generator: Function that returns kernel arguments, defaults to None :type workspace_generator: Callable - :param workspace_count: Number of workspaces (arguments) to loop through, looping through enough workspaces will keep the L2 cache cold + :param workspace_count: Number of workspaces (arguments) to loop through, looping through enough + workspaces will keep the L2 cache cold :type workspace_count: int, optional :param use_cuda_graphs: Whether to use cuda graphs, defaults to False :type use_cuda_graphs: bool, optional @@ -251,31 +253,29 @@ def workspace_generator(): :return: The benchmark time in microseconds :rtype: float """ - if stream is None: - stream = cuda_driver.CUstream( - cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) + stream = cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) if workspace_count < 1: raise ValueError("workspace_count must be at least 1") float("nan") - if workspace_generator == None: + if workspace_generator is None: # If no workspace generator is provided, we need a single workspace if workspace_count != 1: - raise ValueError( - "Need a single workspace if not providing a generator") + raise ValueError("Need a single workspace if not providing a generator") # If no workspace generator is provided, we need a kernel_argument - if kernel_arguments == None: - raise ValueError( - "Please pass a kernel argument if not providing a generator") - workspace_generator = lambda: kernel_arguments + if kernel_arguments is None: + raise ValueError("Please pass a kernel argument if not providing a generator") + + def workspace_generator(): + return kernel_arguments workspaces = [workspace_generator() for _ in range(workspace_count)] for workspace in workspaces: - if type(workspace) != JitArguments: + if not isinstance(workspace, JitArguments): raise TypeError( "workspace_generator and/or kernel_arguments should use JitArguments type" ) @@ -288,11 +288,9 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): return workspace_index # Create CUDA events for timing - err, start_event = cuda_driver.cuEventCreate( - cuda_driver.CUevent_flags.CU_EVENT_DEFAULT) + err, start_event = cuda_driver.cuEventCreate(cuda_driver.CUevent_flags.CU_EVENT_DEFAULT) _cuda_success(err, "Error on creating event") - err, end_event = cuda_driver.cuEventCreate( - cuda_driver.CUevent_flags.CU_EVENT_DEFAULT) + err, end_event = cuda_driver.cuEventCreate(cuda_driver.CUevent_flags.CU_EVENT_DEFAULT) _cuda_success(err, "Error on creating event") elapsed_time = float("nan") @@ -305,8 +303,7 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): cutlass.base_dsl.jit_executor.JitExecutor, ) if not isinstance(callable, compiled_types): - raise TypeError( - "Function must be precompiled to be used with CUDA Graphs") + raise TypeError("Function must be precompiled to be used with CUDA Graphs") # Check if the stream is a non-default stream if int(stream) == int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT): @@ -318,8 +315,8 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): # Capture warmup graph err = cuda_runtime.cudaStreamBeginCapture( - stream, - cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) + stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal + ) _cuda_success(err, "Error on stream capture") workspace_index = _loop_and_call_kernel(warmup_iterations) @@ -337,8 +334,8 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): # Capture profiling graph err = cuda_runtime.cudaStreamBeginCapture( - stream, - cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) + stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal + ) _cuda_success(err, "Error on stream capture") _loop_and_call_kernel(iterations, workspace_index) err, gprofile = cuda_runtime.cudaStreamEndCapture(stream) @@ -369,8 +366,7 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): _cuda_success(err, "Error on synchronizing event") # Get elapsed time - err, elapsed_time = cuda_driver.cuEventElapsedTime( - start_event, end_event) + err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event) _cuda_success(err, "Error on querying event") # Destroy graphs @@ -398,9 +394,10 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): else: if int(stream) != int( - cuda_driver.CUstream_flags.CU_STREAM_DEFAULT - ) and not _does_kernel_use_stream(callable, stream, *workspaces[0].args, - **workspaces[0].kwargs): + cuda_driver.CUstream_flags.CU_STREAM_DEFAULT + ) and not _does_kernel_use_stream( + callable, stream, *workspaces[0].args, **workspaces[0].kwargs + ): raise ValueError( "CUDA stream passed to benchmark does not match the stream the kernel was launched in" ) @@ -418,8 +415,7 @@ def _loop_and_call_kernel(iterations: int, workspace_index: int = 0): # Synchronize end event err = cuda_driver.cuEventSynchronize(end_event) _cuda_success(err, "Error on synchronizing event") - err, elapsed_time = cuda_driver.cuEventElapsedTime( - start_event, end_event) + err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event) _cuda_success(err, "Error on querying event") # Destroy events