diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py index 0204d18f..e8d6206a 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py @@ -975,7 +975,7 @@ class BlockwiseGemmKernel: # Specialized Schedule warp # if warp_idx == self.sched_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_sched_warps) + cute.arch.setmaxregister_decrease(self.num_regs_sched_warps) # # Persistent tile scheduling loop # @@ -1008,10 +1008,7 @@ class BlockwiseGemmKernel: ) # fence view async shared - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1023,7 +1020,7 @@ class BlockwiseGemmKernel: # Specialized TMA load warp # if warp_idx == self.tma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Persistent tile scheduling loop # @@ -1126,10 +1123,7 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1142,7 +1136,7 @@ class BlockwiseGemmKernel: # Specialized Scale load warp # if warp_idx == self.scale_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Persistent tile scheduling loop # @@ -1301,10 +1295,7 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1317,7 +1308,7 @@ class BlockwiseGemmKernel: # Specialized MMA warp # if warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Bar sync for retrieve tensor memory ptr from shared mem # @@ -1459,10 +1450,7 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1475,7 +1463,7 @@ class BlockwiseGemmKernel: # Specialized acc update warps # if warp_idx <= self.acc_update_warp_id[-1]: - cute.arch.warpgroup_reg_alloc(self.num_regs_acc_update_warps) + cute.arch.setmaxregister_increase(self.num_regs_acc_update_warps) # # Bar sync for retrieve tensor memory ptr from shared memory # @@ -1696,10 +1684,7 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1707,7 +1692,7 @@ class BlockwiseGemmKernel: # Specialized epilogue warps # if warp_idx <= self.epilog_warp_id[-1] and warp_idx >= self.epilog_warp_id[0]: - cute.arch.warpgroup_reg_alloc(self.num_regs_epilogue_warps) + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) # # Alloc tensor memory buffer # @@ -1866,10 +1851,7 @@ class BlockwiseGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # @@ -1899,10 +1881,7 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py index ec61d4ef..233dda4a 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py @@ -996,7 +996,7 @@ class BlockwiseContiguousGroupedGemmKernel: # Specialized Schedule warp # if warp_idx == self.sched_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_sched_warps) + cute.arch.setmaxregister_decrease(self.num_regs_sched_warps) # # Persistent tile scheduling loop # @@ -1034,10 +1034,7 @@ class BlockwiseContiguousGroupedGemmKernel: ) # fence view async shared - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1051,7 +1048,7 @@ class BlockwiseContiguousGroupedGemmKernel: # Specialized TMA load warp # if warp_idx == self.tma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Persistent tile scheduling loop # @@ -1153,10 +1150,7 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1168,7 +1162,7 @@ class BlockwiseContiguousGroupedGemmKernel: # Specialized Scale load warp # if warp_idx == self.scale_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Persistent tile scheduling loop # @@ -1328,10 +1322,7 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1344,7 +1335,7 @@ class BlockwiseContiguousGroupedGemmKernel: # Specialized MMA warp # if warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Bar sync for retrieve tensor memory ptr from shared mem # @@ -1488,10 +1479,7 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1504,7 +1492,7 @@ class BlockwiseContiguousGroupedGemmKernel: # Specialized acc update warps # if warp_idx <= self.acc_update_warp_id[-1]: - cute.arch.warpgroup_reg_alloc(self.num_regs_acc_update_warps) + cute.arch.setmaxregister_increase(self.num_regs_acc_update_warps) # # Bar sync for retrieve tensor memory ptr from shared memory # @@ -1727,10 +1715,7 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1738,7 +1723,7 @@ class BlockwiseContiguousGroupedGemmKernel: # Specialized epilogue warps # if warp_idx <= self.epilog_warp_id[-1] and warp_idx >= self.epilog_warp_id[0]: - cute.arch.warpgroup_reg_alloc(self.num_regs_epilogue_warps) + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) # # Alloc tensor memory buffer # @@ -1899,10 +1884,7 @@ class BlockwiseContiguousGroupedGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # @@ -1932,10 +1914,7 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py index 0cf8e988..a0b655ae 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py @@ -995,7 +995,7 @@ class BlockwiseMaskedGroupedGemmKernel: # Specialized Schedule warp # if warp_idx == self.sched_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_sched_warps) + cute.arch.setmaxregister_decrease(self.num_regs_sched_warps) # # Persistent tile scheduling loop # @@ -1041,10 +1041,7 @@ class BlockwiseMaskedGroupedGemmKernel: ) # fence view async shared - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1056,7 +1053,7 @@ class BlockwiseMaskedGroupedGemmKernel: # Specialized TMA load warp # if warp_idx == self.tma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Persistent tile scheduling loop # @@ -1159,10 +1156,7 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1175,7 +1169,7 @@ class BlockwiseMaskedGroupedGemmKernel: # Specialized Scale load warp # if warp_idx == self.scale_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Persistent tile scheduling loop # @@ -1334,10 +1328,7 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1350,7 +1341,7 @@ class BlockwiseMaskedGroupedGemmKernel: # Specialized MMA warp # if warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # # Bar sync for retrieve tensor memory ptr from shared mem # @@ -1492,10 +1483,7 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1508,7 +1496,7 @@ class BlockwiseMaskedGroupedGemmKernel: # Specialized acc update warps # if warp_idx <= self.acc_update_warp_id[-1]: - cute.arch.warpgroup_reg_alloc(self.num_regs_acc_update_warps) + cute.arch.setmaxregister_increase(self.num_regs_acc_update_warps) # # Bar sync for retrieve tensor memory ptr from shared memory # @@ -1729,10 +1717,7 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1740,7 +1725,7 @@ class BlockwiseMaskedGroupedGemmKernel: # Specialized epilogue warps # if warp_idx <= self.epilog_warp_id[-1] and warp_idx >= self.epilog_warp_id[0]: - cute.arch.warpgroup_reg_alloc(self.num_regs_epilogue_warps) + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) # # Alloc tensor memory buffer # @@ -1899,10 +1884,7 @@ class BlockwiseMaskedGroupedGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # @@ -1932,10 +1914,7 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() diff --git a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py index 19f56213..a98b739a 100644 --- a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py @@ -1453,10 +1453,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_prefetch.py b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_prefetch.py index 7193d7e6..009e6c7b 100644 --- a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_prefetch.py +++ b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_prefetch.py @@ -1552,10 +1552,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm.py b/examples/python/CuTeDSL/blackwell/dense_gemm.py index 01aa0078..1e347bd1 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm.py @@ -1030,10 +1030,7 @@ class DenseGemmKernel: c_buffer = subtile_idx % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") pipeline.sync(barrier_id=1) # TMA store C to global memory diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py index 61bcf01f..2b7d7d84 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py @@ -1176,10 +1176,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel: tSR_sC[(None, None, None, c_pipeline_consumer_state.index)], tSR_rC, ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") c_pipeline.consumer_release(c_pipeline_consumer_state) # Advance pipeline states @@ -1206,10 +1203,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel: tiled_copy_r2s, tRS_rD, tRS_sD[(None, None, None, d_buffer)] ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilog_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py index df393021..e54b1412 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py @@ -1125,10 +1125,7 @@ class PersistentDenseGemmKernel: c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilog_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_prefetch.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_prefetch.py index 604c7c6c..1892a842 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_prefetch.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_prefetch.py @@ -1208,10 +1208,7 @@ class PersistentDenseGemmKernel: c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilog_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py index 228ec0e4..3952cd46 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py @@ -983,10 +983,7 @@ class DenseGemmKernel: c_buffer = subtile_idx % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") pipeline.sync(barrier_id=1) # TMA store C to global memory diff --git a/examples/python/CuTeDSL/blackwell/fmha.py b/examples/python/CuTeDSL/blackwell/fmha.py index 8fb65d8a..f52d9538 100644 --- a/examples/python/CuTeDSL/blackwell/fmha.py +++ b/examples/python/CuTeDSL/blackwell/fmha.py @@ -802,13 +802,13 @@ class BlackwellFusedMultiHeadAttentionForward: # EMPTY # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.empty_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + cute.arch.setmaxregister_decrease(self.num_regs_other) # /////////////////////////////////////////////////////////////////////////////// # LOAD # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.load_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + cute.arch.setmaxregister_decrease(self.num_regs_other) tile_sched = fmha_utils.create_fmha_static_tile_scheduler( tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() @@ -997,7 +997,7 @@ class BlackwellFusedMultiHeadAttentionForward: # MMA # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + cute.arch.setmaxregister_decrease(self.num_regs_other) # Alloc tmem buffer tmem_alloc_cols = Int32(self.tmem_alloc_cols) @@ -1271,7 +1271,7 @@ class BlackwellFusedMultiHeadAttentionForward: # Epilogue # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.epilogue_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + cute.arch.setmaxregister_decrease(self.num_regs_other) tile_sched = fmha_utils.create_fmha_static_tile_scheduler( tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() ) @@ -1354,7 +1354,7 @@ class BlackwellFusedMultiHeadAttentionForward: # /////////////////////////////////////////////////////////////////////////////// if warp_idx < self.softmax1_warp_ids[0]: # increase register after decreasing - cute.arch.warpgroup_reg_alloc(self.num_regs_softmax) + cute.arch.setmaxregister_increase(self.num_regs_softmax) self.softmax( stage=0, @@ -1384,7 +1384,7 @@ class BlackwellFusedMultiHeadAttentionForward: and warp_idx >= self.softmax1_warp_ids[0] ): # increase register after decreasing - cute.arch.warpgroup_reg_alloc(self.num_regs_softmax) + cute.arch.setmaxregister_increase(self.num_regs_softmax) self.softmax( stage=1, @@ -1410,7 +1410,7 @@ class BlackwellFusedMultiHeadAttentionForward: # Correction # /////////////////////////////////////////////////////////////////////////////// if warp_idx >= self.correction_warp_ids[0] and warp_idx < self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_correction) + cute.arch.setmaxregister_decrease(self.num_regs_correction) cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) tScS = qk_thr_mma.partition_C(cS) @@ -2327,10 +2327,7 @@ class BlackwellFusedMultiHeadAttentionForward: mLSE[row_idx + cuseqlen_q, blk_coord[2]] = lse # fence view async shared - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") def run( diff --git a/examples/python/CuTeDSL/blackwell/fmha_bwd.py b/examples/python/CuTeDSL/blackwell/fmha_bwd.py index 70b912ec..f40af955 100644 --- a/examples/python/CuTeDSL/blackwell/fmha_bwd.py +++ b/examples/python/CuTeDSL/blackwell/fmha_bwd.py @@ -1043,7 +1043,7 @@ class BlackwellFusedMultiHeadAttentionBackward: # LOAD # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.load_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_load) + cute.arch.setmaxregister_decrease(self.num_regs_load) self.load( K_in, @@ -1081,7 +1081,7 @@ class BlackwellFusedMultiHeadAttentionBackward: # MMA # /////////////////////////////////////////////////////////////////////////////// elif warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_mma) + cute.arch.setmaxregister_decrease(self.num_regs_mma) self.mma( KQ_tiled_mma, @@ -1125,7 +1125,7 @@ class BlackwellFusedMultiHeadAttentionBackward: warp_idx >= self.compute_warp_id[0] and warp_idx <= self.compute_warp_id[-1] ): - cute.arch.warpgroup_reg_alloc(self.num_regs_compute) + cute.arch.setmaxregister_increase(self.num_regs_compute) self.compute( tSTtST, @@ -1176,7 +1176,7 @@ class BlackwellFusedMultiHeadAttentionBackward: warp_idx >= self.reduce_warp_id[0] and warp_idx <= self.reduce_warp_id[-1] ): - cute.arch.warpgroup_reg_alloc(self.num_regs_reduce) + cute.arch.setmaxregister_increase(self.num_regs_reduce) self.reduce( dSK_tiled_mma, @@ -1193,7 +1193,7 @@ class BlackwellFusedMultiHeadAttentionBackward: ) else: - cute.arch.warpgroup_reg_dealloc(self.num_regs_empty) + cute.arch.setmaxregister_decrease(self.num_regs_empty) @cute.kernel def convert( @@ -2161,10 +2161,7 @@ class BlackwellFusedMultiHeadAttentionBackward: cute.autovec_copy(tTR_rdST, sdS_slice) # Notify for dS - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") compute_mma_dS_pipeline.producer_commit(compute_mma_dS_producer_state) compute_mma_dS_producer_state.advance() @@ -2282,10 +2279,7 @@ class BlackwellFusedMultiHeadAttentionBackward: ) # Wait for the stores to all be visible to the TMA - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.reduce_sync_barrier.arrive_and_wait() if warp_idx == 0: diff --git a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py index 97baad03..7cc59105 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py @@ -1610,10 +1610,7 @@ class Sm100GroupedBlockScaledGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/blackwell/grouped_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_gemm.py index c846647d..bae57f8a 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_gemm.py @@ -1261,10 +1261,7 @@ class GroupedGemmKernel: tRS_sC[(None, None, None, epi_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # # store C to global memory with TMA diff --git a/examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py index 7aeef3b2..b37181cb 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py @@ -1175,7 +1175,7 @@ class GroupedMixedInputGemmKernel: # Schedule warp if warp_idx == self.schedule_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_schedule_warp) + cute.arch.setmaxregister_decrease(self.num_regs_schedule_warp) # Persistent tile scheduling loop tile_sched = utils.StaticPersistentRuntimeTileScheduler.create( tile_sched_params, @@ -1222,10 +1222,7 @@ class GroupedMixedInputGemmKernel: - (cta_tile_offset_n * self.cta_tile_shape_mnk[1]) ) # Fence and barrier to ensure tile info store has finished - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.sched_sync_barrier.arrive_and_wait() # Commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1237,7 +1234,7 @@ class GroupedMixedInputGemmKernel: # Specialized TMA load warp for A/B tensor if warp_idx == self.tma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_tma_warps) + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) # Persistent tile scheduling loop tile_info_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_tile_info_stage @@ -1246,10 +1243,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() a_load2trans_producer_state = pipeline.make_pipeline_state( @@ -1333,10 +1327,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() # Wait A/B buffer empty @@ -1345,7 +1336,7 @@ class GroupedMixedInputGemmKernel: # Specialized TMA load for scale tensor if warp_idx == self.scale_tma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_tma_warps) + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): # Persistent tile scheduling loop tile_info_consumer_state = pipeline.make_pipeline_state( @@ -1355,10 +1346,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() scale_load2trans_producer_state = pipeline.make_pipeline_state( @@ -1425,10 +1413,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() # Wait scale buffer empty @@ -1436,7 +1421,7 @@ class GroupedMixedInputGemmKernel: # Specialized transform warps if warp_idx >= self.transform_warp_id[0]: - cute.arch.warpgroup_reg_alloc(self.num_regs_transform_warps) + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) transform_local_tidx = tidx - 32 * self.transform_warp_id[0] # Partition tensors for transform input and output and set up the copy atom # used for loading and storing transformed A tensor @@ -1508,10 +1493,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() a_load2trans_consumer_state = pipeline.make_pipeline_state( @@ -1669,10 +1651,7 @@ class GroupedMixedInputGemmKernel: ): cute.arch.fence_view_async_tmem_store() else: - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") if cutlass.const_expr( self.scale_mode == TransformMode.ConvertScale ): @@ -1697,10 +1676,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() # Wait a_transform buffer empty @@ -1708,7 +1684,7 @@ class GroupedMixedInputGemmKernel: # Specialized MMA warp if warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_mma_warp) + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) tCtAcc_base = accumulators # Persistent tile scheduling loop tile_info_consumer_state = pipeline.make_pipeline_state( @@ -1718,10 +1694,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() trans2mma_consumer_state = pipeline.make_pipeline_state( @@ -1799,10 +1772,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() # Wait for accumulator buffer empty @@ -1810,7 +1780,7 @@ class GroupedMixedInputGemmKernel: # Specialized epilogue warps if warp_idx < self.mma_warp_id: - cute.arch.warpgroup_reg_alloc(self.num_regs_epilogue_warps) + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) epi_tidx = tidx tCtAcc_base = accumulators # Partition for epilogue @@ -1860,10 +1830,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() num_prev_subtiles = cutlass.Int32(0) @@ -1938,10 +1905,7 @@ class GroupedMixedInputGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # TMA store C to global memory if warp_idx == self.epilog_warp_id[0]: @@ -1986,10 +1950,7 @@ class GroupedMixedInputGemmKernel: work_tile = self.make_work_tile_info( sTile_info[(None, tile_info_consumer_state.index)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() diff --git a/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py index eb1954db..c72c24c8 100644 --- a/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py +++ b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py @@ -839,7 +839,7 @@ class SSDKernel: # Specialized TMA load Delta/CumsumDelta/X warp if warp_idx == self.tma_deltas_x_d_warp_id: # Dealloc regs for pre-inter/pre-intra warps - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # ((ATOM_V, REST_V), INPUT_STAGE) # ((ATOM_V, REST_V), 1, 1, C, EH, B) @@ -1003,7 +1003,7 @@ class SSDKernel: # Specialized TMA load B/C warp elif warp_idx == self.tma_b_c_warp_id: # Dealloc regs for pre-inter/pre-intra warps - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # ((ATOM_V, REST_V), INPUT_STAGE) # ((ATOM_V, REST_V), 1, 1, C, G, B) @@ -1107,7 +1107,7 @@ class SSDKernel: # Specialized MMA Intra warp elif warp_idx == self.mma_intra_warp_id: # Dealloc regs for pre-inter/pre-intra warps - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # Make shared/tmem fragments for INTRA_MMA1 B/C/ACC # (MMA, MMA_N, MMA_K, INPUT_STAGE) @@ -1355,7 +1355,7 @@ class SSDKernel: # Specialized MMA Inter warp elif warp_idx == self.mma_inter_warp_id: # Dealloc regs for pre-inter/pre-intra warps - cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + cute.arch.setmaxregister_decrease(self.num_regs_uniform_warps) # Make shared/tmem fragments for INTER_MMA1 X/B/ACC # (MMA, MMA_N, MMA_K, INPUT_STAGE) @@ -1515,7 +1515,7 @@ class SSDKernel: or warp_idx == self.pre_inter_warp_id[3] ): # Alloc regs in pre_inter warps - cute.arch.warpgroup_reg_alloc(self.num_regs_pre_inter_warps) + cute.arch.setmaxregister_increase(self.num_regs_pre_inter_warps) # Make tiledCopy and partition smem/register tensor for smem load Bt # ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N, INPUT_STAGE) @@ -1670,10 +1670,7 @@ class SSDKernel: cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord]) # Fence for shared memory - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # Async arrive INTER2_P buffer full inter2_p_pipeline.producer_commit(inter2_p_producer_state) # Advance INTER2_P producer state @@ -1703,10 +1700,7 @@ class SSDKernel: ] # Fence for shared memory - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # Combine B/Delta/DeltaA/last_column tScaledB = self.pre_inter_scale_bt_with_delta( @@ -1722,10 +1716,7 @@ class SSDKernel: cute.copy(tiled_r2s_b, tBrB_r2s, tBsB_r2s[inter1_b_coord]) # Fence for shared memory - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # Async arrive B/Delta/B_TMEM buffer empty/empty/full b_pipeline.consumer_release( @@ -1774,10 +1765,7 @@ class SSDKernel: cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord]) # Fence for shared memory - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # Async arrive INTER1_ACC/INTER2_P buffer empty/full inter1_acc_pipeline.consumer_release(inter1_acc_consumer_state) @@ -1811,10 +1799,7 @@ class SSDKernel: # Store last INTER2_P (State) from smem to gmem # Wait for all previous stores to smem to be done - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.pre_inter_sync_barrier.arrive_and_wait() if local_warp_idx == 0: @@ -1849,7 +1834,7 @@ class SSDKernel: or warp_idx == self.pre_intra_warp_id[3] ): # Alloc regs in pre_inter warps - cute.arch.warpgroup_reg_alloc(self.num_regs_pre_intra_warps) + cute.arch.setmaxregister_increase(self.num_regs_pre_intra_warps) # Make tmem fragment for INTRA1_ACC # (MMA, MMA_M, MMA_N, INTRA1_ACC_STAGE) @@ -2022,7 +2007,7 @@ class SSDKernel: # Specialized Epilogue warp else: # Dealloc regs for pre-inter/pre-intra warps - cute.arch.warpgroup_reg_dealloc(self.num_regs_epilogue_warps) + cute.arch.setmaxregister_decrease(self.num_regs_epilogue_warps) # (L, D, INPUT_STAGE) sDeltaA = self.epilog_make_delta(smem_cumsum_delta) @@ -2324,10 +2309,7 @@ class SSDKernel: ) # Fence for R2S store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # Sync before TMA store self.epilog_sync_barrier.arrive_and_wait() diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py index 62095af7..a1b14194 100644 --- a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py @@ -717,7 +717,7 @@ class MixedInputFusedMultiHeadAttentionDecode: elif warp_idx == self.tma_kv_warp_id: # Free registers if cutlass.const_expr(self.use_reg_reconfig): - cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Apply block tiler and slice gK = cute.local_tile(mK, tiler=(blk_tile_s, blk_tile_d), coord=(None, 0, coord_hb)) # (TILE_S, TILE_D, #TILE_S) @@ -782,7 +782,7 @@ class MixedInputFusedMultiHeadAttentionDecode: elif warp_idx == self.tma_qo_warp_id: # Free registers if cutlass.const_expr(self.use_reg_reconfig): - cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Apply block tiler and slice gQ = cute.local_tile(mQ, tiler=(blk_tile_h, blk_tile_d), coord=(coord_hr, 0, coord_hb)) # (TILE_H, TILE_D) @@ -852,7 +852,7 @@ class MixedInputFusedMultiHeadAttentionDecode: elif warpgroup_idx in self.cvt_warpgroup_ids: # Free registers if cutlass.const_expr(self.use_reg_reconfig): - cute.arch.warpgroup_reg_dealloc(self.cvt_regs) + cute.arch.setmaxregister_decrease(self.cvt_regs) # Initialize for dual convert if necessary convert_warpgroups = 1 @@ -986,7 +986,7 @@ class MixedInputFusedMultiHeadAttentionDecode: elif warp_idx == self.mma_kq_warp_id: # Free registers if cutlass.const_expr(self.use_reg_reconfig): - cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Setup mma descriptors tBsQ_desc = thrblk_mma_kq.make_fragment_B(tBsQ) @@ -1037,7 +1037,7 @@ class MixedInputFusedMultiHeadAttentionDecode: elif warp_idx == self.mma_vp_warp_id: # Free registers if cutlass.const_expr(self.use_reg_reconfig): - cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Setup mma descriptors tiled_mma_vp.set(tcgen05.Field.ACCUMULATE, True) @@ -1095,7 +1095,7 @@ class MixedInputFusedMultiHeadAttentionDecode: elif warpgroup_idx == self.softmax_warpgroup_id: # Alloc registers if cutlass.const_expr(self.use_reg_reconfig): - cute.arch.warpgroup_reg_alloc(self.softmax_regs) + cute.arch.setmaxregister_increase(self.softmax_regs) # Construct tiled copies tmem_op_width = 32 diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py index 10506046..3117bc56 100644 --- a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py @@ -837,7 +837,7 @@ class MixedInputFusedMultiHeadAttentionPrefill: # Load # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.load_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + cute.arch.setmaxregister_decrease(self.num_regs_other) while work_tile.is_valid_tile: curr_block_coord = work_tile.tile_idx mma_block_coord = ( @@ -923,7 +923,7 @@ class MixedInputFusedMultiHeadAttentionPrefill: # MMA # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + cute.arch.setmaxregister_decrease(self.num_regs_other) tmem.wait_for_alloc() while work_tile.is_valid_tile: curr_block_coord = work_tile.tile_idx @@ -1023,7 +1023,7 @@ class MixedInputFusedMultiHeadAttentionPrefill: # Softmax # /////////////////////////////////////////////////////////////////////////////// if warp_idx < self.mma_warp_id and warp_idx >= self.softmax_warp_ids[0]: - cute.arch.warpgroup_reg_alloc(self.num_regs_softmax) + cute.arch.setmaxregister_increase(self.num_regs_softmax) tmem.allocate(self.num_tmem_alloc_cols) tmem.wait_for_alloc() tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) @@ -1117,7 +1117,7 @@ class MixedInputFusedMultiHeadAttentionPrefill: # Trans # /////////////////////////////////////////////////////////////////////////////// if warp_idx < self.softmax_warp_ids[0]: - cute.arch.warpgroup_reg_dealloc(self.num_regs_transform) + cute.arch.setmaxregister_decrease(self.num_regs_transform) qk_thr_mma_leader_cta = qk_tiled_mma.get_slice(0) pv_thr_mma_leader_cta = pv_tiled_mma.get_slice(0) sScaleK_ = qk_thr_mma_leader_cta.partition_B(sScaleK_s2r_view) @@ -1181,7 +1181,7 @@ class MixedInputFusedMultiHeadAttentionPrefill: # Empty # /////////////////////////////////////////////////////////////////////////////// if warp_idx > self.load_warp_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + cute.arch.setmaxregister_decrease(self.num_regs_other) return diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py index 606e58ab..f10e50d1 100644 --- a/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py @@ -1326,10 +1326,7 @@ class MixedInputGemmKernel: ): cute.arch.fence_view_async_tmem_store() else: - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # Signal the completion of transformation trans2mma_pipeline.producer_commit(trans2mma_producer_state) # Signal the completion of using A and scale tensors @@ -1557,10 +1554,7 @@ class MixedInputGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() # TMA store C to global memory if warp_idx == self.epilog_warp_id[0]: diff --git a/examples/python/CuTeDSL/blackwell/mla.py b/examples/python/CuTeDSL/blackwell/mla.py index d3443103..3d4f20c9 100644 --- a/examples/python/CuTeDSL/blackwell/mla.py +++ b/examples/python/CuTeDSL/blackwell/mla.py @@ -1045,7 +1045,7 @@ class BlackwellMultiHeadLatentAttentionForward: ) # Load page table when isasync is true if warp_idx == self.load_pt_warp_id: - cute.arch.warpgroup_reg_dealloc(self.other_reg_num) + cute.arch.setmaxregister_decrease(self.other_reg_num) load_pt_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.load_pt_stage ) @@ -1084,7 +1084,7 @@ class BlackwellMultiHeadLatentAttentionForward: warp_idx == self.load_cp_async_warp_ids[0] or warp_idx == self.load_cp_async_warp_ids[1] ): - cute.arch.warpgroup_reg_dealloc(self.other_reg_num) + cute.arch.setmaxregister_decrease(self.other_reg_num) load_pt_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.load_pt_stage ) @@ -1161,9 +1161,9 @@ class BlackwellMultiHeadLatentAttentionForward: warp_idx >= self.empty_warp_ids[0] and warp_idx <= self.empty_warp_ids[-1] ): - cute.arch.warpgroup_reg_dealloc(self.other_reg_num) + cute.arch.setmaxregister_decrease(self.other_reg_num) if warp_idx == self.load_tma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.other_reg_num) + cute.arch.setmaxregister_decrease(self.other_reg_num) load_q_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.load_q_stage ) @@ -1232,7 +1232,7 @@ class BlackwellMultiHeadLatentAttentionForward: # MMA warp # /////////////////////////////////////////////////////////////////////////////// if warp_idx == self.mma_warp_id: - cute.arch.warpgroup_reg_dealloc(self.other_reg_num) + cute.arch.setmaxregister_decrease(self.other_reg_num) # Alloc tensor memory buffer cute.arch.alloc_tmem( cute.arch.SM100_TMEM_CAPACITY_COLUMNS, @@ -1339,7 +1339,7 @@ class BlackwellMultiHeadLatentAttentionForward: warp_idx >= self.compute_warp_ids[0] and warp_idx <= self.compute_warp_ids[-1] ): - cute.arch.warpgroup_reg_alloc(self.softmax_reg_num) + cute.arch.setmaxregister_increase(self.softmax_reg_num) mma_s_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.mma_s_stage ) @@ -1413,7 +1413,7 @@ class BlackwellMultiHeadLatentAttentionForward: warp_idx >= self.correction_warp_ids[0] and warp_idx <= self.correction_warp_ids[-1] ): - cute.arch.warpgroup_reg_alloc(self.correction_reg_num) + cute.arch.setmaxregister_increase(self.correction_reg_num) p_cor_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.p_cor_stage ) diff --git a/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py b/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py index 491f7cf6..74d2afa3 100644 --- a/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py +++ b/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py @@ -649,7 +649,7 @@ class Sm120GemmKernel: # MMA warp group if warp_idx < self.num_mma_warps: - cute.arch.warpgroup_reg_alloc(self.mma_register_requirement) + cute.arch.setmaxregister_increase(self.mma_register_requirement) num_k_blocks = cute.size(tCrA, mode=[2]) @@ -873,10 +873,7 @@ class Sm120GemmKernel: tRS_sD[(None, None, None, epi_buffer)], ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # barrier for sync self.epilog_sync_barrier.arrive_and_wait() @@ -901,7 +898,7 @@ class Sm120GemmKernel: # End of MMA warp group # Start of DMA warp group elif warp_idx == self.num_mma_warps: - cute.arch.warpgroup_reg_dealloc(self.load_register_requirement) + cute.arch.setmaxregister_decrease(self.load_register_requirement) while work_tile.is_valid_tile: tile_coord_mnl = work_tile.tile_idx diff --git a/examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py b/examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py index 1995bf95..01a7534b 100644 --- a/examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py +++ b/examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py @@ -1298,10 +1298,7 @@ class PersistentDenseGemmKernel: c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilog_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py b/examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py index c5ea8419..6fe01aaa 100644 --- a/examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py +++ b/examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py @@ -1377,10 +1377,7 @@ class PersistentDenseGemmKernel: c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilogue_sync_barrier.arrive_and_wait() # diff --git a/examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py b/examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py index 3104b0e8..f196768d 100644 --- a/examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py +++ b/examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py @@ -1224,10 +1224,7 @@ class PersistentDenseGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilog_threads = 32 * len(self.epilog_warp_id) cute.arch.barrier( barrier_id=self.epilog_sync_bar_id, diff --git a/examples/python/CuTeDSL/hopper/dense_gemm.py b/examples/python/CuTeDSL/hopper/dense_gemm.py index 16410b03..c00ead6f 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm.py @@ -1006,10 +1006,7 @@ class HopperWgmmaGemmKernel: tiled_copy_r2s, tRS_rD_out, tRS_sD[(None, None, None, epi_buffer)] ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") # barrier for sync pipeline.sync(barrier_id=1) diff --git a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py index b6b58ae6..5b999c0b 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py @@ -723,7 +723,7 @@ class HopperWgmmaGemmPersistentKernel: is_dma_warp_group = warp_group_idx < self.num_dma_warp_groups if is_dma_warp_group: - cute.arch.warpgroup_reg_dealloc(self.load_register_requirement) + cute.arch.setmaxregister_decrease(self.load_register_requirement) if warp_idx == self.load_warp_id: tile_sched = utils.StaticPersistentTileScheduler.create( @@ -783,7 +783,7 @@ class HopperWgmmaGemmPersistentKernel: # MMA warp group if not is_dma_warp_group: - cute.arch.warpgroup_reg_alloc(self.mma_register_requirement) + cute.arch.setmaxregister_increase(self.mma_register_requirement) tile_sched = utils.StaticPersistentTileScheduler.create( tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() ) @@ -952,10 +952,7 @@ class HopperWgmmaGemmPersistentKernel: tRS_sD[(None, None, None, epi_buffer)], ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") self.epilog_sync_barrier.arrive_and_wait() gmem_coord = epi_tile_layout.get_hier_coord(epi_idx) diff --git a/examples/python/CuTeDSL/hopper/fmha.py b/examples/python/CuTeDSL/hopper/fmha.py index 49776bca..9efce75b 100644 --- a/examples/python/CuTeDSL/hopper/fmha.py +++ b/examples/python/CuTeDSL/hopper/fmha.py @@ -661,7 +661,7 @@ class HopperFusedMultiHeadAttentionForward: cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_o) if warp_group_idx == self.load_warp_group_id: - cute.arch.warpgroup_reg_dealloc(self.num_regs_load) + cute.arch.setmaxregister_decrease(self.num_regs_load) tile_sched = fmha_utils.create_fmha_static_tile_scheduler( tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() @@ -784,7 +784,7 @@ class HopperFusedMultiHeadAttentionForward: warp_group_idx == self.compute_epilogue_0_warp_group_id or warp_group_idx == self.compute_epilogue_1_warp_group_id ): - cute.arch.warpgroup_reg_alloc(self.num_regs_mma) + cute.arch.setmaxregister_increase(self.num_regs_mma) tile_sched = fmha_utils.create_fmha_static_tile_scheduler( tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() @@ -1164,10 +1164,7 @@ class HopperFusedMultiHeadAttentionForward: tRS_sD[(None, None, None, epi_buffer, warp_group_idx - 1)], ) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") pipeline.arrive_and_wait( barrier_id=warp_group_idx, num_threads=self.num_threads_per_warp_group, diff --git a/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py index dec19c90..9741c08a 100644 --- a/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py +++ b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py @@ -243,10 +243,7 @@ class ClcDynamicPersistentTileScheduler: result_addr: 16-byte response data (simulating shared memory access) """ m_idx, n_idx, l_idx, vld = cute.arch.clc_response(result_addr, loc=loc, ip=ip) - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") cta_idx_in_cluster, cta_idy_in_cluster, _ = self.cta_id_in_cluster cur_tile_coord = (m_idx + cta_idx_in_cluster, n_idx + cta_idy_in_cluster, l_idx) return WorkTileInfo(cur_tile_coord, vld) diff --git a/python/CuTeDSL/cutlass/utils/gemm/sm100.py b/python/CuTeDSL/cutlass/utils/gemm/sm100.py index 1089a91a..d23fcfba 100644 --- a/python/CuTeDSL/cutlass/utils/gemm/sm100.py +++ b/python/CuTeDSL/cutlass/utils/gemm/sm100.py @@ -280,10 +280,7 @@ def epilogue_tma_store( c_buffer = (num_prev_subtiles + subtile_idx) % gemm_kernel.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilog_sync_barrier.arrive_and_wait() # @@ -675,10 +672,7 @@ def epilogue_tma_store_release_flag( c_buffer = (num_prev_subtiles + subtile_idx) % gemm_kernel.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy( - cute.arch.ProxyKind.async_shared, - space=cute.arch.SharedSpace.shared_cta, - ) + cute.arch.fence_proxy("async.shared", space="cta") epilog_sync_barrier.arrive_and_wait() #