leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 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/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: