leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 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.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(