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

v4.1 release update v2. (#2481)

Browse Source
This commit is contained in:
Junkai-Wu
2025-07-22 10:03:55 +08:00
committed by GitHub
parent 9baa06dd57
commit fd6cfe1ed0
179 changed files with 7878 additions and 1286 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
examples/python/CuTeDSL/ampere/elementwise_add.py
View File

@@ -90,18 +90,17 @@ If you already know the TV layout you want to use for your tiled copy, CuTe DSL
# Tile input tensor to thread blocks: ((TileM,TileN),(RestM,RestN))
gA = cute.zipped_divide(mA, tiler_mn)
where `tiler_mn` is the tile size per thread block and `tv_layout` is the TV layout which maps
thread index and inter-thread index of data array per thread to logical coordinates of elements in
input and output tensors.
Then we can build tiled copy for input and output tensors with `cute.make_tiled_copy` utility.
Then we can build tiled copy for input and output tensors with `cute.make_tiled_copy_tv` utility, which
infers the tiler and tv layout for the tiled copy automatically, where `tiler` is the tile size per thread
block and `tv_layout` is the TV layout which maps thread index and inter-thread index of data array per
thread to logical coordinates of elements in input and output tensors.
.. code-block:: python
blkA = gA[((None, None), bidx)] # (TileM,TileN)
copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type)
tiled_copy_A = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn)
tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
# get slice of tiled_copy_A for current thread
thr_copy_A = tiled_copy_A.get_slice(tidx)
@@ -140,8 +139,8 @@ def elementwise_add_kernel(
gC: cute.Tensor,
cC: cute.Tensor, # coordinate tensor
shape: cute.Shape,
tv_layout: cute.Layout,
tiler_mn: cute.Shape,
thr_layout: cute.Layout,
val_layout: cute.Layout,
):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
@@ -165,9 +164,9 @@ def elementwise_add_kernel(
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(copy_atom_load, tv_layout, tiler_mn)
tiled_copy_B = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn)
tiled_copy_C = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn)
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)
@@ -254,7 +253,7 @@ def elementwise_add(mA, mB, mC, copy_bits: cutlass.Constexpr = 128):
cC = cute.zipped_divide(idC, tiler=tiler_mn)
print(f"[DSL INFO] coord tensor = {cC.type}")
elementwise_add_kernel(gA, gB, gC, cC, mC.shape, tv_layout, tiler_mn).launch(
elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch(
grid=[cute.size(gC, mode=[1]), 1, 1],
block=[cute.size(tv_layout, mode=[0]), 1, 1],
)
@@ -362,7 +361,7 @@ def run_elementwise_add(
workspace_generator=generate_tensors,
workspace_count=10,
warmup_iterations=warmup_iterations,
profiling_iterations=iterations,
iterations=iterations,
)
# Print execution results

2
examples/python/CuTeDSL/ampere/elementwise_apply.py
View File

@@ -353,7 +353,7 @@ def run_elementwise_apply_and_verify(
current_stream,
),
warmup_iterations=warmup_iterations,
profiling_iterations=iterations,
iterations=iterations,
use_cuda_graphs=True,
stream=current_stream,
)

212
examples/python/CuTeDSL/ampere/flash_attention_v2.py
View File

