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

5
examples/python/CuTeDSL/hopper/dense_gemm.py
View File

@@ -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)

9
examples/python/CuTeDSL/hopper/dense_gemm_persistent.py
View File

@@ -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)

9
examples/python/CuTeDSL/hopper/fmha.py
View File

@@ -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,