leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 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/blockwise_gemm/blockwise_gemm.py
+3 -7
View File
@@ -549,7 +549,7 @@ class BlockwiseGemmKernel:
cutlass.Int64, self.num_tile_stage * 2
]
epi_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
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[
@@ -614,7 +614,6 @@ class BlockwiseGemmKernel:
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,
)
@@ -682,9 +681,6 @@ class BlockwiseGemmKernel:
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
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
@@ -771,11 +767,11 @@ class BlockwiseGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.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
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
+3 -7
View File
@@ -568,7 +568,7 @@ class BlockwiseContiguousGroupedGemmKernel:
cutlass.Int64, self.num_tile_stage * 2
]
epi_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
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[
@@ -634,7 +634,6 @@ class BlockwiseContiguousGroupedGemmKernel:
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,
)
@@ -703,9 +702,6 @@ class BlockwiseContiguousGroupedGemmKernel:
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
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
@@ -792,11 +788,11 @@ class BlockwiseContiguousGroupedGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.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
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
+3 -7
View File
@@ -567,7 +567,7 @@ class BlockwiseMaskedGroupedGemmKernel:
cutlass.Int64, self.num_tile_stage * 2
]
epi_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
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[
@@ -633,7 +633,6 @@ class BlockwiseMaskedGroupedGemmKernel:
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,
)
@@ -702,9 +701,6 @@ class BlockwiseMaskedGroupedGemmKernel:
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
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
@@ -791,11 +787,11 @@ class BlockwiseMaskedGroupedGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.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
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py
+3 -3
View File
@@ -615,7 +615,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
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[
@@ -794,11 +794,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.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
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py
+3 -3
View File
@@ -551,7 +551,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
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[
@@ -737,11 +737,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.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
examples/python/CuTeDSL/blackwell/dense_gemm.py
+3 -3
View File
@@ -517,7 +517,7 @@ class DenseGemmKernel:
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()
@@ -564,10 +564,10 @@ class DenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
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
examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py
+3 -3
View File
@@ -493,7 +493,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sD: cute.struct.Align[
@@ -674,11 +674,11 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
)
# 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_ids[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
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
+3 -3
View File
@@ -593,7 +593,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()
@@ -644,11 +644,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
examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py
+3 -3
View File
@@ -614,7 +614,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
clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
clc_response: cute.struct.MemRange[cutlass.Int32, 4]
@@ -686,11 +686,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
examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py
+3 -3
View File
@@ -514,7 +514,7 @@ class DenseGemmKernel:
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()
@@ -562,10 +562,10 @@ class DenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
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
examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py
+3 -3
View File
@@ -573,7 +573,7 @@ class DenseGemmEFC:
# Barriers used by the supplemental load tensor pipeline.
c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (MMA, MMA_M, MMA_K, STAGE)
sA: cute.struct.Align[
@@ -651,11 +651,11 @@ class DenseGemmEFC:
)
# 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=self.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
examples/python/CuTeDSL/blackwell/fmha.py
+2 -2
View File
@@ -998,7 +998,7 @@ class BlackwellFusedMultiHeadAttentionForward:
# Alloc tmem buffer
tmem_alloc_cols = Int32(self.tmem_alloc_cols)
cute.arch.alloc_tmem(tmem_alloc_cols, storage.tmem_holding_buf)
cute.arch.alloc_tmem(tmem_alloc_cols, storage.tmem_holding_buf.ptr)
self.tmem_alloc_barrier.arrive_and_wait()
tile_sched = fmha_utils.create_fmha_static_tile_scheduler(
tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
@@ -1260,7 +1260,7 @@ class BlackwellFusedMultiHeadAttentionForward:
tmem_ptr = cute.arch.retrieve_tmem_ptr(
Float32,
alignment=16,
ptr_to_buffer_holding_addr=storage.tmem_holding_buf,
ptr_to_buffer_holding_addr=storage.tmem_holding_buf.ptr,
)
cute.arch.dealloc_tmem(tmem_ptr, tmem_alloc_cols)
examples/python/CuTeDSL/blackwell/fmha_bwd.py
+1 -2
View File
@@ -708,7 +708,6 @@ class BlackwellFusedMultiHeadAttentionBackward:
grid=bwd_grid,
block=[self.threads_per_cta, 1, 1],
cluster=[1, 1, 1],
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -913,7 +912,7 @@ class BlackwellFusedMultiHeadAttentionBackward:
)
sLSE = storage.sLSE.get_tensor(LSE_smem_layout)
sSum_OdO = storage.sSum_OdO.get_tensor(sum_OdO_smem_layout)
tmem_holding_buf = storage.tmem_holding_buf
tmem_holding_buf = storage.tmem_holding_buf.ptr
sQT_ptr = cute.recast_ptr(sQ.iterator, QT_smem_layout_staged.inner)
sQT = cute.make_tensor(sQT_ptr, QT_smem_layout_staged.outer)
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)
examples/python/CuTeDSL/blackwell/grouped_gemm.py
+3 -3
View File
@@ -425,7 +425,7 @@ class GroupedGemmKernel:
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[
@@ -590,11 +590,11 @@ class GroupedGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.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
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py
+1 -1
View File
@@ -820,7 +820,7 @@ class SSDKernel:
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
smem_storage.tmem_holding_buf,
smem_storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
)
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py
-2
View File
@@ -603,8 +603,6 @@ class MixedInputFusedMultiHeadAttentionDecode:
p_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2)
o_pipeline_ptr = smem.allocate_array(Int64, self.o_stages * 2)
assert smem._allocated_bytes <= self.mbarrier_reserved_bytes
# Declare named barriers
softmax_nbar_id = 1
mma_kq_nbar_id = 2
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py
+3 -3
View File
@@ -403,7 +403,7 @@ class MixedInputFusedMultiHeadAttentionPrefillD256:
s_corr_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2]
sum_mbar_ptr: cute.struct.MemRange[Int64, 2]
mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2]
tmem_dealloc_mbar_ptr: Int64
tmem_dealloc_mbar: Int64
tmem_holding_buf: Int32
self.shared_storage = SharedStorage
@@ -654,11 +654,11 @@ class MixedInputFusedMultiHeadAttentionPrefillD256:
)
# 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.correction_warp_ids[0],
is_two_cta=True,
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
pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py
+3 -3
View File
@@ -390,7 +390,7 @@ class MixedInputFusedMultiHeadAttentionPrefillD512:
p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2]
mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2]
swap_mbar_ptr: cute.struct.MemRange[Int64, self.swap_stage * 2]
tmem_dealloc_mbar_ptr: Int64
tmem_dealloc_mbar: Int64
tmem_holding_buf: Int32
self.shared_storage = SharedStorage
@@ -627,11 +627,11 @@ class MixedInputFusedMultiHeadAttentionPrefillD512:
)
# 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.softmax_warp_ids[0],
is_two_cta=True,
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
pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py
+3 -3
View File
@@ -624,7 +624,7 @@ class GroupedMixedInputGemmKernel:
tile_info_empty_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_tile_info_stage
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
self.shared_storage = SharedStorage
@@ -824,11 +824,11 @@ class GroupedMixedInputGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_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
examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py
+3 -3
View File
@@ -518,7 +518,7 @@ class GroupedMixedInputGemmAccScaleKernel:
tile_info_empty_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_tile_info_stage
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
self.shared_storage = SharedStorage
@@ -708,11 +708,11 @@ class GroupedMixedInputGemmAccScaleKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_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
examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py
+3 -3
View File
@@ -618,7 +618,7 @@ class MixedInputGemmKernel:
]
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
# Tensor buffers
# (EPI_TILE_M, EPI_TILE_N, STAGE)
@@ -820,11 +820,11 @@ class MixedInputGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_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
examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py
+3 -5
View File
@@ -568,7 +568,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
cutlass.Int64, self.load_pt_stage * 2
]
# Tmem dealloc cluster barrier
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
# Tmem holding buffer
tmem_holding_buf: cutlass.Int32
@@ -641,7 +641,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=self.cluster_shape_mnk,
smem=SplitKVKernelSharedStorage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -657,7 +656,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
).launch(
grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]),
block=[self.threads_per_warp * self.num_compute_warps, 1, 1],
smem=MAX_SPLITS * self.acc_dtype.width // 8,
stream=stream,
min_blocks_per_mp=1,
)
@@ -838,11 +836,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_bar,
allocator_warp_id=self.mma_warp_id,
is_two_cta=self.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,
)
load_q_pipeline = self.make_and_init_load_qkv_pipeline(
examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py
+3 -5
View File
@@ -661,7 +661,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
]
# Tmem dealloc cluster barrier
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
# Tmem holding buffer
tmem_holding_buf: cutlass.Int32
@@ -707,7 +707,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=self.cluster_shape_mnk,
smem=SplitKVKernelSharedStorage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -723,7 +722,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
).launch(
grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]),
block=[self.threads_per_warp * self.num_compute_warps, 1, 1],
smem=MAX_SPLITS * self.acc_dtype.width // 8,
stream=stream,
min_blocks_per_mp=1,
)
@@ -904,11 +902,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_bar,
allocator_warp_id=self.mma_warp_id,
is_two_cta=self.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,
)
load_q_pipeline = self.make_and_init_load_qkv_pipeline(
examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py
+9 -9
View File
@@ -43,9 +43,9 @@ from cutlass.cute.runtime import from_dlpack
from dataclasses import dataclass, field
"""
This example provides an experimental implementation of the SM103 batched 3xFP4 blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases.
This example provides an experimental implementation of the SM103 batched FP4 Ultra blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases.
A high-performance persistent batched 3xFP4 blockscaled GEMM example for the NVIDIA Blackwell SM103 architecture
A high-performance persistent batched FP4 Ultra blockscaled GEMM example for the NVIDIA Blackwell SM103 architecture
using CUTE DSL.
- Matrix A is MxKxL, L is batch dimension, A can only be row-major("K") for MXF4/NVF4 input type
- Matrix B is NxKxL, L is batch dimension, B can only be row-major("K") for MXF4/NVF4 input type
@@ -166,7 +166,7 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
cluster_shape_mn: Tuple[int, int],
use_tma_store: bool,
):
"""Initializes the configuration for a Blackwell SM103 3xFP4 GEMM kernel.
"""Initializes the configuration for a Blackwell SM103 FP4 Ultra GEMM kernel.
This configuration includes several key aspects:
@@ -603,7 +603,7 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
sf_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_sf_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
# (MMA, MMA_M, MMA_K, STAGE)
sA: cute.struct.Align[
@@ -800,11 +800,11 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
)
# 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
@@ -1810,7 +1810,7 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
mma_tiler_mn: Tuple[int, int],
a_source: tcgen05.OperandSource = tcgen05.OperandSource.SMEM,
) -> cute.TiledMma:
"""Create a blockscaled trivial tiled MMA for SM103 (3xFP4), K fixed to 96.
"""Create a blockscaled trivial tiled MMA for SM103 (FP4 Ultra), K fixed to 96.
Returns a tcgen05 MMA configured for the given (M, N) tiler and CTA group.
@@ -2653,7 +2653,7 @@ def run(
:return: Execution time of the GEMM kernel
:rtype: float
"""
print(f"Running Sm103 Persistent 3xfp4 Dense BlockScaled GEMM test with:")
print(f"Running Sm103 Persistent FP4 Ultra Dense BlockScaled GEMM test with:")
print(f"mnkl: {mnkl}")
print(f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}")
print(f"C dtype: {c_dtype}")
@@ -2954,7 +2954,7 @@ if __name__ == "__main__":
)
parser = argparse.ArgumentParser(
description="Example of Sm103 3xfp4 Dense Persistent BlockScaled GEMM."
description="Example of Sm103 FP4 Ultra Dense Persistent BlockScaled GEMM."
)
parser.add_argument(
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py
+49 -64
View File
@@ -8,13 +8,10 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
import argparse
from typing import Tuple, Type, Callable
from functools import partial, lru_cache
from typing import Tuple
import cutlass
from cutlass import Numeric
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
@@ -69,7 +66,6 @@ def kernel(
a_smem_layout: cute.ComposedLayout,
b_smem_layout: cute.ComposedLayout,
):
# Current thread/warp/block coordinates
tidx, _, _ = cute.arch.thread_idx()
warp_idx = cute.arch.warp_idx()
@@ -103,7 +99,7 @@ def kernel(
num_threads=threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
)
num_tmem_cols = 512
@@ -143,15 +139,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(0)
# (MMA, MMA_M, MMA_K, RestK)
# (MMA, MMA_M, MMA_K)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K, RestK)
# (MMA, MMA_N, MMA_K)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K, STAGE)
# (MMA, MMA_M, MMA_K)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
# (MMA, MMA_N, MMA_K)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
@@ -199,14 +195,14 @@ def kernel(
tmem_thr_copy = tmem_tiled_copy.get_slice(tidx)
# (TmemCpy,NumTmemCpy,NumTiles)
tCtC = tmem_thr_copy.partition_S(tCtAcc_epi)
tDtC = tmem_thr_copy.partition_S(tCtAcc_epi)
# (TmemCpy,NumTmemCpy,NumTiles)
tCgC = tmem_thr_copy.partition_D(gC_epi)
tDgC = tmem_thr_copy.partition_D(gC_epi)
# (TmemCpy,NumTmemCpy)
tCrAcc = cute.make_rmem_tensor(tCgC[None, None, 0].shape, acc_dtype)
tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
# (TmemCpy,NumTmemCpy)
tCrC = cute.make_rmem_tensor(tCgC[None, None, 0].shape, io_dtype)
tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
#
# 2. Main loop
@@ -233,8 +229,6 @@ def kernel(
# Execute one K-block worth of MMA instructions
ab_full = ab_consumer.wait_and_advance()
# tCtAcc += tCrA * tCrB
num_k_blocks = cute.size(tCrA, mode=[2])
for k_block_idx in cutlass.range_constexpr(num_k_blocks):
k_block_coord = (None, None, k_block_idx, ab_full.index)
@@ -265,10 +259,10 @@ def kernel(
# TMEM -> RMEM -> GEMM
# Sub-tiling for better instruction-level parallelism
for i in cutlass.range(cute.size(tCtC, mode=[2])):
cute.copy(tmem_tiled_copy, tCtC[None, None, i], tCrAcc)
for i in cutlass.range(cute.size(tDtC, mode=[2])):
cute.copy(tmem_tiled_copy, tDtC[None, None, i], tCrAcc)
tCrC.store(tCrAcc.load().to(io_dtype))
cute.autovec_copy(tCrC, tCgC[None, None, i])
cute.autovec_copy(tCrC, tDgC[None, None, i])
acc_full.release()
# Deallocate TMEM
@@ -350,44 +344,10 @@ def host_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
)
@lru_cache(maxsize=1)
def prepare_run(
callable: Callable,
m: int,
n: int,
k: int,
a_dtype: Type[Numeric],
b_dtype: Type[Numeric],
c_dtype: Type[Numeric],
) -> tuple[Callable, tuple]:
import cutlass.torch as cutlass_torch
a, b, c = cutlass_torch.prepare_tensors_for_gemm(
(m, n, k), a_dtype, b_dtype, c_dtype
)
a_ = (
from_dlpack(a, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
b_ = (
from_dlpack(b, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
c_ = (
from_dlpack(c, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=n)
)
compiled_fn = cute.compile(callable, a_, b_, c_, options="--generate-line-info")
return partial(compiled_fn, a_, b_, c_), (a, b, c)
def run_dense_gemm(
mnk: Tuple[int, int, int],
tolerance: float,
) -> None:
):
global torch, cutlass_torch
import torch
import cutlass.torch as cutlass_torch
@@ -402,23 +362,48 @@ def run_dense_gemm(
m, n, k = mnk
torch.manual_seed(1111)
run_fn, (a, b, c) = prepare_run(
host_function, m, n, k, io_dtype, io_dtype, io_dtype
# Make K-major tensors (torch tensors are row-major)
def make_tensors(mn, k, dtype):
shape = (mn, k)
return (
torch.empty(*shape, dtype=torch.int32)
.random_(-2, 2)
.to(dtype=dtype, device="cuda")
)
a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
a_tensor = (
from_dlpack(a, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
b_tensor = (
from_dlpack(b, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
c_tensor = (
from_dlpack(c, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=n)
)
# Entry point to the host JIT function
run_fn()
host_function(a_tensor, b_tensor, c_tensor, no_cache=True)
# Compute reference result and verify
ref = torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))
ref = (torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))).cpu()
torch.testing.assert_close(
c, ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
)
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> list[int]:
def parse_comma_separated_ints(s: str):
try:
return [int(x.strip()) for x in s.split(",")]
except ValueError:
@@ -443,14 +428,14 @@ if __name__ == "__main__":
parser.add_argument(
"--tolerance", type=float, default=1e-01, help="Tolerance for validation"
)
args = parser.parse_args()
if len(args.mnk) != 3:
parser.error("--mnk must contain exactly 3 values")
if args.mnk[0] % mma_tiler_mnk[0] != 0 or args.mnk[1] % mma_tiler_mnk[1] != 0:
parser.error("m n must be divisible by mma_tiler_mn")
run_dense_gemm(args.mnk, args.tolerance)
run_dense_gemm(
args.mnk,
args.tolerance,
)
print("PASS")
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py
+22 -32
View File
@@ -65,8 +65,7 @@ Constraints for this example:
io_dtype = cutlass.Float16
acc_dtype = cutlass.Float32
use_2cta_instrs = True
cluster_shape_mnk = (2, 1, 1) if use_2cta_instrs else (1, 1, 1)
cluster_shape_mnk = (2, 1, 1)
mma_inst_shape_mnk = (256, 256, 16)
mma_tiler_mnk = (256, 256, 64)
threads_per_cta = 128
@@ -96,7 +95,6 @@ def kernel(
b_smem_layout: cute.ComposedLayout,
cta_layout_vmnk: cute.Layout,
):
# Current thread/warp/block coordinates
tidx, _, _ = cute.arch.thread_idx()
warp_idx = cute.arch.warp_idx()
@@ -174,15 +172,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0])
# (MMA, MMA_M, MMA_K, RestK)
# (MMA, MMA_M, MMA_K)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K, RestK)
# (MMA, MMA_N, MMA_K)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K, STAGE)
# (MMA, MMA_M, MMA_K)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
# (MMA, MMA_N, MMA_K)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
@@ -217,10 +215,10 @@ def kernel(
num_threads=threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
num_tmem_cols = 512
tmem.allocate(num_tmem_cols)
@@ -232,7 +230,7 @@ def kernel(
# Swap the pointer in tCtAcc
tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc.layout)
subtile_cnt = 1 if mma_tiler_mnk[0] == 64 else 4
subtile_cnt = 4
# (EpiTile)
epi_tiler = (
(cute.size(tCtAcc, mode=[0, 0]), cute.size(tCtAcc, mode=[0, 1]) // subtile_cnt),
@@ -244,24 +242,21 @@ def kernel(
# Every thread loads 64 x fp32
tmem_atom = cute.make_copy_atom(
tcgen05.Ld16x256bOp(tcgen05.Repetition.x8)
if mma_tiler_mnk[0] == 64
else tcgen05.Ld32x32bOp(tcgen05.Repetition.x64),
tcgen05.Ld32x32bOp(tcgen05.Repetition.x64),
cutlass.Float32,
)
tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_atom, tCtAcc_epi[None, 0])
tmem_thr_copy = tmem_tiled_copy.get_slice(tidx)
# (TmemCpy,NumTmemCpy,NumTiles)
tCtC = tmem_thr_copy.partition_S(tCtAcc_epi)
tDtC = tmem_thr_copy.partition_S(tCtAcc_epi)
# (TmemCpy,NumTmemCpy,NumTiles)
tCgC = tmem_thr_copy.partition_D(gC_epi)
tDgC = tmem_thr_copy.partition_D(gC_epi)
# (TmemCpy,NumTmemCpy)
tCrAcc = cute.make_rmem_tensor(tCgC[None, None, 0].shape, acc_dtype)
tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
# (TmemCpy,NumTmemCpy)
tCrC = cute.make_rmem_tensor(tCgC[None, None, 0].shape, io_dtype)
tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
#
# 2. Main loop
@@ -271,8 +266,8 @@ def kernel(
if warp_idx == 0:
# Wait for a empty accumulator buffer
if is_leader_cta:
acc_producer.acquire()
for k_tile in cutlass.range(num_k_tiles, prefetch_stages=ab_stages - 2):
acc_producer.acquire_and_advance()
for _ in cutlass.range(num_k_tiles, prefetch_stages=ab_stages - 2):
# Issue TMA loads
ab_empty = ab_producer.acquire_and_advance()
cute.copy(
@@ -310,7 +305,6 @@ def kernel(
# Signal that the accumulator is fully computed
if is_leader_cta:
acc_producer.commit()
acc_producer.advance()
#
# 3. Epilogue
@@ -321,13 +315,12 @@ def kernel(
# Wait for the accumulator buffer to be full
acc_full = acc_consumer.wait_and_advance()
# TMEM -> RMEM -> GEMM
# Sub-tiling for better instruction-level parallelism
for i in cutlass.range(cute.size(tCtC, mode=[2])):
cute.copy(tmem_tiled_copy, tCtC[None, None, i], tCrAcc)
for i in cutlass.range(cute.size(tDtC, mode=[2])):
cute.copy(tmem_tiled_copy, tDtC[None, None, i], tCrAcc)
tCrC.store(tCrAcc.load().to(io_dtype))
cute.autovec_copy(tCrC, tCgC[None, None, i])
cute.autovec_copy(tCrC, tDgC[None, None, i])
acc_full.release()
# Ensure used buffers are properly synchronized before producer exit.
@@ -353,7 +346,7 @@ def host_function(
io_dtype,
acc_dtype,
mma_inst_shape_mnk,
tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE,
tcgen05.CtaGroup.TWO,
tcgen05.OperandSource.SMEM,
tcgen05.OperandMajorMode.K,
tcgen05.OperandMajorMode.K,
@@ -381,16 +374,14 @@ def host_function(
cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,))
# Construct TMA load atoms
op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(
tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
)
op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
a_tma_atom, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
op,
a,
a_smem_layout_one_stage,
mma_tiler_mnk,
tiled_mma,
cta_layout_vmnk.shape,
cta_layout_vmnk.shape, # take the layout and extract the shape internally
)
b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
op,
@@ -403,8 +394,7 @@ def host_function(
grid_shape = cute.round_up(
cute.ceil_div(
(*c.layout.shape, 1),
(mma_tiler_mnk[0] // (2 if use_2cta_instrs else 1), *mma_tiler_mnk[1:]),
(*c.layout.shape, 1), (mma_tiler_mnk[0] // 2, *mma_tiler_mnk[1:])
),
cluster_shape_mnk,
)
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_4.py
+1 -1
View File
@@ -981,7 +981,7 @@ def run_dense_gemm(
import cutlass.torch as cutlass_torch
print("===================================================================")
print("Running Blackwell fp16 GEMM example 4 (with MIX CGA support):")
print("Running Blackwell fp16 GEMM example 4 (with MIX cluster size support):")
print(f" mnk: {mnk}")
print(f" tolerance: {tolerance}")
print(f" Preferred cluster shape: {preferred_cluster_shape_mnk}")
examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py
+1 -1
View File
@@ -500,7 +500,7 @@ class Sm100BlockScaledDenseGemmKernel:
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
)
tmem.allocate(self.num_tmem_alloc_cols)
examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py
+3 -3
View File
@@ -436,7 +436,7 @@ class Sm100BlockScaledDenseGemmKernel:
class SharedStorage:
ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
acc_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()
@@ -638,10 +638,10 @@ class Sm100BlockScaledDenseGemmKernel:
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()