@@ -32,13 +32,13 @@ from typing import Type, Union, Callable
import torch
import cuda.bindings.driver as cuda
import cutlass.cute.testing as testing
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, warp
import cutlass.torch as cutlass_torch
from cutlass.cute.runtime import from_dlpack
import cutlass.utils.ampere_helpers as sm80_utils
import cutlass.utils as utils
"""
A flash attention v2 forward pass example for NVIDIA Ampere SM80 architecture using CUTE DSL.
@@ -163,7 +163,7 @@ class FlashAttentionForwardAmpere:
# Check if block size setting is out of shared memory capacity
# Shared memory usage: Q tile + (K tile + V tile) where K and V use the same tile size
smem_usage = (m_block_size * head_dim + n_block_size * head_dim * 2) * 2
smem_capacity = sm80_utils.SMEM_CAPACITY["sm80"]
smem_capacity = utils.get_smem_capacity_in_bytes("sm_80")
if smem_usage > smem_capacity:
return False
@@ -469,21 +469,9 @@ class FlashAttentionForwardAmpere:
warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4),
self._dtype,
)
smem_tiled_copy_Q = cute.make_tiled_copy(
smem_copy_atom_Q,
layout_tv=tiled_mma.tv_layout_A_tiled,
tiler_mn=(tiled_mma.get_tile_size(0), tiled_mma.get_tile_size(2)),
)
smem_tiled_copy_K = cute.make_tiled_copy(
smem_copy_atom_K,
layout_tv=tiled_mma.tv_layout_B_tiled,
tiler_mn=(tiled_mma.get_tile_size(1), tiled_mma.get_tile_size(2)),
)
smem_tiled_copy_V = cute.make_tiled_copy(
smem_copy_atom_V,
layout_tv=tiled_mma.tv_layout_B_tiled,
tiler_mn=(tiled_mma.get_tile_size(1), tiled_mma.get_tile_size(2)),
)
smem_tiled_copy_Q = cute.make_tiled_copy_A(smem_copy_atom_Q, tiled_mma)
smem_tiled_copy_K = cute.make_tiled_copy_B(smem_copy_atom_K, tiled_mma)
smem_tiled_copy_V = cute.make_tiled_copy_B(smem_copy_atom_V, tiled_mma)
smem_thr_copy_Q = smem_tiled_copy_Q.get_slice(tidx)
smem_thr_copy_K = smem_tiled_copy_K.get_slice(tidx)
@@ -702,11 +690,7 @@ class FlashAttentionForwardAmpere:
cute.nvgpu.CopyUniversalOp(), self._dtype
)
# tiled copy atom for O
smem_tiled_copy_O = cute.make_tiled_copy(
smem_copy_atom_O,
layout_tv=tiled_mma.tv_layout_C_tiled,
tiler_mn=(tiled_mma.get_tile_size(0), tiled_mma.get_tile_size(1)),
)
smem_tiled_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma)
smem_thr_copy_O = smem_tiled_copy_O.get_slice(tidx)
taccOrO = smem_thr_copy_O.retile(rO)
taccOsO = smem_thr_copy_O.partition_D(sO)
@@ -1178,7 +1162,7 @@ class FlashAttentionForwardAmpere:
return cute.arch.exp2(x)
def run_flash_attention_fwd(
def run(
dtype: Type[cutlass.Numeric],
batch_size: int,
seqlen_q: int,
@@ -1193,6 +1177,8 @@ def run_flash_attention_fwd(
warmup_iterations: int = 0,
iterations: int = 1,
skip_ref_check: bool = False,
use_cold_l2: bool = False,
**kwargs,
):
# Skip unsupported testcase
if not FlashAttentionForwardAmpere.can_implement(
@@ -1207,6 +1193,23 @@ def run_flash_attention_fwd(
f"Unsupported testcase {dtype}, {head_dim}, {m_block_size}, {n_block_size}, {num_threads}, {is_causal}"
)
print(f"Running Ampere SM80 FlashAttentionForward test with:")
print(f" dtype: {dtype}")
print(f" batch_size: {batch_size}")
print(f" seqlen_q: {seqlen_q}")
print(f" seqlen_k: {seqlen_k}")
print(f" num_head: {num_head}")
print(f" head_dim: {head_dim}")
print(f" softmax_scale: {softmax_scale}")
print(f" m_block_size: {m_block_size}")
print(f" n_block_size: {n_block_size}")
print(f" num_threads: {num_threads}")
print(f" is_causal: {is_causal}")
print(f" warmup_iterations: {warmup_iterations}")
print(f" iterations: {iterations}")
print(f" skip_ref_check: {skip_ref_check}")
print(f" use_cold_l2: {use_cold_l2}")
# Create tensor Q/K/V/O
def create_tensor(
batch_size: int,
@@ -1217,22 +1220,28 @@ def run_flash_attention_fwd(
) -> cute.Tensor:
# (batch_size, seqlen, num_head, head_dim)
shape = (batch_size, seqlen, num_head, head_dim)
return (
torch.empty(*shape, dtype=torch.int32).random_(-2, 2).to(dtype=dtype).cuda()
torch_tensor = (
torch.empty(*shape, dtype=torch.int32)
.random_(-2, 2)
.to(dtype=cutlass_torch.dtype(dtype))
.cuda()
)
# assume input is 16B aligned.
cute_tensor = (
from_dlpack(torch_tensor, assumed_align=16)
.mark_layout_dynamic(leading_dim=3)
.mark_compact_shape_dynamic(
mode=3,
stride_order=torch_tensor.dim_order(),
divisibility=(128 // dtype.width),
)
)
return cute_tensor, torch_tensor
q = create_tensor(
batch_size, seqlen_q, num_head, head_dim, cutlass_torch.dtype(dtype)
)
k = create_tensor(
batch_size, seqlen_k, num_head, head_dim, cutlass_torch.dtype(dtype)
)
v = create_tensor(
batch_size, seqlen_k, num_head, head_dim, cutlass_torch.dtype(dtype)
)
o = create_tensor(
batch_size, seqlen_q, num_head, head_dim, cutlass_torch.dtype(dtype)
)
q, q_torch = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype)
k, k_torch = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype)
v, v_torch = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype)
o, o_torch = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype)
fa2_fwd = FlashAttentionForwardAmpere(
head_dim,
@@ -1241,78 +1250,63 @@ def run_flash_attention_fwd(
num_threads,
is_causal,
)
# assume input is 16B align.
q_tensor = (
from_dlpack(q, assumed_align=16)
.mark_layout_dynamic(leading_dim=3)
.mark_compact_shape_dynamic(
mode=3, stride_order=q.dim_order(), divisibility=(128 // dtype.width)
)
)
k_tensor = (
from_dlpack(k, assumed_align=16)
.mark_layout_dynamic(leading_dim=3)
.mark_compact_shape_dynamic(
mode=3, stride_order=k.dim_order(), divisibility=(128 // dtype.width)
)
)
v_tensor = (
from_dlpack(v, assumed_align=16)
.mark_layout_dynamic(leading_dim=3)
.mark_compact_shape_dynamic(
mode=3, stride_order=v.dim_order(), divisibility=(128 // dtype.width)
)
)
o_tensor = (
from_dlpack(o, assumed_align=16)
.mark_layout_dynamic(leading_dim=3)
.mark_compact_shape_dynamic(
mode=3, stride_order=o.dim_order(), divisibility=(128 // dtype.width)
)
)
# Get current CUDA stream from PyTorch
torch_stream = torch.cuda.current_stream()
# Get the raw stream pointer as a CUstream
current_stream = cuda.CUstream(torch_stream.cuda_stream)
# compile the fa2 forward pass
compiled_fa2_fwd = cute.compile(
fa2_fwd, q_tensor, k_tensor, v_tensor, o_tensor, softmax_scale, current_stream
compiled_fa2_fwd = cute.compile(fa2_fwd, q, k, v, o, softmax_scale, current_stream)
if not skip_ref_check:
compiled_fa2_fwd(q, k, v, o, softmax_scale, current_stream)
torch.cuda.synchronize()
q_ref = q_torch.permute(0, 2, 1, 3)
k_ref = k_torch.permute(0, 2, 1, 3)
v_ref = v_torch.permute(0, 2, 1, 3)
torch.backends.cuda.enable_flash_sdp(enabled=True)
ref_o = torch.nn.functional.scaled_dot_product_attention(
q_ref, k_ref, v_ref, scale=softmax_scale, is_causal=is_causal
).permute(0, 2, 1, 3)
torch.testing.assert_close(o_torch.cpu(), ref_o.cpu(), atol=1e-02, rtol=1e-04)
print("Results verified successfully!")
def generate_tensors():
q_workspace, _ = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype)
k_workspace, _ = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype)
v_workspace, _ = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype)
o_workspace, _ = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype)
return testing.JitArguments(
q_workspace,
k_workspace,
v_workspace,
o_workspace,
softmax_scale,
current_stream,
)
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
q_torch.numel() * q_torch.element_size()
+ k_torch.numel() * k_torch.element_size()
+ v_torch.numel() * v_torch.element_size()
+ o_torch.numel() * o_torch.element_size()
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
)
avg_time_us = testing.benchmark(
compiled_fa2_fwd,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
stream=current_stream,
warmup_iterations=warmup_iterations,
iterations=iterations,
)
# warmup
for _ in range(warmup_iterations):
compiled_fa2_fwd(
q_tensor,
k_tensor,
v_tensor,
o_tensor,
softmax_scale,
current_stream,
)
# run the compiled fa2 forward pass
for _ in range(iterations):
compiled_fa2_fwd(
q_tensor,
k_tensor,
v_tensor,
o_tensor,
softmax_scale,
current_stream,
)
torch.cuda.synchronize()
if skip_ref_check:
return
# reference implementation
q_ref = q.permute(0, 2, 1, 3)
k_ref = k.permute(0, 2, 1, 3)
v_ref = v.permute(0, 2, 1, 3)
torch.backends.cuda.enable_flash_sdp(enabled=True)
ref_o = torch.nn.functional.scaled_dot_product_attention(
q_ref, k_ref, v_ref, scale=softmax_scale, is_causal=is_causal
).permute(0, 2, 1, 3)
torch.testing.assert_close(o.cpu(), ref_o.cpu(), atol=1e-02, rtol=1e-04)
return avg_time_us # Return execution time in microseconds
if __name__ == "__main__":
parser = argparse.ArgumentParser(
@@ -1334,9 +1328,15 @@ if __name__ == "__main__":
parser.add_argument(
"--skip_ref_check", action="store_true", help="Skip reference check"
)
parser.add_argument(
"--use_cold_l2",
action="store_true",
default=False,
help="Use circular buffer tensor sets to ensure L2 cold cache",
)
args = parser.parse_args()
run_flash_attention_fwd(
run(
args.dtype,
args.batch_size,
args.seqlen_q,
@@ -1348,6 +1348,10 @@ if __name__ == "__main__":
args.n_block_size,
args.num_threads,
args.is_causal,
args.warmup_iterations,
args.iterations,
args.skip_ref_check,
args.use_cold_l2,
)
print("PASS")

131
examples/python/CuTeDSL/ampere/sgemm.py
View File

@@ -634,16 +634,50 @@ class SGemm:
return
def main(
def run(
mnk: Tuple[int, int, int],
a_major: str,
b_major: str,
c_major: str,
problem_shape: Tuple[int, int, int],
static_shape: bool = False,
warmup_iterations: int = 2,
iterations: int = 100,
skip_ref_check: bool = False,
use_cold_l2: bool = False,
**kwargs,
):
M, N, K = problem_shape
"""Execute SIMT GEMM operation and benchmark performance.
:param mnk: GEMM problem size (M, N, K, L)
:type mnk: Tuple[int, int, int, int]
:param a_major: Memory layout of tensor A
:type a_major: str
:param b_major: Memory layout of tensor B
:type b_major: str
:param c_major: Memory layout of tensor C
:type c_major: str
:param static_shape: Whether to use static shape optimization, defaults to False
:type static_shape: bool, optional
:param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 2
:type warmup_iterations: int, optional
:param iterations: Number of benchmark iterations to run, defaults to 100
:type iterations: int, optional
:param skip_ref_check: Skip validation against reference implementation, defaults to False
:type skip_ref_check: bool, optional
:param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False
:type use_cold_l2: bool, optional
:return: Execution time of the GEMM kernel in microseconds
:rtype: float
"""
print(f"Running Ampere SIMT GEMM example:")
print(f"mnk: {mnk}")
print(f"A major: {a_major}, B major: {b_major}, C major: {c_major}")
print(f"Static shape: {static_shape}")
print(f"Warmup iterations: {warmup_iterations}")
print(f"Iterations: {iterations}")
print(f"Skip reference checking: {skip_ref_check}")
print(f"Use cold L2: {use_cold_l2}")
M, N, K = mnk
# Create and permute tensor A/B/C
def create_and_permute_tensor(mode0, mode1, is_mode0_major, dtype):
@@ -710,20 +744,6 @@ def main(
print("Executing GEMM kernel...")
avg_time_us = testing.benchmark(
gemm,
kernel_arguments=testing.JitArguments(
a_tensor, b_tensor, c_tensor, current_stream
),
warmup_iterations=warmup_iterations,
profiling_iterations=iterations,
use_cuda_graphs=False,
stream=current_stream,
)
# Print execution results
print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
if not skip_ref_check:
gemm(a_tensor, b_tensor, c_tensor)
torch.cuda.synchronize()
@@ -732,6 +752,71 @@ def main(
torch.testing.assert_close(c.cpu(), ref.cpu(), atol=1e-03, rtol=1e-05)
print("Results verified successfully!")
def generate_tensors():
# Create new tensors for each workspace to ensure cold L2 cache
a_workspace = create_and_permute_tensor(M, K, a_major == "m", torch.float32)
b_workspace = create_and_permute_tensor(N, K, b_major == "n", torch.float32)
c_workspace = create_and_permute_tensor(M, N, c_major == "m", torch.float32)
if static_shape:
a_tensor_workspace = (
from_dlpack(a_workspace, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if a_major == "k" else 0))
.mark_compact_shape_dynamic(
mode=(1 if a_major == "k" else 0),
divisibility=divisibility_a,
)
)
else:
a_tensor_workspace = from_dlpack(a_workspace, assumed_align=16)
b_tensor_workspace = (
from_dlpack(b_workspace, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if b_major == "k" else 0))
.mark_compact_shape_dynamic(
mode=(1 if b_major == "k" else 0),
divisibility=divisibility_b,
)
)
c_tensor_workspace = (
from_dlpack(c_workspace, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if c_major == "n" else 0))
.mark_compact_shape_dynamic(
mode=(1 if c_major == "n" else 0),
divisibility=divisibility_c,
)
)
return testing.JitArguments(
a_tensor_workspace, b_tensor_workspace, c_tensor_workspace, current_stream
)
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
a.numel() * a.element_size()
+ b.numel() * b.element_size()
+ c.numel() * c.element_size()
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
)
avg_time_us = testing.benchmark(
gemm,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
stream=current_stream,
warmup_iterations=warmup_iterations,
iterations=iterations,
)
# Print execution results
print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
return avg_time_us # Return execution time in microseconds
if __name__ == "__main__":
@@ -753,19 +838,27 @@ if __name__ == "__main__":
parser.add_argument("--warmup_iterations", default=2, type=int)
parser.add_argument("--iterations", default=100, type=int)
parser.add_argument("--skip_ref_check", action="store_true")
parser.add_argument(
"--use_cold_l2",
action="store_true",
default=False,
help="Use circular buffer tensor sets to ensure L2 cold cache",
)
args = parser.parse_args()
print("Running SIMT GEMM example:")
torch.manual_seed(1024)
main(
run(
args.mnk,
args.a_major,
args.b_major,
args.c_major,
args.mnk,
args.static_shape,
args.warmup_iterations,
args.iterations,
args.skip_ref_check,
args.use_cold_l2,
)
print("PASS")

150
examples/python/CuTeDSL/ampere/tensorop_gemm.py
View File

@@ -51,7 +51,7 @@ This GEMM kernel supports the following features:
- Utilizes Ampere's tensor cores for matrix multiply-accumulate (MMA) operations
- Threadblock rasterization to improve data re-use
- Supports multi-stage pipeline to overlap computation and memory access
- Implements shared memory buffering for epilogue to increase coalesed global memory access
- Implements shared memory buffering for epilogue to increase coalesced global memory access
This GEMM works as follows:
1. Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using asynchronous copies.
@@ -214,7 +214,7 @@ class TensorOpGemm:
atom_async_copy, mB.element_type, self.b_major_mode, ab_copy_bits
)
# Creates a synchonous copy atom and thread layouts for the epilogue
# Creates a synchronous copy atom and thread layouts for the epilogue
c_copy_bits = 128
atom_sync_copy = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
@@ -550,16 +550,8 @@ class TensorOpGemm:
# Creates the tiled copy so that it matches the thread-value layout
# expected by the tiled mma
tiled_copy_s2r_A = cute.make_tiled_copy(
atom_copy_s2r_A,
layout_tv=tiled_mma.tv_layout_A_tiled,
tiler_mn=(tiled_mma.get_tile_size(0), tiled_mma.get_tile_size(2)),
)
tiled_copy_s2r_B = cute.make_tiled_copy(
atom_copy_s2r_B,
layout_tv=tiled_mma.tv_layout_B_tiled,
tiler_mn=(tiled_mma.get_tile_size(1), tiled_mma.get_tile_size(2)),
)
tiled_copy_s2r_A = cute.make_tiled_copy_A(atom_copy_s2r_A, tiled_mma)
tiled_copy_s2r_B = cute.make_tiled_copy_B(atom_copy_s2r_B, tiled_mma)
thr_copy_ldmatrix_A = tiled_copy_s2r_A.get_slice(tidx)
thr_copy_ldmatrix_B = tiled_copy_s2r_B.get_slice(tidx)
@@ -836,8 +828,7 @@ class TensorOpGemm:
if major_mode == utils.LayoutEnum.ROW_MAJOR
else cute.make_layout((copy_elems, 1))
)
tiler_mn, layout_tv = cute.make_layout_tv(thread_layout, value_layout)
return cute.make_tiled_copy(atom_copy, layout_tv, tiler_mn)
return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout)
def raster_tile(self, i, j, f):
new_i = i // f
@@ -845,20 +836,33 @@ class TensorOpGemm:
return (new_i, new_j)
def run_tensor_op_gemm(
def run(
a_major: str,
b_major: str,
c_major: str,
ab_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
acc_dtype: Type[cutlass.Numeric],
problem_shape: Tuple[int, int, int, int],
mnkl: Tuple[int, int, int, int],
atom_layout_mnk: Tuple[int, int, int],
warmup_iterations: int = 2,
iterations: int = 100,
skip_ref_check: bool = False,
use_cold_l2: bool = False,
**kwargs,
):
M, N, K, L = problem_shape
print(f"Running Ampere tensor core GEMM example:")
print(f"mnkl: {mnkl}")
print(
f"A dtype: {ab_dtype}, B dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}"
)
print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}")
print(f"Atoms layout: {atom_layout_mnk}")
print(f"Warmup iterations: {warmup_iterations}")
print(f"Iterations: {iterations}")
print(f"Skip reference checking: {skip_ref_check}")
print(f"Use cold L2: {use_cold_l2}")
M, N, K, L = mnkl
# Create and permute tensor A/B/C
def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype):
@@ -866,23 +870,28 @@ def run_tensor_op_gemm(
# else: (l, mode0, mode1) -> (mode0, mode1, l)
shape = (l, mode1, mode0) if is_mode0_major else (l, mode0, mode1)
permute_order = (2, 1, 0) if is_mode0_major else (1, 2, 0)
return (
torch_tensor = (
torch.empty(*shape, dtype=torch.int32)
.random_(-2, 2)
.to(dtype=dtype)
.to(dtype=cutlass_torch.dtype(dtype))
.permute(permute_order)
.cuda()
)
# assume input is 16B aligned
cute_tensor = (
from_dlpack(torch_tensor, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if not is_mode0_major else 0))
.mark_compact_shape_dynamic(
mode=(1 if not is_mode0_major else 0),
stride_order=(2, 0, 1) if not is_mode0_major else (2, 1, 0),
divisibility=(128 // dtype.width),
)
)
return cute_tensor, torch_tensor
a = create_and_permute_tensor(
L, M, K, a_major == "m", cutlass_torch.dtype(ab_dtype)
)
b = create_and_permute_tensor(
L, N, K, b_major == "n", cutlass_torch.dtype(ab_dtype)
)
c = create_and_permute_tensor(L, M, N, c_major == "m", cutlass_torch.dtype(c_dtype))
ref = torch.einsum("mkl,nkl->mnl", a, b).to(cutlass_torch.dtype(c_dtype))
mA, a_torch = create_and_permute_tensor(L, M, K, a_major == "m", ab_dtype)
mB, b_torch = create_and_permute_tensor(L, N, K, b_major == "n", ab_dtype)
mC, c_torch = create_and_permute_tensor(L, M, N, c_major == "m", c_dtype)
tensor_op_gemm = TensorOpGemm(
ab_dtype,
@@ -891,56 +900,49 @@ def run_tensor_op_gemm(
atom_layout_mnk,
)
# assume input is 16B aligned
a_tensor = (
from_dlpack(a, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if a_major == "k" else 0))
.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=(128 // ab_dtype.width),
)
)
b_tensor = (
from_dlpack(b, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if b_major == "k" else 0))
.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=(128 // ab_dtype.width),
)
)
c_tensor = (
from_dlpack(c, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if c_major == "n" else 0))
.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=(128 // c_dtype.width),
)
)
print("Compiling kernel with cute.compile ...")
gemm = cute.compile(tensor_op_gemm, a_tensor, b_tensor, c_tensor)
compiled_gemm = cute.compile(tensor_op_gemm, mA, mB, mC)
print("Executing GEMM kernel...")
if not skip_ref_check:
ref = torch.einsum(
"mkl,nkl->mnl",
a_torch.to(dtype=torch.float32),
b_torch.to(dtype=torch.float32),
).to(cutlass_torch.dtype(c_dtype))
compiled_gemm(mA, mB, mC)
print("Verifying results...")
torch.testing.assert_close(c_torch.cpu(), ref.cpu(), atol=1e-03, rtol=1e-05)
print("Results verified successfully!")
def generate_tensors():
a_workspace, _ = create_and_permute_tensor(L, M, K, a_major == "m", ab_dtype)
b_workspace, _ = create_and_permute_tensor(L, N, K, b_major == "n", ab_dtype)
c_workspace, _ = create_and_permute_tensor(L, M, N, c_major == "m", c_dtype)
return testing.JitArguments(a_workspace, b_workspace, c_workspace)
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
a_torch.numel() * a_torch.element_size()
+ b_torch.numel() * b_torch.element_size()
+ c_torch.numel() * c_torch.element_size()
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
)
avg_time_us = testing.benchmark(
gemm,
kernel_arguments=testing.JitArguments(a_tensor, b_tensor, c_tensor),
compiled_gemm,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
warmup_iterations=warmup_iterations,
profiling_iterations=iterations,
iterations=iterations,
use_cuda_graphs=False,
)
print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
if not skip_ref_check:
gemm(a_tensor, b_tensor, c_tensor)
print("Verifying results...")
torch.testing.assert_close(c.cpu(), ref.cpu(), atol=1e-03, rtol=1e-05)
print("Results verified successfully!")
return avg_time_us # Return execution time in microseconds
if __name__ == "__main__":
@@ -985,10 +987,15 @@ if __name__ == "__main__":
parser.add_argument("--warmup_iterations", default=2, type=int)
parser.add_argument("--iterations", default=100, type=int)
parser.add_argument("--skip_ref_check", action="store_true")
parser.add_argument(
"--use_cold_l2",
action="store_true",
default=False,
help="Use circular buffer tensor sets to ensure L2 cold cache",
)
args = parser.parse_args()
print("Running Ampere tensor core GEMM example:")
run_tensor_op_gemm(
run(
args.a_major,
args.b_major,
args.c_major,
@@ -1000,5 +1007,6 @@ if __name__ == "__main__":
args.warmup_iterations,
args.iterations,
args.skip_ref_check,
args.use_cold_l2,
)
print("PASS")