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

v4.5 dev update. (#3153)

Browse Source
This commit is contained in:
Junkai-Wu
2026-04-07 12:16:05 -04:00
committed by GitHub
parent 418d38a5de
commit a221da7ccf
265 changed files with 4913 additions and 1478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py
+6 -7
View File
@@ -567,7 +567,7 @@ class Sm100GroupedBlockScaledGemmKernel:
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[
@@ -641,7 +641,6 @@ class Sm100GroupedBlockScaledGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -737,8 +736,8 @@ class Sm100GroupedBlockScaledGemmKernel:
+ Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8
)
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_ptr = storage.tmem_holding_buf.ptr
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
@@ -1249,7 +1248,7 @@ class Sm100GroupedBlockScaledGemmKernel:
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
ptr_to_buffer_holding_addr=tmem_holding_buf_ptr,
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
@@ -1446,7 +1445,7 @@ class Sm100GroupedBlockScaledGemmKernel:
if warp_idx == self.epilog_warp_id[0]:
cute.arch.alloc_tmem(
self.num_tmem_alloc_cols,
tmem_holding_buf,
tmem_holding_buf_ptr,
is_two_cta=use_2cta_instrs,
)
@@ -1461,7 +1460,7 @@ class Sm100GroupedBlockScaledGemmKernel:
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
ptr_to_buffer_holding_addr=tmem_holding_buf_ptr,
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)