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

Update blackwell tutorial to be compatible with 4.5-dev version (#3130)

Browse Source 
* Update blackwell tutorial to be compatible with 4.5-dev version

* update example for reverted changes

* add more example fix
This commit is contained in:
Longsheng Du
2026-04-09 14:40:33 +08:00
committed by GitHub
parent bd01dd3651
commit 08185b9c3e
12 changed files with 29 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/python/CuTeDSL/distributed/distributed_all_gather_gemm_blackwell.py
View File

@@ -756,7 +756,7 @@ class PersistentDenseGemmKernel:
acc_full_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_acc_stage * 2
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
@@ -806,11 +806,11 @@ class PersistentDenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init

6
examples/python/CuTeDSL/distributed/distributed_gemm_all_reduce_blackwell.py
View File

@@ -672,7 +672,7 @@ class PersistentDenseGemmKernel:
acc_full_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_acc_stage * 2
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
@@ -723,11 +723,11 @@ class PersistentDenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilogue_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init

6
examples/python/CuTeDSL/distributed/distributed_gemm_reduce_scatter_blackwell.py
View File

@@ -541,7 +541,7 @@ class PersistentDenseGemmKernel:
ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -660,8 +660,8 @@ class PersistentDenseGemmKernel:
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr
tmem_holding_buf = storage.tmem_holding_buf
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar.ptr
tmem_holding_buf = storage.tmem_holding_buf.ptr
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)