From 94b022f91fa786609d3a7b7752444d150bbe35a9 Mon Sep 17 00:00:00 2001 From: Yanqin Zhai Date: Fri, 14 Aug 2026 18:53:28 -0700 Subject: [PATCH] 1 --- .../frost/kernel_templates/_tile_helpers.py | 44 ++++++++++++ .../sm100_block_scale_matmul_1ctamma.py | 9 ++- ...sm100_block_scale_matmul_1ctamma_static.py | 9 ++- .../sm100_block_scale_matmul_2ctamma.py | 11 +-- ...sm100_block_scale_matmul_2ctamma_static.py | 11 +-- .../kernel_templates/sm100_matmul_1ctamma.py | 6 +- .../sm100_matmul_1ctamma_static.py | 6 +- .../kernel_templates/sm100_matmul_2ctamma.py | 8 ++- .../sm100_matmul_2ctamma_static.py | 8 ++- .../sm100_matmul_mainloop_1ctamma.py | 6 +- .../sm100_matmul_mainloop_2ctamma.py | 8 ++- ..._grouped_block_scale_matmul_fwd_1ctamma.py | 9 ++- ..._grouped_block_scale_matmul_fwd_2ctamma.py | 11 +-- .../sm100_moe_grouped_matmul_fwd_1ctamma.py | 6 +- .../sm100_moe_grouped_matmul_fwd_2ctamma.py | 8 ++- .../sm103_block_scale_matmul_1ctamma.py | 11 +-- .../sm103_block_scale_matmul_2ctamma.py | 13 ++-- .../sm107_block_scale_matmul_1ctamma.py | 9 ++- .../sm107_block_scale_matmul_2ctamma.py | 11 +-- ..._grouped_block_scale_matmul_fwd_1ctamma.py | 9 ++- ..._grouped_block_scale_matmul_fwd_2ctamma.py | 11 +-- .../gemm/frost/test_block_scale_matmul.py | 32 ++++++--- test/python/gemm/frost/test_matmul.py | 68 +++++++++++++++++++ 23 files changed, 249 insertions(+), 75 deletions(-) diff --git a/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py b/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py index bffc67a9f..4496c7b26 100644 --- a/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py +++ b/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py @@ -129,3 +129,47 @@ def copy_tensormap_to_workspace(src_desc_ptr, dst_i64_ptr) -> None: src_words = cute.make_ptr(cutlass.Int64, src_desc_ptr.toint(), mem_space=cute.AddressSpace.generic) for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): dst_i64_ptr.subview(i).store((src_words + i).load()) + + +def tcgen05_alloc(tmem_ptr, num_cols, *, is_exclusive=False, group=None): + if is_exclusive: + nvvm.tcgen05_alloc(tmem_ptr, num_cols, is_exclusive=True, group=group) + else: + nvvm.tcgen05_alloc(tmem_ptr, num_cols, group=group) + + +def tcgen05_dealloc(tmem_ptr, num_cols, *, is_exclusive=False, group=None): + if is_exclusive: + nvvm.tcgen05_dealloc(tmem_ptr, num_cols, is_exclusive=True, group=group) + else: + nvvm.tcgen05_dealloc(tmem_ptr, num_cols, group=group) + + +def tcgen05_mma_block_scale(mma_kind, cta_group, d, a, b, idesc, *, enable_input_d, scale_a, scale_b, scale_vec_size, b_collector_op=None): + if b_collector_op is None: + nvvm.tcgen05_mma_block_scale( + mma_kind, + cta_group, + d, + a, + b, + idesc, + enable_input_d=enable_input_d, + scale_a=scale_a, + scale_b=scale_b, + scale_vec_size=scale_vec_size, + ) + else: + nvvm.tcgen05_mma_block_scale( + mma_kind, + cta_group, + d, + a, + b, + idesc, + enable_input_d=enable_input_d, + scale_a=scale_a, + scale_b=scale_b, + scale_vec_size=scale_vec_size, + b_collector_op=b_collector_op, + ) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py index bfbd1378c..ee7c5c44a 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py @@ -25,6 +25,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -608,7 +611,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -762,7 +765,7 @@ def _kernel( # shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, acc_tmem_ptrs[g][mi], @@ -831,7 +834,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py index 15a125691..640cbdfc8 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py @@ -27,6 +27,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -513,7 +516,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -668,7 +671,7 @@ def _kernel( # shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, acc_tmem_ptrs[g][mi], @@ -721,7 +724,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py index 4314cb39b..1f33dd1c0 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py @@ -24,6 +24,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -635,7 +638,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -797,7 +800,7 @@ def _kernel( # are shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, acc_tmem_ptrs[g][mi], @@ -874,7 +877,7 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -914,7 +917,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py index 5906c2117..8d9b50a5b 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py @@ -25,6 +25,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -535,7 +538,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -697,7 +700,7 @@ def _kernel( # are shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, acc_tmem_ptrs[g][mi], @@ -759,7 +762,7 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -777,7 +780,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py index 2dc69605a..835373875 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py @@ -28,6 +28,8 @@ from cudnn.gemm.frost.kernel_templates._tile_helpers import ( epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -543,7 +545,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -681,7 +683,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py index 8f018f03d..d32b5c6ae 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py @@ -26,6 +26,8 @@ from cudnn.gemm.frost.kernel_templates._tile_helpers import ( epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -440,7 +442,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -559,7 +561,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py index d69dd8256..65ad6b0eb 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py @@ -25,6 +25,8 @@ from cudnn.gemm.frost.kernel_templates._tile_helpers import ( epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -563,7 +565,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -717,7 +719,7 @@ def _kernel( pass nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -756,7 +758,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py index 8c2c2e7c9..987cd5181 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py @@ -25,6 +25,8 @@ from cudnn.gemm.frost.kernel_templates._tile_helpers import ( epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -461,7 +463,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -599,7 +601,7 @@ def _kernel( pass nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -616,7 +618,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py index c57d96fcb..96f1dc1c5 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py @@ -26,6 +26,8 @@ from cudnn.gemm.frost.kernel_templates._tile_helpers import ( epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -556,7 +558,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -707,7 +709,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py index 5270ee597..f89a9f569 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py @@ -27,6 +27,8 @@ from cudnn.gemm.frost.kernel_templates._tile_helpers import ( epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -652,7 +654,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -815,7 +817,7 @@ def _kernel( pass nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -854,7 +856,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py index 404b91ffa..c59ccd773 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py @@ -26,6 +26,9 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, TENSOR_MAP_QWORDS, ) import cutlass.experimental.cuda.tensor_map as _tma @@ -657,7 +660,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -833,7 +836,7 @@ def _kernel( # are shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, acc_tmem_ptrs[g][mi], @@ -891,7 +894,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py index e7af83c91..cd3074233 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py @@ -31,6 +31,9 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, TENSOR_MAP_QWORDS, ) import cutlass.experimental.cuda.tensor_map as _tma @@ -686,7 +689,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -862,7 +865,7 @@ def _kernel( # are shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, acc_tmem_ptrs[g][mi], @@ -923,7 +926,7 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -959,7 +962,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py index 3ae4d6bc9..9208d47f1 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py @@ -33,6 +33,8 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, TENSOR_MAP_QWORDS, ) import cutlass.experimental.cuda.tensor_map as _tma @@ -516,7 +518,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -662,7 +664,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py index cc64ecc4a..6d858a125 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py @@ -34,6 +34,8 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, TENSOR_MAP_QWORDS, ) import cutlass.experimental.cuda.tensor_map as _tma @@ -546,7 +548,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -699,7 +701,7 @@ def _kernel( pass nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -734,7 +736,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py index 65402c61f..22dc5a5ae 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py @@ -46,6 +46,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -694,7 +697,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(mma_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -851,7 +854,7 @@ def _kernel( # the same stride; both read one name so they cannot drift. desc_a = _sm103_make_circular_mma_desc(desc_a_circ[_ai][mi][_kc], _ph, desc_a_next[_ai][mi][_kn]) if is_mma_leader: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, acc_tmem_ptrs[g][mi], @@ -962,7 +965,7 @@ def _kernel( # the same stride; both read one name so they cannot drift. desc_a = _sm103_make_circular_mma_desc(desc_a_circ[_ai][mi][_kc], _ph, desc_a_next[_ai][mi][_kn]) if is_mma_leader: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, acc_tmem_ptrs[g][mi], @@ -1032,7 +1035,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py index 1974b5e11..03ae73089 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py @@ -46,6 +46,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -713,7 +716,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(mma_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -883,7 +886,7 @@ def _kernel( # destination uses the same name so they cannot drift. desc_a = _sm103_make_circular_mma_desc(desc_a_circ[_ai][mi][_kc], _ph, desc_a_next[_ai][mi][_kn]) if is_mma_leader: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, acc_tmem_ptrs[g][mi], @@ -996,7 +999,7 @@ def _kernel( # destination uses the same name so they cannot drift. desc_a = _sm103_make_circular_mma_desc(desc_a_circ[_ai][mi][_kc], _ph, desc_a_next[_ai][mi][_kn]) if is_mma_leader: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, acc_tmem_ptrs[g][mi], @@ -1073,7 +1076,7 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -1113,7 +1116,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py index 4569c960d..7a58499fe 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py @@ -43,6 +43,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -625,7 +628,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -818,7 +821,7 @@ def _kernel( # shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, acc_tmem_ptrs[g][mi], @@ -887,7 +890,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py index 873f5c19c..4c20d844b 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py @@ -42,6 +42,9 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( l2_swizzle_tile as _l2_swizzle_tile, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, ) import cutlass.experimental.cuda.tensor_map as _tma import cutlass._mlir_helpers.vector as _cvec @@ -652,7 +655,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -854,7 +857,7 @@ def _kernel( # are shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, acc_tmem_ptrs[g][mi], @@ -931,7 +934,7 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -971,7 +974,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py index 8f9bcbfbf..33ee8800f 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py @@ -45,6 +45,9 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, TENSOR_MAP_QWORDS, ) import cutlass.experimental.cuda.tensor_map as _tma @@ -675,7 +678,7 @@ def _kernel( if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -890,7 +893,7 @@ def _kernel( # are shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, acc_tmem_ptrs[g][mi], @@ -948,7 +951,7 @@ def _kernel( nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py index a7e34b247..b91e571cf 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py @@ -50,6 +50,9 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + tcgen05_alloc as _tcgen05_alloc, + tcgen05_dealloc as _tcgen05_dealloc, + tcgen05_mma_block_scale as _tcgen05_mma_block_scale, TENSOR_MAP_QWORDS, ) import cutlass.experimental.cuda.tensor_map as _tma @@ -704,7 +707,7 @@ def _kernel( ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) - nvvm.tcgen05_alloc( + _tcgen05_alloc( tmem_ptr_i32, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -919,7 +922,7 @@ def _kernel( # are shared; A's SF word block follows the M block. desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) if elect_one: - nvvm.tcgen05_mma_block_scale( + _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, acc_tmem_ptrs[g][mi], @@ -980,7 +983,7 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.mbarrier_arrive(peer_mbar, scope=nvvm.MemScope.CLUSTER, relaxed=True) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, @@ -1016,7 +1019,7 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) - nvvm.tcgen05_dealloc( + _tcgen05_dealloc( alloc_ptr, cutlass.Int32(num_tmem_alloc_cols), is_exclusive=tmem_alloc_exclusive, diff --git a/test/python/gemm/frost/test_block_scale_matmul.py b/test/python/gemm/frost/test_block_scale_matmul.py index f514e00f1..3ba949fd3 100644 --- a/test/python/gemm/frost/test_block_scale_matmul.py +++ b/test/python/gemm/frost/test_block_scale_matmul.py @@ -1908,10 +1908,14 @@ def test_sm107_templates_reject_older_blackwell(monkeypatch): ("mxfp8", 128, False, 1, 2, 1, "cutlass.Float8E4M3FN"), ], ) -def test_render_sm107_tile_constants(combo, cta_n, omma, k_mode, scales_per_inst, word_atoms, idesc_dtype): +def test_render_sm107_tile_constants(_pretend_sm107, combo, cta_n, omma, k_mode, scales_per_inst, word_atoms, idesc_dtype): """One MMA spans a 64-byte K, so it eats twice sm100's scales; when that outgrows a 4-scale utccp atom the SF word spans word_atoms of them. fp4 - rides the OMMA descriptor, mxfp8 the MX one; both take the real dtype.""" + rides the OMMA descriptor, mxfp8 the MX one; both take the real dtype. + + The render sizes TMEM from the LIVE arch, and the cta_n=256 tile's SFB span + needs 520 of SM 10.7's 576 columns — so this has to pretend, or it renders + against a 512-column part and raises.""" chain = analyze(_bs_chain(combo)) cfg = by_name(f"CONFIG_sm107_128x{cta_n}x128_128x{cta_n}x64_cluster1x1") txt = C._render_block_scale_tile_constants(cfg, chain, 1) @@ -2194,13 +2198,26 @@ def test_fp4_scale_dtype_and_block_are_orthogonal(sf_dt, sf_name, block_size): assert (bs.a_dtype, bs.sf_dtype, bs.block_size) == ("fp4_e2m1", sf_name, block_size) +# Block-scale cases that only some GPUs decode: SM 10.7 added the E5M3 scale +# format (either K-block) and E4M3 at block 32. Keyed by (SF dtype, K-block). +_GPU_GATED_FP4_CASES = {("fp8_e5m3", 16), ("fp8_e5m3", 32), ("fp8_e4m3", 32)} +_DTYPE_GATED_SF_DTYPES = {"fp8_e5m3"} +_GPU_GATED_RANGES = ((107, 110),) + + @_GPU @pytest.mark.parametrize("sf_dt,sf_name", [(cudnn.data_type.FP8_E4M3, "fp8_e4m3"), (cudnn.data_type.FP8_E8M0, "fp8_e8m0")]) @pytest.mark.parametrize("block_size", [16, 32]) -@pytest.mark.parametrize("config_name", ["CONFIG_sm100_128x128x128_128x128x32_cluster1x1", _SM107_128]) +@pytest.mark.parametrize("config_name", ["CONFIG_sm100_128x128x128_128x128x32_cluster1x1", pytest.param(_SM107_128, marks=requires_sm107)]) def test_fp4_all_scale_block_corners_numerics(config_name, sf_dt, sf_name, block_size): """Numerics for the whole non-E5M3 fp4 matrix, including the two corners the - nvfp4 / mxfp4 pair leaves out: e4m3 at block 32 and e8m0 at block 16.""" + nvfp4 / mxfp4 pair leaves out: e4m3 at block 32 and e8m0 at block 16. + + e4m3 at block 32 is one of the GPU-gated cases — it is a 10.7 addition on + EVERY pipeline, so it runs here only on a 10.7 part.""" + if (sf_name, block_size) in _GPU_GATED_FP4_CASES and not any(lo <= _SM < hi for lo, hi in _GPU_GATED_RANGES): + spans = " or ".join(f"{lo} <= SM < {hi}" for lo, hi in _GPU_GATED_RANGES) + pytest.skip(f"fp4+{sf_name} at block {block_size} decodes only on {spans}, have sm_{_SM}") dev = "cuda" torch.manual_seed(0) M, N, K = 256, 256, 512 @@ -2231,13 +2248,6 @@ def test_fp4_all_scale_block_corners_numerics(config_name, sf_dt, sf_name, block torch.testing.assert_close(c[0], (a_s @ b_s.t()).to(torch.float16), atol=2e-1, rtol=2e-2) -# Block-scale cases that only some GPUs decode: SM 10.7 added the E5M3 scale -# format (either K-block) and E4M3 at block 32. Keyed by (SF dtype, K-block). -_GPU_GATED_FP4_CASES = {("fp8_e5m3", 16), ("fp8_e5m3", 32), ("fp8_e4m3", 32)} -_DTYPE_GATED_SF_DTYPES = {"fp8_e5m3"} -_GPU_GATED_RANGES = ((107, 110),) - - def test_gpu_gated_cases_are_narrowed_everywhere(): """The load-bearing invariant behind putting the GPU-gated fp4 cases in the ordinary case sets: EVERY one of them, on EVERY pipeline that carries it, diff --git a/test/python/gemm/frost/test_matmul.py b/test/python/gemm/frost/test_matmul.py index 71f5668f9..3c4f9c622 100644 --- a/test/python/gemm/frost/test_matmul.py +++ b/test/python/gemm/frost/test_matmul.py @@ -13,7 +13,9 @@ from __future__ import annotations import os +import pathlib import sys +import textwrap import pytest import torch @@ -2402,3 +2404,69 @@ def test_tma_store_gate_follows_the_mma_block_height(name: str) -> None: g.matmul(A=A, B=B, name="mm").set_output(True) cfg = by_name(name) assert _use_tma_store_epi(analyze(g), cfg, 16, 1) == (cfg.mma_inst_m == 128) + + +# --- cutlass-dsl version-gated kwargs ---------------------------------------- + +# `is_exclusive` (the >512-column TMEM grant) and `b_collector_op` (B-operand +# collector reuse) only reached the cutlass-dsl `nvvm.*` wrappers in 4.8. Those +# wrappers take no **kwargs, so NAMING one on an older DSL is a TypeError at JIT +# regardless of the value -- `is_exclusive=False` is just as fatal as True. + +_VERSION_GATED_KWARGS = { + "tcgen05_alloc": "is_exclusive", + "tcgen05_dealloc": "is_exclusive", + "tcgen05_mma_block_scale": "b_collector_op", +} + + +def _template_dir(): + # kernel_templates has no __init__.py (it is exec'd per render), so go + # through the package that does. + return pathlib.Path(cudnn.gemm.frost.__file__).parent / "kernel_templates" + + +def test_templates_route_version_gated_kwargs_through_the_guarded_wrappers(): + """Every template must reach these ops through `_tile_helpers`, which emits + the kwarg only on the branch that wants it. Calling `nvvm.` directly and + passing the inert False/None compiles on an internal wheel and fails every + single JIT on the public one -- which is how the whole gemm suite went red.""" + import ast + + offenders = [] + for path in sorted(_template_dir().glob("sm*.py")): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "nvvm" + and node.func.attr in _VERSION_GATED_KWARGS + ): + offenders.append(f"{path.name}:{node.lineno} nvvm.{node.func.attr}(...)") + assert not offenders, "call the _tile_helpers wrapper instead of nvvm directly:\n " + "\n ".join(offenders) + + +def test_the_guarded_wrappers_keep_the_kwarg_off_the_default_branch(): + """...and the wrappers themselves only name it under the flag. Pinned as + source structure because the failure mode is a TypeError at trace time on a + DSL we cannot install here, so no runtime assertion can see it.""" + import ast + import inspect + + import cudnn.gemm.frost.kernel_templates._tile_helpers as helpers + + for fn_name, kwarg in _VERSION_GATED_KWARGS.items(): + fn = getattr(helpers, fn_name) + tree = ast.parse(textwrap.dedent(inspect.getsource(fn))) + calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and isinstance(n.func.value, ast.Name) and n.func.value.id == "nvvm" + ] + assert len(calls) == 2, f"{fn_name} should have exactly one guarded and one plain nvvm call, got {len(calls)}" + named = [c for c in calls if any(k.arg == kwarg for k in c.keywords)] + assert len(named) == 1, f"{fn_name}: exactly one branch may name {kwarg!r}, got {len(named)}" + # ...and the other branch must be reachable without the newer DSL. + assert len(calls) - len(named) == 1, f"{fn_name}: no branch left that omits {kwarg!r}"