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

v4.5 tag update (#3202)

Browse Source 
* Python DSL examples reorganization.

* v4.5 tag update.
This commit is contained in:
Junkai-Wu
2026-05-06 08:55:27 +08:00
committed by GitHub
parent f74fea9ce3
commit cb37157db5
351 changed files with 36688 additions and 8117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

155
examples/python/CuTeDSL/cute_ext/ampere/memcpy_simt_universal_copy.py Normal file
View File

@@ -0,0 +1,155 @@
# Copyright (c) 2025 - 2026 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 torch
import pytest
from cutlass import cute
from cutlass.cute import experimental as cute_ext
from cutlass.cute.runtime import from_dlpack
import cutlass.utils as utils
@cute.experimental.kernel
def memcpy_simt_universal_copy_kernel(
mA: cute.Tensor, mD: cute.Tensor, addend: cute.Float16
):
tile_mn = cute.core._pack_shape((128, 64))
gA = cute.zipped_divide(mA, tile_mn)
gD = cute.zipped_divide(mD, tile_mn)
cta_m, cta_n, cta_l = cute.arch.block_idx()
tid_x, _, _ = cute.arch.thread_idx()
gA_tile = gA[(None, None), (cta_m, cta_n, cta_l)]
gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)]
buffer = cute_ext.allocate(
cute.Float16,
cute.AddressSpace.rmem,
cute.make_layout(((8, 1), (1, 8)), stride=((1, 8), (1, 8))),
alignment=16,
)
tCgA = cute_ext.partition(
gA_tile,
tid_x,
layout_tv=cute.make_layout(((16, 8), (8, 1)), stride=((8, 128), (1, 1024))),
tiler=cute.core._pack_tile((128, 8)),
)
tCgD = cute_ext.partition(
gD_tile,
tid_x,
layout_tv=cute.make_layout(((16, 8), (8, 1)), stride=((8, 128), (1, 1024))),
tiler=cute.core._pack_tile((128, 8)),
)
# cute_ext.copy() automatically computes predicates based on the shape of
# the tensor passed to the @cute.experimental.kernel argument
cute_ext.copy(
tCgA,
buffer,
copy_atom=cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
tCgD.element_type,
num_bits_per_copy=128,
),
)
# Update the RMEM tensor in place using elementwise addition.
buffer.store(buffer.load() + addend)
# cute_ext.copy() automatically computes predicates based on the shape of
# the tensor passed to the @cute.experimental.kernel argument
cute_ext.copy(
buffer,
tCgD,
copy_atom=cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
tCgD.element_type,
num_bits_per_copy=128,
),
)
@cute.experimental.jit
def memcpy_simt_universal_copy(
src: cute.Tensor, dst: cute.Tensor, addend: cute.Float16
):
tile_mn = cute.core._pack_shape((128, 64))
div = cute.tiled_divide(src, tile_mn)
grid = (div.shape[1], div.shape[2], div.shape[3])
memcpy_simt_universal_copy_kernel(src, dst, addend).launch(
grid=grid,
block=(128, 1, 1),
smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_80")),
)
def run_simt_universal_memcpy(M, N, L):
src = torch.randn(L, N, M).permute(2, 1, 0).to(torch.float16).cuda()
dst = torch.randn(L, N, M).permute(2, 1, 0).to(torch.float16).cuda()
mA = (
from_dlpack(src, assumed_align=16)
.mark_layout_dynamic(leading_dim=0)
.mark_compact_shape_dynamic(
mode=0, stride_order=src.dim_order(), divisibility=8
)
)
mD = (
from_dlpack(dst, assumed_align=16)
.mark_layout_dynamic(leading_dim=0)
.mark_compact_shape_dynamic(
mode=0, stride_order=dst.dim_order(), divisibility=8
)
)
addend = 5.0
memcpy_simt_universal_copy(
mA,
mD,
cute.Float16(addend),
no_cache=True,
)
torch.testing.assert_close(src.cpu() + addend, dst.cpu())
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Example memory copy example using CuTe auto predication features."
)
parser.add_argument("--mnl", default=[136, 7, 9], nargs="+", type=int)
args = parser.parse_args()
M, N, L = tuple(args.mnl)
run_simt_universal_memcpy(M, N, L)
print("PASS")

1021
examples/python/CuTeDSL/cute_ext/blackwell/dense_block_scaled_gemm.py Normal file
View File

File diff suppressed because it is too large Load Diff

1389
examples/python/CuTeDSL/cute_ext/blackwell/dense_gemm.py Normal file
View File

File diff suppressed because it is too large Load Diff

519
examples/python/CuTeDSL/cute_ext/blackwell/dense_gemm_2sm.py Normal file
View File

