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/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