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

v4.3 update. (#2709)

Browse Source 
* v4.3 update.

* Update the cute_dsl_api changelog's doc link

* Update version to 4.3.0

* Update the example link

* Update doc to encourage user to install DSL from requirements.txt

---------

Co-authored-by: Larry Wu <larwu@nvidia.com>
This commit is contained in:
Junkai-Wu
2025-10-21 14:26:30 -04:00
committed by GitHub
co-authored by Larry Wu
parent e6e2cc29f5
commit b1d6e2c9b3
244 changed files with 59272 additions and 10455 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py
+2927
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
+3056
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
+3033
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py
+143 -130
View File
@@ -27,7 +27,7 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
from typing import Optional, Type, Tuple, Union
from typing import Type, Tuple, Union
import cuda.bindings.driver as cuda
import torch
@@ -85,7 +85,7 @@ Input arguments to this example is shown below:
.. code-block:: bash
python examples/blackwell/dense_blockscaled_gemm_persistent.py \
python examples/blackwell/dense_blockscaled_gemm_persistent.py \
--ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \
--c_dtype Float16 \
--mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \
@@ -95,7 +95,7 @@ To collect performance with NCU profiler:
.. code-block:: bash
ncu python examples/blackwell/dense_blockscaled_gemm_persistent.py \
ncu python examples/blackwell/dense_blockscaled_gemm_persistent.py \
--ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \
--c_dtype Float16 \
--mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \
@@ -108,7 +108,7 @@ Constraints:
see detailed valid dtype combinations in below Sm100BlockScaledPersistentDenseGemmKernel class documentation
* A/B tensor must have the same data type, mixed data type is not supported (e.g., mxf8 x mxf4)
* Mma tiler M must be 128 or 256(use_2cta_instrs)
* Mma tiler N must be 128 or 256
* Mma tiler N must be 64/128/192/256
* Cluster shape M/N must be positive and power of 2, total cluster size <= 16
* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs)
* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned,
@@ -144,7 +144,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
- Float8E4M3FN/Float8E5M2
:note: Constraints:
- MMA tiler M must be 128 or 256 (use_2cta_instrs)
- MMA tiler N must be 128/256
- MMA tiler N must be 64/128/192/256
- Cluster shape M must be multiple of 2 if Mma tiler M is 256
- Cluster shape M/N must be positive and power of 2, total cluster size <= 16
- Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
@@ -209,9 +209,18 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# Set barrier id for cta sync, epilogue sync and tmem ptr sync
self.cta_sync_bar_id = 0
self.epilog_sync_bar_id = 1
self.tmem_ptr_sync_bar_id = 2
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
SM100_TMEM_CAPACITY_COLUMNS = 512
self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
@@ -228,21 +237,17 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
- Computing epilogue subtile
- Setting up A/B/SFA/SFB/C stage counts in shared memory
- Computing A/B/SFA/SFB/C shared memory layout
- Computing tensor memory allocation columns
"""
# Compute mma instruction shapes
mma_inst_bits_k = 256
# (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K)
self.mma_inst_shape_mnk = (
self.mma_inst_shape_mn = (
self.mma_tiler[0],
self.mma_tiler[1],
mma_inst_bits_k // self.a_dtype.width,
)
# (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
self.mma_inst_shape_mnk_sfb = (
self.mma_inst_shape_mnk[0] // (2 if self.use_2cta_instrs else 1),
cute.round_up(self.mma_inst_shape_mnk[1], 128),
self.mma_inst_shape_mnk[2],
self.mma_inst_shape_mn_sfb = (
self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
cute.round_up(self.mma_inst_shape_mn[1], 128),
)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
@@ -252,7 +257,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.sf_dtype,
self.sf_vec_size,
self.cta_group,
self.mma_inst_shape_mnk[:2],
self.mma_inst_shape_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
@@ -262,20 +267,21 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.sf_dtype,
self.sf_vec_size,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mnk_sfb[:2],
self.mma_inst_shape_mn_sfb,
)
# Compute mma/cluster/tile shapes
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = 4
self.mma_tiler = (
self.mma_inst_shape_mnk[0],
self.mma_inst_shape_mnk[1],
self.mma_inst_shape_mnk[2] * mma_inst_tile_k,
self.mma_inst_shape_mn[0],
self.mma_inst_shape_mn[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.mma_tiler_sfb = (
self.mma_inst_shape_mnk_sfb[0],
self.mma_inst_shape_mnk_sfb[1],
self.mma_inst_shape_mnk_sfb[2] * mma_inst_tile_k,
self.mma_inst_shape_mn_sfb[0],
self.mma_inst_shape_mn_sfb[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
@@ -314,9 +320,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
tiled_mma,
self.mma_tiler,
self.a_dtype,
self.a_major_mode,
self.b_dtype,
self.b_major_mode,
self.epi_tile,
self.c_dtype,
self.c_layout,
@@ -431,7 +435,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.sf_dtype,
self.sf_vec_size,
self.cta_group,
self.mma_inst_shape_mnk[:2],
self.mma_inst_shape_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
@@ -441,7 +445,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.sf_dtype,
self.sf_vec_size,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mnk_sfb[:2],
self.mma_inst_shape_mn_sfb,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
@@ -507,6 +511,31 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
internal_type=cutlass.Int16,
)
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
x = tma_tensor_sfb.stride[0][1]
y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4)
new_shape = (
(
tma_tensor_sfb.shape[0][0],
((2, 2), y)
),
tma_tensor_sfb.shape[1],
tma_tensor_sfb.shape[2]
)
# Use right multiplication for ScaledBasis (3 * x instead of x * 3)
x_times_3 = 3 * x
new_stride = (
(
tma_tensor_sfb.stride[0][0],
((x, x), x_times_3)
),
tma_tensor_sfb.stride[1],
tma_tensor_sfb.stride[2]
)
tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride)
tma_tensor_sfb = cute.make_tensor(tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout)
a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
@@ -628,7 +657,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
mSFA_mkl: cute.Tensor,
tma_atom_sfb: cute.CopyAtom,
mSFB_nkl: cute.Tensor,
tma_atom_c: Optional[cute.CopyAtom],
tma_atom_c: cute.CopyAtom,
mC_mnl: cute.Tensor,
cluster_layout_vmnk: cute.Layout,
cluster_layout_sfb_vmnk: cute.Layout,
@@ -636,7 +665,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None],
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
epi_tile: cute.Tile,
tile_sched_params: utils.PersistentTileSchedulerParams,
epilogue_op: cutlass.Constexpr,
@@ -684,9 +713,6 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
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
@@ -719,14 +745,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
# Tensor memory dealloc barrier init
if use_2cta_instrs:
if warp_idx == self.tma_warp_id:
num_tmem_dealloc_threads = 32
with cute.arch.elect_one():
cute.arch.mbarrier_init(
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
)
cute.arch.mbarrier_init_fence()
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
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,
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
@@ -790,7 +815,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
# (bN, bK, RestN, RestK, RestL)
gSFB_nkl = cute.local_tile(
mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
mSFB_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
@@ -894,9 +921,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
cute.arch.barrier(
barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta
)
self.cta_sync_barrier.arrive_and_wait()
#
# Specialized TMA load warp
@@ -915,7 +940,6 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
@@ -940,9 +964,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
tAgSFA_slice = tAgSFA[
(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
]
slice_n = mma_tile_coord_mnl[1]
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
slice_n = mma_tile_coord_mnl[1] // 2
# ((atom_v, rest_v), RestK)
tBgSFB_slice = tBgSFB[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
(None, slice_n, None, mma_tile_coord_mnl[2])
]
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
@@ -1017,21 +1045,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
#
# Bar sync for retrieve tensor memory ptr from shared mem
#
tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id))
cute.arch.barrier(
barrier_id=self.tmem_ptr_sync_bar_id,
number_of_threads=tmem_ptr_read_threads,
)
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor
#
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
# Make accumulator tmem tensor
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
@@ -1067,12 +1087,16 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
#
# Partition for S2T copy of SFA/SFB
#
tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = (
self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
)
tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = (
self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
)
(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t,
tCtSFA_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
(
tiled_copy_s2t_sfb,
tCsSFB_compact_s2t,
tCtSFB_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
#
# Persistent tile scheduling loop
@@ -1116,6 +1140,30 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
if is_leader_cta:
acc_pipeline.producer_acquire(acc_producer_state)
tCtSFB_mma = tCtSFB
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
# If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB)
offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
shifted_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ offset,
dtype=self.sf_dtype,
)
tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
# Move in increments of 64 columns of SFB
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
shifted_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ offset,
dtype=self.sf_dtype,
)
tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
#
# Reset the ACCUMULATE field for each tile
#
@@ -1170,7 +1218,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB[sf_kblock_coord].iterator,
tCtSFB_mma[sf_kblock_coord].iterator,
)
cute.gemm(
@@ -1220,30 +1268,17 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
#
# Alloc tensor memory buffer
#
if warp_idx == self.epilog_warp_id[0]:
cute.arch.alloc_tmem(
self.num_tmem_alloc_cols,
tmem_holding_buf,
is_two_cta=use_2cta_instrs,
)
tmem.allocate(self.num_tmem_alloc_cols)
#
# Bar sync for retrieve tensor memory ptr from shared memory
#
tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id))
cute.arch.barrier(
barrier_id=self.tmem_ptr_sync_bar_id,
number_of_threads=tmem_ptr_read_threads,
)
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator tensor
#
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
)
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
@@ -1251,20 +1286,24 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Partition for epilogue
#
epi_tidx = tidx
tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = (
self.epilog_tmem_copy_and_partition(
epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs
)
(
tiled_copy_t2r,
tTR_tAcc_base,
tTR_rAcc,
) = self.epilog_tmem_copy_and_partition(
epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs
)
tTR_rC = cute.make_fragment(tTR_rAcc.shape, self.c_dtype)
tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
tiled_copy_t2r, tTR_rC, epi_tidx, sC
)
tma_atom_c, bSG_sC, bSG_gC_partitioned = (
self.epilog_gmem_copy_and_partition(
epi_tidx, tma_atom_c, tCgC, epi_tile, sC
)
(
tma_atom_c,
bSG_sC,
bSG_gC_partitioned,
) = self.epilog_gmem_copy_and_partition(
epi_tidx, tma_atom_c, tCgC, epi_tile, sC
)
#
@@ -1283,7 +1322,6 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread,
32 * len(self.epilog_warp_id),
32 * len(self.epilog_warp_id),
)
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage,
@@ -1291,7 +1329,6 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
@@ -1360,11 +1397,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
cute.arch.ProxyKind.async_shared,
space=cute.arch.SharedSpace.shared_cta,
)
epilog_threads = 32 * len(self.epilog_warp_id)
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id,
number_of_threads=epilog_threads,
)
self.epilog_sync_barrier.arrive_and_wait()
#
# TMA store C to global memory
@@ -1378,10 +1411,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Fence and barrier to make sure shared memory store is visible to TMA store
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id,
number_of_threads=epilog_threads,
)
self.epilog_sync_barrier.arrive_and_wait()
#
# Async arrive accumulator buffer empty
@@ -1399,21 +1429,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
#
# Dealloc the tensor memory buffer
#
if warp_idx == self.epilog_warp_id[0]:
cute.arch.relinquish_tmem_alloc_permit(is_two_cta=use_2cta_instrs)
epilog_threads = 32 * len(self.epilog_warp_id)
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads
)
if warp_idx == self.epilog_warp_id[0]:
if use_2cta_instrs:
cute.arch.mbarrier_arrive(
tmem_dealloc_mbar_ptr, cta_rank_in_cluster ^ 1
)
cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0)
cute.arch.dealloc_tmem(
acc_tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=use_2cta_instrs
)
tmem.relinquish_alloc_permit()
self.epilog_sync_barrier.arrive_and_wait()
tmem.free(acc_tmem_ptr)
#
# Wait for C store complete
#
@@ -1520,7 +1538,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
# (T2R, T2R_M, T2R_N)
tTR_rAcc = cute.make_fragment(
tTR_rAcc = cute.make_rmem_tensor(
tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
)
return tiled_copy_t2r, tTR_tAcc, tTR_rAcc
@@ -1614,9 +1632,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
tiled_mma: cute.TiledMma,
mma_tiler_mnk: Tuple[int, int, int],
a_dtype: Type[cutlass.Numeric],
a_major_mode: tcgen05.OperandMajorMode,
b_dtype: Type[cutlass.Numeric],
b_major_mode: tcgen05.OperandMajorMode,
epi_tile: cute.Tile,
c_dtype: Type[cutlass.Numeric],
c_layout: utils.LayoutEnum,
@@ -1633,12 +1649,8 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
:type mma_tiler_mnk: tuple[int, int, int]
:param a_dtype: Data type of operand A.
:type a_dtype: type[cutlass.Numeric]
:param a_major_mode: Major mode of operand A.
:type a_major_mode: tcgen05.OperandMajorMode
:param b_dtype: Data type of operand B.
:type b_dtype: type[cutlass.Numeric]
:param b_major_mode: Major mode of operand B.
:type b_major_mode: tcgen05.OperandMajorMode
:param epi_tile: The epilogue tile shape.
:type epi_tile: cute.Tile
:param c_dtype: Data type of operand C (output).
@@ -1830,7 +1842,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
c_major: str,
) -> bool:
"""
Check if the dtypes and sf_vec_size are valid combinations
Check if layouts and dtypes are valid combinations
:param ab_dtype: The data type of the A and B operands
:type ab_dtype: Type[cutlass.Numeric]
@@ -1870,9 +1882,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
"""
is_valid = True
# Skip invalid mma tile shape
if not mma_tiler_mn[0] in [128, 256]:
if mma_tiler_mn[0] not in [128, 256]:
is_valid = False
if not mma_tiler_mn[1] in [128, 256]:
if mma_tiler_mn[1] not in [64, 128, 192, 256]:
is_valid = False
# Skip illegal cluster shape
if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0:
@@ -2088,7 +2100,7 @@ def run(
:return: Execution time of the GEMM kernel
:rtype: float
"""
print(f"Running Sm100 Persistent Dense BlockScaled GEMM test with:")
print("Running Sm100 Persistent 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}")
@@ -2143,21 +2155,21 @@ def run(
c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16
)
# Mark tensor to be byte aligned
# Mark tensor with element divisibility for 16B alignment
a_tensor.mark_compact_shape_dynamic(
mode=1 if a_major == "k" else 0,
stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0),
divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1,
divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16,
)
b_tensor.mark_compact_shape_dynamic(
mode=1 if b_major == "k" else 0,
stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0),
divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1,
divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16,
)
c_tensor.mark_compact_shape_dynamic(
mode=1 if c_major == "n" else 0,
stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0),
divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1,
divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16,
)
# Create scale factor tensor SFA/SFB
@@ -2374,6 +2386,7 @@ def run(
return exec_time # Return execution time in microseconds
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
examples/python/CuTeDSL/blackwell/dense_gemm.py
+539 -643
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py
+2218
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
+626 -863
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py
+514 -611
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/fmha.py
+582 -439
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/fmha_bwd.py
+3569
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py
+3222
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/grouped_gemm.py
+102 -167
View File
@@ -38,6 +38,7 @@ import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.torch as cutlass_torch
@@ -152,12 +153,24 @@ class GroupedGemmKernel:
self.threads_per_cta = 32 * len(
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# Set barrier id for cta sync, epilog sync, tmem ptr sync and tensormap update sync
self.cta_sync_bar_id = 0
self.epilog_sync_bar_id = 1
self.tmem_ptr_sync_bar_id = 2
# Barrier ID used by MMA/TMA warps to signal A/B tensormap initialization completion
self.tensormap_ab_init_bar_id = 4
# Set barrier for cta sync, epilog sync, tmem ptr sync and tensormap update sync
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
)
# Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion
self.tensormap_ab_init_barrier = pipeline.NamedBarrier(
barrier_id=4,
num_threads=32 * (len(self.epilog_warp_id) + 1),
)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
self.num_tma_load_bytes = 0
@@ -251,14 +264,6 @@ class GroupedGemmKernel:
self.num_epi_stage,
)
tensor_smem_bytes = self._get_tensor_smem_bytes(
self.a_smem_layout_staged,
self.a_dtype,
self.b_smem_layout_staged,
self.b_dtype,
self.epi_smem_layout_staged,
self.c_dtype,
)
mbar_smem_bytes = self._get_mbar_smem_bytes(
num_acc_stage=self.num_acc_stage,
num_ab_stage=self.num_ab_stage,
@@ -390,15 +395,12 @@ class GroupedGemmKernel:
# Setup TMA store for C
tma_atom_c = None
tma_tensor_c = None
c_cta_v_layout = cute.composition(
cute.make_identity_layout(initial_c.shape), self.epi_tile
)
epi_smem_layout = cute.slice_(self.epi_smem_layout_staged, (None, None, 0))
tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(),
initial_c,
epi_smem_layout,
c_cta_v_layout,
self.epi_tile,
)
self.tile_sched_params, grid = self._compute_grid(
@@ -558,8 +560,6 @@ class GroupedGemmKernel:
ab_empty_mbar_ptr = storage.ab_empty_mbar_ptr.data_ptr()
acc_full_mbar_ptr = storage.acc_full_mbar_ptr.data_ptr()
acc_empty_mbar_ptr = storage.acc_empty_mbar_ptr.data_ptr()
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr
tmem_holding_buf = storage.tmem_holding_buf
# init barrier for loading A, B with TMA
if warp_idx == self.epilog_warp_id[0]:
@@ -579,13 +579,13 @@ class GroupedGemmKernel:
acc_empty_mbar_ptr + acc_stage, 8 if use_2cta_instrs else 4
)
# Tensor memory dealloc barrier init
if use_2cta_instrs:
if warp_idx == self.tma_warp_id:
num_tmem_dealloc_threads = 32
with cute.arch.elect_one():
cute.arch.mbarrier_init(
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
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,
)
cute.arch.mbarrier_init_fence()
# Cluster arrive after barrier init
@@ -721,9 +721,7 @@ class GroupedGemmKernel:
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
cute.arch.barrier(
barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta
)
self.cta_sync_barrier.arrive_and_wait()
#
# Get tensormap buffer address
@@ -826,10 +824,7 @@ class GroupedGemmKernel:
# wait tensormap initialization complete before update
if tensormap_init_done == False:
if cutlass.const_expr(self.delegate_tensormap_ab_init):
cute.arch.barrier(
barrier_id=self.tensormap_ab_init_bar_id,
number_of_threads=64,
)
self.tensormap_ab_init_barrier.arrive_and_wait()
tensormap_manager.fence_tensormap_initialization()
tensormap_init_done = True
@@ -951,33 +946,13 @@ class GroupedGemmKernel:
# Specialized MMA warp
#
if warp_idx == self.mma_warp_id:
# initialize tensormap A, B for TMA warp
if cutlass.const_expr(self.delegate_tensormap_ab_init):
tensormap_manager.init_tensormap_from_atom(
tma_atom_a, tensormap_a_init_ptr, self.mma_warp_id
)
tensormap_manager.init_tensormap_from_atom(
tma_atom_b, tensormap_b_init_ptr, self.mma_warp_id
)
# signal tensormap initialization has finished
cute.arch.barrier(
barrier_id=self.tensormap_ab_init_bar_id, number_of_threads=64
)
# Bar sync for retrieve tmem ptr from shared mem
tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id))
cute.arch.barrier(
barrier_id=self.tmem_ptr_sync_bar_id,
number_of_threads=tmem_ptr_read_threads,
)
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator tensor
#
tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
)
tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
@@ -1019,22 +994,22 @@ class GroupedGemmKernel:
# Peek (try_wait) AB buffer full for k_tile = 0
mma_rd_k_tile = cutlass.Int32(0)
smem_rd_buffer = (num_prev_k_blk + mma_rd_k_tile) % self.num_ab_stage
need_check_rd_buffer_full = (
mma_rd_k_tile < cur_k_tile_cnt and is_leader_cta
)
mma_rd_ab_full_phase = (
(num_prev_k_blk + mma_rd_k_tile) // self.num_ab_stage % 2
)
peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
need_check_rd_buffer_full,
ab_full_mbar_ptr + smem_rd_buffer,
mma_rd_ab_full_phase,
)
#
# Wait for accumulator buffer empty
#
if is_leader_cta:
need_check_rd_buffer_full = (
mma_rd_k_tile < cur_k_tile_cnt and is_leader_cta
)
mma_rd_ab_full_phase = (
(num_prev_k_blk + mma_rd_k_tile) // self.num_ab_stage % 2
)
peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
need_check_rd_buffer_full,
ab_full_mbar_ptr + smem_rd_buffer,
mma_rd_ab_full_phase,
)
#
# Wait for accumulator buffer empty
#
acc_empty_phase = (
tile_sched.num_tiles_executed // self.num_acc_stage % 2 ^ 1
)
@@ -1042,25 +1017,24 @@ class GroupedGemmKernel:
acc_empty_mbar_ptr + acc_buf_idx, acc_empty_phase
)
#
# Reset the ACCUMULATE field for each tile
#
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
#
# Reset the ACCUMULATE field for each tile
#
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
#
# Mma mainloop
#
for k_tile in range(cur_k_tile_cnt):
mma_rd_k_tile_next = cutlass.Int32(k_tile + 1)
smem_rd_buffer_next = (
num_prev_k_blk + mma_rd_k_tile_next
) % self.num_ab_stage
mma_rd_ab_full_phase_next = (
mma_rd_ab_full_phase ^ 1
if smem_rd_buffer_next == 0
else mma_rd_ab_full_phase
)
if is_leader_cta:
#
# Mma mainloop
#
for k_tile in range(cur_k_tile_cnt):
mma_rd_k_tile_next = cutlass.Int32(k_tile + 1)
smem_rd_buffer_next = (
num_prev_k_blk + mma_rd_k_tile_next
) % self.num_ab_stage
mma_rd_ab_full_phase_next = (
mma_rd_ab_full_phase ^ 1
if smem_rd_buffer_next == 0
else mma_rd_ab_full_phase
)
# Wait for AB buffer full
if peek_ab_full_status == 0:
cute.arch.mbarrier_wait(
@@ -1090,25 +1064,24 @@ class GroupedGemmKernel:
self.cta_group,
)
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
need_check_rd_buffer_full = (
mma_rd_k_tile_next < cur_k_tile_cnt and is_leader_cta
)
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
need_check_rd_buffer_full = (
mma_rd_k_tile_next < cur_k_tile_cnt and is_leader_cta
)
peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
need_check_rd_buffer_full,
ab_full_mbar_ptr + smem_rd_buffer_next,
mma_rd_ab_full_phase_next,
)
peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
need_check_rd_buffer_full,
ab_full_mbar_ptr + smem_rd_buffer_next,
mma_rd_ab_full_phase_next,
)
mma_rd_k_tile = mma_rd_k_tile_next
smem_rd_buffer = smem_rd_buffer_next
mma_rd_ab_full_phase = mma_rd_ab_full_phase_next
mma_rd_k_tile = mma_rd_k_tile_next
smem_rd_buffer = smem_rd_buffer_next
mma_rd_ab_full_phase = mma_rd_ab_full_phase_next
#
# Async arrive accumulator buffer full
#
if is_leader_cta:
#
# Async arrive accumulator buffer full
#
with cute.arch.elect_one():
tcgen05.commit(
acc_full_mbar_ptr + acc_buf_idx,
@@ -1126,6 +1099,16 @@ class GroupedGemmKernel:
# Specialized epilogue warps
#
if warp_idx < self.mma_warp_id:
# initialize tensormap A, B for TMA warp
if cutlass.const_expr(self.delegate_tensormap_ab_init):
tensormap_manager.init_tensormap_from_atom(
tma_atom_a, tensormap_a_init_ptr, self.epilog_warp_id[0]
)
tensormap_manager.init_tensormap_from_atom(
tma_atom_b, tensormap_b_init_ptr, self.epilog_warp_id[0]
)
# signal tensormap initialization has finished
self.tensormap_ab_init_barrier.arrive_and_wait()
# initialize tensorap for C
tensormap_manager.init_tensormap_from_atom(
tma_atom_c,
@@ -1133,30 +1116,17 @@ class GroupedGemmKernel:
self.epilog_warp_id[0],
)
# Alloc tensor memory buffer
if warp_idx == self.epilog_warp_id[0]:
cute.arch.alloc_tmem(
self.num_tmem_alloc_cols,
tmem_holding_buf,
is_two_cta=use_2cta_instrs,
)
tmem.allocate(self.num_tmem_alloc_cols)
#
# Bar sync for retrieve tensor memory ptr from shared memory
#
tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id))
cute.arch.barrier(
barrier_id=self.tmem_ptr_sync_bar_id,
number_of_threads=tmem_ptr_read_threads,
)
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator tensor
#
tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
)
tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
@@ -1172,7 +1142,7 @@ class GroupedGemmKernel:
epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs
)
tTR_rC = cute.make_fragment(tTR_rAcc.shape, self.c_dtype)
tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
tiled_copy_t2r, tTR_rC, epi_tidx, sC
)
@@ -1303,11 +1273,7 @@ class GroupedGemmKernel:
cute.arch.ProxyKind.async_shared,
space=cute.arch.SharedSpace.shared_cta,
)
epilog_threads = 32 * len(self.epilog_warp_id)
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id,
number_of_threads=epilog_threads,
)
self.epilog_sync_barrier.arrive_and_wait()
#
# store C to global memory with TMA
#
@@ -1325,10 +1291,7 @@ class GroupedGemmKernel:
cute.arch.cp_async_bulk_wait_group(
self.num_epi_stage - 1, read=True
)
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id,
number_of_threads=epilog_threads,
)
self.epilog_sync_barrier.arrive_and_wait()
#
# Async arrive accumulator buffer empty
#
@@ -1348,21 +1311,9 @@ class GroupedGemmKernel:
#
# Dealloc the tensor memory buffer
#
if warp_idx == self.epilog_warp_id[0]:
cute.arch.relinquish_tmem_alloc_permit(is_two_cta=use_2cta_instrs)
epilog_threads = 32 * len(self.epilog_warp_id)
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads
)
if warp_idx == self.epilog_warp_id[0]:
if use_2cta_instrs:
cute.arch.mbarrier_arrive(
tmem_dealloc_mbar_ptr, cta_rank_in_cluster ^ 1
)
cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0)
cute.arch.dealloc_tmem(
tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=use_2cta_instrs
)
tmem.relinquish_alloc_permit()
self.epilog_sync_barrier.arrive_and_wait()
tmem.free(tmem_ptr)
#
# Wait a/b buffer empty
@@ -1417,7 +1368,7 @@ class GroupedGemmKernel:
)
strides_tensor_gmem = strides_abc[(group_idx, tensor_index, None)]
strides_tensor_reg = cute.make_fragment(
strides_tensor_reg = cute.make_rmem_tensor(
cute.make_layout(2),
strides_abc.element_type,
)
@@ -1507,7 +1458,7 @@ class GroupedGemmKernel:
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
# (T2R, T2R_M, T2R_N)
tTR_rAcc = cute.make_fragment(
tTR_rAcc = cute.make_rmem_tensor(
tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
)
return tiled_copy_t2r, tTR_tAcc, tTR_rAcc
@@ -1761,23 +1712,6 @@ class GroupedGemmKernel:
else:
raise ValueError(f"Invalid tensormap update mode: {tensormap_update_mode}")
@staticmethod
def _get_tensor_smem_bytes(
a_smem_layout_staged: cute.Layout,
a_dtype: Type[cutlass.Numeric],
b_smem_layout_staged: cute.Layout,
b_dtype: Type[cutlass.Numeric],
epi_smem_layout_staged: cute.Layout,
c_dtype: Type[cutlass.Numeric],
) -> int:
"""Compute the total SMEM consumption for tensor A, B and C."""
ab_bytes = cute.size_in_bytes(
a_dtype, a_smem_layout_staged
) + cute.size_in_bytes(b_dtype, b_smem_layout_staged)
epi_bytes = cute.size_in_bytes(c_dtype, epi_smem_layout_staged)
return ab_bytes + epi_bytes
@staticmethod
def _compute_num_tmem_alloc_cols(
tiled_mma: cute.TiledMma,
@@ -1821,7 +1755,7 @@ def create_tensor_and_stride(
is_dynamic_layout: bool = True,
torch_tensor_cpu: torch.Tensor = None,
) -> tuple[int, torch.Tensor, cute.Tensor, torch.Tensor, tuple[int, int]]:
"""Create a GPU tensor from scratch or based on an existing CPU tensor.
"""Create GPU tensor from either a new or existing CPU tensor.
:param torch_tensor_cpu: Optional existing CPU tensor to reuse. If None, creates a new one.
:type torch_tensor_cpu: torch.Tensor, optional
@@ -1967,7 +1901,7 @@ def run(
:return: Execution time of the GEMM kernel in microseconds
:rtype: float
"""
print(f"Running Blackwell Grouped GEMM test with:")
print("Running Blackwell Grouped GEMM test with:")
print(f"{num_groups} groups")
for i, (m, n, k, l) in enumerate(problem_sizes_mnkl):
print(f"Group {i}: {m}x{n}x{k}x{l}")
@@ -2102,6 +2036,7 @@ def run(
is_dynamic_layout=False,
assumed_align=16,
)
# layout (num_groups, 3, 2):(6, 2, 1)
tensor_of_strides_abc, tensor_of_strides_abc_torch = cutlass_torch.cute_tensor_like(
torch.tensor(strides_abc, dtype=torch.int32),
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py
+71 -92
View File
@@ -86,9 +86,9 @@ class SSDKernel:
cutlass.BFloat16,
}, "Do not support other I/O types."
assert acc_dtype in {cutlass.Float32}, "Do not support other ACC types."
assert cumsum_delta_dtype in {
cutlass.Float32
}, "Do not support other cumsum types."
assert cumsum_delta_dtype in {cutlass.Float32}, (
"Do not support other cumsum types."
)
assert not (not has_d and d_has_hdim), "D cannot have Hdim if has_d is False"
# Hardcode default setting
@@ -129,10 +129,18 @@ class SSDKernel:
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
# Named barriers
self.pre_inter_sync_bar_id = 1
self.epilog_sync_bar_id = 2
self.pre_intra_sync_bar_id = 3
self.tmem_dealloc_sync_bar_id = 4
self.pre_inter_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=len(self.pre_inter_warp_id) * 32,
)
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
num_threads=len(self.epilog_warp_id) * 32,
)
self.tmem_dealloc_sync_barrier = pipeline.NamedBarrier(
barrier_id=3,
num_threads=self.threads_per_cta,
)
# Number of registers used by each warp
self.num_regs_uniform_warps = 24
@@ -467,15 +475,12 @@ class SSDKernel:
)
# TMA store for y
y_cta_v_layout = cute.composition(
cute.make_identity_layout(y.shape), self.epi_tile
)
y_smem_layout = cute.slice_(self.y_smem_layout, (None, None, 0))
tma_atom_y, tma_tensor_y = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(),
y,
y_smem_layout,
y_cta_v_layout,
self.epi_tile,
)
# TMA store for fstate(p)
@@ -512,7 +517,9 @@ class SSDKernel:
d_empty: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
# Intra1 acc stage barriers
intra1_acc_full: cute.struct.MemRange[cutlass.Int64, self.intra1_acc_stages] # type: ignore
intra1_acc_empty: cute.struct.MemRange[cutlass.Int64, self.intra1_acc_stages] # type: ignore
intra1_acc_empty: cute.struct.MemRange[
cutlass.Int64, self.intra1_acc_stages
] # type: ignore
# Internal stage barriers
intra2_q_full: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
intra2_q_empty: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
@@ -811,23 +818,22 @@ class SSDKernel:
if cute.size(self.cluster_shape_mnk) > 1:
cute.arch.cluster_wait()
# Alloc tmem buffer
if warp_idx == self.epilog_warp_id[0]:
cute.arch.alloc_tmem(
self.num_tmem_cols_total,
smem_storage.tmem_holding_buf,
is_two_cta=self.use_2cta_instrs,
)
tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=0,
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
smem_storage.tmem_holding_buf,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
)
tmem.allocate(self.num_tmem_cols_total)
# Bar sync before retrieving tmem ptr from shared mem
cute.arch.barrier()
# Barrier before retrieve tensor memory ptr from shared memory
tmem.wait_for_alloc()
# Retrieve tmem ptr
tmem_ptr_base = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=smem_storage.tmem_holding_buf,
)
tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype)
# Specialized TMA load Delta/CumsumDelta/X warp
if warp_idx == self.tma_deltas_x_d_warp_id:
@@ -1579,7 +1585,7 @@ class SSDKernel:
) = self.pre_inter_tmem_load_and_partition_p(local_tidx, tInter1, smem_pt)
# Make fragment for register to hold P after post-processing (in acc dtype)
tState = cute.make_fragment(tTR_rP.shape, self.acc_dtype)
tState = cute.make_rmem_tensor(tTR_rP.shape, self.acc_dtype)
# Make tiledCopy and partition smem/register tensor for smem store INTER2_P
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N)
@@ -1621,7 +1627,7 @@ class SSDKernel:
tma_p_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.internal_stages,
producer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
),
)
@@ -1808,10 +1814,7 @@ class SSDKernel:
cute.arch.ProxyKind.async_shared,
space=cute.arch.SharedSpace.shared_cta,
)
cute.arch.barrier(
barrier_id=self.pre_inter_sync_bar_id,
number_of_threads=len(self.pre_inter_warp_id) * 32,
)
self.pre_inter_sync_barrier.arrive_and_wait()
if local_warp_idx == 0:
# TMA store P
@@ -1824,10 +1827,7 @@ class SSDKernel:
tma_p_pipeline.producer_commit()
tma_p_pipeline.producer_acquire()
cute.arch.barrier(
barrier_id=self.pre_inter_sync_bar_id,
number_of_threads=len(self.pre_inter_warp_id) * 32,
)
self.pre_inter_sync_barrier.arrive_and_wait()
tma_p_pipeline.producer_tail()
# Advance to next tile
@@ -2085,7 +2085,7 @@ class SSDKernel:
local_tidx, smem_y, tiled_t2r_inter2
)
tRS_rCompute = cute.make_fragment(tRS_rY.shape, self.acc_dtype)
tRS_rCompute = cute.make_rmem_tensor(tRS_rY.shape, self.acc_dtype)
tiled_s2r_x = None
tSR_sX = None
@@ -2128,7 +2128,7 @@ class SSDKernel:
tma_y_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.output_stages,
producer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
),
)
@@ -2328,10 +2328,7 @@ class SSDKernel:
space=cute.arch.SharedSpace.shared_cta,
)
# Sync before TMA store
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id,
number_of_threads=len(self.epilog_warp_id) * 32,
)
self.epilog_sync_barrier.arrive_and_wait()
# Async arrive Delta/INTRA2_ACC/INTER2_ACC buffer empty
if (
@@ -2366,10 +2363,7 @@ class SSDKernel:
# Wait for TMA store
tma_y_pipeline.producer_acquire()
# Sync before smem store
cute.arch.barrier(
barrier_id=self.epilog_sync_bar_id,
number_of_threads=len(self.epilog_warp_id) * 32,
)
self.epilog_sync_barrier.arrive_and_wait()
# Advance deltas/intra2_acc/inter2_acc consumer states
deltas_consumer_state.advance()
@@ -2406,22 +2400,12 @@ class SSDKernel:
# Producer tail for TMA store Y
tma_y_pipeline.producer_tail()
# Release tensor memory allocation lock
tmem.relinquish_alloc_permit()
# Sync before deallocating tmem
self.tmem_dealloc_sync_barrier.arrive_and_wait()
# Dealloc tmem buffer
if warp_idx == self.epilog_warp_id[0]:
cute.arch.barrier(
barrier_id=self.tmem_dealloc_sync_bar_id,
number_of_threads=self.threads_per_cta,
)
cute.arch.dealloc_tmem(
tmem_ptr_base,
self.num_tmem_cols_total,
is_two_cta=self.use_2cta_instrs,
)
else:
cute.arch.barrier_arrive(
barrier_id=self.tmem_dealloc_sync_bar_id,
number_of_threads=self.threads_per_cta,
)
tmem.free(tmem_ptr_base)
return
@@ -2597,7 +2581,7 @@ class SSDKernel:
len([self.mma_intra_warp_id, self.mma_inter_warp_id]),
)
x_consumer_group_async = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
)
return pipeline.PipelineTmaMultiConsumersAsync.create(
num_stages=self.input_stages,
@@ -2616,7 +2600,7 @@ class SSDKernel:
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
)
b_consumer_group_async = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
)
return pipeline.PipelineTmaMultiConsumersAsync.create(
num_stages=self.input_stages,
@@ -2651,9 +2635,6 @@ class SSDKernel:
len(
[*self.pre_inter_warp_id, *self.pre_intra_warp_id, *self.epilog_warp_id]
),
len(
[*self.pre_inter_warp_id, *self.pre_intra_warp_id, *self.epilog_warp_id]
),
)
return pipeline.PipelineTmaAsync.create(
@@ -2672,9 +2653,7 @@ class SSDKernel:
pipeline.Agent.Thread, len([self.tma_deltas_x_d_warp_id])
)
d_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread,
len(self.epilog_warp_id),
len(self.epilog_warp_id),
pipeline.Agent.Thread, len(self.epilog_warp_id)
)
return pipeline.PipelineTmaAsync.create(
@@ -2690,7 +2669,7 @@ class SSDKernel:
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
)
intra1_acc_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.pre_intra_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.pre_intra_warp_id)
)
return pipeline.PipelineUmmaAsync.create(
num_stages=self.intra1_acc_stages,
@@ -2701,7 +2680,7 @@ class SSDKernel:
def make_and_init_intra2_q_pipeline(self, intra2_q_full_mbar_ptr):
intra2_q_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.pre_intra_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.pre_intra_warp_id)
)
intra2_q_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
@@ -2718,7 +2697,7 @@ class SSDKernel:
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
)
intra2_acc_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
)
return pipeline.PipelineUmmaAsync.create(
num_stages=self.internal_stages,
@@ -2729,7 +2708,7 @@ class SSDKernel:
def make_and_init_inter1_b_pipeline(self, inter1_b_full_mbar_ptr):
inter1_b_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
)
inter1_b_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
@@ -2746,7 +2725,7 @@ class SSDKernel:
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
)
inter1_acc_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
)
return pipeline.PipelineUmmaAsync.create(
num_stages=self.internal_stages,
@@ -2757,7 +2736,7 @@ class SSDKernel:
def make_and_init_inter2_p_pipeline(self, inter2_p_full_mbar_ptr):
inter2_p_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
)
inter2_p_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
@@ -2774,7 +2753,7 @@ class SSDKernel:
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
)
inter2_acc_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id), 128
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
)
return pipeline.PipelineUmmaAsync.create(
num_stages=self.internal_stages,
@@ -3035,7 +3014,7 @@ class SSDKernel:
# Partition tmem/register tensor for tensor memory store INTRA2_Q
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, ...)
tRT_rQ = cute.make_fragment(
tRT_rQ = cute.make_rmem_tensor(
cute.slice_(thr_r2t_q.partition_S(tCrQ).shape, (None, None, None, None, 0)),
dtype,
)
@@ -3049,10 +3028,10 @@ class SSDKernel:
self, tTR_rQ, tQrDeltaA_Row, tQrDeltaA_Col, tQrDelta, tCoord, tRT_rQ
):
# Make tmp acc type fragments
tCrDeltaA_Row = cute.make_fragment(tQrDeltaA_Row.shape, self.acc_dtype)
tCrDeltaA_Col = cute.make_fragment(tQrDeltaA_Col.shape, self.acc_dtype)
tCrDelta = cute.make_fragment(tQrDelta.shape, self.acc_dtype)
tCompute = cute.make_fragment(tRT_rQ.shape, self.acc_dtype)
tCrDeltaA_Row = cute.make_rmem_tensor(tQrDeltaA_Row.shape, self.acc_dtype)
tCrDeltaA_Col = cute.make_rmem_tensor(tQrDeltaA_Col.shape, self.acc_dtype)
tCrDelta = cute.make_rmem_tensor(tQrDelta.shape, self.acc_dtype)
tCompute = cute.make_rmem_tensor(tRT_rQ.shape, self.acc_dtype)
# Combine tTR_rQ/tCrDeltaA_Row/tCrDeltaA_Col/tCrDelta
tCrDeltaA_Row.store(tQrDeltaA_Row.load().to(self.acc_dtype))
@@ -3127,7 +3106,7 @@ class SSDKernel:
tBsB_s2r = thr_s2r_b.partition_S(smem_bt)
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N)
tBrB_s2r = cute.make_fragment(
tBrB_s2r = cute.make_rmem_tensor(
cute.slice_(tBsB_s2r.shape, (None, None, None, 0)),
dtype,
)
@@ -3167,7 +3146,7 @@ class SSDKernel:
# Make register fragments for smem load/store of Delta/DeltaA
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N)
tBrDelta_s2r = cute.make_fragment(tBsDelta_s2r[smem_tile_coord].shape, dtype)
tBrDelta_s2r = cute.make_rmem_tensor(tBsDelta_s2r[smem_tile_coord].shape, dtype)
return s2r_atom_delta, tBsDelta_s2r, tBrDelta_s2r
def pre_inter_tmem_load_and_partition_p(self, local_tidx, tInter1, smem_pt):
@@ -3195,7 +3174,7 @@ class SSDKernel:
tTR_s = thr_t2r.partition_D(smem_tensor)
# Make register fragments for tmem load INTER1_ACC
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
tTR_r = cute.make_fragment(
tTR_r = cute.make_rmem_tensor(
tTR_s.shape,
dtype,
)
@@ -3213,7 +3192,7 @@ class SSDKernel:
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N, INTERNAL_STAGE)
tRS_sP = thr_r2s_p.partition_D(smem_pt)
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N)
tRS_rP = cute.make_fragment(
tRS_rP = cute.make_rmem_tensor(
cute.slice_(tRS_sP.shape, (None, None, None, 0)), self.io_dtype
)
return tiled_r2s_p, tRS_rP, tRS_sP
@@ -3239,10 +3218,10 @@ class SSDKernel:
def pre_inter_scale_bt_with_delta(
self, tBrB_s2r, tBrDelta_s2r, tBrDeltaA_s2r, last_column
):
tCompute = cute.make_fragment(tBrB_s2r.shape, self.acc_dtype)
tBrB_Compute = cute.make_fragment(tBrB_s2r.shape, self.acc_dtype)
tBrDelta_Compute = cute.make_fragment(tBrDelta_s2r.shape, self.acc_dtype)
tBrDeltaA_Compute = cute.make_fragment(tBrDeltaA_s2r.shape, self.acc_dtype)
tCompute = cute.make_rmem_tensor(tBrB_s2r.shape, self.acc_dtype)
tBrB_Compute = cute.make_rmem_tensor(tBrB_s2r.shape, self.acc_dtype)
tBrDelta_Compute = cute.make_rmem_tensor(tBrDelta_s2r.shape, self.acc_dtype)
tBrDeltaA_Compute = cute.make_rmem_tensor(tBrDeltaA_s2r.shape, self.acc_dtype)
tBrB_Compute.store(tBrB_s2r.load().to(self.acc_dtype))
tBrDelta_Compute.store(tBrDelta_s2r.load().to(self.acc_dtype))
@@ -3323,7 +3302,7 @@ class SSDKernel:
# (R2S_ATOM, R2S_M, R2S_N, EPI_M, EPI_N, INPUT_STAGES)
tSR_sX = thr_s2r_x.partition_S(cute.flat_divide(smem_xt, epi_tile))
# (R2S_ATOM, R2S_M, R2S_N)
tSR_rX = cute.make_fragment(
tSR_rX = cute.make_rmem_tensor(
cute.slice_(tSR_sX.shape, (None, None, None, 0, 0, 0)), dtype
)
return tiled_s2r_x, tSR_sX, tSR_rX
@@ -3360,7 +3339,7 @@ def run(
has_d = fuse_scale_d != "none"
d_has_hdim = fuse_scale_d == "vector"
print(f"Running B100 Mamba2 SSD with:")
print("Running B100 Mamba2 SSD with:")
print(f"GBEHCDLN: {gbehcdln}")
print(
f"Input/Output dtype: {io_dtype}, Intermediate delta dtype: {cumsum_delta_dtype}, Acc dtype: {acc_dtype}"
@@ -3405,7 +3384,7 @@ def run(
# Build torch_dtype torch tensor
torch_dtype = cutlass_torch.dtype(dtype)
dst_tensor = ref_tensor.to(torch_dtype).cuda()
dst_tensor = ref_tensor.to(dtype=torch_dtype).cuda()
cute_tensor = from_dlpack(dst_tensor, assumed_align=16)
for mode in dynamic_modes:
cute_tensor = cute_tensor.mark_compact_shape_dynamic(
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd_reference.py
+2 -2
View File
@@ -212,7 +212,7 @@ def analyze_relative_diffs(actual, expected):
)
# Print max relative difference info
print(f"Maximum relative difference:")
print("Maximum relative difference:")
print(f"Position: {max_rel_diff_pos}")
print(f"Value: {max_rel_diff:.6e}")
print(f"Actual value: {actual.flatten()[max_rel_diff_pos]}")
@@ -236,7 +236,7 @@ def analyze_relative_diffs(actual, expected):
print(f"Elements with rtol <= {rtol:.0e}: {count} ({percentage:.2f}%)")
else:
print(
f"Elements with {rtol_levels[i-1]:.0e} < rtol <= {rtol:.0e}: {count} ({percentage:.2f}%)"
f"Elements with {rtol_levels[i - 1]:.0e} < rtol <= {rtol:.0e}: {count} ({percentage:.2f}%)"
)
# Print elements exceeding the largest rtol
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd_tile_scheduler.py
+1 -2
View File
@@ -29,7 +29,6 @@
from typing import Tuple
from cutlass.cutlass_dsl import (
Boolean,
Integer,
Int32,
min,
@@ -121,8 +120,8 @@ class Mamba2SSDTileScheduler:
)
# called by host
@dsl_user_op
@staticmethod
@dsl_user_op
def create(
params: Mamba2SSDTileSchedulerParams,
block_idx: Tuple[Integer, Integer, Integer],
examples/python/CuTeDSL/blackwell/mixed_input_gemm.py
+3113
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/mla.py
+5193
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/programmatic_dependent_launch.py
+381
View File
@@ -0,0 +1,381 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
from cutlass.cute.runtime import from_dlpack
def supports_pdl():
return torch.cuda.get_device_capability()[0] >= 9
"""
This example demonstrates the use of Programmatic Dependent Launch (PDL) using
CuTe DSL.
PDL is a mechanism which allows for overlapping execution of back-to-back kernels
within the same stream.
For example, consider the following two elementwise add operations, where the second
operation's first operand is the result of the first operation. While performing
``w = u + v`` we will load u and v, add them, and then store the result. Once we
have finished loading data, we are no longer utilizing the read bandwidth.
To effectively utilize the read bandwidth, we can start loading ``x``
immediately upon finishing reading. This is what PDL enables us to do.
.. code-block:: bash
w = u + v
y = w + x
To enable PDL, we need to do two things:
1. Insert the ``griddepcontrol.launch_dependents`` and ``griddepcontrol.wait`` instructions in the kernel.
2. Set the PDL launch attribute when launching the kernel.
The ``griddepcontrol.launch_dependents`` and ``griddepcontrol.wait``
instructions enable fine-grained control over kernel execution in PDL.
Once all thread blocks execute the ``griddepcontrol.launch_dependents``
instruction, the dependent kernels can opportunistically be early-launched.
``griddepcontrol.wait`` functions as a synchronization barrier - any warp
executing this instruction will block until the previous kernel finishes
execution. This allows precise control over data dependencies between kernels.
The following diagram shows the overlapping execution of two dependent kernels.
We call the instructions before ``griddepcontrol.wait`` as prologue (``P0``),
which may include barrier initialization and loading of independent data, etc.
We call the instructions after ``griddepcontrol.launch_dependents`` as epilogue
(``P2``), which may include math operations, data stores, etc. PDL enables
these prologue and epilogue phases to execute concurrently across dependent
kernels, improving GPU resource utilization. This is particularly beneficial
when prologue and epilogue are bound by different resources (e.g., memory
bandwidth vs compute throughput).
# P0: Prologue, P1: Main compute, P2: Epilogue
P0 P1 P2
K1: |=====|+++++|-----|
<-----> K2 can start early
(K1's P2 overlaps with K2's P0)
P0 P1 P2
K2: |=====| |+++++|-----|
^
|
wait for K1 to complete
Time ------------------------------------------------------>
We could run this example with and without PDL:
.. code-block:: bash
python examples/blackwell/programmatic_dependent_launch.py --benchmark
python examples/blackwell/programmatic_dependent_launch.py --benchmark --use_pdl
From the benchmark results, you can see some speedups for the PDL version in most cases, benefiting from
the overlapping execution of consecutive kernels. Moreover, you can use nsys to observe the overlapping execution.
.. code-block:: bash
nsys profile python examples/blackwell/programmatic_dependent_launch.py --benchmark --use_pdl
Note, PDL feature is supported on Hopper and later GPUs.
See [the programming guide](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization)
and the [PTX documentation](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-griddepcontrol)
for more details.
"""
@cute.kernel
def elementwise_add_kernel(
gA: cute.Tensor,
gB: cute.Tensor,
gC: cute.Tensor,
cC: cute.Tensor, # coordinate tensor
shape: cute.Shape,
thr_layout: cute.Layout,
val_layout: cute.Layout,
use_pdl: cutlass.Constexpr = True,
is_first_kernel: cutlass.Constexpr = True,
):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
blk_coord = ((None, None), bidx)
blkA = gA[blk_coord] # (TileM,TileN)
blkB = gB[blk_coord] # (TileM,TileN)
blkC = gC[blk_coord] # (TileM,TileN)
blkCrd = cC[blk_coord] # (TileM, TileN)
copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type)
copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gC.element_type)
tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
tiled_copy_B = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
tiled_copy_C = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout)
thr_copy_A = tiled_copy_A.get_slice(tidx)
thr_copy_B = tiled_copy_B.get_slice(tidx)
thr_copy_C = tiled_copy_C.get_slice(tidx)
thrA = thr_copy_A.partition_S(blkA)
thrB = thr_copy_B.partition_S(blkB)
thrC = thr_copy_C.partition_S(blkC)
frgA = cute.make_fragment_like(thrA)
frgB = cute.make_fragment_like(thrB)
frgC = cute.make_fragment_like(thrC)
thrCrd = thr_copy_C.partition_S(blkCrd)
frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
for i in range(cute.size(frgPred)):
val = cute.elem_less(thrCrd[i], shape)
frgPred[i] = val
# Note: when not using cuda-graph, the kernel execution may be blocked by the host overhead.
# In this case we won't see overlapping even when pdl is enabled.
# In this example, we add a loop (10 times) for all the copy and compute operations in the following code
# to make kernel running longer and make pdl benefits observable for both cuda-graph enabled and disabled cases.
if not use_pdl:
for _ in range(10):
cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)
cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)
else:
if is_first_kernel:
for _ in range(10):
cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)
cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)
# Here we add the launch dependents instruction for the first kernel as a hint to the runtime to early-launch
# the next kernel. If the next kernel becomes concurrent, we will have overlap where the second kernel
# can start reading x to ensure an E2E speedup. Note the placement of launch dependents has no implication
# on correctness, only performance.
cute.arch.griddepcontrol_launch_dependents()
else:
# In this example, the second kernel's second operand ``gB`` has no dependencies, its loading can overlap
# with the computation of ``gC`` from the first kernel.
for _ in range(10):
cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)
# For the second kernel, its first operand ``gA`` is dependent on the previous kernel, we must call
# griddepcontrol.wait to assure correctness. This instruction will block until the prior kernels finishes
# and its memory operations are visible. Since gA is written by the prior kernel, this will block until gA
# is visible to our kernel. Without it, we would have undefined behavior due to a race condition.
cute.arch.griddepcontrol_wait()
for _ in range(10):
cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)
for _ in range(10):
result = frgA.load() + frgB.load()
frgC.store(result)
cute.copy(copy_atom_store, frgC, thrC, pred=frgPred)
@cute.jit
def elementwise_add(
mA,
mB,
mC,
stream: cuda.CUstream,
use_pdl: cutlass.Constexpr = True,
is_first_kernel: cutlass.Constexpr = True,
):
dtype = mA.element_type
# copy_bits for a thread is 128 bits, and we use 128 // dtype.width to get the vector size
vector_size = 128 // dtype.width
thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN))
gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN))
gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN))
idC = cute.make_identity_tensor(mC.shape)
cC = cute.zipped_divide(idC, tiler=tiler_mn)
elementwise_add_kernel(
gA, gB, gC, cC, mC.shape, thr_layout, val_layout, use_pdl, is_first_kernel
).launch(
grid=[cute.size(gC, mode=[1]), 1, 1],
block=[cute.size(tv_layout, mode=[0]), 1, 1],
# set cluster to enable cuLaunchKernelEx API for additional launch attributes setting
cluster=(1, 1, 1),
stream=stream,
# Currently, pdl launch attribute is set in compile phase,
# so we need to recompile the function if we change the value of use_pdl for multiple runs.
use_pdl=use_pdl,
)
def run_pdl_example(
M,
N,
skip_ref_check=False,
benchmark=True,
warmup_iterations=5,
iterations=10,
use_pdl=True,
):
if not torch.cuda.is_available():
raise RuntimeError("Blackwell/Hopper GPU is required to run this example!")
print("\nRunning Elementwise Add test with:")
print(f"Tensor dimensions: [{M}, {N}]")
print(f"Use PDL: {use_pdl}")
u = torch.randn(M, N, dtype=torch.float32, device="cuda")
v = torch.randn(M, N, dtype=torch.float32, device="cuda")
w = torch.randn(M, N, dtype=torch.float32, device="cuda")
x = torch.randn(M, N, dtype=torch.float32, device="cuda")
y = torch.empty(M, N, dtype=torch.float32, device="cuda")
u_tensor = from_dlpack(u).mark_layout_dynamic()
v_tensor = from_dlpack(v).mark_layout_dynamic()
w_tensor = from_dlpack(w).mark_layout_dynamic()
x_tensor = from_dlpack(x).mark_layout_dynamic()
y_tensor = from_dlpack(y).mark_layout_dynamic()
stream = torch.cuda.Stream()
current_stream = cuda.CUstream(stream.cuda_stream)
# Since use_pdl and is_first_kernel are cutlass.Constexpr, we need to compile for
# the first and second kernel separately.
compiled_func_first_kernel = cute.compile(
elementwise_add,
u_tensor,
v_tensor,
w_tensor,
current_stream,
use_pdl,
is_first_kernel=True,
)
compiled_func_second_kernel = cute.compile(
elementwise_add,
w_tensor,
x_tensor,
y_tensor,
current_stream,
use_pdl,
is_first_kernel=False,
)
# launch and run the two consecutive kernels in a same stream.
# Here, we simply use default stream.
def run_func(current_stream, u_tensor, v_tensor, w_tensor, x_tensor, y_tensor):
# Run first operation: w_tensor = u_tensor + v_tensor
compiled_func_first_kernel(
u_tensor,
v_tensor,
w_tensor,
current_stream,
)
# Run second operation: y_tensor = w_tensor + x_tensor
# its first operand ``w_tensor`` is the result of the first operation,
# they use the same memory space.
compiled_func_second_kernel(
w_tensor,
x_tensor,
y_tensor,
current_stream,
)
if not skip_ref_check:
run_func(current_stream, u_tensor, v_tensor, w_tensor, x_tensor, y_tensor)
print("Verifying results...")
torch.testing.assert_close(u.cpu() + v.cpu() + x.cpu(), y.cpu())
print("Results verified successfully!")
if not benchmark:
return
def generate_kernel_arguments():
u = torch.randn(M, N, dtype=torch.float32, device="cuda")
v = torch.randn(M, N, dtype=torch.float32, device="cuda")
w = torch.randn(M, N, dtype=torch.float32, device="cuda")
x = torch.randn(M, N, dtype=torch.float32, device="cuda")
y = torch.empty(M, N, dtype=torch.float32, device="cuda")
u_tensor = from_dlpack(u).mark_layout_dynamic()
v_tensor = from_dlpack(v).mark_layout_dynamic()
w_tensor = from_dlpack(w).mark_layout_dynamic()
x_tensor = from_dlpack(x).mark_layout_dynamic()
y_tensor = from_dlpack(y).mark_layout_dynamic()
return testing.JitArguments(
current_stream, u_tensor, v_tensor, w_tensor, x_tensor, y_tensor
)
avg_time_us = testing.benchmark(
run_func,
workspace_generator=generate_kernel_arguments,
workspace_count=10,
warmup_iterations=warmup_iterations,
iterations=iterations,
stream=current_stream,
)
print(f"Execution time: {avg_time_us:.4f} us")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="example of Programmatic Dependent Launch (PDL) using CuTe DSL"
)
parser.add_argument("--M", default=512, type=int)
parser.add_argument("--N", default=512, type=int)
parser.add_argument("--warmup_iterations", default=3, type=int)
parser.add_argument("--iterations", default=10, type=int)
parser.add_argument("--skip_ref_check", action="store_true")
parser.add_argument("--benchmark", action="store_true")
parser.add_argument("--use_pdl", action="store_true")
args = parser.parse_args()
if supports_pdl():
run_pdl_example(
args.M,
args.N,
skip_ref_check=args.skip_ref_check,
benchmark=args.benchmark,
warmup_iterations=args.warmup_iterations,
iterations=args.iterations,
use_pdl=args.use_pdl,
)
print("\nPASS")
else:
print(
"PDL is not supported on this device, it requires Hopper or newer generations"
)
examples/python/CuTeDSL/blackwell/tutorial_gemm/README.md
+25
View File
@@ -0,0 +1,25 @@
# CUTLASS Tutorial Examples for Blackwell GEMM
This folder contains tutorial examples demonstrating how to write performant GEMM (General Matrix Multiplication) kernels using Tensor Cores on NVIDIA Blackwell GPUs.
## Overview
The examples showcase different scenarios and optimization techniques for implementing GEMM operations:
- Basic FP16 GEMM implementation
- Software Pipeline optimizations
- Tensor Core utilization
- Thread/warp/block level parallelism
## Examples
### tutorial_fp16_gemm_0.py
A basic example showing:
- FP16 GEMM implementation using Tensor Cores
- TMA (Tensor Memory Access) for efficient data loading
- SMEM (Shared Memory) layouts and access patterns
- Usage of ``cutlass.range(..., prefetch_stages=...)`` to replace boilerplate code for multi-stage software pipeline
With some minor optimization tricks
- Tiling Epilogue to avoid bursty write out and reduce register pressure
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py
+444
View File
@@ -0,0 +1,444 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
import argparse
import torch
from typing import Tuple
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.torch as cutlass_torch
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
"""
The first tutorial GEMM demonstrating a simple kernel implementation in CuTeDSL
This dense GEMM kernel is implemented in just over 200 lines of code.
With large tile sizes, it can achieve very high performance on 8k×8k×8k problem sizes.
It can serve as a starting point to help users quickly experiment
with optimizations for challenges that may arise with other problem sizes.
To run this example:
.. code-block:: bash
python examples/blackwell/tutorial_fp16_gemm_0.py \
--mnk 8192,8192,8192 \
--tolerance 1e-01
Constraints for this example:
* The problem size of m and n must be divisible by the tile size m & n (128, 256)
"""
io_dtype = cutlass.Float16
acc_dtype = cutlass.Float32
mma_inst_shape_mnk = (128, 256, 16)
mma_tiler_mnk = (128, 256, 64)
threads_per_cta = 128
# Pipeline stage configuration
ab_stages = 4
acc_stage = 1
@cute.struct
class SharedStorage:
ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stage * 2]
tmem_holding_buf: cutlass.Int32
@cute.kernel
def kernel(
tiled_mma: cute.TiledMma,
tma_atom_a: cute.CopyAtom,
mA_mkl: cute.Tensor,
tma_atom_b: cute.CopyAtom,
mB_nkl: cute.Tensor,
mC_mnl: cute.Tensor,
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()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
bidx, bidy, _ = cute.arch.block_idx()
mma_coord_mnk = (bidx, bidy, None)
#
# 1. Prepare args
#
# Allocate SMEM
smem = cutlass.utils.SmemAllocator()
storage = smem.allocate(SharedStorage)
sA = smem.allocate_tensor(
element_type=io_dtype,
layout=a_smem_layout.outer,
byte_alignment=128,
swizzle=a_smem_layout.inner,
)
sB = smem.allocate_tensor(
element_type=io_dtype,
layout=b_smem_layout.outer,
byte_alignment=128,
swizzle=b_smem_layout.inner,
)
# Allocate all TMEM columns
tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=tmem_alloc_barrier,
)
num_tmem_cols = 512
tmem.allocate(num_tmem_cols)
# Prefetch tma descriptor
if warp_idx == 0:
cpasync.prefetch_descriptor(tma_atom_a)
cpasync.prefetch_descriptor(tma_atom_b)
# Pipeline configuration
num_tma_copy_bytes = cute.size_in_bytes(
io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2])
) + cute.size_in_bytes(io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]))
ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
num_stages=ab_stages,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
tx_count=num_tma_copy_bytes,
barrier_storage=storage.ab_mbar_ptr.data_ptr(),
).make_participants()
acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
num_stages=acc_stage,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread, threads_per_cta
),
barrier_storage=storage.acc_mbar_ptr.data_ptr(),
).make_participants()
# Partition tensors for MMA and make fragments
# (bM, bK, RestK)
gA = cute.local_tile(mA_mkl, mma_tiler_mnk, mma_coord_mnk, proj=(1, None, 1))
# (bN, bK, RestK)
gB = cute.local_tile(mB_nkl, mma_tiler_mnk, mma_coord_mnk, proj=(None, 1, 1))
# (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)
tCgA = thr_mma.partition_A(gA)
# (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)
tCrA = tiled_mma.make_fragment_A(sA)
# (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])
# (MMA, MMA_M, MMA_N)
tCtAcc = tiled_mma.make_fragment_C(acc_shape)
# Partition tensors for TMA; This requires the tensors partitioned for MMA
tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
tma_atom_a,
0,
cute.make_layout(1),
cute.group_modes(sA, 0, 3),
cute.group_modes(tCgA, 0, 3),
)
tBsB, tBgB = cute.nvgpu.cpasync.tma_partition(
tma_atom_b,
0,
cute.make_layout(1),
cute.group_modes(sB, 0, 3),
cute.group_modes(tCgB, 0, 3),
)
# CTA-wide sync before retrieving the pointer to the start of the allocated TMEM
# Only warp 0 does the allocation so we need to sync before retrieving the TMEM start address
tmem.wait_for_alloc()
tmem_ptr = tmem.retrieve_ptr(acc_dtype)
# Swap the pointer in tCtAcc
tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc.layout)
subtile_cnt = 4
# (EpiTile)
epi_tiler = (
(cute.size(tCtAcc, mode=[0, 0]), cute.size(tCtAcc, mode=[0, 1]) // subtile_cnt),
)
# (EpiTile, NumTiles)
tCtAcc_epi = cute.zipped_divide(tCtAcc, epi_tiler)
# (EpiTile, NumTiles)
gC_epi = cute.zipped_divide(tCgC, epi_tiler)
# Every thread loads 32x128 bits
tmem_atom = cute.make_copy_atom(
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)
tDtC = tmem_thr_copy.partition_S(tCtAcc_epi)
# (TmemCpy,NumTmemCpy,NumTiles)
tDgC = tmem_thr_copy.partition_D(gC_epi)
# (TmemCpy,NumTmemCpy)
tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
# (TmemCpy,NumTmemCpy)
tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
#
# 2. Main loop
#
num_k_tiles = cute.size(gA, mode=[2])
if warp_idx == 0:
# Wait for a empty accumulator buffer
acc_empty = acc_producer.acquire_and_advance()
for k_tile_idx in cutlass.range(num_k_tiles, prefetch_stages=ab_stages - 2):
# Issue TMA loads
ab_empty = ab_producer.acquire_and_advance()
cute.copy(
tma_atom_a,
tAgA[(None, ab_empty.count)],
tAsA[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_b,
tBgB[(None, ab_empty.count)],
tBsB[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
# Execute one K-block worth of MMA instructions
ab_full = ab_consumer.wait_and_advance()
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)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[k_block_coord],
tCrB[k_block_coord],
tCtAcc,
)
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Signal that the A/B buffers have been consumed and are ready for the next load
ab_full.release()
# Signal that the accumulator is fully computed
acc_empty.commit()
#
# 3. Epilogue
#
# Release TMEM allocation lock
tmem.relinquish_alloc_permit()
# 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(tDtC, mode=[2])):
cute.copy(tmem_tiled_copy, tDtC[None, None, i], tCrAcc)
tCrC.store(tCrAcc.load().to(io_dtype))
cute.autovec_copy(tCrC, tDgC[None, None, i])
acc_full.release()
# Deallocate TMEM
pipeline.sync(barrier_id=1)
tmem.free(tmem_ptr)
@cute.jit
def host_function(
a: cute.Tensor,
b: cute.Tensor,
c: cute.Tensor,
):
# Construct tiled MMA
op = tcgen05.MmaF16BF16Op(
io_dtype,
acc_dtype,
mma_inst_shape_mnk,
tcgen05.CtaGroup.ONE,
tcgen05.OperandSource.SMEM,
tcgen05.OperandMajorMode.K,
tcgen05.OperandMajorMode.K,
)
tiled_mma = cute.make_tiled_mma(op)
# Construct SMEM layouts for A and B
a_smem_layout = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
a.element_type,
ab_stages,
)
b_smem_layout = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
b.element_type,
ab_stages,
)
a_smem_layout_one_stage = cute.select(a_smem_layout, mode=[0, 1, 2])
b_smem_layout_one_stage = cute.select(b_smem_layout, mode=[0, 1, 2])
# Construct TMA load atoms
op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
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,
)
b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
op,
b,
b_smem_layout_one_stage,
mma_tiler_mnk,
tiled_mma,
)
# Pretty prints kernel attributes useful for debugging
# print(f"a = {cute.pretty_str(a)}")
# print(f"b = {cute.pretty_str(b)}")
# print(f"c = {cute.pretty_str(c)}")
# print(f"tiled_mma = {cute.pretty_str(tiled_mma)}")
# print(f"a_tma_atom = {cute.pretty_str(a_tma_atom)}")
# print(f"b_tma_atom = {cute.pretty_str(b_tma_atom)}")
# print(f"a_tma_tensor = {cute.pretty_str(a_tma_tensor)}")
# print(f"b_tma_tensor = {cute.pretty_str(b_tma_tensor)}")
# Launch the kernel
grid_shape = cute.ceil_div((*c.layout.shape, 1), mma_tiler_mnk[:2])
kernel(
tiled_mma,
a_tma_atom,
a_tma_tensor,
b_tma_atom,
b_tma_tensor,
c,
a_smem_layout,
b_smem_layout,
).launch(
grid=grid_shape,
block=(threads_per_cta, 1, 1),
)
def run_dense_gemm(
mnk: Tuple[int, int, int],
tolerance: float,
):
print("===================================================================")
print("Running Blackwell fp16 GEMM example 0 with:")
print(f" mnk: {mnk}")
print(f" tolerance: {tolerance}")
print("===================================================================")
print()
m, n, k = mnk
torch.manual_seed(1111)
# 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
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))).cpu()
torch.testing.assert_close(
c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
)
if __name__ == "__main__":
def parse_comma_separated_ints(s: str):
try:
return [int(x.strip()) for x in s.split(",")]
except ValueError:
raise argparse.ArgumentTypeError(
"Invalid format. Expected comma-separated integers."
)
if not torch.cuda.is_available():
raise RuntimeError("A GPU is required to run this example")
parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 0")
parser.add_argument(
"--mnk",
type=parse_comma_separated_ints,
default=[8192, 8192, 8192],
help="MNK dimensions (comma-separated)",
)
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,
)
print("PASS")