@@ -0,0 +1,519 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 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.
"""
2SM Dense GEMM example using cute_ext decorators.
"""
import torch
import math
import cutlass
from cutlass import cute
from cutlass.cute import experimental as cute_ext
from cutlass.cute.runtime import from_dlpack
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils as utils
from cutlass.base_dsl.typing import Numeric
from typing import Type
def create_gemm_tensors_torch(
M,
N,
K,
majors: tuple[
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
],
dtypes: tuple[torch.dtype, torch.dtype, torch.dtype],
):
A = None
B = None
D = None
if majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.MN:
A = torch.empty(K, M).random_(-4, 4).permute(1, 0).to(dtypes[0]).cuda()
elif majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.K:
A = torch.empty(M, K).random_(-4, 4).permute(0, 1).to(dtypes[0]).cuda()
if majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.MN:
B = torch.empty(K, N).random_(-4, 4).permute(1, 0).to(dtypes[1]).cuda()
elif majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.K:
B = torch.empty(N, K).random_(-4, 4).permute(0, 1).to(dtypes[1]).cuda()
if majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.MN:
D = torch.empty(N, M).random_(-4, 4).permute(1, 0).to(dtypes[2]).cuda()
elif majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.K:
D = torch.empty(M, N).random_(-4, 4).permute(0, 1).to(dtypes[2]).cuda()
return A, B, D
def get_gemm_tensors(
M,
N,
K,
majors: tuple[
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
],
dtypes: tuple[torch.dtype, torch.dtype, torch.dtype],
):
A, B, D = create_gemm_tensors_torch(M, N, K, majors, dtypes)
A_cute = from_dlpack(A, assumed_align=16).mark_layout_dynamic(
leading_dim=1 if majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0
)
B_cute = from_dlpack(B, assumed_align=16).mark_layout_dynamic(
leading_dim=1 if majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0
)
D_cute = from_dlpack(D, assumed_align=16).mark_layout_dynamic(
leading_dim=1 if majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0
)
return A, B, D, A_cute, B_cute, D_cute
def sm100_4x4x1_kernel_builder(
use_tma_multicast: bool,
use_2cta_instrs: bool,
acc_dtype: Type[Numeric],
M: int,
N: int,
):
CLUSTER_SHAPE = (2, 1, 1)
GRID_SHAPE = (
math.ceil(M / 128),
math.ceil(N / 256),
1,
) # TODO (xpbowler): remove hard-code
NUM_WARPS_PER_CTA = 6
TMA_STORE_PIPE_DEPTH = 4
MAINLOOP_STAGE_DEPTH = 4 # pipeline depth of TMA->MMA
# pipeline depth of mainloop->epilogue. only useful if using persistent CTA
EPILOGUE_STAGE_DEPTH = 1
# m256n256k16 2SM MMA / m128n256k16 1SM MMA
mma_inst_shape_mnk = (256, 256, 16) if use_2cta_instrs else (128, 256, 16)
@cute_ext.kernel
def kernel(
mA: cute.Tensor,
mB: cute.Tensor,
mD: cute.Tensor,
):
d_layout = utils.LayoutEnum.from_tensor(mD)
d_dtype = mD.element_type
ab_dtype = mA.element_type
mma_inst_shape_m, mma_inst_shape_n, mma_inst_shape_k = mma_inst_shape_mnk
if cutlass.const_expr(use_2cta_instrs):
cta_group = cute.nvgpu.tcgen05.CtaGroup.TWO
else:
cta_group = cute.nvgpu.tcgen05.CtaGroup.ONE
tiled_mma = sm100_utils.make_trivial_tiled_mma(
ab_dtype,
utils.LayoutEnum.from_tensor(mA).mma_major_mode(),
utils.LayoutEnum.from_tensor(mB).mma_major_mode(),
acc_dtype,
cta_group,
(mma_inst_shape_m, mma_inst_shape_n),
)
mma_inst_tile_k = (
4 # 4 MMAs per MMA tile K. For 16b types, tcgen05.mma has K=16.
)
mma_inst_tile_m = mma_inst_tile_n = 1 # 1 MMAs per MMA tile M/N
bM = mma_inst_shape_m * mma_inst_tile_m
bN = mma_inst_shape_n * mma_inst_tile_n
bK = mma_inst_shape_k * mma_inst_tile_k
mnk_tiler = (bM, bN, bK)
cta_m, cta_n, _ = cute.arch.block_idx()
tid_x, _, _ = cute.arch.thread_idx()
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout(CLUSTER_SHAPE),
cute.core._pack_shape((cute.size(tiled_mma.thr_id.shape),)),
)
cluster_layout_v_size = cute.size(cluster_layout_vmnk.shape[0])
mma_coord_vmnk = (
cta_m % cluster_layout_v_size,
cta_m // cluster_layout_v_size,
cta_n,
)
gA = cute.zipped_divide(mA, (bM, bK)) # ((bM, bK), (M/bM, K/bK))
gA_tma = cute.zipped_divide(
mA, (bM // cluster_layout_v_size, bK)
) # ((bM/2, bK), (2*M/bM, K/bK))
tAgA = gA_tma[(None, None), (cta_m, None)] # ((bM/2, bK), (1, K/bK))
gB_tma = cute.zipped_divide(
mB, (bN // cluster_layout_v_size, bK)
) # ((bN/2, bK), (2*M/bM, K/bK))
# ((bN/2, bK), (1, K/bK))
tBgB = gB_tma[
(None, None),
(cluster_layout_v_size * cta_n + cta_m % cluster_layout_v_size, None),
]
gD_tma = cute.zipped_divide(
mD, (bM // cluster_layout_v_size, bN)
) # ((bM/2, bN), (2*M/bM, N/bN))
tDgD = gD_tma[(None, None), (cta_m, cta_n)] # ((bM/2, bN), (1, 1))
a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
mnk_tiler,
ab_dtype,
MAINLOOP_STAGE_DEPTH,
)
b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
mnk_tiler,
ab_dtype,
MAINLOOP_STAGE_DEPTH,
)
cta_tile_shape_mnk = cute.shape_div(mnk_tiler, (cluster_layout_v_size, 1, 1))
epi_tile = sm100_utils.compute_epilogue_tile_shape(
cta_tile_shape_mnk,
use_2cta_instrs,
d_layout,
d_dtype,
)
sc_smem_layout_staged = sm100_utils.make_smem_layout_epi(
d_dtype,
d_layout,
epi_tile,
TMA_STORE_PIPE_DEPTH,
)
tmem_layout = cute_ext.make_tmem_layout_acc(
tiled_mma, mnk_tiler, EPILOGUE_STAGE_DEPTH
)
bufferA = cute_ext.allocate(
ab_dtype,
cute.AddressSpace.smem,
a_smem_layout_staged,
alignment=1024,
)
bufferB = cute_ext.allocate(
ab_dtype,
cute.AddressSpace.smem,
b_smem_layout_staged,
alignment=1024,
)
bufferAcc = cute_ext.allocate(
acc_dtype,
cute.AddressSpace.tmem,
tmem_layout,
alignment=16,
is2cta=use_2cta_instrs,
)
bufferC = cute_ext.allocate(
d_dtype,
cute.AddressSpace.smem,
sc_smem_layout_staged,
alignment=1024,
)
copy_atom_t2r = sm100_utils.get_tmem_load_op(
cta_tile_shape_mnk,
d_layout,
d_dtype,
acc_dtype,
epi_tile,
use_2cta_instrs,
)
# Take only one stage of the TMEM buffer for the epilogue
accumulators = cute.zipped_divide(bufferAcc, ((epi_tile), 1))
acc_epi_div = accumulators[((None, None), 0), 0]
# Create the TMEM copy atom based on the size of transfer within one iteration of epilogue
tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div)
# Calculate the per thread destination size per iteration for output of TMEM and input of SMEM
gC_mnl_epi = cute.flat_divide(tDgD, epi_tile)
acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout(
tiled_copy_t2r, gC_mnl_epi, tid_x
)
bufferRAcc = cute_ext.allocate(
acc_dtype,
cute.AddressSpace.rmem,
acc_d_rmem_layout,
alignment=32,
)
bufferRD = cute_ext.allocate(
d_dtype,
cute.AddressSpace.rmem,
acc_d_rmem_layout,
alignment=32,
)
tma_mcast_proj_A = 2
tma_mcast_proj_B = 1
mma_operation_type = tma_operation_type = None
acc_pipe = mainloop_pipe = None
if cutlass.const_expr(use_2cta_instrs):
mma_operation_type = cute_ext.OperationTypeEnum.SM100_MMA_2SM_SS
if cutlass.const_expr(use_tma_multicast):
tma_operation_type = (
cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM_MULTICAST
)
else:
tma_operation_type = cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM
else:
mma_operation_type = cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS
if cutlass.const_expr(use_tma_multicast):
tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD_MULTICAST
else:
tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD
# MMA <-> TMEM load pipeline
# if 2CTA MMA, warpgroup from both peer and leader CTA consumer.release
acc_pipe_consumer_arv_count = 256 if use_2cta_instrs else 128
acc_pipe = cute_ext.UMMAtoAsyncPipeline.create(
num_stages=EPILOGUE_STAGE_DEPTH,
mma_operation_type=mma_operation_type,
consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R,
consumer_arv_count=acc_pipe_consumer_arv_count,
cluster_layout_vmnk=cluster_layout_vmnk,
)
if cutlass.const_expr(use_tma_multicast):
# TMA load <-> MMA pipeline
mainloop_pipe = cute_ext.TMAToUMMAPipeline.create_with_mask(
num_stages=MAINLOOP_STAGE_DEPTH,
tma_operation_type=tma_operation_type,
mma_operation_type=mma_operation_type,
cluster_layout_vmnk=cluster_layout_vmnk,
)
else:
mainloop_pipe = cute_ext.TMAToUMMAPipeline.create(
num_stages=MAINLOOP_STAGE_DEPTH,
mma_operation_type=mma_operation_type,
tma_operation_type=tma_operation_type,
cluster_layout_vmnk=cluster_layout_vmnk,
)
tma_store_warp_id = 0
mma_warp_id = 4
tma_load_warp_id = 5
is_tma_thr = warp_idx == tma_load_warp_id
is_mma_thr = warp_idx == mma_warp_id
is_epi_thr = warp_idx < 4
is_leader_cta = mma_coord_vmnk[0] == 0
# SMEM -> GMEM
tma_store_pipe = cute_ext.TMAStorePipeline(
stages=TMA_STORE_PIPE_DEPTH,
arv_count=128,
barrier_id=1,
tma_warp_id=tma_store_warp_id,
)
k_tile_count = cute.size(gA, mode=[1, 1])
if is_tma_thr:
for k_tile in cutlass.range(0, k_tile_count, 1, unroll=1):
gA_k = tAgA[None, None, k_tile]
gB_k = tBgB[None, None, k_tile]
producer_stage_token, idx = (
mainloop_pipe.producer_acquire_and_get_stage()
)
mbar = cute_ext.get_mbarrier(producer_stage_token)
bufferA_sliced = bufferA[None, None, None, idx]
bufferB_sliced = bufferB[None, None, None, idx]
a_cta_v_map = cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A")
b_cta_v_map = cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B")
if cutlass.const_expr(use_tma_multicast):
cute_ext.tma_load_multicast(
gA_k,
bufferA_sliced,
mbar,
vmnk_layout=cluster_layout_vmnk,
cta_v_map=a_cta_v_map,
tma_operation_type=tma_operation_type,
multicast_mode=tma_mcast_proj_A,
)
cute_ext.tma_load_multicast(
gB_k,
bufferB_sliced,
mbar,
vmnk_layout=cluster_layout_vmnk,
cta_v_map=b_cta_v_map,
tma_operation_type=tma_operation_type,
multicast_mode=tma_mcast_proj_B,
)
else:
cute_ext.tma_load(
gA_k,
bufferA_sliced,
mbar,
cta_v_map=a_cta_v_map,
tma_operation_type=tma_operation_type,
)
cute_ext.tma_load(
gB_k,
bufferB_sliced,
mbar,
cta_v_map=b_cta_v_map,
tma_operation_type=tma_operation_type,
)
if is_leader_cta:
mainloop_pipe.producer_commit()
mainloop_pipe.producer_state = cute_ext.pipeline_advance_iterator(
mainloop_pipe.raw_pipeline, mainloop_pipe.producer_state
)
if is_mma_thr and is_leader_cta:
producer_stage_token, idx = acc_pipe.producer_acquire_and_get_stage()
accumulators_sliced = bufferAcc[None, None, None, idx]
mma_atom = cute.make_mma_atom(tiled_mma.op)
mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False)
for k_tile in cutlass.range(0, k_tile_count, 1, unroll=1):
_, mainloop_idx = mainloop_pipe.consumer_wait_and_get_stage()
bufferA_sliced_stage = cute.core.slice_(
bufferA, (None, None, None, mainloop_idx)
)
bufferB_sliced_stage = cute.core.slice_(
bufferB, (None, None, None, mainloop_idx)
)
for k_block in cutlass.range(mma_inst_tile_k, unroll_full=True):
cute_ext.dot(
mma_atom,
cute.append_ones(
bufferA_sliced_stage[None, None, k_block], up_to_rank=3
),
cute.append_ones(
bufferB_sliced_stage[None, None, k_block], up_to_rank=3
),
accumulators_sliced,
)
mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True)
mainloop_pipe.consumer_release_and_advance()
acc_pipe.producer_commit_and_advance()
if is_epi_thr:
_, idx = acc_pipe.consumer_wait_and_get_stage()
accumulators_sliced = bufferAcc[(None, None), 0, 0, idx]
acc_epi_div_tiled = cute.flat_divide(accumulators_sliced, epi_tile)
tiled_copy_r2s = cute.make_tiled_copy_D(
cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), d_dtype),
tiled_copy_t2r,
)
c_cta_v_map = cute_ext.get_cta_v_map_c(mD, epi_tile)
subtile_cnt = cute.size(acc_epi_div_tiled.shape, mode=[3])
for mn in range(subtile_cnt):
# TMEM -> RMEM
cute_ext.partition_and_copy(
tiled_copy_t2r.get_slice(tid_x),
acc_epi_div_tiled[None, None, 0, mn],
bufferRAcc,
)
# RMEM -> RMEM
bufferRD.store(bufferRAcc.load().to(d_dtype))
tma_store_pipe.acquire_sync()
store_idx = tma_store_pipe.get_index()
# RMEM -> SMEM
cute_ext.partition_and_copy(
tiled_copy_r2s.get_slice(tid_x),
bufferRD,
bufferC[None, None, store_idx],
)
tma_store_pipe.commit_sync()
if warp_idx == tma_store_warp_id:
cute_ext.tma_store(
bufferC[None, None, store_idx],
gC_mnl_epi[None, None, 0, mn],
cta_v_map=c_cta_v_map,
)
tma_store_pipe.release_advance()
tma_store_pipe.tail()
acc_pipe.consumer_release_and_advance()
# Return a callable that launches the kernel with proper grid/block/cluster
@cute_ext.jit
def launch_kernel(mA: cute.Tensor, mB: cute.Tensor, mD: cute.Tensor):
kernel(mA, mB, mD).launch(
grid=GRID_SHAPE,
block=(32 * NUM_WARPS_PER_CTA, 1, 1),
cluster=CLUSTER_SHAPE,
smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")),
)
return launch_kernel
if __name__ == "__main__":
M = 256
N = 256
K = 64
use_tma_multicast = True
use_2cta_instrs = True
acc_dtype = cutlass.Float32
majors = (
cute.nvgpu.tcgen05.OperandMajorMode.K,
cute.nvgpu.tcgen05.OperandMajorMode.K,
cute.nvgpu.tcgen05.OperandMajorMode.K,
)
dtypes = (torch.float16, torch.float16, torch.float16)
A_torch, B_torch, D_torch, A_cute, B_cute, D_cute = get_gemm_tensors(
M, N, K, majors, dtypes
)
kernel_launcher = sm100_4x4x1_kernel_builder(
use_tma_multicast, use_2cta_instrs, acc_dtype, M, N
)
compiled_kernel = cute_ext.compile(kernel_launcher, A_cute, B_cute, D_cute)
compiled_kernel(A_cute, B_cute, D_cute)
# Reference check (may fail on simulator/unsupported GPU)
try:
ref = torch.mm(A_torch.float(), B_torch.float().T)
torch.testing.assert_close(D_torch.float(), ref, atol=1e-2, rtol=1e-2)
print("PASS")
except RuntimeError as e:
if "no kernel image is available" in str(e):
print("SKIP: Reference check skipped - GPU not supported by PyTorch")
else:
raise

