leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 30 Packages Projects Releases Wiki Activity

Update nvvm API call from nvvm enum to str (#2985)

Browse Source
This commit is contained in:
Xiao Song
2026-01-27 17:28:29 +08:00
committed by GitHub
parent 7a14467776
commit acb45938e9
29 changed files with 132 additions and 327 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py
+13 -34
View File
@@ -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()
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
+13 -34
View File
@@ -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()
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
+13 -34
View File
@@ -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()
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py
+1 -4
View File
@@ -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()
#
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_prefetch.py
+1 -4
View File
@@ -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()
#
examples/python/CuTeDSL/blackwell/dense_gemm.py
+1 -4
View File
@@ -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
examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py
+2 -8
View File
@@ -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()
#
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
+1 -4
View File
@@ -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()
#
examples/python/CuTeDSL/blackwell/dense_gemm_persistent_prefetch.py
+1 -4
View File
@@ -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()
#
examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py
+1 -4
View File
@@ -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
examples/python/CuTeDSL/blackwell/fmha.py
+8 -11
View File
@@ -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(
examples/python/CuTeDSL/blackwell/fmha_bwd.py
+7 -13
View File
@@ -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:
examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py
+1 -4
View File
@@ -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()
#
examples/python/CuTeDSL/blackwell/grouped_gemm.py
+1 -4
View File
@@ -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
examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py
+19 -58
View File
@@ -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()
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py
+13 -31
View File
@@ -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()
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py
+6 -6
View File
@@ -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
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py
+5 -5
View File
@@ -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
examples/python/CuTeDSL/blackwell/mixed_input_gemm.py
+2 -8
View File
@@ -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]:
examples/python/CuTeDSL/blackwell/mla.py
+7 -7
View File
@@ -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
)