1803
examples/python/CuTeDSL/cute_ext/blackwell/dense_gemm_cute_pipeline.py Executable file
View File

File diff suppressed because it is too large Load Diff

724
examples/python/CuTeDSL/cute_ext/blackwell/dense_gemm_ptr_array.py Executable file
View File

@@ -0,0 +1,724 @@
# Copyright (c) 2025 - 2026 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 torch
from typing import Type, Tuple, List
import cutlass
from cutlass.cute import experimental as cute_ext
from cutlass.base_dsl.typing import Numeric
from cutlass import cute as cute
from cutlass import utils
from cutlass import torch as cutlass_torch
from cutlass.cute.experimental.host_runtime import QueryDeviceWorkspaceFunc
from cutlass.cute.runtime import from_dlpack
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.cute.testing as testing
class DenseGemmPtrArrayKernel:
def __init__(
self,
mn_tiler: tuple[int, int],
mma_dtype: tuple[Type[Numeric], Type[Numeric], Type[Numeric]],
tmem_output_dtype: Type[Numeric],
batch_count: int, # Number of batches, each batch will have its own pointer for the matrix
A_shape: tuple, # Shape of the matrix A
A_stride: tuple, # Stride of the matrix A
B_shape: tuple, # Shape of the matrix B
B_stride: tuple, # Stride of the matrix B
D_shape: tuple, # Shape of the matrix D
D_stride: tuple, # Stride of the matrix D
epilogue_op=lambda x: x,
):
self.mn_tiler = mn_tiler
self.ab_dtype, self.acc_dtype, self.d_dtype = mma_dtype
self.tmem_output_dtype = tmem_output_dtype
self.use_2cta_instrs = False
self.TMA_STORE_STAGE = 4
self.epilogue_op = epilogue_op
self.batch_count = batch_count
self.A_shape = A_shape
self.A_stride = A_stride
self.B_shape = B_shape
self.B_stride = B_stride
self.D_shape = D_shape
self.D_stride = D_stride
"""
Helper function to convert an int64 to a cute.ptr of a certain type.
The cute.ptr is always located in Gmem.
This is used to load the pointers for A/B/D from the Ptr array.
"""
@cute.experimental.jit
def _get_pointer(self, address_as_int, cute_type):
cute_ptr = cute.make_ptr(
cute_type,
address_as_int,
mem_space=cute.AddressSpace.gmem,
assumed_align=16,
)
return cute_ptr
@cute.experimental.jit
def __call__(
self, mA_tensor: cute.Tensor, mB_tensor: cute.Tensor, mD_tensor: cute.Tensor
):
# Get the pointer to the first batch of D
d_ptr = self._get_pointer(mD_tensor[0], self.d_dtype)
d_ptr_base_tensor = cute.make_tensor(
d_ptr, layout=cute.make_layout(self.D_shape, stride=self.D_stride)
)
tile_mn = cute.core._pack_shape((*self.mn_tiler, 1))
div = cute.tiled_divide(d_ptr_base_tensor, tile_mn)
grid = (div.shape[1], div.shape[2], div.shape[3])
self.kernel(mA_tensor, mB_tensor, mD_tensor).launch(
grid=grid,
block=(192, 1, 1),
cluster=(1, 1, 1),
smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")),
)
@cute.experimental.kernel
def kernel(
self,
mA_tensor: cute.Tensor,
mB_tensor: cute.Tensor,
mD_tensor: cute.Tensor,
):
# Get pointers for the first batch to perform shape and stage calculations
A_0_ptr = self._get_pointer(mA_tensor[0], self.ab_dtype)
B_0_ptr = self._get_pointer(mB_tensor[0], self.ab_dtype)
mA = cute.make_tensor(
A_0_ptr, layout=cute.make_layout(self.A_shape, stride=self.A_stride)
)
mB = cute.make_tensor(
B_0_ptr, layout=cute.make_layout(self.B_shape, stride=self.B_stride)
)
tiled_mma = sm100_utils.make_trivial_tiled_mma(
self.ab_dtype,
self.ab_dtype,
utils.LayoutEnum.from_tensor(mA).mma_major_mode(),
utils.LayoutEnum.from_tensor(mB).mma_major_mode(),
self.acc_dtype,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mn_tiler,
)
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = 4
mnk_tiler = (
self.mn_tiler[0],
self.mn_tiler[1],
mma_inst_shape_k * mma_inst_tile_k,
)
tiler_mk = (mnk_tiler[0], mnk_tiler[2])
tiler_nk = (mnk_tiler[1], mnk_tiler[2])
gA = cute.zipped_divide(mA, tiler_mk)
gB = cute.zipped_divide(mB, tiler_nk)
mainloop_stage = 2
acc_stage = 2
cta_m, cta_n, cta_l = cute.arch.block_idx()
gA_tile = gA[(None, None), (cta_m, None, cta_l)]
gB_tile = gB[(None, None), (cta_n, None, cta_l)]
# Compute A/B/C shared memory layout
a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
mnk_tiler,
self.ab_dtype,
mainloop_stage,
)
b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
mnk_tiler,
self.ab_dtype,
mainloop_stage,
)
cta_tile_shape_mnk = cute.shape_div(
mnk_tiler, (cute.size(tiled_mma.thr_id.shape), 1, 1)
)
# UMMA ACC TMEM Layout
tmem_layout = cute_ext.make_tmem_layout_acc(tiled_mma, mnk_tiler, acc_stage)
# Allocate UMMA Buffers
bufferA = cute_ext.allocate(
self.ab_dtype,
cute.AddressSpace.smem,
a_smem_layout_staged,
alignment=1024,
)
bufferB = cute_ext.allocate(
self.ab_dtype,
cute.AddressSpace.smem,
b_smem_layout_staged,
alignment=1024,
)
bufferAcc = cute_ext.allocate(
self.acc_dtype,
cute.AddressSpace.tmem,
tmem_layout,
alignment=16,
)
# TMA -> UMMA
mainloop_pipe = cute_ext.TMAToUMMAPipeline.create(
num_stages=mainloop_stage,
mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS,
)
# UMMA -> TMEM
acc_pipe = cute_ext.UMMAtoAsyncPipeline.create(
num_stages=acc_stage,
mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS,
consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R,
consumer_arv_count=128,
)
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
# warp assignment: [0]-tma_store, [0-3]-epi, [4]-mma, [5]-tma_load
tma_store_warp_id = 0
mma_warp_id = 4
tma_load_warp_id = 5
is_tma_thr = warp_idx == tma_load_warp_id
is_mma_thr = warp_idx == mma_warp_id
is_epi_thr = warp_idx < 4
# SMEM -> GMEM
tma_store_pipe = cute_ext.TMAStorePipeline(
stages=self.TMA_STORE_STAGE,
arv_count=128,
barrier_id=1,
tma_warp_id=tma_store_warp_id,
)
k_tile_size = cute.size(gA, mode=[1, 1])
# Outer loop over batches and perform GEMM for each batch as usual
# This is a dynamic for loop that lowers to an scf.for
# Note that the tensor loading is done in the if `thread` warp specialized
# sections. This is essential to ensure proper synchronization of tma loads
# and tma updates across batches.
for batch_idx in range(0, self.batch_count):
# Load pointers for the current batch
ptr_A = self._get_pointer(mA_tensor[batch_idx], self.ab_dtype)
ptr_B = self._get_pointer(mB_tensor[batch_idx], self.ab_dtype)
ptr_D = self._get_pointer(mD_tensor[batch_idx], self.d_dtype)
gALayout = cute.zipped_divide(mA, tiler_mk)
k_tile_size = cute.size(gALayout, mode=[1, 1])
if is_tma_thr:
mA = cute.make_tensor(
ptr_A, layout=cute.make_layout(self.A_shape, stride=self.A_stride)
)
mB = cute.make_tensor(
ptr_B, layout=cute.make_layout(self.B_shape, stride=self.B_stride)
)
gA = cute.zipped_divide(mA, tiler_mk)
gB = cute.zipped_divide(mB, tiler_nk)
gA_tile = gA[(None, None), (cta_m, None, cta_l)]
gB_tile = gB[(None, None), (cta_n, None, cta_l)]
for k in cutlass.range(0, k_tile_size, 1, unroll=1):
gA_k = gA_tile[None, None, k]
gB_k = gB_tile[None, None, k]
# Scoped state management - pipeline object manages state internally
(
producer_stage_token,
idx,
) = mainloop_pipe.producer_acquire_and_get_stage()
mbar = cute_ext.get_mbarrier(producer_stage_token)
## producer_body begin ##
bufferA_sliced = bufferA[None, None, None, idx]
bufferB_sliced = bufferB[None, None, None, idx]
a_cta_v_map = cute_ext.get_cta_v_map_ab(
mA, mnk_tiler, tiled_mma, "A"
)
b_cta_v_map = cute_ext.get_cta_v_map_ab(
mB, mnk_tiler, tiled_mma, "B"
)
cute_ext.tma_load(
gA_k,
bufferA_sliced,
mbar,
cta_v_map=a_cta_v_map,
)
cute_ext.tma_load(
gB_k,
bufferB_sliced,
mbar,
cta_v_map=b_cta_v_map,
)
## producer_body end ##
mainloop_pipe.producer_commit_and_advance()
# MMA section remains same as a regular GEMM
if is_mma_thr:
producer_stage_token, idx = acc_pipe.producer_acquire_and_get_stage()
accumulators_sliced = bufferAcc[None, None, None, idx]
(updated_a_pipe, _updated_b_pipe) = cute_ext.mainloop_mma(
tiled_mma,
bufferA,
bufferB,
accumulators_sliced,
0,
k_tile_size,
mma_inst_tile_k,
mainloop_pipe,
mainloop_pipe,
)
mainloop_pipe = updated_a_pipe
acc_pipe.producer_commit_and_advance()
if is_epi_thr:
# Load the D tensor in the warp specialized section
mD = cute.make_tensor(
ptr_D, layout=cute.make_layout(self.D_shape, stride=self.D_stride)
)
_, idx = acc_pipe.consumer_wait_and_get_stage()
accumulators_sliced = bufferAcc[(None, None), 0, 0, idx]
cta_d_tile_coord = (cta_m, cta_n, cta_l)
tma_store_pipe = cute_ext.epilogue_tma_store(
cta_tile_shape_mnk,
self.use_2cta_instrs,
accumulators_sliced,
mD,
cta_d_tile_coord,
tma_store_pipe,
tma_store_warp_id,
self.epilogue_op,
)
acc_pipe.consumer_release_and_advance()
def create_tensors(l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype):
torch.manual_seed(1111)
a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype)
b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype)
d_torch_cpu = cutlass_torch.matrix(l, m, n, d_major == "m", d_dtype)
a_tensor, a_torch_gpu = cutlass_torch.cute_tensor_like(
a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
)
b_tensor, b_torch_gpu = cutlass_torch.cute_tensor_like(
b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
)
d_tensor, d_torch_gpu = cutlass_torch.cute_tensor_like(
d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16
)
return (
a_tensor,
b_tensor,
d_tensor,
a_torch_cpu,
b_torch_cpu,
d_torch_cpu,
a_torch_gpu,
b_torch_gpu,
d_torch_gpu,
)
# Helper creates a cute.Tensor from a List of device pointers
def make_tensor_of_ptrs(torch_tensor_array: List):
tensor_of_ptrs_torch = torch.tensor(
[t.data_ptr() for t in torch_tensor_array],
dtype=torch.int64,
device="cuda",
requires_grad=False,
)
tensor_of_ptrs_cute, backing_torch_tensor = cutlass_torch.cute_tensor_like(
tensor_of_ptrs_torch,
cutlass.Int64,
is_dynamic_layout=False,
assumed_align=16,
)
return tensor_of_ptrs_cute, backing_torch_tensor
def create_tensors_for_ptr_array(
l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype
):
# Store torch gpu pointers
As_torch_gpu = []
Bs_torch_gpu = []
Ds_torch_gpu = []
# Store cute tensors
A_cutes = []
B_cutes = []
D_cutes = []
for batch_idx in range(l):
torch.manual_seed(111 + batch_idx)
(
A_tensor,
B_tensor,
D_tensor,
A_torch_cpu,
B_torch_cpu,
D_torch_cpu,
A_torch_gpu,
B_torch_gpu,
D_torch_gpu,
) = create_tensors(
1, # outer loop creates a new tensor for each batch
m,
n,
k,
a_major,
b_major,
d_major,
ab_dtype,
d_dtype,
)
A_cutes.append(A_tensor)
B_cutes.append(B_tensor)
D_cutes.append(D_tensor)
As_torch_gpu.append(A_torch_gpu)
Bs_torch_gpu.append(B_torch_gpu)
Ds_torch_gpu.append(D_torch_gpu)
# Create cute tensors of pointers
a_tensor, a_backing_torch_tensor = make_tensor_of_ptrs(As_torch_gpu)
b_tensor, b_backing_torch_tensor = make_tensor_of_ptrs(Bs_torch_gpu)
d_tensor, d_backing_torch_tensor = make_tensor_of_ptrs(Ds_torch_gpu)
return (
a_tensor,
b_tensor,
d_tensor,
a_backing_torch_tensor,
b_backing_torch_tensor,
d_backing_torch_tensor,
A_cutes,
B_cutes,
D_cutes,
As_torch_gpu,
Bs_torch_gpu,
Ds_torch_gpu,
)
def compare(a_torch_cpu, b_torch_cpu, d_torch_gpu, d_dtype, tolerance):
ref = torch.einsum("mkl,nkl->mnl", a_torch_cpu, b_torch_cpu)
_, ref_torch_gpu = cutlass_torch.cute_tensor_like(
ref, d_dtype, is_dynamic_layout=True, assumed_align=16
)
ref_result = ref_torch_gpu.cpu()
torch.testing.assert_close(
d_torch_gpu.cpu(), ref_result, atol=tolerance, rtol=1e-05
)
def run(
mnkl: Tuple[int, int, int, int],
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
ab_dtype: Type[Numeric],
c_dtype: Type[Numeric],
acc_dtype: Type[Numeric],
a_major: str,
b_major: str,
c_major: str,
warmup_iterations: int = 0,
iterations: int = 1,
use_cold_l2: bool = False,
tolerance: float = 1e-02,
skip_ref_check: bool = False,
**kwargs,
):
"""Execute a Pointer array batched dense GEMM operation on Blackwell architecture with performance benchmarking.
The main difference between this and a regular bathced GEMM is that the inputs to the kernel are arrays of pointers.
Every batch of each operand (A/B/D) has its own pointer. Thus, the size of the array of pointers is the batch size.
These pointers NEED NOT be stored contiguously in memory.
This example also demonstrates how cute_ext.tma_load/cute_ext.tma_store performs automatic device side TMA updates.
Note that the dimensions of the operand for each batch are the same across all batches. That is, all batches of A have the same shape and stride, same for B and D.
This function prepares input tensors, configures and launches the GEMM kernel,
optionally performs reference validation, and benchmarks the execution performance.
:param mnkl: Problem size (M, N, K, L)
:type mnkl: Tuple[int, int, int, int]
:param mma_tiler_mn: MMA tiling size.
:type mma_tiler_mn: Tuple[int, int]
:param cluster_shape_mn: Cluster shape.
:type cluster_shape_mn: Tuple[int, int]
:param ab_dtype: Data type for input tensors A and B
:type ab_dtype: Type[Numeric]
:param d_dtype: Data type for output tensor D
:type d_dtype: Type[Numeric]
"""
print("Running Blackwell Dense GEMM test with:")
print(f"mnkl: {mnkl}")
print(f"AB dtype: {ab_dtype}, D dtype: {c_dtype}, Acc dtype: {acc_dtype}")
print(f"Matrix majors - A: {a_major}, B: {b_major}, D: {c_major}")
print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}")
print(f"Tolerance: {tolerance}")
print(f"Warmup iterations: {warmup_iterations}")
print(f"Iterations: {iterations}")
print(f"Skip reference checking: {skip_ref_check}")
print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
m, n, k, l = mnkl
ab_dtype = ab_dtype
d_major = c_major
d_dtype = c_dtype
sm100_utils.check_gemm_tma_alignment(
m,
n,
k,
ab_dtype,
ab_dtype,
d_dtype,
a_major,
b_major,
d_major,
output_tensor_name="D",
)
# a_tensor, b_tensor, d_tensor are cute Tensors where each element is an Int64 pointer to global memory
# A_cutes, B_cutes, D_cutes are lists of cute Tensors for each batch of A/B/D
(
a_tensor,
b_tensor,
d_tensor,
a_backing_torch_tensor,
b_backing_torch_tensor,
d_backing_torch_tensor,
A_cutes,
B_cutes,
D_cutes,
As_torch_gpu,
Bs_torch_gpu,
Ds_torch_gpu,
) = create_tensors_for_ptr_array(
l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype
)
ptr_array_dense_gemm = DenseGemmPtrArrayKernel(
mn_tiler=mma_tiler_mn,
mma_dtype=(ab_dtype, acc_dtype, d_dtype),
tmem_output_dtype=d_dtype,
batch_count=l,
A_shape=A_cutes[0].shape,
A_stride=A_cutes[0].stride,
B_shape=B_cutes[0].shape,
B_stride=B_cutes[0].stride,
D_shape=D_cutes[0].shape,
D_stride=D_cutes[0].stride,
)
compiled_dense_gemm = cute_ext.compile(
ptr_array_dense_gemm, a_tensor, b_tensor, d_tensor
)
query = compiled_dense_gemm.get_aux_func(
QueryDeviceWorkspaceFunc, kernel=ptr_array_dense_gemm.kernel
)
req = query(a_tensor, b_tensor, d_tensor)
workspace = torch.empty(req.size_in_bytes, dtype=torch.uint8, device="cuda")
workspace_cute = from_dlpack(workspace)
compiled_dense_gemm(a_tensor, b_tensor, d_tensor, workspace_cute)
if not skip_ref_check:
for batch_idx in range(l):
compare(
As_torch_gpu[batch_idx].cpu(),
Bs_torch_gpu[batch_idx].cpu(),
Ds_torch_gpu[batch_idx],
d_dtype,
tolerance,
)
print("check reference: PASS")
def generate_tensors():
(
a_tensor,
b_tensor,
d_tensor,
a_backing_torch_tensor,
b_backing_torch_tensor,
d_backing_torch_tensor,
A_cutes,
B_cutes,
D_cutes,
As_torch_gpu,
Bs_torch_gpu,
Ds_torch_gpu,
) = create_tensors_for_ptr_array(
l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype
)
ws = torch.empty(req.size_in_bytes, dtype=torch.uint8, device="cuda")
ws_cute = from_dlpack(ws)
args = testing.JitArguments(a_tensor, b_tensor, d_tensor, ws_cute)
args.add_to_scope([A_cutes, B_cutes, D_cutes])
return args
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
sum(
As_torch_gpu[batch_idx].numel() * As_torch_gpu[batch_idx].element_size()
for batch_idx in range(l)
)
+ sum(
Bs_torch_gpu[batch_idx].numel() * Bs_torch_gpu[batch_idx].element_size()
for batch_idx in range(l)
)
+ sum(
Ds_torch_gpu[batch_idx].numel() * Ds_torch_gpu[batch_idx].element_size()
for batch_idx in range(l)
)
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
)
exec_time = testing.benchmark(
compiled_dense_gemm,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
warmup_iterations=warmup_iterations,
iterations=iterations,
)
return exec_time
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
try:
return tuple(int(x.strip()) for x in s.split(","))
except ValueError:
raise argparse.ArgumentTypeError(
"Invalid format. Expected comma-separated integers."
)
parser = argparse.ArgumentParser(description="Example of Dense GEMM on Blackwell.")
parser.add_argument(
"--mnkl",
type=parse_comma_separated_ints,
default=(256, 256, 512, 1),
help="mnkl dimensions (comma-separated)",
)
parser.add_argument(
"--mma_tiler_mn",
type=parse_comma_separated_ints,
default=(128, 128),
help="Mma tile shape (comma-separated)",
)
parser.add_argument(
"--cluster_shape_mn",
type=parse_comma_separated_ints,
default=(1, 1),
help="Cluster shape (comma-separated)",
)
parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float32)
parser.add_argument("--d_dtype", type=cutlass.dtype, default=cutlass.Float32)
parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32)
parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k")
parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k")
parser.add_argument("--d_major", choices=["n", "m"], type=str, default="n")
parser.add_argument(
"--warmup_iterations", type=int, default=0, help="Warmup iterations"
)
parser.add_argument(
"--iterations", type=int, default=1, help="Number of iterations"
)
parser.add_argument("--use_cold_l2", action="store_true", help="Use cold L2")
parser.add_argument(
"--tolerance", type=float, default=1e-02, help="Tolerance for validation"
)
parser.add_argument(
"--skip_ref_check", action="store_true", help="Skip reference checking"
)
args = parser.parse_args()
if len(args.mnkl) != 4:
parser.error("--mnkl must contain exactly 4 values")
if len(args.mma_tiler_mn) != 2:
parser.error("--mma_tiler_mn must contain exactly 2 values")
exec_time = run(
args.mnkl,
args.mma_tiler_mn,
args.cluster_shape_mn,
args.ab_dtype,
args.d_dtype,
args.acc_dtype,
args.a_major,
args.b_major,
args.d_major,
args.warmup_iterations,
args.iterations,
args.use_cold_l2,
args.tolerance,
args.skip_ref_check,
)
print(f"Execution time: {exec_time} microseconds per iteration")