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/ampere/call_bypass_dlpack.py
+6 -6
View File
@@ -89,7 +89,7 @@ def tensor_op_gemm_wrapper(
k: cutlass.Int32,
l: cutlass.Int32,
):
print(f"\n[DSL INFO] Input Parameters:")
print("\n[DSL INFO] Input Parameters:")
print(f"[DSL INFO] mnkl: {(m, n, k, l)}")
# Assume alignment of shape to call tensorop_gemm example
@@ -111,7 +111,7 @@ def tensor_op_gemm_wrapper(
tensor_op_gemm = TensorOpGemm(
a_ptr.value_type, c_ptr.value_type, cutlass.Float32, (2, 2, 1)
)
print(f"\n[DSL INFO] Created TensorOpGemm instance")
print("\n[DSL INFO] Created TensorOpGemm instance")
print(f"[DSL INFO] Input dtype: {a_ptr.value_type}")
print(f"[DSL INFO] Output dtype: {c_ptr.value_type}")
print(f"[DSL INFO] Accumulation dtype: {cutlass.Float32}")
@@ -119,11 +119,11 @@ def tensor_op_gemm_wrapper(
# No need to compile inside jit function
tensor_op_gemm(mA, mB, mC)
print(f"\n[DSL INFO] Executed TensorOpGemm")
print("\n[DSL INFO] Executed TensorOpGemm")
def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
print(f"\nRunning TensorOpGemm test with:")
print("\nRunning TensorOpGemm test with:")
print(f"Tensor dimensions: {mnkl}")
# (M,K,L)
@@ -139,7 +139,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
mnkl[3], mnkl[0], mnkl[1], dtype=torch.float16, device="cuda"
).permute(1, 2, 0)
print(f"Input tensor shapes:")
print("Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
@@ -158,7 +158,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
ref = torch.einsum("mkl,nkl->mnl", a, b)
torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
print(f"\n[DSL INFO] Results verified successfully!")
print("\n[DSL INFO] Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3, :3]}")
examples/python/CuTeDSL/ampere/call_from_jit.py
+7 -7
View File
@@ -169,7 +169,7 @@ def tensor_op_gemm_wrapper(
acc_dtype: Type[cutlass.Numeric],
atom_layout_mnk: cutlass.Constexpr[tuple[int, int, int]],
):
print(f"\n[DSL INFO] Input Parameters:")
print("\n[DSL INFO] Input Parameters:")
print(f"[DSL INFO] mnkl: {mnkl}")
print(f"[DSL INFO] buffer_a: {buffer_a}")
print(f"[DSL INFO] buffer_b: {buffer_b}")
@@ -181,7 +181,7 @@ def tensor_op_gemm_wrapper(
mB = buffer_b.to_tensor(cute.select(mnkl, mode=[3, 1, 2]))
mC = buffer_c.to_tensor(cute.select(mnkl, mode=[3, 0, 1]))
print(f"\n[DSL INFO] Created Tensors:")
print("\n[DSL INFO] Created Tensors:")
print(f"[DSL INFO] mA = {mA}")
print(f"[DSL INFO] mB = {mB}")
print(f"[DSL INFO] mC = {mC}")
@@ -192,7 +192,7 @@ def tensor_op_gemm_wrapper(
acc_dtype,
atom_layout_mnk,
)
print(f"\n[DSL INFO] Created TensorOpGemm instance")
print("\n[DSL INFO] Created TensorOpGemm instance")
print(f"[DSL INFO] Input dtype: {buffer_a.ptr.value_type}")
print(f"[DSL INFO] Output dtype: {buffer_c.ptr.value_type}")
print(f"[DSL INFO] Accumulation dtype: {acc_dtype}")
@@ -200,11 +200,11 @@ def tensor_op_gemm_wrapper(
# No need to compile inside jit function
tensor_op_gemm(mA, mB, mC)
print(f"\n[DSL INFO] Executed TensorOpGemm")
print("\n[DSL INFO] Executed TensorOpGemm")
def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
print(f"\nRunning TensorOpGemm test with:")
print("\nRunning TensorOpGemm test with:")
print(f"Tensor dimensions: {mnkl}")
ab_dtype = cutlass.Float16
@@ -220,7 +220,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
mnkl[3], mnkl[0], mnkl[1], dtype=torch_dtype(c_dtype), device="cuda"
)
print(f"Input tensor shapes:")
print("Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
@@ -251,7 +251,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
ref = torch.einsum("lmk,lnk->lmn", a, b)
torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
print(f"\n[DSL INFO] Results verified successfully!")
print("\n[DSL INFO] Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3, :3]}")
examples/python/CuTeDSL/ampere/elementwise_add.py
+13 -12
View File
@@ -28,11 +28,10 @@
import argparse
import torch
import time
from typing import Type
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
@@ -154,7 +153,7 @@ def elementwise_add_kernel(
blkCrd = cC[blk_coord] # (TileM, TileN)
# Note: these prints only run at compile/jit time
print(f"[DSL INFO] Sliced Tensors per thread block:")
print("[DSL INFO] Sliced Tensors per thread block:")
print(f"[DSL INFO] blkA = {blkA.type}")
print(f"[DSL INFO] blkB = {blkB.type}")
print(f"[DSL INFO] blkC = {blkC.type}")
@@ -182,9 +181,9 @@ def elementwise_add_kernel(
frgC = cute.make_fragment_like(thrC)
thrCrd = thr_copy_C.partition_S(blkCrd)
frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
print(f"[DSL INFO] Sliced Tensors per thread:")
print("[DSL INFO] Sliced Tensors per thread:")
print(f"[DSL INFO] thrA = {thrA.type}")
print(f"[DSL INFO] thrB = {thrB.type}")
print(f"[DSL INFO] thrC = {thrC.type}")
@@ -233,18 +232,18 @@ def elementwise_add(mA, mB, mC, copy_bits: cutlass.Constexpr = 128):
val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
print(f"[DSL INFO] Input Tensors:")
print("[DSL INFO] Input Tensors:")
print(f"[DSL INFO] mA = {mA.type}")
print(f"[DSL INFO] mB = {mB.type}")
print(f"[DSL INFO] Tiling Parameters:")
print("[DSL INFO] Tiling Parameters:")
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
print(f"[DSL INFO] tv_layout = {tv_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))
print(f"[DSL INFO] Tiled Tensors:")
print("[DSL INFO] Tiled Tensors:")
print(f"[DSL INFO] gA = {gA.type}")
print(f"[DSL INFO] gB = {gB.type}")
print(f"[DSL INFO] gC = {gC.type}")
@@ -271,7 +270,7 @@ def run_elementwise_add(
warmup_iterations=2,
iterations=200,
):
print(f"\nRunning Elementwise Add test with:")
print("\nRunning Elementwise Add test with:")
print(f"Tensor dimensions: [{M}, {N}]")
print(f"Input and Output Data type: {dtype}")
@@ -285,7 +284,7 @@ def run_elementwise_add(
c = torch.zeros_like(a)
print(f"Input tensor shapes:")
print("Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
@@ -307,7 +306,9 @@ def run_elementwise_add(
print("Compiling kernel with cute.compile ...")
start_time = time.time()
compiled_func = cute.compile(elementwise_add, a_tensor, b_tensor, c_tensor)
compiled_func = cute.compile(
elementwise_add, a_tensor, b_tensor, c_tensor, options="--generate-line-info"
)
compilation_time = time.time() - start_time
print(f"Compilation time: {compilation_time:.4f} seconds")
@@ -386,7 +387,7 @@ if __name__ == "__main__":
args = parser.parse_args()
if not torch.cuda.is_available():
raise RuntimeError(f"Ampere GPU is required to run this example!")
raise RuntimeError("Ampere GPU is required to run this example!")
run_elementwise_add(
args.M,
examples/python/CuTeDSL/ampere/elementwise_apply.py
+140 -137
View File
@@ -30,17 +30,18 @@
import argparse
import operator
import time
from typing import Type, List
from functools import partial
from typing import List, Type
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.torch as cutlass_torch
import torch
from cutlass.cute.runtime import from_dlpack
import cutlass
"""
An Elementwise Apply Example using CuTe DSL.
@@ -78,103 +79,83 @@ while maintaining high performance through efficient memory access patterns.
@cute.kernel
def elementwise_apply_kernel(
op: cutlass.Constexpr,
inputs: List[cute.Tensor],
gC: cute.Tensor,
mInputs: List[cute.Tensor],
mC: cute.Tensor,
cC: cute.Tensor, # coordinate tensor
shape: cute.Shape,
tv_layout: cute.Layout, # (tid, vid) -> logic coord
):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bidx, bidy, _ = cute.arch.block_idx()
###############################################################################
# Slice to local tile of thread block
###############################################################################
blk_crd = ((None, None), (bidx, bidy))
# slice for CTAs
cta_coord = ((None, None), bidx)
# logical coord -> address
# Leverage the meta-programming capability of the DSL to slice the tensors for each input
# All for loops below on input tensors would be fully unrolled automatically at compile time
ctaInputs = [t[cta_coord] for t in inputs] # (TileM, TileN)
ctaC = gC[cta_coord] # (TileM, TileN)
ctaCrd = cC[cta_coord] # (TileM, TileN)
# logical coord -> memory address
gInputs = [t[blk_crd] for t in mInputs] # (TileM, TileN)
gC = mC[blk_crd] # (TileM, TileN)
gCrd = cC[blk_crd] # (TileM, TileN)
print(f"[DSL INFO] Sliced Tensors per thread block:")
for i in cutlass.range_constexpr(len(ctaInputs)):
print(f"[DSL INFO] ctaInputs{i} = {ctaInputs[i].type}")
print(f"[DSL INFO] ctaC = {ctaC.type}")
print(f"[DSL INFO] ctaCrd = {ctaCrd.type}")
print("[DSL INFO] Sliced Tensors per thread block:")
for i in cutlass.range_constexpr(len(gInputs)):
print(f"[DSL INFO] ctaInputs{i} = {gInputs[i].type}")
print(f"[DSL INFO] gC = {gC.type}")
print(f"[DSL INFO] gCrd = {gCrd.type}")
# compose with CTA TV layout
# (tid, vid) -> address
tidfrgInputs = [cute.composition(t, tv_layout) for t in ctaInputs]
tidfrgC = cute.composition(ctaC, tv_layout)
tidfrgCrd = cute.composition(ctaCrd, tv_layout)
# print(f"{tv_layout = }")
# print(f"{tidfrgAB[0] = }")
###############################################################################
# Compose with thread block TV layout to map thread & value indices to memory address
###############################################################################
# (tid, vid) -> memory address
tidfrgInputs = [cute.composition(t, tv_layout) for t in gInputs]
tidfrgC = cute.composition(gC, tv_layout)
tidfrgCrd = cute.composition(gCrd, tv_layout)
thr_coord = (tidx, (None, None))
# repeat None like vid to remove hierarchy of layout
thr_crd = (tidx, cute.repeat_like(None, tidfrgInputs[0][1]))
# slice for threads
###############################################################################
# Slice to local tile of thread
###############################################################################
# vid -> address
thrInputs = [t[thr_coord] for t in tidfrgInputs] # (V)
thrC = tidfrgC[thr_coord] # (V)
thrCrd = tidfrgCrd[thr_coord]
thrInputs = [t[thr_crd] for t in tidfrgInputs] # (V)
thrC = tidfrgC[thr_crd] # (V)
thrCrd = tidfrgCrd[thr_crd]
print(f"[DSL INFO] Sliced Tensors per thread:")
print("[DSL INFO] Sliced Tensors per thread:")
for i in cutlass.range_constexpr(len(thrInputs)):
print(f"[DSL INFO] thrInputs{i} = {thrInputs[i].type}")
print(f"[DSL INFO] thrC = {thrC.type}")
print(f"[DSL INFO] thrCrd = {thrCrd.type}")
# allocate fragments for gmem->rmem
frgInputs = [cute.make_fragment_like(t, t.element_type) for t in thrInputs]
frgC = cute.make_fragment_like(thrC, gC.element_type)
frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
###############################################################################
# Compute predicate for out of boundary checks
###############################################################################
frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
print(f"[DSL INFO] frgPred = {frgPred.type}")
for i in cutlass.range(cute.size(frgPred), unroll=1):
for i in cutlass.range_constexpr(cute.size(frgPred)):
frgPred[i] = cute.elem_less(thrCrd[i], shape)
# if tidx == 0 and bidx == 0:
# cute.print_tensor(frgPred)
##########################################################
# Move data to reg address space
# Load data and compute result
##########################################################
# declare the atoms which will be used later for memory copy
# Compile time validation: expect same element type for all input tensors so as to reuse the copy atom for load
assert all(t.element_type == inputs[0].element_type for t in inputs)
copy_atom_load = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
inputs[0].element_type,
num_bits_per_copy=inputs[0].element_type.width,
)
copy_atom_store = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
gC.element_type,
num_bits_per_copy=gC.element_type.width,
)
for thrInput, frgInput in zip(thrInputs, frgInputs):
cute.copy(copy_atom_load, thrInput, frgInput, pred=frgPred)
# Load data before use. The compiler will optimize the copy and load
# operations to convert some memory ld/st into register uses.
result = op(*[frgInput.load() for frgInput in frgInputs])
# Save the results back to registers. Here we reuse b's registers.
frgC.store(result)
# Copy the results back to c
cute.copy(copy_atom_store, frgC, thrC, pred=frgPred)
result = op(*[thrInput.load() for thrInput in thrInputs])
thrC.store(result)
@cute.jit
def elementwise_apply(
op: cutlass.Constexpr,
a: cute.Tensor,
b: cute.Tensor,
result: cute.Tensor,
stream: cuda.CUstream,
op: cutlass.Constexpr, inputs, result: cute.Tensor, stream: cuda.CUstream
):
"""CUDA kernel applying binary operator on each element of two n-D input tensors in
CuTe Python and store to result tensor.
@@ -232,51 +213,71 @@ def elementwise_apply(
# Opt-3: SOL with 2D thread tile
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (16, 128) logical tile
# * TV layout map to (64, 256) logical tile
# * tidx maps to mode-1 and input layout is contiguous on mode-1 for coalesced load-store
thr_layout = cute.make_layout((4, 32), stride=(32, 1))
val_layout = cute.make_layout((4, 4), stride=(4, 1))
# Use 128bit(16B) load as canonicalized form of val_layout then recast to target element-type
coalesced_ldst_bytes = 16
# Compile time validation: expect same element type for all input tensors
assert all(t.element_type == inputs[0].element_type for t in inputs)
dtype = inputs[0].element_type
thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0))
val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0))
val_layout = cute.recast_layout(dtype.width, 8, val_layout)
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
print(f"[DSL INFO] Input Tensors:")
print(f"[DSL INFO] a = {a.type}")
print(f"[DSL INFO] b = {b.type}")
print(f"[DSL INFO] result = {result.type}")
print("[DSL INFO] Input Tensors:")
for i, t in enumerate(inputs):
print(f"[DSL INFO] inputs{i} = {t}")
print(f"[DSL INFO] result = {result}")
print(f"[DSL INFO] Tiling Parameters:")
print("[DSL INFO] Tiling Parameters:")
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
print(f"[DSL INFO] tv_layout = {tv_layout}")
gA = cute.zipped_divide(a, tiler_mn) # ((TileM, TileN), (RestM, RestN))
gB = cute.zipped_divide(b, tiler_mn) # ((TileM, TileN), (RestM, RestN))
gC = cute.zipped_divide(result, tiler_mn) # ((TileM, TileN), (RestM, RestN))
print("[DSL INFO] Tiled Tensors:")
mInputs = [cute.zipped_divide(input, tiler_mn) for input in inputs]
# ((TileM, TileN), (RestM, RestN))
mC = cute.zipped_divide(result, tiler_mn)
print(f"[DSL INFO] Tiled Tensors:")
print(f"[DSL INFO] gA = {gA.type}")
print(f"[DSL INFO] gB = {gB.type}")
print(f"[DSL INFO] gC = {gC.type}")
# (RestM, RestN) -> (RestN, RestM)
remap_block = cute.make_ordered_layout(
cute.select(mInputs[0].shape[1], mode=[1, 0]), order=(1, 0)
)
for i, t in enumerate(mInputs):
print(f"[DSL INFO] gInputs{i} = {mInputs[i]}")
mInputs[i] = cute.composition(t, (None, remap_block))
print(f"[DSL INFO] gInputs{i} (remapped) = {mInputs[i]}")
mC = cute.composition(mC, (None, remap_block))
print(f"[DSL INFO] gC = {mC}")
idC = cute.make_identity_tensor(result.shape)
cC = cute.zipped_divide(idC, tiler=tiler_mn)
print(f"[DSL INFO] coord tensor = {cC.type}")
print(f"[DSL INFO] coord tensor = {cC}")
# Launch the kernel asynchronously
# Async token(s) can also be specified as dependencies
elementwise_apply_kernel(
op,
[gA, gB], # Group input tensors into a list as a single argument
gC,
cC,
result.shape,
tv_layout,
).launch(
grid=[cute.size(gC, mode=[1]), 1, 1],
# Group input tensors into a list as a single argument
elementwise_apply_kernel(op, mInputs, mC, cC, result.shape, tv_layout).launch(
# Compute production at each mode of mC.shape[1] to get multi-dimensional grid size
grid=cute.product_each(mC.shape[1]),
block=[cute.size(tv_layout, mode=[0]), 1, 1],
stream=stream,
)
def run_elementwise_apply_and_verify(
@cutlass.dsl_user_op
def leaky_relu(x, alpha, *, loc=None, ip=None):
return cute.where(x > 0, x, alpha * x, loc=loc, ip=ip)
def leaky_relu_ref(x, alpha):
return torch.where(x > 0, x, alpha * x)
def run_and_verify(
op,
M,
N,
@@ -287,14 +288,23 @@ def run_elementwise_apply_and_verify(
iterations=100,
):
if not torch.cuda.is_available():
raise RuntimeError(f"Ampere GPU is required to run this example!")
raise RuntimeError("NVIDIA GPU is required to run this example!")
if op == "leaky_relu":
op = partial(leaky_relu, alpha=0.01)
ref_op = partial(leaky_relu_ref, alpha=0.01)
num_inputs = 1
else:
op = getattr(operator, op)
ref_op = op
num_inputs = 2
# Create non default CUDA stream from PyTorch
torch_stream = torch.cuda.Stream()
# Get the raw stream pointer as a CUstream
current_stream = cuda.CUstream(torch_stream.cuda_stream)
print(f"\nRunning Elementwise Apply test with:")
print("\nRunning Elementwise Apply test with:")
print(f"Tensor dimensions: [{M}, {N}]")
print(f"Input and Output Data type: {dtype}")
print(f"Warmup iterations: {warmup_iterations}")
@@ -303,85 +313,78 @@ def run_elementwise_apply_and_verify(
torch_dtype = cutlass_torch.dtype(dtype)
# Allocate tensors with random values.
a = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
b = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
c = torch.zeros_like(a)
inputs = [
torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
for _ in range(num_inputs)
]
c = torch.zeros_like(inputs[0])
print(f"Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print("Input tensor shapes:")
for i in range(num_inputs):
print(f"inputs[{i}]: {inputs[i].shape}, dtype: {inputs[i].dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
epsilon = 1.2
if op in (operator.truediv, operator.floordiv):
b = torch.where(b == 0, torch.tensor(epsilon), b)
inputs[1] = torch.where(inputs[1] == 0, torch.tensor(epsilon), inputs[1])
print("Executing elementwise apply kernel...")
inputs_ = [from_dlpack(t, assumed_align=16) for t in inputs]
c_ = from_dlpack(c, assumed_align=16).mark_layout_dynamic()
print("Compiling kernel with cute.compile ...")
start_time = time.time()
compiled_fn = cute.compile[cute.GenerateLineInfo(True)](
elementwise_apply, op, inputs_, c_, current_stream
)
compilation_time = time.time() - start_time
print(f"Compilation time: {compilation_time:.4f} seconds")
if not skip_ref_check:
elementwise_apply(
op,
from_dlpack(a),
from_dlpack(b),
from_dlpack(c).mark_layout_dynamic(),
current_stream,
)
print("Executing elementwise apply kernel...")
compiled_fn(inputs_, c_, current_stream)
print("Verifying results...")
torch.testing.assert_close(op(a, b), c)
torch.testing.assert_close(ref_op(*inputs), c)
print("Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3]}")
if not benchmark:
return
compiled_func = cute.compile(
elementwise_apply,
op,
from_dlpack(a),
from_dlpack(b),
from_dlpack(c).mark_layout_dynamic(),
current_stream,
)
# When compiled we inlined op in the kernel, so we do not pass it when benchmarking
print("Benchmarking elementwise apply kernel...")
avg_time_us = testing.benchmark(
compiled_func,
kernel_arguments=testing.JitArguments(
from_dlpack(a),
from_dlpack(b),
from_dlpack(c).mark_layout_dynamic(),
current_stream,
),
compiled_fn,
kernel_arguments=testing.JitArguments(inputs_, c_, current_stream),
warmup_iterations=warmup_iterations,
iterations=iterations,
use_cuda_graphs=True,
stream=current_stream,
)
avg_time = avg_time_us / 1e3
num_elements = sum(input.numel() for input in inputs) + c.numel()
# Print execution results
print(f"Kernel execution time: {avg_time:.4f} ms")
print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
print(
f"Achieved memory throughput: {(3 * a.numel() * dtype.width // 8) / (avg_time / 1000) / 1e9:.2f} GB/s"
f"Achieved memory throughput: {(num_elements * dtype.width // 8) / (avg_time_us * 1000):.2f} GB/s"
)
print(f"First few elements of result: \n{c[:3, :3]}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="example of elementwise apply to demonstrate building elementwise kernels"
description="Demonstration of building customizable elementwise CUDA kernels using the CuTe DSL"
)
parser.add_argument("--M", default=128, type=int)
parser.add_argument("--N", default=128, type=int)
parser.add_argument("--M", default=4096, type=int)
parser.add_argument("--N", default=4096, type=int)
parser.add_argument("--op", default="add", type=str)
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("--benchmark", action="store_true")
args = parser.parse_args()
run_elementwise_apply_and_verify(
getattr(operator, args.op),
run_and_verify(
args.op,
args.M,
args.N,
dtype=cutlass.Float32,
examples/python/CuTeDSL/ampere/flash_attention_v2.py
+18 -12
View File
@@ -28,7 +28,7 @@
import argparse
from types import SimpleNamespace
from typing import Type, Union, Callable
from typing import Type, Callable
import torch
import cuda.bindings.driver as cuda
@@ -38,6 +38,7 @@ 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.pipeline as pipeline
import cutlass.utils as utils
"""
@@ -126,6 +127,10 @@ class FlashAttentionForwardAmpere:
self._num_threads = num_threads
self._is_causal = is_causal
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1, num_threads=num_threads
)
@staticmethod
def can_implement(
dtype, head_dim, m_block_size, n_block_size, num_threads, is_causal
@@ -450,7 +455,7 @@ class FlashAttentionForwardAmpere:
acc_shape_O = thr_mma.partition_shape_C(
(self._m_block_size, self._head_dim_padded)
)
acc_O = cute.make_fragment(acc_shape_O, cutlass.Float32)
acc_O = cute.make_rmem_tensor(acc_shape_O, cutlass.Float32)
acc_O.fill(0.0)
# ///////////////////////////////////////////////////////////////////////////////
@@ -506,7 +511,7 @@ class FlashAttentionForwardAmpere:
tKVcKV = gmem_thr_copy_QKV.partition_S(cKV)
# Allocate predicate tensors for m and n, here we only allocate the tile of k, and do special process for mn.
# This is to reduce register pressure and gets 2-3% performance gain compared with allocating the whole tile.
tQpQ = cute.make_fragment(
tQpQ = cute.make_rmem_tensor(
cute.make_layout(
(
tQsQ.shape[0][1],
@@ -517,7 +522,7 @@ class FlashAttentionForwardAmpere:
),
cutlass.Boolean,
)
tKVpKV = cute.make_fragment(
tKVpKV = cute.make_rmem_tensor(
cute.make_layout(
(
tKsK.shape[0][1],
@@ -571,11 +576,11 @@ class FlashAttentionForwardAmpere:
# Softmax intermediate result: row_max and row_sum
# ///////////////////////////////////////////////////////////////////////////////
# shape: (atom_v_m * rest_m)
row_max = cute.make_fragment(
row_max = cute.make_rmem_tensor(
(acc_O.shape[0][0] * acc_O.shape[1]), cutlass.Float32
)
# shape: (atom_v_m * rest_m)
row_sum = cute.make_fragment(
row_sum = cute.make_rmem_tensor(
(acc_O.shape[0][0] * acc_O.shape[1]), cutlass.Float32
)
row_max.fill(-cutlass.Float32.inf)
@@ -710,7 +715,7 @@ class FlashAttentionForwardAmpere:
tOgO = gmem_thr_copy_O.partition_D(gO)
tOrO = cute.make_fragment_like(tOgO, self._dtype)
# sync before all smem stores are done.
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# load acc O from smem to rmem for wider vectorization
cute.copy(
gmem_tiled_copy_O,
@@ -724,7 +729,7 @@ class FlashAttentionForwardAmpere:
(m_block, 0),
)
tOcO = gmem_thr_copy_O.partition_D(cO)
tOpO = cute.make_fragment(
tOpO = cute.make_rmem_tensor(
cute.make_layout(
(tOgO.shape[0][1], tOgO.shape[1], tOgO.shape[2]),
stride=(tOgO.shape[2], 0, 1),
@@ -778,12 +783,12 @@ class FlashAttentionForwardAmpere:
acc_shape_S = mma_params.thr_mma.partition_shape_C(
(self._m_block_size, self._n_block_size)
)
acc_S = cute.make_fragment(acc_shape_S, cutlass.Float32)
acc_S = cute.make_rmem_tensor(acc_shape_S, cutlass.Float32)
acc_S.fill(0.0)
# wait for smem tile QK before mma calculation for S
cute.arch.cp_async_wait_group(0)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# load smem tile V for O, special process for the first tile to avoid loading nan.
# The `if` here is a constexpr, won't be generated in the IR.
if is_first_n_block:
@@ -847,7 +852,7 @@ class FlashAttentionForwardAmpere:
# wait for smem tile V for O
cute.arch.cp_async_wait_group(0)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
if basic_params.n_block > 0:
cute.copy(
@@ -1170,7 +1175,7 @@ def run(
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("Running Ampere SM80 FlashAttentionForward test with:")
print(f" dtype: {dtype}")
print(f" batch_size: {batch_size}")
print(f" seqlen_q: {seqlen_q}")
@@ -1285,6 +1290,7 @@ def run(
return avg_time_us # Return execution time in microseconds
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="example of flash attention v2 with CuTe on GPU"
examples/python/CuTeDSL/ampere/hstu_attention.py
+1149
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/ampere/inline_ptx.py
+245
View File
@@ -0,0 +1,245 @@
# 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.
from functools import partial
from typing import Union
import torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
from cutlass._mlir.dialects import llvm
from cutlass.cute.typing import Boolean, Int32, Int, Constexpr
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass.cute.arch.nvvm_wrappers import FULL_MASK, WARP_SIZE
"""
A simple example to show how to wrap PTX instructions by using inline_asm op in llvm dialect.
Situations like:
1. Instructions that are not already exposed by CuTe DSL via `nvvm` module
2. Sequences of instructions that the compiler otherwise does not generate optimally
motivate developers to inline PTX themselves.
In this example, we inline the vote.sync.ballot.b32, vote.sync.any.pred, vote.sync.all.pred,
vote.sync.uni.pred, and use the corresponding ops in nvvm_wrappers.py for the test.
You can refer to the documentation of `inline_asm op in llvm dialect <https://mlir.llvm.org/docs/Dialects/LLVM/#llvminline_asm-llvminlineasmop>`_
and `vote.sync <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`_
for more details.
To run this example:
.. code-block:: bash
python examples/ampere/inline_ptx.py
The example will run the vote kernel with inline PTX and nvvm dialect separately.
The results from inline PTX and nvvm dialect will be verified correspondingly.
"""
@dsl_user_op
def ptx_vote_sync_op(
pred: Boolean, kind: str, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Union[Int32, Boolean]:
return_type = Boolean
return_type_str = "pred"
return return_type(
llvm.inline_asm(
T.bool(),
[
Boolean(pred).ir_value(loc=loc, ip=ip),
Int32(mask).ir_value(loc=loc, ip=ip),
],
f"""{{\n\t
.reg .pred ps;\n\t
.reg .pred pd;\n\t
setp.ne.b32 ps, $1, 0;\n\t
vote.sync.{kind}.{return_type_str} pd, ps, $2;\n\t
selp.b32 $0, 1, 0, pd;\n\t
}}""",
"=r,r,i",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
)
ptx_vote_any_sync = partial(ptx_vote_sync_op, kind="any")
ptx_vote_all_sync = partial(ptx_vote_sync_op, kind="all")
ptx_vote_uni_sync = partial(ptx_vote_sync_op, kind="uni")
@dsl_user_op
def ptx_vote_ballot_sync(
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Union[Int32, Boolean]:
return_type = Int32
return_type_str = "b32"
return return_type(
llvm.inline_asm(
T.i32(),
[
Boolean(pred).ir_value(loc=loc, ip=ip),
Int32(mask).ir_value(loc=loc, ip=ip),
],
f"""{{\n\t
.reg .pred p;\n\t
setp.ne.b32 p, $1, 0;\n\t
vote.sync.ballot.{return_type_str} $0, p, $2;\n\t
}}""",
"=r,r,i",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
)
@cute.kernel
def vote_kernel(
mBallot: cute.Tensor,
mAny: cute.Tensor,
mAll: cute.Tensor,
mUni: cute.Tensor,
use_inline_ptx: Constexpr[bool],
):
tidx, _, _ = cute.arch.thread_idx()
vote_ballot = (
ptx_vote_ballot_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_ballot_sync(tidx < 10)
)
vote_any = (
ptx_vote_any_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_any_sync(tidx < 10)
)
vote_all = (
ptx_vote_all_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_all_sync(tidx < 10)
)
vote_uni = (
ptx_vote_uni_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_uni_sync(tidx < 10)
)
mBallot[tidx] = vote_ballot
mAny[tidx] = vote_any
mAll[tidx] = vote_all
mUni[tidx] = vote_uni
@cute.jit
def vote(
mBallot: cute.Tensor,
mAny: cute.Tensor,
mAll: cute.Tensor,
mUni: cute.Tensor,
use_inline_ptx: Constexpr[bool],
):
vote_kernel(
mBallot,
mAny,
mAll,
mUni,
use_inline_ptx,
).launch(
grid=[1, 1, 1],
block=[cute.size(WARP_SIZE, mode=[0]), 1, 1],
)
def run():
ballot_ptx = torch.randint(
0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32
)
any_ptx = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
all_ptx = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
uni_ptx = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
mBallotPTX = from_dlpack(ballot_ptx).mark_layout_dynamic()
mAnyPTX = from_dlpack(any_ptx).mark_layout_dynamic()
mAllPTX = from_dlpack(all_ptx).mark_layout_dynamic()
mUniPTX = from_dlpack(uni_ptx).mark_layout_dynamic()
# get the results from ptx
vote(mBallotPTX, mAnyPTX, mAllPTX, mUniPTX, use_inline_ptx=True)
ballot_nvvm = torch.randint(
0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32
)
any_nvvm = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
all_nvvm = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
uni_nvvm = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
mBallotNVVM = from_dlpack(ballot_nvvm).mark_layout_dynamic()
mAnyNVVM = from_dlpack(any_nvvm).mark_layout_dynamic()
mAllNVVM = from_dlpack(all_nvvm).mark_layout_dynamic()
mUniNVVM = from_dlpack(uni_nvvm).mark_layout_dynamic()
# get the results from nvvm
vote(mBallotNVVM, mAnyNVVM, mAllNVVM, mUniNVVM, use_inline_ptx=False)
print("Verifying ballot results...")
torch.testing.assert_close(ballot_ptx, ballot_nvvm)
print("Verifying any results...")
torch.testing.assert_close(any_ptx, any_nvvm)
print(torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE))))
assert torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE)))
print("Verifying all results...")
torch.testing.assert_close(all_ptx, all_nvvm)
assert torch.all(all_ptx == all(i < 10 for i in range(WARP_SIZE)))
print("Verifying uni results...")
torch.testing.assert_close(uni_ptx, uni_nvvm)
assert torch.all(uni_ptx == (len(set(i < 10 for i in range(WARP_SIZE))) == 1))
print("Results verified successfully!")
if __name__ == "__main__":
run()
examples/python/CuTeDSL/ampere/sgemm.py
+29 -28
View File
@@ -36,7 +36,7 @@ import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.torch as cutlass_torch
import cutlass.pipeline as pipeline
import cutlass.utils as utils
from cutlass.cute.runtime import from_dlpack
@@ -103,6 +103,9 @@ class SGemm:
assert self._bM % 16 == 0, "multiple of 16 required for tile dimension M"
assert self._bN % 16 == 0, "multiple of 16 required for tile dimension N"
assert self._num_stages >= 3, "num_stages must be greater than or equal to 3"
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1, num_threads=num_threads
)
@cute.jit
def __call__(
@@ -166,9 +169,8 @@ class SGemm:
mA.element_type,
num_bits_per_copy=mB.element_type.width,
)
if cutlass.const_expr(self.a_major_mode == utils.LayoutEnum.COL_MAJOR):
num_vectorized = 4 if (mA.layout.max_alignment % 16 == 0) else 1
num_vectorized = 4 if (mA.layout[0].max_alignment % 16 == 0) else 1
atom_async_copy_A = cute.make_copy_atom(
cute.nvgpu.cpasync.CopyG2SOp(),
mA.element_type,
@@ -182,7 +184,7 @@ class SGemm:
vA = cute.make_layout((num_vectorized, 1))
if cutlass.const_expr(self.b_major_mode == utils.LayoutEnum.COL_MAJOR):
num_vectorized = 4 if (mB.layout.max_alignment % 16 == 0) else 1
num_vectorized = 4 if (mB.layout[0].max_alignment % 16 == 0) else 1
atom_async_copy_B = cute.make_copy_atom(
cute.nvgpu.cpasync.CopyG2SOp(),
mA.element_type,
@@ -294,7 +296,7 @@ class SGemm:
# tile (instead of the last one) irregular in shape when k is irregular.
# We first handle the irregular tile to avoid checking for this
# condition within the mainloop.
residue_k = mA.shape[1] - cutlass.Int32(self._bK) * gA.shape[2]
residue_k = mA.shape[1] - self._bK * gA.shape[2]
gA = cute.domain_offset((0, residue_k, 0), gA)
gB = cute.domain_offset((0, residue_k, 0), gB)
@@ -342,7 +344,7 @@ class SGemm:
tAcA = thr_copy_A.partition_S(cA)
tBcB = thr_copy_B.partition_S(cB)
# Allocate predicate tensors for m and n
tApA = cute.make_fragment(
tApA = cute.make_rmem_tensor(
cute.make_layout(
(
tAsA.shape[0][1],
@@ -353,7 +355,7 @@ class SGemm:
),
cutlass.Boolean,
)
tBpB = cute.make_fragment(
tBpB = cute.make_rmem_tensor(
cute.make_layout(
(
tBsB.shape[0][1],
@@ -365,7 +367,7 @@ class SGemm:
cutlass.Boolean,
)
# Allocate predicate tensors for m, n and k for residue k-tile
tApA_residue_k = cute.make_fragment(
tApA_residue_k = cute.make_rmem_tensor(
cute.make_layout(
(
tAsA.shape[0][1],
@@ -380,7 +382,7 @@ class SGemm:
),
cutlass.Boolean,
)
tBpB_residue_k = cute.make_fragment(
tBpB_residue_k = cute.make_rmem_tensor(
cute.make_layout(
(
tBsB.shape[0][1],
@@ -508,7 +510,7 @@ class SGemm:
if k_block_max > 1:
# Wait until our first prefetched tile is loaded in
cute.arch.cp_async_wait_group(k_pipe_max - 2)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# Prefetch the first rmem from the first k-tile
cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0])
@@ -545,7 +547,7 @@ class SGemm:
tCsA_p = tCsA[None, None, None, smem_pipe_read]
tCsB_p = tCsB[None, None, None, smem_pipe_read]
cute.arch.cp_async_wait_group(k_pipe_max - 2)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# Load A, B from shared memory to registers for k_block + 1
k_block_next = (k_block + 1) % k_block_max # static
@@ -611,13 +613,13 @@ class SGemm:
# them without vectorization.
# ///////////////////////////////////////////////////////////////////////////////
cute.arch.cp_async_wait_group(0)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
tCrC.store(epilogue_op(tCrC.load()))
# predicate
cC = cute.make_identity_tensor(gC.shape)
tCpC = thr_mma.partition_C(cC)
predC = cute.make_fragment(tCrC.layout, cutlass.Boolean)
predC = cute.make_rmem_tensor(tCrC.layout, cutlass.Boolean)
residue_m = mC.shape[0] - cutlass.Int32(self._bM) * bidx
residue_n = mC.shape[1] - cutlass.Int32(self._bN) * bidy
for i in range(cute.size(tCrC.shape)):
@@ -664,7 +666,7 @@ def run(
:return: Execution time of the GEMM kernel in microseconds
:rtype: float
"""
print(f"Running Ampere SIMT GEMM example:")
print("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}")
@@ -697,14 +699,17 @@ def run(
divisibility_b = b.shape[1] if b_major == "k" else b.shape[0]
divisibility_c = c.shape[1] if c_major == "n" else c.shape[0]
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),
divisibility=divisibility_a,
if static_shape:
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),
divisibility=divisibility_a,
)
)
)
else:
a_tensor = from_dlpack(a, assumed_align=16)
b_tensor = (
from_dlpack(b, assumed_align=16)
@@ -733,12 +738,8 @@ def run(
print("Compiling kernel with cute.compile ...")
start_time = time.time()
compiled_fn = cute.compile(
sgemm,
a_tensor,
b_tensor,
c_tensor,
stream=current_stream,
compiled_fn = cute.compile[cute.GenerateLineInfo](
sgemm, a_tensor, b_tensor, c_tensor, stream=current_stream
)
compilation_time = time.time() - start_time
print(f"Compilation time: {compilation_time:.4f} seconds")
@@ -833,7 +834,7 @@ if __name__ == "__main__":
parser.add_argument(
"--mnk", type=parse_comma_separated_ints, default=(256, 256, 64)
)
parser.add_argument("--a_major", choices=["k", "m"], default="k")
parser.add_argument("--a_major", choices=["k", "m"], default="m")
parser.add_argument("--b_major", choices=["k", "n"], default="k")
parser.add_argument("--c_major", choices=["n", "m"], default="n")
parser.add_argument("--warmup_iterations", default=2, type=int)
examples/python/CuTeDSL/ampere/smem_allocator.py
+24 -15
View File
@@ -69,7 +69,7 @@ class complex:
class SharedStorage:
# struct elements with natural alignment
a: cute.struct.MemRange[cutlass.Float32, 32] # array
b: cutlass.Int64 # saclar
b: cutlass.Int64 # scalar
c: complex # nested struct
# struct elements with strict alignment
x: cute.struct.Align[
@@ -90,10 +90,17 @@ def kernel(
dst_c: cute.Tensor,
):
# Note: SMEM_SIZE bytes (specified in kernel().launch(smem=...)) can be reserved for developer to utilize
# Note: alignment of inital allocator base ptr is 1024
# Note: alignment of initial allocator base ptr is 1024
allocator = cutlass.utils.SmemAllocator()
# base ptr of allocator points at: SMEM_ADDR_START (the starting address of available shared memory)
# -- Allocate a scalar
int_ptr = allocator.allocate(cutlass.Int32)
# base ptr of allocator now points at: SMEM_ADDR_AFTER_INT = SMEM_ADDR_START + aligned_size(int)
assert int_ptr.dtype == cutlass.Int32, "Expected Int32, but got {}".format(
int_ptr.dtype
)
# -- Allocate a struct --
# Note: when specified alignment, max(alignment, alignof(struct)) will be applied
# reserves the section of struct in smem, elements in the struct can be accessed by ptr
@@ -153,7 +160,7 @@ def kernel(
@cute.jit
def run_allocation_kernel(
def host(
const_a: cutlass.Constexpr,
dst_a: cute.Tensor,
const_b: cutlass.Constexpr,
@@ -161,22 +168,18 @@ def run_allocation_kernel(
const_c: cutlass.Constexpr,
dst_c: cute.Tensor,
):
# additional size for the example, 64(section) + 112(array) + 128(tensor) < 384
addtional_bytes = 384
# Note: launch shared memory size is: SMEM_SIZE = 512 + 384 = 896 bytes
# Note: Shared Memory size is automatically calculated now
kernel(const_a, dst_a, const_b, dst_b, const_c, dst_c).launch(
grid=(1, 1, 1),
block=(1, 1, 1),
smem=SharedStorage.size_in_bytes() + addtional_bytes,
grid=(1, 1, 1), block=(1, 1, 1)
)
def veify_allocation_kernel(const_a, const_b, const_c):
def run_and_verify(const_a, const_b, const_c):
dst_a = torch.zeros((8, 4), dtype=torch.float32, device="cuda")
dst_b = torch.zeros((8, 2), dtype=torch.float32, device="cuda")
dst_c = torch.zeros((16, 2), dtype=torch.float32, device="cuda")
run_allocation_kernel(
host(
const_a,
from_dlpack(dst_a),
const_b,
@@ -185,9 +188,15 @@ def veify_allocation_kernel(const_a, const_b, const_c):
from_dlpack(dst_c),
)
np.testing.assert_equal(const_a, dst_a.detach().cpu().numpy()[0])
np.testing.assert_equal(const_b, dst_b.detach().cpu().numpy()[0])
np.testing.assert_equal(const_c, dst_c.detach().cpu().numpy()[0])
assert const_a == dst_a.cpu()[0, 0], (
f"Expected {const_a}, but got {dst_a.cpu()[0, 0]}"
)
assert const_b == dst_b.cpu()[0, 0], (
f"Expected {const_b}, but got {dst_b.cpu()[0, 0]}"
)
assert const_c == dst_c.cpu()[0, 0], (
f"Expected {const_c}, but got {dst_c.cpu()[0, 0]}"
)
if __name__ == "__main__":
@@ -197,4 +206,4 @@ if __name__ == "__main__":
const_a = 0.5
const_b = 1.0
const_c = 2.0
veify_allocation_kernel(const_a, const_b, const_c)
run_and_verify(const_a, const_b, const_c)
examples/python/CuTeDSL/ampere/tensorop_gemm.py
+11 -12
View File
@@ -28,10 +28,8 @@
import argparse
import math
import time
from typing import Tuple, Type
import cuda.bindings.driver as cuda
import torch
import cutlass
@@ -121,12 +119,12 @@ class TensorOpGemm:
self.mma_inst_shape = (16, 8, 16)
mmaM, mmaN, mmaK = self.mma_inst_shape
assert (
self.bM % (atom_lay_M * mmaM) == 0
), "bM must be divisible by MMA instruction"
assert (
self.bN % (atom_lay_N * mmaN) == 0
), "bN must be divisible by MMA instruction"
assert self.bM % (atom_lay_M * mmaM) == 0, (
"bM must be divisible by MMA instruction"
)
assert self.bN % (atom_lay_N * mmaN) == 0, (
"bN must be divisible by MMA instruction"
)
assert atom_lay_K == 1, "this example does not support atom layout K > 1"
assert self.bK % mmaK == 0, "bK must be divisible by MMA instruction"
assert self.num_stages >= 3, "num_stages must be greater than or equal to 3"
@@ -428,7 +426,7 @@ class TensorOpGemm:
# at the granularity of a copy atom, so the predicate tensor does not
# need separate booleans for individual elements within a copy
# atom (for example, the elements of tAgA.shape[0][0].)
tApA = cute.make_fragment(
tApA = cute.make_rmem_tensor(
cute.make_layout(
(
tAgA.shape[0][1],
@@ -439,7 +437,7 @@ class TensorOpGemm:
),
cutlass.Boolean,
)
tBpB = cute.make_fragment(
tBpB = cute.make_rmem_tensor(
cute.make_layout(
(
tBsB.shape[0][1],
@@ -707,7 +705,7 @@ class TensorOpGemm:
cute.autovec_copy(tCsC_epilogue, tCrC_epilogue)
# Create predication tensor for m
tCpC = cute.make_fragment(
tCpC = cute.make_rmem_tensor(
cute.make_layout(
(
tCgC_epilogue.shape[0][1],
@@ -851,7 +849,7 @@ def run(
use_cold_l2: bool = False,
**kwargs,
):
print(f"Running Ampere tensor core GEMM example:")
print("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}"
@@ -944,6 +942,7 @@ def run(
return avg_time_us # Return execution time in microseconds
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
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")
examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py
+1326
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/cute/ffi/CMakeLists.txt
+9 -2
View File
@@ -30,11 +30,12 @@ cmake_minimum_required(VERSION 3.15)
project(tensor)
# Find Python
find_package(Python COMPONENTS Interpreter Development REQUIRED)
find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
# Get Python site-packages directory using Python
execute_process(
COMMAND ${Python_EXECUTABLE} -c "import site; print(site.getsitepackages()[0])"
COMMAND ${Python3_EXECUTABLE} -c "import site; print(site.getsitepackages()[0])"
OUTPUT_VARIABLE Python_SITE_PACKAGES
OUTPUT_STRIP_TRAILING_WHITESPACE
)
@@ -45,7 +46,13 @@ message(STATUS "Python site-packages directory: ${Python_SITE_PACKAGES}")
list(APPEND CMAKE_PREFIX_PATH ${Python_SITE_PACKAGES}/nanobind/cmake)
# Find nanobind
find_package(nanobind REQUIRED)
find_package(nanobind)
if(NOT nanobind_FOUND)
message(FATAL_ERROR
"nanobind not found!\n"
"Please install nanobind with: pip install nanobind\n"
)
endif()
# Add the module
nanobind_add_module(tensor tensor.cpp)
examples/python/CuTeDSL/cute/ffi/jit_argument.py
+18 -6
View File
@@ -54,7 +54,6 @@ import cutlass.cute as cute
from cutlass._mlir import ir
from cutlass._mlir.dialects import llvm
import cutlass._mlir.extras.types as T
class ExampleTensorValue(ir.Value):
@@ -244,7 +243,7 @@ import tempfile
import torch
def run_test(tmpdir=None):
def run_test(tmpdir=None, cmake_args=""):
# Skip cleanup if user provides tmpdir
cleanup = tmpdir is None
# Initialize temporary build directory
@@ -253,7 +252,8 @@ def run_test(tmpdir=None):
try:
current_dir = os.path.dirname(os.path.abspath(__file__))
subprocess.run(["cmake", "-B", tmpdir, current_dir], check=True)
cmake_args = cmake_args.split()
subprocess.run(["cmake", "-B", tmpdir, current_dir] + cmake_args, check=True)
subprocess.run(["cmake", "--build", tmpdir], check=True)
sys.path.append(tmpdir)
@@ -284,7 +284,10 @@ def run_test(tmpdir=None):
# Execute compiled function
compiled_func(tensor)
except Exception as e:
print(e)
import traceback
traceback.print_exception(type(e), e, e.__traceback__)
raise e
finally:
if cleanup:
# Clean up the temporary directory
@@ -298,8 +301,17 @@ if __name__ == "__main__":
description="Set temporary directory for building C modules"
)
parser.add_argument(
"--tmp-dir", type=str, help="Temporary directory path for building C modules"
"--tmp-dir",
type=str,
default=None,
help="Temporary directory path for building C modules",
)
parser.add_argument(
"--cmake-args",
type=str,
default="",
help="Extra CMake arguments for building C modules",
)
args = parser.parse_args()
run_test(args.tmp_dir)
run_test(tmpdir=args.tmp_dir, cmake_args=args.cmake_args)
examples/python/CuTeDSL/cute/torch_fake_tensor.py
+77
View File
@@ -0,0 +1,77 @@
# 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 torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
"""Example demonstrating how to use CuTe with PyTorch's FakeTensor mode.
This example shows how to:
1. Use PyTorch's FakeTensor mode to compile a CuTe function without real data
2. Execute the compiled function on real data later
FakeTensor mode allows compiling code without allocating real memory, which is useful
for ahead-of-time compilation scenarios. The compiled function can then be executed
on real tensors that match the expected shapes and dtypes.
Primary goals of this example are to demonstrate: How to use PyTorch's FakeTensor mode with CuTe
to enable ahead-of-time compilation without real data allocation.
The example:
1. Creates a fake tensor in PyTorch using FakeTensor mode
2. Compiles a CuTe function using the fake tensor without allocating real memory
3. Creates a real tensor with matching shape and dtype
4. Executes the compiled function on the real tensor
To run this example:
.. code-block:: bash
python examples/cute/torch_fake_tensor.py
"""
@cute.jit
def print_tensor(t: cute.Tensor):
cute.print_tensor(t)
if __name__ == "__main__":
from torch._subclasses.fake_tensor import FakeTensorMode
shape = (3, 4)
with FakeTensorMode():
fake_tensor = torch.zeros(shape, dtype=torch.float32)
compiled_fn = cute.compile(print_tensor, from_dlpack(fake_tensor))
real_tensor = torch.randn(shape, dtype=torch.float32)
compiled_fn(from_dlpack(real_tensor))
examples/python/CuTeDSL/hopper/dense_gemm.py
+150 -128
View File
@@ -91,10 +91,11 @@ To collect performance with NCU profiler:
--a_major k --b_major k --c_major n
Constraints:
* Supported input data types: fp16, fp8 (e4m3fn, e5m2)
* Supported input data types: fp16, fp8 (e4m3fn, e5m2), int8, uint8
* For fp16 types, A and B must have the same data type
* For fp8 types, A and B can have different types (e4m3fn or e5m2) but both must be 8-bit
* Fp8 types only support k-major layout
* For fp8 types, A and B can have different types (e4m3fn or e5m2)
* For 8-bit integer types, A and B can have different types (int8 or uint8)
* 8-bit types (e4m3fn, e5m2, int8, uint8) only support k-major layout
* CTA tile shape M must be 64/128
* CTA tile shape N must be 64/128/256
* Cluster shape M/N must be positive and power of 2, total cluster size <= 4
@@ -212,17 +213,19 @@ class HopperWgmmaGemmKernel:
:param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
:type cluster_shape_mn: Tuple[int, int]
:note: Data type requirements:
- For 16-bit types: A and B must have the same data type
- For 8-bit types: A and B can have different types (Float8E4M3FN/Float8E5M2) as long as both are 8-bit
- Float8 types only support k-major layout
:note: Supported data types:
:note: Supported A/B data types:
- Float16
A and B must have the same data type
- Float8E4M3FN/Float8E5M2
A and B can have different types (Float8E4M3FN/Float8E5M2)
only support k-major layout
- Int8/Uint8
A and B can have different types (Int8/Uint8)
only support k-major layout
:note: Supported accumulation types:
- Float32 (for all floating point inputs)
- Float32/Float16 (for all floating point inputs)
- Int32 (for Int8/Uint8 inputs)
:note: Constraints:
- CTA tile M must be 64/128
@@ -339,7 +342,7 @@ class HopperWgmmaGemmKernel:
self.is_b_mcast = self.num_mcast_ctas_b > 1
is_cooperative = self.atom_layout_mnk == (2, 1, 1)
self.epi_tile = self._sm90_compute_tile_shape_or_override(
self.epi_tile = sm90_utils.compute_tile_shape_or_override(
self.tile_shape_mnk, self.c_dtype, is_cooperative=is_cooperative
)
@@ -411,7 +414,7 @@ class HopperWgmmaGemmKernel:
f"Type width mismatch: {self.a_dtype.width} != {self.b_dtype.width}"
)
if cutlass.const_expr(self.a_dtype.width != 16 and self.a_dtype.width != 8):
raise TypeError(f"a_dtype should be float16 or float8")
raise TypeError("a_dtype should be float16 or float8")
self._setup_attributes()
@@ -708,7 +711,7 @@ class HopperWgmmaGemmKernel:
tCrB = tiled_mma.make_fragment_B(tCsB)
acc_shape = tCgC.shape
accumulators = cute.make_fragment(acc_shape, self.acc_dtype)
accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype)
# ///////////////////////////////////////////////////////////////////////////////
# Cluster wait
@@ -960,7 +963,7 @@ class HopperWgmmaGemmKernel:
# Allocate D registers.
rD_shape = cute.shape(thr_copy_r2s.partition_S(sC))
tRS_rD_layout = cute.make_layout(rD_shape[:3])
tRS_rD = cute.make_fragment_like(tRS_rD_layout, self.acc_dtype)
tRS_rD = cute.make_rmem_tensor_like(tRS_rD_layout, self.acc_dtype)
size_tRS_rD = cute.size(tRS_rD)
sepi_for_tma_partition = cute.group_modes(sC, 0, 2)
@@ -982,7 +985,7 @@ class HopperWgmmaGemmKernel:
# Initialize tma store c_pipeline
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, self.threads_per_cta, self.threads_per_cta
pipeline.Agent.Thread, self.threads_per_cta
)
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.epi_stage,
@@ -995,7 +998,7 @@ class HopperWgmmaGemmKernel:
tRS_rD[epi_v] = tRS_rAcc[epi_idx * size_tRS_rD + epi_v]
# Type conversion
tRS_rD_out = cute.make_fragment_like(tRS_rD_layout, self.c_dtype)
tRS_rD_out = cute.make_rmem_tensor_like(tRS_rD_layout, self.c_dtype)
acc_vec = tRS_rD.load()
tRS_rD_out.store(acc_vec.to(self.c_dtype))
@@ -1010,7 +1013,7 @@ class HopperWgmmaGemmKernel:
space=cute.arch.SharedSpace.shared_cta,
)
# barrier for sync
cute.arch.barrier()
pipeline.sync(barrier_id=1)
gmem_coord = epi_tile_layout.get_hier_coord(epi_idx)
# Copy from shared memory to global memory
@@ -1023,7 +1026,7 @@ class HopperWgmmaGemmKernel:
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
cute.arch.barrier()
pipeline.sync(barrier_id=1)
if warp_idx == 0:
c_pipeline.producer_tail()
@@ -1073,39 +1076,6 @@ class HopperWgmmaGemmKernel:
) // ab_bytes_per_stage
return ab_stage, epi_stage
@staticmethod
def _sm90_compute_tile_shape_or_override(
tile_shape_mnk: tuple[int, int, int],
element_type: type[cutlass.Numeric],
is_cooperative: bool = False,
epi_tile_override: tuple[int, int] | None = None,
) -> tuple[int, int]:
"""Compute the epilogue tile shape or use override if provided.
:param tile_shape_mnk: CTA tile shape (M,N,K)
:type tile_shape_mnk: Tuple[int, int, int]
:param element_type: Data type of elements
:type element_type: type[cutlass.Numeric]
:param is_cooperative: Whether to use cooperative approach
:type is_cooperative: bool
:param epi_tile_override: Optional override for epilogue tile shape
:type epi_tile_override: Tuple[int, int] or None
:return: Computed epilogue tile shape
:rtype: Tuple[int, int]
"""
if epi_tile_override is not None:
return epi_tile_override
if is_cooperative:
tile_m = min(128, cute.size(tile_shape_mnk, mode=[0]))
tile_n = min(32, cute.size(tile_shape_mnk, mode=[1]))
return (tile_m, tile_n)
else:
n_perf = 64 if element_type.width == 8 else 32
tile_m = min(64, cute.size(tile_shape_mnk, mode=[0]))
tile_n = min(n_perf, cute.size(tile_shape_mnk, mode=[1]))
return (tile_m, tile_n)
@staticmethod
def _make_smem_layouts(
tile_shape_mnk: tuple[int, int, int],
@@ -1145,60 +1115,25 @@ class HopperWgmmaGemmKernel:
:return: Tuple of shared memory layouts for A, B, and C
:rtype: Tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]
"""
a_smem_shape = cute.slice_(tile_shape_mnk, (None, 0, None))
a_is_k_major = (
a_layout.sm90_mma_major_mode() == cute.nvgpu.warpgroup.OperandMajorMode.K
)
b_is_k_major = (
b_layout.sm90_mma_major_mode() == cute.nvgpu.warpgroup.OperandMajorMode.K
)
a_major_mode_size = tile_shape_mnk[2 if a_is_k_major else 0]
a_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom(
sm90_utils.get_smem_layout_atom(
a_layout,
a_dtype,
a_major_mode_size,
),
a_smem_layout_staged = sm90_utils.make_smem_layout_a(
a_layout,
tile_shape_mnk,
a_dtype,
)
a_smem_layout_staged = cute.tile_to_shape(
a_smem_layout_atom,
cute.append(a_smem_shape, ab_stage),
order=(0, 1, 2) if a_is_k_major else (1, 0, 2),
ab_stage,
)
b_smem_shape = cute.slice_(tile_shape_mnk, (0, None, None))
b_major_mode_size = tile_shape_mnk[2 if b_is_k_major else 1]
b_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom(
sm90_utils.get_smem_layout_atom(
b_layout,
b_dtype,
b_major_mode_size,
),
b_smem_layout_staged = sm90_utils.make_smem_layout_b(
b_layout,
tile_shape_mnk,
b_dtype,
)
b_smem_layout_staged = cute.tile_to_shape(
b_smem_layout_atom,
cute.append(b_smem_shape, ab_stage),
order=(0, 1, 2) if b_is_k_major else (1, 0, 2),
ab_stage,
)
c_smem_shape = epi_tile
c_major_mode_size = epi_tile[1] if c_layout.is_n_major_c() else epi_tile[0]
c_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom(
sm90_utils.get_smem_layout_atom(
c_layout,
c_dtype,
c_major_mode_size,
),
epi_smem_layout_staged = sm90_utils.make_smem_layout_epi(
c_dtype,
)
epi_smem_layout_staged = cute.tile_to_shape(
c_smem_layout_atom,
cute.append(c_smem_shape, epi_stage),
order=(1, 0, 2) if c_layout.is_m_major_c() else (0, 1, 2),
c_layout,
epi_tile,
epi_stage,
)
return a_smem_layout_staged, b_smem_layout_staged, epi_smem_layout_staged
@@ -1248,14 +1183,11 @@ class HopperWgmmaGemmKernel:
:rtype: Tuple[cute.CopyAtom, cute.Tensor]
"""
epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
c_cta_v_layout = cute.composition(
cute.make_identity_layout(tensor_c.shape), epi_tile
)
tma_atom_c, tma_tensor_c = cute.nvgpu.cpasync.make_tiled_tma_atom(
cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(),
tensor_c,
epi_smem_layout,
c_cta_v_layout,
epi_tile,
)
return tma_atom_c, tma_tensor_c
@@ -1326,44 +1258,130 @@ class HopperWgmmaGemmKernel:
:rtype: bool
"""
is_valid = True
# tested a_dtype
if a_dtype not in {
valid_ab_dtypes = {
cutlass.Float16,
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
}:
cutlass.Uint8,
cutlass.Int8,
}
if a_dtype not in valid_ab_dtypes:
is_valid = False
# tested b_dtype
if b_dtype not in {
cutlass.Float16,
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
}:
is_valid = False
# tested acc_dtype
if acc_dtype not in {cutlass.Float32, cutlass.Float16}:
is_valid = False
# tested c_dtype
if c_dtype not in {
cutlass.Float32,
cutlass.Float16,
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
}:
if b_dtype not in valid_ab_dtypes:
is_valid = False
# make sure a_dtype == b_dtype for Float16
if a_dtype.width == 16 and a_dtype != b_dtype:
is_valid = False
# make sure a_dtype.width == b_dtype.width (i.e, Float8E4M3FN or Float8E5M2)
if a_dtype.width != b_dtype.width:
is_valid = False
if not a_dtype.is_same_kind(b_dtype):
is_valid = False
# for Float8 types, this implementation only supports k-major layout
# for 8-bit types, this implementation only supports k-major layout
if (a_dtype.width == 8 and a_major != "k") or (
b_dtype.width == 8 and b_major != "k"
):
is_valid = False
# Define compatibility mapping between accumulator type and AB type
acc_ab_compatibility = {
cutlass.Float32: {
cutlass.Float16,
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
},
cutlass.Float16: {
cutlass.Float16,
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
},
cutlass.Int32: {cutlass.Uint8, cutlass.Int8},
}
# Check compatibility between accumulator type and A type
if a_dtype not in acc_ab_compatibility[acc_dtype]:
is_valid = False
# Define compatibility mapping between accumulator type and C type
acc_c_compatibility = {
cutlass.Float32: {
cutlass.Float32,
cutlass.Float16,
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
},
cutlass.Float16: {
cutlass.Float32,
cutlass.Float16,
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
},
cutlass.Int32: {
cutlass.Float32,
cutlass.Float16,
cutlass.Int32,
cutlass.Int8,
cutlass.Uint8,
},
}
# Check compatibility between accumulator type and C type
if c_dtype not in acc_c_compatibility[acc_dtype]:
is_valid = False
return is_valid
@staticmethod
def is_valid_tensor_alignment(
m: int,
n: int,
k: int,
l: int,
ab_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
) -> bool:
"""
Check if the tensor alignment is valid
:param m: The number of rows in the A tensor
:type m: int
:param n: The number of columns in the B tensor
:type n: int
:param k: The number of columns in the A tensor
:type k: int
:param l: The number of columns in the C tensor
:type l: int
:param ab_dtype: The data type of the A and B operands
:type ab_dtype: Type[cutlass.Numeric]
:param c_dtype: The data type of the output tensor
:type c_dtype: Type[cutlass.Numeric]
:param a_major: The major axis of the A tensor
:type a_major: str
:param b_major: The major axis of the B tensor
:type b_major: str
:param c_major: The major axis of the C tensor
:type c_major: str
:return: True if the problem shape is valid, False otherwise
:rtype: bool
"""
is_valid = True
def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
major_mode_idx = 0 if is_mode0_major else 1
num_major_elements = tensor_shape[major_mode_idx]
num_contiguous_elements = 16 * 8 // dtype.width
return num_major_elements % num_contiguous_elements == 0
if (
not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l))
or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l))
or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l))
):
is_valid = False
return is_valid
@@ -1418,7 +1436,7 @@ def run(
:rtype: float
"""
print(f"Running Hopper Dense GEMM with:")
print("Running Hopper Dense GEMM with:")
print(f"mnkl: {mnkl}")
print(
f"A dtype: {a_dtype}, B dtype: {b_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}"
@@ -1434,15 +1452,19 @@ def run(
# Unpack parameters
m, n, k, l = mnkl
# Skip unsupported types
if not HopperWgmmaGemmKernel.is_valid_dtypes(
a_dtype, b_dtype, acc_dtype, c_dtype, a_major, b_major
):
raise TypeError(
f"Skipping due to unsupported combination of types and majors: {a_dtype}, {b_dtype}, {acc_dtype}, {c_dtype}, {a_major=}, {b_major=}"
f"unsupported combination of types and majors: A {a_dtype}, B {b_dtype}, Acc {acc_dtype}, C {c_dtype}, {a_major=}, {b_major=}"
)
if not HopperWgmmaGemmKernel.is_valid_tensor_alignment(
m, n, k, l, a_dtype, c_dtype, a_major, b_major, c_major
):
raise TypeError(
"the contiguous dimension of A/B/C tensors is not 16 bytes aligned"
)
# Prepare pytorch tensors: A, B (random from 0 to 2) and C (all zero)
if not torch.cuda.is_available():
raise RuntimeError("GPU is required to run this example!")
examples/python/CuTeDSL/hopper/dense_gemm_persistent.py
+1619
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/hopper/fmha.py
+2537
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/notebooks/async_pipeline.ipynb
+599
View File
@@ -0,0 +1,599 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"import cutlass\n",
"import cutlass.cute as cute\n",
"from cutlass.cute.runtime import from_dlpack"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<style>\n",
"div.mermaid > svg {\n",
" width: 50% !important;\n",
" height: auto !important;\n",
"}\n",
"</style>\n",
"\n",
"# Tutorial: Warp Specialization with Async Pipeline in CuTe DSL\n",
"\n",
"This tutorial explores advanced CUDA programming techniques for implementing efficient producer-consumer \n",
"patterns using asynchronous communication primitives in the CuTe Domain Specific Language (DSL).\n",
"\n",
"## Foundation: Inter-Warp Communication Basics\n",
"\n",
"### Understanding CUDA Warps and Shared Memory\n",
"\n",
"A **warp** is the fundamental execution unit in CUDA, consisting of 32 threads that execute instructions in Single Instruction, \n",
"Multiple Thread (SIMT) fashion on a Streaming Multiprocessor (SM). Understanding warp-level programming is crucial for \n",
"achieving optimal GPU performance.\n",
"\n",
"**Key Concepts:**\n",
"- Warps execute in lockstep, making them ideal for SIMD operations\n",
"- Multiple warps within a thread block (CTA) can cooperate through shared memory\n",
"- Shared memory provides low-latency, high-bandwidth communication between threads\n",
"\n",
"### Shared Memory Architecture\n",
"\n",
"**Shared memory** serves as a programmer-managed cache with several important characteristics:\n",
"\n",
"- **Speed**: ~100x faster than global memory access\n",
"- **Scope**: Accessible by all threads within the same thread block\n",
"- **Organization**: Divided into banks (typically 32) to enable parallel access\n",
"- **Conflicts**: Bank conflicts occur when multiple threads access the same bank simultaneously\n",
"\n",
"### Traditional Synchronous Communication\n",
"\n",
"The conventional approach for inter-warp communication relies on explicit synchronization barriers. The following sequence diagram \n",
"illustrates the typical producer-consumer pattern:\n",
"\n",
"```mermaid\n",
"sequenceDiagram\n",
" participant W0 as Producer Warp\n",
" participant SMEM as Shared Memory\n",
" participant W1 as Consumer Warp\n",
" \n",
" W0->>SMEM: Write data\n",
" critical Synchronization Barrier\n",
" W0-->W1: __syncthreads()\n",
" SMEM->>W1: Read data\n",
" W0-->W1: __syncthreads()\n",
" end\n",
"```\n",
"\n",
"**Limitations of Synchronous Communication:**\n",
"- All warps must wait at synchronization points\n",
"- No opportunity for overlapped computation\n",
"- Reduced overall throughput due to forced serialization"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def synced_producer_consumer(SharedStorage: cutlass.Constexpr, res: cute.Tensor):\n",
" warp_idx = cute.arch.warp_idx()\n",
" warp_idx = cute.arch.make_warp_uniform(warp_idx)\n",
"\n",
" smem = cutlass.utils.SmemAllocator()\n",
" storage = smem.allocate(SharedStorage, 64)\n",
"\n",
" staging_smem = storage.staging_buffer.get_tensor(cute.make_layout(1))\n",
" staging_smem.fill(0)\n",
" cute.arch.sync_threads()\n",
"\n",
" for i in cutlass.range(cute.size(res)):\n",
" if warp_idx == 0:\n",
" staging_smem[0] = i * 1.0\n",
" # mark enter of critical region\n",
" cute.arch.sync_threads()\n",
" if warp_idx == 1:\n",
" res[i] = staging_smem[0]\n",
" # mark exit of critical region\n",
" cute.arch.sync_threads()\n",
"\n",
"\n",
"@cute.jit\n",
"def run_synced_producer_consumer(res: cute.Tensor):\n",
" @cute.struct\n",
" class SharedStorage:\n",
" staging_buffer: cute.struct.Align[\n",
" cute.struct.MemRange[cutlass.Float32, 1], 1024\n",
" ]\n",
"\n",
" synced_producer_consumer(SharedStorage, res).launch(\n",
" grid=(1, 1, 1), block=(64, 1, 1), smem=SharedStorage.size_in_bytes()\n",
" )\n",
"\n",
"\n",
"res = torch.zeros((8,), device=\"cuda\")\n",
"run_synced_producer_consumer(from_dlpack(res))"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"tensor([0., 1., 2., 3., 4., 5., 6., 7.], device='cuda:0')"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res"
]
},
{
"cell_type": "markdown",
"metadata": {
"editable": true,
"slideshow": {
"slide_type": ""
},
"tags": []
},
"source": [
"<style>\n",
"div.mermaid > svg {\n",
" width: 50% !important;\n",
" height: auto !important;\n",
"}\n",
"</style>\n",
"\n",
"## Asynchronous Communication: Breaking the Synchronization Bottleneck\n",
"\n",
"### The Problem with Synchronous Patterns\n",
"\n",
"The previous example demonstrates traditional synchronized communication between warps. While functional, this approach \n",
"has significant performance limitations:\n",
"\n",
"**Critical Section Analysis:**\n",
"- **First `__syncthreads()`**: Ensures data is written and ready for consumption\n",
"- **Second `__syncthreads()`**: Guarantees data has been consumed and memory can be safely overwritten\n",
"\n",
"**Performance Impact:**\n",
"- All warps are forced into lockstep execution\n",
"- No computational overlap between producer and consumer operations\n",
"- Wasted cycles as warps wait at synchronization barriers\n",
"\n",
"### Hopper Architecture: Enabling Asynchronous Primitives\n",
"\n",
"Starting with the Hopper architecture, CUDA introduced sophisticated asynchronous communication primitives that enable \n",
"**warp specialization**—allowing different warps to perform distinct, specialized roles while maintaining loose coupling.\n",
"\n",
"**Key Benefits:**\n",
"- **Overlapped Execution**: Producer and consumer warps can perform computations concurrently\n",
"- **Reduced Latency**: Eliminates unnecessary synchronization stalls\n",
"- **Better Resource Utilization**: Maximizes SM occupancy and throughput\n",
"\n",
"### Async Pipeline Communication Pattern\n",
"\n",
"The async pipeline abstraction provides a elegant solution for producer-consumer communication without rigid synchronization constraints:\n",
"\n",
"```mermaid\n",
"sequenceDiagram\n",
" participant W0 as Producer Warp\n",
" participant Pipeline as Async Pipeline\n",
" participant SMEM as Shared Memory \n",
" participant W1 as Consumer Warp\n",
" \n",
" W0->>Pipeline: Acquire (request write slot)\n",
" activate W1\n",
" Pipeline-->>W0: Grant access\n",
" deactivate W1\n",
" \n",
" W1->>Pipeline: Wait (for data availability)\n",
" activate Pipeline\n",
" \n",
" W0->>SMEM: Write data\n",
" W0->>Pipeline: Commit (signal data ready)\n",
" \n",
" Pipeline-->>W1: Data available\n",
" deactivate Pipeline\n",
" \n",
" activate W0\n",
" SMEM->>W1: Read data\n",
" deactivate W0\n",
" W1->>Pipeline: Release (mark slot available)\n",
"```\n",
"\n",
"**Async Pipeline Advantages:**\n",
"- **Non-blocking Operations**: Warps can perform other work while waiting\n",
"- **Fine-grained Control**: Explicit control over data readiness and consumption\n",
"- **Scalable**: Supports multiple producer-consumer pairs efficiently"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Async Pipeline API Reference\n",
"\n",
"The `PipelineAsync` abstraction in CuTe DSL provides a comprehensive set of primitives for implementing efficient producer-consumer patterns:\n",
"\n",
"#### Producer Operations\n",
"- **`PipelineProducer.acquire()`**: Blocks until a write slot becomes available (released by consumer)\n",
" - Returns with a handle pointing to a available slot immediately if there is\n",
" - Enables backpressure control to prevent buffer overflow\n",
" - **`PipelineProducer.acquire_and_advance()`** additionally moves the producer's write index to the next buffer slot\n",
"\n",
"- **`PipelineProducer.commit(PipelineProducer.ImmutableProducerHandle)`** / **`PipelineProducer.ImmutableProducerHandle.commit()`**: Signals that data has been written to the handle-pointed slot and is ready for consumption\n",
" - Triggers waiting consumers\n",
" - Maintains data consistency guarantees\n",
" - If no assigned handle, **`PipelineConsumerHandle.release()`** tracks its internal maintained handle (pointed to the last one it acquires)\n",
"\n",
"#### Consumer Operations \n",
"- **`PipelineConsumer.wait()`**: Blocks until data becomes available for reading\n",
" - Returns with a handle pointing to a committed slot when producer commits new data\n",
" - Supports timeout and polling variants\n",
" - **`PipelineConsumer.wait_and_advance()`** additionally moves the consumer's read index to the next buffer slot\n",
"\n",
"- **`PipelineConsumerHandle.release(PipelineConsumer.ImmutableConsumerHandle)`** / **`PipelineConsumer.ImmutableConsumerHandle.release()`**: Marks data as consumed and the handle-pointed slot as consumed and available for reuse\n",
" - Enables producers to acquire released slots\n",
" - Critical for preventing deadlock in circular buffers\n",
" - If no assigned handle, **`PipelineConsumerHandle.release()`** tracks its internal maintained handle (pointed to the last one it waits for)\n",
"\n",
"#### Disclaimer\n",
"\n",
"The `pipeline` APIs provided abstractions for developers to manage synchornization between warps, thread-blocks, etc. It doesn't provide deadlock-free guarantee. It's still developer's responsibility to write correct code to avoid deadlock.\n",
"\n",
"#### Performance Characteristics\n",
"\n",
"**Computational Overlap**: This asynchronous communication pattern enables limited but significant computational overlap:\n",
"- **Producer**: Can perform preprocessing, data transformation, or prefetching while consumer processes previous data\n",
"- **Consumer**: Can execute post-processing, result computation, or output operations while producer prepares next data\n",
"\n",
"**Memory Efficiency**: Explicit slot management ensures optimal memory utilization without unnecessary copying or buffering."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def async_pipeline_kernel(res: cute.Tensor):\n",
" warp_idx = cute.arch.warp_idx()\n",
" warp_idx = cute.arch.make_warp_uniform(warp_idx)\n",
"\n",
" @cute.struct\n",
" class SharedStorage:\n",
" tma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]\n",
" staging_buffer: cute.struct.Align[\n",
" cute.struct.MemRange[cutlass.Float32, 1], 1024\n",
" ]\n",
"\n",
" smem = cutlass.utils.SmemAllocator()\n",
" storage = smem.allocate(SharedStorage, 64)\n",
"\n",
" # Warp 0\n",
" producer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
" # Warp 1\n",
" consumer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
"\n",
" pipeline = cutlass.pipeline.PipelineAsync.create(\n",
" num_stages=1,\n",
" producer_group=producer_group,\n",
" consumer_group=consumer_group,\n",
" barrier_storage=storage.tma_mbar_ptr.data_ptr(),\n",
" )\n",
"\n",
" staging_smem = storage.staging_buffer.get_tensor(cute.make_layout(1))\n",
" staging_smem.fill(0)\n",
" cute.arch.sync_threads()\n",
"\n",
" producer, consumer = pipeline.make_participants()\n",
"\n",
" # Producer warp\n",
" if warp_idx == 0:\n",
" for i in cutlass.range(cute.size(res)):\n",
" # Producer: Wait for data buffer is available\n",
" handle = producer.acquire_and_advance()\n",
" # Producer: Write data to shared memory\n",
" staging_smem[handle.index] = 1.0 * i\n",
" # Producer: Signal data is ready for consumption\n",
" handle.commit()\n",
" producer.tail()\n",
"\n",
" # Consumer warp\n",
" if warp_idx == 1:\n",
" for i in cutlass.range(cute.size(res)):\n",
" # Consumer: Wait for producer to signal when data is available for use\n",
" handle = consumer.wait_and_advance()\n",
" # Conumer: consumes data\n",
" res[i] = staging_smem[handle.index]\n",
" # Conumer: Signal data buffer is ready for write\n",
" handle.release()\n",
"\n",
"\n",
"@cute.jit\n",
"def async_pipeline(res: cute.Tensor):\n",
" # Launch kernel with two warps: producer and consumer\n",
" async_pipeline_kernel(res).launch(grid=(1, 1, 1), block=(64, 1, 1))\n",
"\n",
"\n",
"res = torch.zeros((8,), device=\"cuda\")\n",
"async_pipeline(from_dlpack(res))"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"tensor([0., 1., 2., 3., 4., 5., 6., 7.], device='cuda:0')"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<style>\n",
"div.mermaid > svg {\n",
" width: 50% !important;\n",
" height: auto !important;\n",
"}\n",
"</style>\n",
"\n",
"## Advanced Pattern: Staged Async Pipeline with Circular Buffering\n",
"\n",
"### Limitations of Single-Stage Pipelines\n",
"\n",
"While async communication provides significant improvements over synchronous patterns, single-stage pipelines \n",
"still exhibit serialization bottlenecks:\n",
"\n",
"**Dependency Chain Analysis:**\n",
"```mermaid\n",
"sequenceDiagram\n",
" participant W0 as Producer\n",
" participant Pipeline as Pipeline\n",
" participant W1 as Consumer\n",
" \n",
" W0->>Pipeline: Acquire\n",
" Note over W0,W1: Producer waits here\n",
" W1->>Pipeline: Release\n",
" Pipeline-->>W0: Granted\n",
"```\n",
"\n",
"**Performance Bottleneck**: The producer must wait for the consumer to complete processing and release the buffer \n",
"before acquiring the next write slot. This creates a serialization point that limits overall throughput.\n",
"\n",
"### Multi-Stage Pipeline Architecture\n",
"\n",
"The **staged async pipeline** implements a circular buffer managed by an array of synchronization barriers, \n",
"enabling much higher degrees of parallelism:\n",
"\n",
"#### Core Concepts\n",
"\n",
"**Circular Buffer Management:**\n",
"- **Multiple Stages**: Support for N concurrent buffer slots (typically 2-8 stages)\n",
"- **Independent Indexing**: Producer and consumer maintain separate advancement indices\n",
"- **Barrier Array**: Each stage has an associated memory barrier for fine-grained synchronization\n",
"\n",
"#### Enhanced API Operations\n",
"\n",
"- **`PipelineProducer.advance()`**: Moves the producer's write index to the next buffer slot\n",
" - Enables round-robin buffer allocation\n",
" - Allows producer to continue without waiting for all previous data to be consumed\n",
" - Can be conducted implicitly when calling **`PipelineProducer.require_and_advance()`**\n",
"\n",
"- **`PipelineConsumer.advance()`**: Moves the consumer's read index to the next buffer slot\n",
" - Maintains proper ordering of data consumption\n",
" - Signals availability of processed slots\n",
" - Can be conducted implicitly when calling **`PipelineConsumer.wait_and_advance()`**\n",
"\n",
"- **`PipelineProducer.ImmutableResourceHandle.index`** / **`PipelineConsumer.ImmutableResourceHandle.index`**: Returns pointed buffer slot index\n",
" - Used for addressing specific staging buffer locations\n",
" - Enables direct slot-based data access\n",
"\n",
"### Circular Buffer State Visualization\n",
"\n",
"```\n",
"Legend:\n",
" W: Currently being written (producer active)\n",
" D: Data ready for consumption \n",
" R: Currently being read (consumer active)\n",
" X: Empty slot available for writing\n",
" \n",
" Advance Direction\n",
" <-------------------\n",
"\n",
" Producer Consumer\n",
" | ^\n",
" V |\n",
" +-----------------+\n",
" --|X|X|W|D|D|D|D|R|X|<-.\n",
" / +-----------------+ \\\n",
" | |\n",
" `------------------------' \n",
"```\n",
"\n",
"**Key Advantages:**\n",
"- **Increased Throughput**: Producer can stay ahead of consumer by multiple stages\n",
"- **Latency Hiding**: Consumer processing latency is hidden by buffered data\n",
"- **Better Resource Utilization**: Both warps can maintain high activity levels\n",
"- **Scalable Design**: Buffer depth can be tuned based on workload characteristics\n",
"\n",
"The following implementation demonstrates efficient multi-stage pipeline communication with proper circular buffer management:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"@cute.kernel\n",
"def async_pipeline_staged_kernel(\n",
" SharedStorage: cutlass.Constexpr, res: cute.Tensor, staging: cute.Tensor\n",
"):\n",
" stages = cute.size(staging)\n",
"\n",
" warp_idx = cute.arch.warp_idx()\n",
" warp_idx = cute.arch.make_warp_uniform(warp_idx)\n",
"\n",
" smem = cutlass.utils.SmemAllocator()\n",
" storage = smem.allocate(SharedStorage, 64)\n",
"\n",
" # Warp 0\n",
" producer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
" # Warp 1\n",
" consumer_group = cutlass.pipeline.CooperativeGroup(\n",
" cutlass.pipeline.Agent.Thread, 32\n",
" )\n",
"\n",
" pipeline = cutlass.pipeline.PipelineAsync.create(\n",
" num_stages=stages,\n",
" producer_group=producer_group,\n",
" consumer_group=consumer_group,\n",
" barrier_storage=storage.tma_mbar_ptr.data_ptr(),\n",
" )\n",
"\n",
" staging_smem = storage.staging_buffer.get_tensor(staging.layout)\n",
" staging_smem.fill(0)\n",
" cute.arch.sync_threads()\n",
"\n",
" producer, consumer = pipeline.make_participants()\n",
"\n",
" # Producer warp\n",
" if warp_idx == 0:\n",
" for i in cutlass.range(cute.size(res)):\n",
" handle = producer.acquire_and_advance()\n",
" staging_smem[handle.index] = 1.0 * i\n",
" handle.commit() # or producer.commit(handle)\n",
"\n",
" # prevents CTA0 from retiring until it receives all expected arrives.\n",
" producer.tail()\n",
"\n",
" # Consumer warp\n",
" if warp_idx == 1:\n",
" for i in cutlass.range(cute.size(res)):\n",
" handle = consumer.wait_and_advance()\n",
" res[i] = staging_smem[handle.index]\n",
" handle.release() # or consumer.release(handle)\n",
"\n",
" tidx, _, _ = cute.arch.thread_idx()\n",
" if tidx == 0:\n",
" staging.store(staging_smem.load())\n",
"\n",
"\n",
"@cute.jit\n",
"def async_pipeline_staged(res: cute.Tensor, staging: cute.Tensor):\n",
" stages = cute.size(staging)\n",
"\n",
" @cute.struct\n",
" class SharedStorage:\n",
" tma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, stages * 2]\n",
" staging_buffer: cute.struct.Align[\n",
" cute.struct.MemRange[cutlass.Float32, stages], 1024\n",
" ]\n",
"\n",
" async_pipeline_staged_kernel(SharedStorage, res, staging).launch(\n",
" grid=(1, 1, 1), block=(64, 1, 1), smem=SharedStorage.size_in_bytes()\n",
" )\n",
"\n",
"\n",
"res = torch.zeros((8,), device=\"cuda\")\n",
"staging = torch.zeros((5,), device=\"cuda\")\n",
"async_pipeline_staged(from_dlpack(res), from_dlpack(staging))\n",
"torch.cuda.synchronize()"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(tensor([0., 1., 2., 3., 4., 5., 6., 7.], device='cuda:0'),\n",
" tensor([5., 6., 7., 3., 4.], device='cuda:0'))"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"res, staging"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Try Acquire/Wait\n",
"\n",
"In some circumstances, developers may want to just check status of pipeline state without blocking. This could benefit some cases that we have independent instructions to hide latency of checking pipeline state. We provided `try_aquire` or `try_wait` which are non-blocking APIs. "
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.10"
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"state": {},
"version_major": 2,
"version_minor": 0
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
examples/python/CuTeDSL/notebooks/benchmark_autotune.ipynb
+460
View File
@@ -0,0 +1,460 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"import torch\n",
"\n",
"import cutlass\n",
"import cutlass.cute as cute\n",
"import cutlass.cute.testing as testing\n",
"import cutlass.torch as cutlass_torch"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## The Usage of Benchmark and Autotune Utilities in CuTe DSL\n",
"\n",
"CuTe DSL provides autotune and benchmark utilities to help users evaluate and optimize kernel performance. This notebook demonstrates how to use these tools.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"### Autotune\n",
"\n",
"We provides two kinds of autotune utilities for users: `autotune.jit` decorator and the `tune` function. The former is used as a decorator used on top of `@cute.jit` while the latter is used as an individual function.\n",
"\n",
"#### @autotune.jit\n",
"\n",
"We take the `elementwise_add_kernel` as an example. After writing the jit host function and kernel, we could add the `@autotune_jit` decorator on top of the jit host function to enable autotune. \n",
"```python\n",
"@testing.autotune_jit(\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" update_on_change=[\"M\", \"N\"],\n",
" warmup_iterations=100,\n",
" iterations=100,\n",
")\n",
"```\n",
"\n",
"The `autotune_jit` decorator provides several parameters to control the autotuning process:\n",
"\n",
"- params_dict: A dictionary containing the parameters to be tuned and their possible values\n",
"- update_on_change: A list of argument names that trigger re-tuning when their values change\n",
"- warmup_iterations: Number of warmup iterations before timing\n",
"- iterations: Number of iterations for timing each parameter combination\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"@cute.kernel\n",
"def elementwise_add_kernel(\n",
" gA: cute.Tensor,\n",
" gB: cute.Tensor,\n",
" gC: cute.Tensor,\n",
" cC: cute.Tensor, # coordinate tensor\n",
" shape: cute.Shape,\n",
" thr_layout: cute.Layout,\n",
" val_layout: cute.Layout,\n",
"):\n",
" tidx, _, _ = cute.arch.thread_idx()\n",
" bidx, _, _ = cute.arch.block_idx()\n",
"\n",
" # slice for CTAs\n",
" # logical id -> address\n",
" blk_coord = ((None, None), bidx)\n",
" blkA = gA[blk_coord] # (TileM,TileN)\n",
" blkB = gB[blk_coord] # (TileM,TileN)\n",
" blkC = gC[blk_coord] # (TileM,TileN)\n",
" blkCrd = cC[blk_coord] # (TileM, TileN)\n",
"\n",
" # # declare the atoms which will be used later for memory copy\n",
" copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type)\n",
" copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gC.element_type)\n",
"\n",
" tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)\n",
" tiled_copy_B = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)\n",
" tiled_copy_C = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout)\n",
"\n",
" thr_copy_A = tiled_copy_A.get_slice(tidx)\n",
" thr_copy_B = tiled_copy_B.get_slice(tidx)\n",
" thr_copy_C = tiled_copy_C.get_slice(tidx)\n",
"\n",
" thrA = thr_copy_A.partition_S(blkA)\n",
" thrB = thr_copy_B.partition_S(blkB)\n",
" thrC = thr_copy_C.partition_S(blkC)\n",
"\n",
" # allocate fragments for gmem->rmem\n",
" frgA = cute.make_fragment_like(thrA)\n",
" frgB = cute.make_fragment_like(thrB)\n",
" frgC = cute.make_fragment_like(thrC)\n",
"\n",
" thrCrd = thr_copy_C.partition_S(blkCrd)\n",
" frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)\n",
"\n",
" for i in range(0, cute.size(frgPred), 1):\n",
" val = cute.elem_less(thrCrd[i], shape)\n",
" frgPred[i] = val\n",
"\n",
" ##########################################################\n",
" # Move data to reg address space\n",
" ##########################################################\n",
"\n",
" cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)\n",
" cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)\n",
"\n",
" # Load data before use. The compiler will optimize the copy and load\n",
" # operations to convert some memory ld/st into register uses.\n",
" result = frgA.load() + frgB.load()\n",
"\n",
" # Save the results back to registers. Here we reuse b's registers.\n",
" frgC.store(result)\n",
"\n",
" # Copy the results back to c\n",
" cute.copy(copy_atom_store, frgC, thrC, pred=frgPred)\n",
"\n",
"\n",
"@testing.autotune_jit(\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" update_on_change=[\"M\", \"N\"],\n",
" warmup_iterations=100,\n",
" iterations=100,\n",
")\n",
"@cute.jit\n",
"def elementwise_add_autotune(mA, mB, mC, M, N, copy_bits: cutlass.Constexpr = 128):\n",
" dtype = mA.element_type\n",
" vector_size = copy_bits // dtype.width\n",
"\n",
" thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))\n",
" val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))\n",
" tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)\n",
"\n",
" gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN))\n",
" gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN))\n",
" gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN))\n",
" idC = cute.make_identity_tensor(mC.shape)\n",
" cC = cute.zipped_divide(idC, tiler=tiler_mn)\n",
"\n",
" elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch(\n",
" grid=[cute.size(gC, mode=[1]), 1, 1],\n",
" block=[cute.size(tv_layout, mode=[0]), 1, 1],\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When we run the jit funciton `elementwise_add_autotune`, the CuTe DSL will help us tune the kernels by looping the specified configs and run the kernel with the best config.\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"\n",
"M, N = 1024, 1024\n",
"dtype = cutlass.Float32\n",
"skip_ref_check = False\n",
"\n",
"print(f\"\\nRunning Elementwise Add test with:\")\n",
"print(f\"Tensor dimensions: [{M}, {N}]\")\n",
"print(f\"Input and Output Data type: {dtype}\")\n",
"\n",
"torch_dtype = cutlass_torch.dtype(dtype)\n",
"\n",
"a = torch.randn(M, N, device=torch.device(\"cuda\"), dtype=torch_dtype)\n",
"b = torch.randn(M, N, device=torch.device(\"cuda\"), dtype=torch_dtype)\n",
"\n",
"c = torch.zeros_like(a)\n",
"\n",
"print(f\"Input tensor shapes:\")\n",
"print(f\"a: {a.shape}, dtype: {a.dtype}\")\n",
"print(f\"b: {b.shape}, dtype: {b.dtype}\")\n",
"print(f\"c: {c.shape}, dtype: {c.dtype}\\n\")\n",
"\n",
"elementwise_add_autotune(a, b, c, M, N)\n",
"\n",
"if not skip_ref_check:\n",
" print(\"Verifying results for autotuned function ...\")\n",
" torch.testing.assert_close(a + b, c)\n",
" print(\"Results verified successfully!\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The output is as follows:\n",
"\n",
"```\n",
"Running Elementwise Add test with:\n",
"Tensor dimensions: [1024, 1024]\n",
"Input and Output Data type: Float32\n",
"Input tensor shapes:\n",
"a: torch.Size([1024, 1024]), dtype: torch.float32\n",
"b: torch.Size([1024, 1024]), dtype: torch.float32\n",
"c: torch.Size([1024, 1024]), dtype: torch.float32\n",
"Verifying results for autotuned function ...\n",
"Results verified successfully!\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"To monitor the autotuning process in detail, you can enable logging by setting the environment variable `CUTE_DSL_LOG_AUTOTUNE`. \n",
"```shell\n",
"export CUTE_DSL_LOG_AUTOTUNE=1\n",
"```\n",
"This will display comprehensive information including:\n",
"- Each configuration being evaluated and its corresponding execution time\n",
"- The optimal configuration that was selected\n",
"- Total time spent on tuning\n",
"- Cache hit/miss statistics\n",
"\n",
"\n",
"Below is a sample output showing the autotuning process with different configurations:\n",
"```python\n",
"2025-07-23 06:17:03,978 - cutlass.cute.testing_Autotune - INFO - Tuning configuration: {'copy_bits': 64}\n",
"2025-07-23 06:17:04,519 - cutlass.cute.testing_Autotune - INFO - Execution time: 0.010857919985428453 us\n",
"2025-07-23 06:17:04,519 - cutlass.cute.testing_Autotune - INFO - Tuning configuration: {'copy_bits': 128}\n",
"2025-07-23 06:17:04,683 - cutlass.cute.testing_Autotune - INFO - Execution time: 0.011117440033704042 us\n",
"2025-07-23 06:17:04,683 - cutlass.cute.testing_Autotune - INFO - Best configuration: {'copy_bits': 64}, execution time: 0.010857919985428453 us\n",
"2025-07-23 06:17:04,683 - cutlass.cute.testing_Autotune - INFO - Total tuning time: 0.7053244113922119 s\n",
"...\n",
"2025-07-23 06:17:04,700 - cutlass.cute.testing_Autotune - INFO - Using cached best configuration: {'copy_bits': 64}\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### tune\n",
"\n",
"We also provide a `tune` funtion. The interface of the `tune` function is as follows:\n",
"\n",
"```python\n",
"def tune(\n",
" func: Callable[[Any], Callable[[], Any]],\n",
" params_dict: Dict[str, List[Any]] = None,\n",
" kernel_arguments: JitArguments = JitArguments(),\n",
" warmup_iterations=10,\n",
" iterations=100,\n",
" stream: Optional[cuda_driver.CUstream] = None,\n",
") -> Dict[str, Any]:\n",
"```\n",
"\n",
"The `tune` function takes the following parameters:\n",
"\n",
"- func: A callable that takes configuration parameters and returns a kernel function\n",
"- params_dict: Dictionary mapping parameter names to lists of possible values to tune\n",
"- kernel_arguments: Arguments to pass to the kernel for tuning\n",
"- warmup_iterations: Number of warmup iterations before timing (default: 10)\n",
"- iterations: Number of timing iterations per configuration (default: 100)\n",
"- stream: Optional CUDA stream to use for execution. defaults to default CUDA stream. The stream parameter must match the stream passed to the kernel, mismatched streams will result in an error.\n",
"\n",
"It returns a dictionary containing the best kernel configuration found.\n",
"\n",
"\n",
"Here is an example to use the `tune` function:\n",
"\n",
"1. First remove the `@testing.autotune_jit` decorator from the `elementwise_add_autotune` function:\n",
" ```python\n",
" @testing.autotune_jit(\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" update_on_change=[\"M\", \"N\"], \n",
" warmup_iterations=100,\n",
" iterations=100,\n",
" )\n",
" ```\n",
"\n",
" 2. Define a `tune_func` that:\n",
" - Takes input tensors (a, b, c), dimensions (M, N) and tuning parameter copy_bits\n",
" - Compiles the `elementwise_add_autotune` function using `cute.compile()`\n",
" - Returns a lambda function that executes the compiled kernel\n",
"\n",
" 3. Pass `tune_func` to `testing.tune` function along with:\n",
" - Parameter space to explore (copy_bits values)\n",
" - Kernel arguments wrapped in JitArguments\n",
" - The `tune` function will find optimal parameters automatically\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"def tune_func(a, b, c, M, N, copy_bits=128):\n",
" compiled_func = cute.compile(elementwise_add_autotune, a, b, c, M, N, copy_bits=128)\n",
" return lambda: compiled_func(a, b, c, M, N)\n",
"\n",
"params = testing.tune(\n",
" tune_func,\n",
" params_dict={\"copy_bits\": [64, 128]},\n",
" kernel_arguments=testing.JitArguments(a, b, c, M, N),\n",
")\n",
"print(f\"The best kernel configs found: {params}\")\n",
"\n",
"# run the kernel with the best config\n",
"compiled_func = cute.compile(elementwise_add_autotune, a, b, c, M, N, **params)\n",
"compiled_func(a, b, c, M, N)\n",
" "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### benchmark\n",
"\n",
"In CuTe DSL, the benchmark utility can be used to measure kernel execution time. The interface of benchmark routine is as follows:\n",
"\n",
"```python\n",
"def benchmark(\n",
" callable: Callable,\n",
" *,\n",
" warmup_iterations: int = 10,\n",
" iterations: int = 100,\n",
" stream: Optional[cuda_driver.CUstream] = None,\n",
" kernel_arguments: Optional[JitArguments] = None,\n",
" workspace_generator: Optional[Callable[[], JitArguments]] = None,\n",
" workspace_count: int = 1,\n",
" use_cuda_graphs: bool = False,\n",
") -> float:\n",
"```\n",
"\n",
"The benchmark utility exposes several key configuration parameters to control profiling behavior:\n",
"\n",
"- callable: The function to be benchmarked\n",
"- warmup_iterations: Controls the number of initial warmup iterations before measurement begins (default: 10)\n",
"- iterations: Specifies how many iterations to profile for performance measurement (default: 100)\n",
"- stream: Designates which CUDA stream to execute the kernel on (default: default stream) \n",
"- use_cuda_graphs: Whether enables CUDA graph for the callable function to minimize kernel launch overhead (default: False)\n",
"- workspace_generator: Provides a function that generates fresh kernel arguments each iteration to avoid caching effects\n",
"- workspace_count: Determines how many different workspaces to cycle through during profiling (default: 1)\n",
"\n",
"When benchmarking, there are several key parameters that can be configured:\n",
"\n",
"1. Core parameters:\n",
" - The function to profile (callable)\n",
" - Number of warmup iterations before measurement\n",
" - Number of profiling iterations for measurement\n",
"\n",
"2. Stream configuration:\n",
" - For kernels running in non-default streams, the stream must be specified\n",
" - The stream parameter must match the stream passed to the kernel, mismatched streams will result in an error\n",
"\n",
"3. Cache effects mitigation:\n",
" - To prevent L2 cache effects from skewing results, multiple workspaces can be cycled through\n",
" - This is controlled via workspace_count and workspace_generator parameters\n",
" - Each workspace provides fresh kernel arguments\n",
"\n",
"4. CUDA Graph support:\n",
" - Enables measuring kernel execution time without host overhead\n",
" - Requires the callable to be decorated with @cute.jit\n",
" - Must use a non-default CUDA stream when using graphs\n",
"\n",
"This function will return the execution time of the callable in microseconds. As GPU frequency can vary dynamically, we could fix the SM and memory frequencies to get more stable and reproducible benchmark results. This can be done by setting the GPU clocks using nvidia-smi before running the benchmark. In the next, let's use the benchmark function to get the execution time of the above elementwise_add kernel."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"def generate_kernel_arguments():\n",
" a = torch.randn(\n",
" M, N, device=torch.device(\"cuda\"), dtype=torch_dtype\n",
" )\n",
" b = torch.randn(\n",
" M, N, device=torch.device(\"cuda\"), dtype=torch_dtype\n",
" )\n",
"\n",
" c = torch.zeros_like(a)\n",
"\n",
" return testing.JitArguments(a, b, c, M, N)\n",
"\n",
"avg_time_us = testing.benchmark(\n",
" elementwise_add_autotune,\n",
" workspace_generator=generate_kernel_arguments,\n",
" workspace_count=10,\n",
" warmup_iterations=10,\n",
" iterations=100,\n",
")\n",
"\n",
"# Print execution results\n",
"print(\n",
" f\"Kernel execution time for cute.jit kernel with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms\"\n",
")\n",
"print(\n",
" f\"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s\"\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After running the code, we will get output similar to the following:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```\n",
"Kernel execution time for cute.jit kernel with M=1024, N=1024: 0.0403 ms\n",
"Achieved memory throughput for M=1024, N=1024: 312.37 GB/s\n",
"```"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
examples/python/CuTeDSL/notebooks/composed_layout.ipynb
+225
View File
@@ -0,0 +1,225 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "0c7cf795",
"metadata": {},
"source": [
"# Composed Layout in CuTe\n",
"\n",
"A **Composed Layout** is a powerful abstraction in CuTe that enables complex data transformations through \n",
"the composition of layouts and transformations. It provides a flexible way to manipulate memory layouts \n",
"and coordinate systems.\n",
"\n",
"## Components\n",
"\n",
"A Composed Layout consists of three key components:\n",
"\n",
"1. **Inner Layout/Transformation** (`inner`):\n",
" - Can be a layout, swizzle, or custom transformation function\n",
" - Applies the final transformation to the coordinates\n",
" - Supports arbitrary coordinate manipulations\n",
"\n",
"2. **Offset** (`offset`):\n",
" - Typically represented as an integer tuple\n",
" - Adds a constant displacement to coordinates\n",
" - Enables fine-grained control over data positioning\n",
"\n",
"3. **Outer Layout** (`outer`):\n",
" - The layout visible to the user\n",
" - Defines the initial coordinate transformation\n",
" - Determines the shape and organization of the data structure\n",
"\n",
"## Mathematical Representation\n",
"\n",
"The mathematical composition of these components is defined as:\n",
"\n",
"$\n",
"R(c) := (inner \\circ offset \\circ outer)(c) := inner(offset + outer(c))\n",
"$\n",
"\n",
"Where:\n",
"- $c$ represents the input coordinates\n",
"- $\\circ$ denotes function composition\n",
"- The transformation is applied from right to left\n",
"\n",
"## Usage in Python\n",
"\n",
"To create a Composed Layout in Python, use the `make_composed_layout` function:\n",
"\n",
"```python\n",
"layout = cute.make_composed_layout(inner, offset, outer)\n",
"```\n",
"\n",
"## Key Benefits\n",
"\n",
"1. **Flexibility**: Supports complex transformations that direct composition cannot handle\n",
"2. **Modularity**: Separates different aspects of the transformation\n",
"3. **Performance**: Enables optimized memory access patterns for GPU computations\n",
"4. **Compatibility**: Works with various types of transformations and layouts"
]
},
{
"cell_type": "markdown",
"id": "24448f7d",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"source": [
"## Custom Transformation Example\n",
"\n",
"This example demonstrates how to create a Composed Layout with a custom transformation function. We'll create a simple transformation that:\n",
"\n",
"1. Takes a 2D coordinate input `(x, y)`\n",
"2. Increments the y-coordinate by 1\n",
"3. Combines this with an offset and identity layout\n",
"\n",
"The example shows how to:\n",
"- Define a custom transformation function\n",
"- Create a composed layout with the transformation\n",
"- Apply the layout to coordinates\n",
"- Print the results for verification"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "184f30e6",
"metadata": {},
"outputs": [],
"source": [
"import cutlass\n",
"import cutlass.cute as cute\n",
"from cutlass.cute.runtime import from_dlpack, make_ptr\n",
"\n",
"\n",
"@cute.jit\n",
"def customized_layout():\n",
" def inner(c):\n",
" x, y = c\n",
" return x, y + 1\n",
"\n",
" layout = cute.make_composed_layout(\n",
" inner, (1, 0), cute.make_identity_layout(shape=(8, 4))\n",
" )\n",
" print(layout)\n",
" cute.printf(layout(0))\n",
"\n",
"\n",
"customized_layout()"
]
},
{
"cell_type": "markdown",
"id": "c897187f",
"metadata": {},
"source": [
"## Gather/Scatter Operations with Composed Layout\n",
"\n",
"Gather and Scatter operations are fundamental data access patterns in parallel computing and GPU programming. In CuTe, we can implement these operations elegantly using Composed Layout.\n",
"\n",
"### Gather Operation\n",
"A gather operation collects elements from a source array using an index array (also called an indirection array). It's defined as:\n",
"```python\n",
"output[i] = source[index[i]]\n",
"```\n",
"\n",
"#### Components in CuTe Implementation:\n",
"1. **Offset Tensor**: Contains the indices for gathering (`offset_tensor`)\n",
"2. **Data Pointer**: Points to the source data array (`data_ptr`)\n",
"3. **Shape**: Defines the shape of logic tensor viewed by user (`shape`)\n",
"\n",
"### How it Works\n",
"1. The inner transformation function reads from the offset tensor:\n",
" ```python\n",
" def inner(c):\n",
" return offset_tensor[c] # Returns the gather index\n",
" ```\n",
"2. The composed layout maps input coordinates through the offset tensor:\n",
" ```python\n",
" gather_layout = cute.make_composed_layout(inner, 0, cute.make_layout(shape))\n",
" ```\n",
"3. This creates an indirect access pattern where:\n",
" - Input coordinate `i` → `offset_tensor[i]` → `data_ptr[offset_tensor[i]]`\n",
"\n",
"4. notably, layout operations like slice, partition can still be applied on `outer` layout\n",
"\n",
"### Use Cases\n",
"- **Sparse Operations**: Accessing non-contiguous memory efficiently\n",
"- **Graph Processing**: Following edge connections in graph algorithms\n",
"- **Feature Embedding**: Looking up embeddings for discrete tokens\n",
"- **Irregular Data Access**: Any pattern requiring indirect memory access\n",
"\n",
"### Example Output Interpretation\n",
"The example code prints pairs of numbers `i -> j` where:\n",
"- `i` is the output index\n",
"- `j` is the gathered source index from `offset_tensor`\n",
"\n",
"This demonstrates how the composed layout transforms coordinates for indirect memory access.\n",
"\n",
"Note: Scatter operations (writing to indirect locations) can be implemented similarly by reversing the data flow direction.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d68f9476",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"\n",
"@cute.jit\n",
"def gather_tensor(\n",
" offset_tensor: cute.Tensor, data_ptr: cute.Pointer, shape: cute.Shape\n",
"):\n",
" def inner(c):\n",
" return offset_tensor[c]\n",
"\n",
" gather_layout = cute.make_composed_layout(inner, 0, cute.make_layout(shape))\n",
" for i in cutlass.range_constexpr(cute.size(shape)):\n",
" cute.printf(\"%d -> %d\", i, gather_layout(i))\n",
"\n",
" # TODO: support in future\n",
" # gather_tensor = cute.make_tensor(data_ptr, gather_layout)\n",
" # cute.printf(gather_tensor[0])\n",
"\n",
"\n",
"shape = (16,)\n",
"offset_tensor = torch.randint(0, 256, shape, dtype=torch.int32)\n",
"data_tensor = torch.arange(0, 256, dtype=torch.int32)\n",
"\n",
"\n",
"gather_tensor(\n",
" from_dlpack(offset_tensor),\n",
" make_ptr(cutlass.Int32, data_tensor.data_ptr(), cute.AddressSpace.generic),\n",
" shape,\n",
")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv3_12",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
examples/python/CuTeDSL/notebooks/cuda_graphs.ipynb
+6 -3
View File
@@ -26,10 +26,11 @@
"source": [
"# import torch for CUDA graphs\n",
"import torch\n",
"import cutlass\n",
"import cutlass.cute as cute\n",
"\n",
"# import CUstream type from the cuda driver bindings\n",
"from cuda.bindings.driver import CUstream\n",
"\n",
"# import the current_stream function from torch\n",
"from torch.cuda import current_stream"
]
@@ -61,13 +62,15 @@
" \"\"\"\n",
" cute.printf(\"Hello world\")\n",
"\n",
"\n",
"@cute.jit\n",
"def hello_world(stream : CUstream):\n",
"def hello_world(stream: CUstream):\n",
" \"\"\"\n",
" Host function that launches our (1,1,1), (1,1,1) grid in stream\n",
" \"\"\"\n",
" hello_world_kernel().launch(grid=[1, 1, 1], block=[1, 1, 1], stream=stream)\n",
"\n",
"\n",
"# Grab a stream from PyTorch, this will also initialize our context\n",
"# so we can omit cutlass.cuda.initialize_cuda_context()\n",
"stream = current_stream()\n",
@@ -585,7 +588,7 @@
"\n",
"# Calculate the time spent when launching kernels in a stream\n",
"# Results are in ms\n",
"stream_time = start.elapsed_time(end) \n",
"stream_time = start.elapsed_time(end)\n",
"\n",
"# Warmup our GPU again\n",
"g.replay()\n",
examples/python/CuTeDSL/notebooks/cute_layout_algebra.ipynb
+90 -61
View File
@@ -90,7 +90,9 @@
" \"\"\"\n",
" Demonstrates coalesce operation flattening and combining modes\n",
" \"\"\"\n",
" layout = cute.make_layout((2, (1, 6)), stride=(1, (cutlass.Int32(6), 2))) # Dynamic stride\n",
" layout = cute.make_layout(\n",
" (2, (1, 6)), stride=(1, (cutlass.Int32(6), 2))\n",
" ) # Dynamic stride\n",
" result = cute.coalesce(layout)\n",
"\n",
" print(\">>> Original:\", layout)\n",
@@ -98,6 +100,7 @@
" print(\">>> Coalesced:\", result)\n",
" cute.printf(\">?? Coalesced: {}\", result)\n",
"\n",
"\n",
"coalesce_example()"
]
},
@@ -275,8 +278,7 @@
" 3. for all i, 0 <= i < size(@a layout), @a result(i) == @a layout(i)\n",
" \"\"\"\n",
" layout = cute.make_layout(\n",
" ((2, (3, 4)), (3, 2), 1),\n",
" stride=((4, (8, 24)), (2, 6), 12)\n",
" ((2, (3, 4)), (3, 2), 1), stride=((4, (8, 24)), (2, 6), 12)\n",
" )\n",
" result = cute.coalesce(layout)\n",
"\n",
@@ -288,21 +290,26 @@
" original_size = cute.size(layout)\n",
" coalesced_size = cute.size(result)\n",
" print(f\"Original size: {original_size}, Coalesced size: {coalesced_size}\")\n",
" assert coalesced_size == original_size, \\\n",
" f\"Size mismatch: original {original_size}, coalesced {coalesced_size}\"\n",
" \n",
" assert coalesced_size == original_size, (\n",
" f\"Size mismatch: original {original_size}, coalesced {coalesced_size}\"\n",
" )\n",
"\n",
" print(\">>> 2. Checking depth of coalesced layout <= 1:\")\n",
" depth = cute.depth(result)\n",
" print(f\"Depth of coalesced layout: {depth}\")\n",
" assert depth <= 1, f\"Depth of coalesced layout should be <= 1, got {depth}\"\n",
"\n",
" print(\">>> 3. Checking layout functionality remains the same after the coalesce operation:\")\n",
" print(\n",
" \">>> 3. Checking layout functionality remains the same after the coalesce operation:\"\n",
" )\n",
" for i in cutlass.range_constexpr(original_size):\n",
" original_value = layout(i)\n",
" coalesced_value = result(i)\n",
" print(f\"Index {i}: original {original_value}, coalesced {coalesced_value}\")\n",
" assert coalesced_value == original_value, \\\n",
" assert coalesced_value == original_value, (\n",
" f\"Value mismatch at index {i}: original {original_value}, coalesced {coalesced_value}\"\n",
" )\n",
"\n",
"\n",
"coalesce_post_conditions()"
]
@@ -338,11 +345,12 @@
"\n",
" # Coalesce with mode-wise profile (1,1) = coalesce both modes\n",
" result = cute.coalesce(layout, target_profile=(1, 1))\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Original: \", layout)\n",
" print(\">>> Coalesced Result: \", result)\n",
"\n",
"\n",
"bymode_coalesce_example()"
]
},
@@ -387,18 +395,19 @@
" \"\"\"\n",
" Demonstrates basic layout composition R = A ◦ B\n",
" \"\"\"\n",
" A = cute.make_layout((6, 2), stride=(cutlass.Int32(8), 2)) # Dynamic stride\n",
" A = cute.make_layout((6, 2), stride=(cutlass.Int32(8), 2)) # Dynamic stride\n",
" B = cute.make_layout((4, 3), stride=(3, 1))\n",
" R = cute.composition(A, B)\n",
"\n",
" # Print static and dynamic information\n",
" print(\">>> Layout A:\", A)\n",
" cute.printf(\">?? Layout A: {}\", A)\n",
" print(\">>> Layout B:\", B) \n",
" print(\">>> Layout B:\", B)\n",
" cute.printf(\">?? Layout B: {}\", B)\n",
" print(\">>> Composition R = A ◦ B:\", R)\n",
" cute.printf(\">?? Composition R: {}\", R)\n",
"\n",
"\n",
"composition_example()"
]
},
@@ -438,14 +447,8 @@
" Shows difference between static and dynamic composition results\n",
" \"\"\"\n",
" # Static version - using compile-time values\n",
" A_static = cute.make_layout(\n",
" (10, 2), \n",
" stride=(16, 4)\n",
" )\n",
" B_static = cute.make_layout(\n",
" (5, 4), \n",
" stride=(1, 5)\n",
" )\n",
" A_static = cute.make_layout((10, 2), stride=(16, 4))\n",
" B_static = cute.make_layout((5, 4), stride=(1, 5))\n",
" R_static = cute.composition(A_static, B_static)\n",
"\n",
" # Static print shows compile-time info\n",
@@ -457,20 +460,21 @@
" # Dynamic version - using runtime Int32 values\n",
" A_dynamic = cute.make_layout(\n",
" (cutlass.Int32(10), cutlass.Int32(2)),\n",
" stride=(cutlass.Int32(16), cutlass.Int32(4))\n",
" stride=(cutlass.Int32(16), cutlass.Int32(4)),\n",
" )\n",
" B_dynamic = cute.make_layout(\n",
" (cutlass.Int32(5), cutlass.Int32(4)),\n",
" stride=(cutlass.Int32(1), cutlass.Int32(5))\n",
" stride=(cutlass.Int32(1), cutlass.Int32(5)),\n",
" )\n",
" R_dynamic = cute.composition(A_dynamic, B_dynamic)\n",
" \n",
"\n",
" # Dynamic printf shows runtime values\n",
" cute.printf(\">?? Dynamic composition:\")\n",
" cute.printf(\">?? A_dynamic: {}\", A_dynamic)\n",
" cute.printf(\">?? B_dynamic: {}\", B_dynamic)\n",
" cute.printf(\">?? R_dynamic: {}\", R_dynamic)\n",
"\n",
"\n",
"composition_static_vs_dynamic_layout()"
]
},
@@ -511,12 +515,12 @@
" \"\"\"\n",
" # Define the original layout A\n",
" A = cute.make_layout(\n",
" (cutlass.Int32(12), (cutlass.Int32(4), cutlass.Int32(8))), \n",
" stride=(cutlass.Int32(59), (cutlass.Int32(13), cutlass.Int32(1)))\n",
" (cutlass.Int32(12), (cutlass.Int32(4), cutlass.Int32(8))),\n",
" stride=(cutlass.Int32(59), (cutlass.Int32(13), cutlass.Int32(1))),\n",
" )\n",
"\n",
" # Define the tiler for by-mode composition\n",
" tiler = (3, 8) # Apply 3:1 to mode-0 and 8:1 to mode-1\n",
" tiler = (3, 8) # Apply 3:1 to mode-0 and 8:1 to mode-1\n",
"\n",
" # Apply by-mode composition\n",
" result = cute.composition(A, tiler)\n",
@@ -529,6 +533,7 @@
" print(\">>> By-mode Composition Result:\", result)\n",
" cute.printf(\">?? By-mode Composition Result: {}\", result)\n",
"\n",
"\n",
"bymode_composition_example()"
]
},
@@ -571,19 +576,20 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((4, 2, 3), stride=(2, 1, 8)) # (4,2,3):(2,1,8)\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout(4, stride=2) # Apply to layout 4:2\n",
" \n",
"\n",
" # Apply logical divide\n",
" result = cute.logical_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Logical Divide Result:\", result)\n",
" cute.printf(\">?? Logical Divide Result: {}\", result)\n",
"\n",
"\n",
"logical_divide_1d_example()"
]
},
@@ -620,21 +626,26 @@
" Result Shape : ((TileM,RestM), (TileN,RestN), L, ...)\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply logical divide\n",
" result = cute.logical_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Logical Divide Result:\", result)\n",
" cute.printf(\">?? Logical Divide Result: {}\", result)\n",
"\n",
"\n",
"logical_divide_2d_example()"
]
},
@@ -673,21 +684,26 @@
" Result Shape : ((TileM,TileN), (RestM,RestN,L,...))\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply zipped divide\n",
" result = cute.zipped_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Zipped Divide Result:\", result)\n",
" cute.printf(\">?? Zipped Divide Result: {}\", result)\n",
"\n",
"\n",
"zipped_divide_example()"
]
},
@@ -724,21 +740,26 @@
" Result Shape : ((TileM,TileN), RestM, RestN, L, ...)\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply tiled divide\n",
" result = cute.tiled_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Tiled Divide Result:\", result)\n",
" cute.printf(\">?? Tiled Divide Result: {}\", result)\n",
"\n",
"\n",
"tiled_divide_example()"
]
},
@@ -775,21 +796,26 @@
" Result Shape : (TileM, TileN, RestM, RestN, L, ...)\n",
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((9, (4, 8)), stride=(59, (13, 1))) # (9,(4,8)):(59,(13,1))\n",
" \n",
" layout = cute.make_layout(\n",
" (9, (4, 8)), stride=(59, (13, 1))\n",
" ) # (9,(4,8)):(59,(13,1))\n",
"\n",
" # Define the tiler\n",
" tiler = (cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8))) # Apply to mode-1 layout (2,4):(1,8)\n",
" \n",
" tiler = (\n",
" cute.make_layout(3, stride=3), # Apply to mode-0 layout 3:3\n",
" cute.make_layout((2, 4), stride=(1, 8)),\n",
" ) # Apply to mode-1 layout (2,4):(1,8)\n",
"\n",
" # Apply flat divide\n",
" result = cute.flat_divide(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Flat Divide Result:\", result)\n",
" cute.printf(\">?? Flat Divide Result: {}\", result)\n",
"\n",
"\n",
"flat_divide_example()"
]
},
@@ -834,19 +860,20 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((2, 2), stride=(4, 1)) # (2,2):(4,1)\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout(6, stride=1) # Apply to layout 6:1\n",
" \n",
"\n",
" # Apply logical product\n",
" result = cute.logical_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
" print(\">>> Logical Product Result:\", result)\n",
" cute.printf(\">?? Logical Product Result: {}\", result)\n",
"\n",
"\n",
"logical_product_1d_example()"
]
},
@@ -886,16 +913,16 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((2, 5), stride=(5, 1))\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout((3, 4), stride=(1, 3))\n",
" \n",
"\n",
" # Apply blocked product\n",
" blocked_result = cute.blocked_product(layout, tiler=tiler)\n",
"\n",
" # Apply raked product\n",
" raked_result = cute.raked_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Print results\n",
" print(\">>> Layout:\", layout)\n",
" print(\">>> Tiler :\", tiler)\n",
@@ -904,6 +931,7 @@
" cute.printf(\">?? Blocked Product Result: {}\", blocked_result)\n",
" cute.printf(\">?? Raked Product Result: {}\", raked_result)\n",
"\n",
"\n",
"blocked_raked_product_example()"
]
},
@@ -950,16 +978,16 @@
" \"\"\"\n",
" # Define the original layout\n",
" layout = cute.make_layout((2, 5), stride=(5, 1))\n",
" \n",
"\n",
" # Define the tiler\n",
" tiler = cute.make_layout((3, 4), stride=(1, 3))\n",
"\n",
" # Apply zipped product\n",
" zipped_result = cute.zipped_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Apply tiled product\n",
" tiled_result = cute.tiled_product(layout, tiler=tiler)\n",
" \n",
"\n",
" # Apply flat product\n",
" flat_result = cute.flat_product(layout, tiler=tiler)\n",
"\n",
@@ -973,6 +1001,7 @@
" cute.printf(\">?? Tiled Product Result: {}\", tiled_result)\n",
" cute.printf(\">?? Flat Product Result: {}\", flat_result)\n",
"\n",
"\n",
"zipped_tiled_flat_product_example()"
]
}
examples/python/CuTeDSL/notebooks/data_types.ipynb
+8 -7
View File
@@ -6,8 +6,6 @@
"metadata": {},
"outputs": [],
"source": [
"from typing import List\n",
"\n",
"import cutlass\n",
"import cutlass.cute as cute"
]
@@ -83,12 +81,13 @@
"@cute.jit\n",
"def bar():\n",
" a = cutlass.Float32(3.14)\n",
" print(\"a(static) =\", a) # prints `a(static) = ?`\n",
" cute.printf(\"a(dynamic) = {}\", a) # prints `a(dynamic) = 3.140000`\n",
" print(\"a(static) =\", a) # prints `a(static) = ?`\n",
" cute.printf(\"a(dynamic) = {}\", a) # prints `a(dynamic) = 3.140000`\n",
"\n",
" b = cutlass.Int32(5)\n",
" print(\"b(static) =\", b) # prints `b(static) = 5`\n",
" cute.printf(\"b(dynamic) = {}\", b) # prints `b(dynamic) = 5`\n",
" print(\"b(static) =\", b) # prints `b(static) = 5`\n",
" cute.printf(\"b(dynamic) = {}\", b) # prints `b(dynamic) = 5`\n",
"\n",
"\n",
"bar()"
]
@@ -154,6 +153,7 @@
" f = e.to(cutlass.Int8)\n",
" cute.printf(\"Int32({}) => Int8({}) (truncated due to range limitation)\", e, f)\n",
"\n",
"\n",
"type_conversion()"
]
},
@@ -241,7 +241,8 @@
" not_a = ~a\n",
" cute.printf(\"~a = {}\", not_a)\n",
"\n",
"operator_demo()\n"
"\n",
"operator_demo()"
]
}
],
examples/python/CuTeDSL/notebooks/elementwise_add.ipynb
+752 -372
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/notebooks/hello_world.ipynb
+18 -9
View File
@@ -27,8 +27,8 @@
"metadata": {},
"outputs": [],
"source": [
"import cutlass \n",
"import cutlass.cute as cute "
"import cutlass\n",
"import cutlass.cute as cute"
]
},
{
@@ -80,14 +80,13 @@
"source": [
"@cute.jit\n",
"def hello_world():\n",
"\n",
" # Print hello world from host code\n",
" cute.printf(\"hello world\")\n",
"\n",
" # Launch kernel\n",
" kernel().launch(\n",
" grid=(1, 1, 1), # Single thread block\n",
" block=(32, 1, 1) # One warp (32 threads) per thread block\n",
" grid=(1, 1, 1), # Single thread block\n",
" block=(32, 1, 1), # One warp (32 threads) per thread block\n",
" )"
]
},
@@ -107,7 +106,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 5,
"metadata": {},
"outputs": [
{
@@ -115,17 +114,19 @@
"output_type": "stream",
"text": [
"Running hello_world()...\n",
"hello world\n",
"Compiling...\n",
"hello world\n",
"Hello world\n",
"Compiling with PTX/CUBIN dumped...\n",
"Running compiled version...\n",
"hello world\n"
"hello world\n",
"Hello world\n"
]
}
],
"source": [
"# Initialize CUDA context for launching a kernel with error checking\n",
"# We make context initialization explicit to allow users to control the context creation \n",
"# We make context initialization explicit to allow users to control the context creation\n",
"# and avoid potential issues with multiple contexts\n",
"cutlass.cuda.initialize_cuda_context()\n",
"\n",
@@ -137,6 +138,14 @@
"print(\"Compiling...\")\n",
"hello_world_compiled = cute.compile(hello_world)\n",
"\n",
"# Dump PTX/CUBIN files while compiling\n",
"from cutlass.cute import KeepPTX, KeepCUBIN\n",
"\n",
"print(\"Compiling with PTX/CUBIN dumped...\")\n",
"# Alternatively, compile with string based options like\n",
"# cute.compile(hello_world, options=\"--keep-ptx --keep-cubin\") would also work.\n",
"hello_world_compiled_ptx_on = cute.compile[KeepPTX, KeepCUBIN](hello_world)\n",
"\n",
"# Run the pre-compiled version\n",
"print(\"Running compiled version...\")\n",
"hello_world_compiled()"
examples/python/CuTeDSL/notebooks/print.ipynb
+28 -17
View File
@@ -83,8 +83,8 @@
" print(\">>>\", type(b)) # => <class 'int'>\n",
"\n",
" layout = cute.make_layout((a, b))\n",
" print(\">>>\", layout) # => (?,2):(1,?)\n",
" cute.printf(\">?? {}\", layout) # => (8,2):(1,8)"
" print(\">>>\", layout) # => (?,2):(1,?)\n",
" cute.printf(\">?? {}\", layout) # => (8,2):(1,8)"
]
},
{
@@ -221,6 +221,7 @@
" layout = cute.make_layout((a, b))\n",
" print(f\"layout: {layout}\")\n",
"\n",
"\n",
"print(\"Direct run output:\")\n",
"format_string_example(cutlass.Int32(8), 2)"
]
@@ -246,23 +247,26 @@
"source": [
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_basic(x : cute.Tensor):\n",
"def print_tensor_basic(x: cute.Tensor):\n",
" # Print the tensor\n",
" print(\"Basic output:\")\n",
" cute.print_tensor(x)\n",
" \n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_verbose(x : cute.Tensor):\n",
"def print_tensor_verbose(x: cute.Tensor):\n",
" # Print the tensor with verbose mode\n",
" print(\"Verbose output:\")\n",
" cute.print_tensor(x, verbose=True)\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_slice(x : cute.Tensor, coord : tuple):\n",
"def print_tensor_slice(x: cute.Tensor, coord: tuple):\n",
" # slice a 2D tensor from the 3D tensor\n",
" sliced_data = cute.slice_(x, coord)\n",
" y = cute.make_fragment(sliced_data.layout, sliced_data.element_type)\n",
" y = cute.make_rmem_tensor(sliced_data.layout, sliced_data.element_type)\n",
" # Convert to TensorSSA format by loading the sliced data into the fragment\n",
" y.store(sliced_data.load())\n",
" print(\"Slice output:\")\n",
@@ -302,12 +306,13 @@
"source": [
"def tensor_print_example1():\n",
" shape = (4, 3, 2)\n",
" \n",
"\n",
" # Creates [0,...,23] and reshape to (4, 3, 2)\n",
" data = np.arange(24, dtype=np.float32).reshape(*shape) \n",
" \n",
" data = np.arange(24, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_basic(from_dlpack(data))\n",
"\n",
"\n",
"tensor_print_example1()"
]
},
@@ -348,12 +353,13 @@
"source": [
"def tensor_print_example2():\n",
" shape = (4, 3)\n",
" \n",
"\n",
" # Creates [0,...,11] and reshape to (4, 3)\n",
" data = np.arange(12, dtype=np.float32).reshape(*shape) \n",
" \n",
" data = np.arange(12, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_verbose(from_dlpack(data))\n",
"\n",
"\n",
"tensor_print_example2()"
]
},
@@ -390,13 +396,14 @@
"source": [
"def tensor_print_example3():\n",
" shape = (4, 3)\n",
" \n",
"\n",
" # Creates [0,...,11] and reshape to (4, 3)\n",
" data = np.arange(12, dtype=np.float32).reshape(*shape) \n",
" \n",
" data = np.arange(12, dtype=np.float32).reshape(*shape)\n",
"\n",
" print_tensor_slice(from_dlpack(data), (None, 0))\n",
" print_tensor_slice(from_dlpack(data), (1, None))\n",
"\n",
"\n",
"tensor_print_example3()"
]
},
@@ -418,9 +425,10 @@
" print(src)\n",
" cute.print_tensor(src)\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_host(src: cute.Tensor):\n",
" print_tensor_gpu(src).launch(grid=(1,1,1), block=(1,1,1))"
" print_tensor_gpu(src).launch(grid=(1, 1, 1), block=(1, 1, 1))"
]
},
{
@@ -449,11 +457,14 @@
],
"source": [
"import torch\n",
"\n",
"\n",
"def tensor_print_example4():\n",
" a = torch.randn(4, 3, device=\"cuda\")\n",
" cutlass.cuda.initialize_cuda_context()\n",
" print_tensor_host(from_dlpack(a))\n",
"\n",
"\n",
"tensor_print_example4()"
]
},
examples/python/CuTeDSL/notebooks/tensor.ipynb
+63 -113
View File
@@ -2,7 +2,7 @@
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -43,7 +43,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -69,24 +69,9 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor(raw_ptr(0x000000000736b0c0: f32, generic, align<4>) o (8,5):(5,1), data=\n",
" [[ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" ...\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ],\n",
" [ 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, ]])\n"
]
}
],
"outputs": [],
"source": [
"import torch\n",
"\n",
@@ -115,12 +100,13 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"\n",
"@cute.jit\n",
"def print_tensor_dlpack(src: cute.Tensor):\n",
" print(src)\n",
@@ -129,25 +115,9 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor<ptr<f32, generic> o (8,5):(5,1)>\n",
"tensor(raw_ptr(0x0000000007559340: f32, generic, align<4>) o (8,5):(5,1), data=\n",
" [[-1.151769, 1.019397, -0.371175, -0.717776, 0.502176, ],\n",
" [ 0.114282, 0.900084, 0.320770, 1.564574, -0.632329, ],\n",
" [-0.570140, 0.178112, -0.423079, 1.936198, 0.003355, ],\n",
" ...\n",
" [-2.425393, -0.275528, 1.267157, -0.811101, -0.985456, ],\n",
" [ 0.777889, -2.114074, 0.357184, -0.321312, -0.938138, ],\n",
" [ 1.959564, 1.797602, 0.116901, 0.306198, -1.837295, ]])\n"
]
}
],
"outputs": [],
"source": [
"a = torch.randn(8, 5, dtype=torch_dtype(cutlass.Float32))\n",
"\n",
@@ -156,25 +126,9 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor<ptr<f32, generic> o (8,8):(8,1)>\n",
"tensor(raw_ptr(0x0000000007979da0: f32, generic, align<4>) o (8,8):(8,1), data=\n",
" [[ 0.122739, -0.605744, -1.442022, ..., -0.356501, -0.993329, -0.091110, ],\n",
" [ 0.278448, 0.318482, -0.276867, ..., 1.542181, -1.701539, -0.309454, ],\n",
" [ 0.563565, -0.753936, 0.131214, ..., 0.437912, -0.482277, -0.051540, ],\n",
" ...\n",
" [-1.974096, -0.177881, 0.426807, ..., -1.579115, -0.304974, 0.451164, ],\n",
" [ 0.149851, -0.704689, -0.295063, ..., -0.653001, 0.008871, 0.903916, ],\n",
" [ 1.188619, 1.519662, 1.270734, ..., 0.404082, 0.173200, 0.093476, ]])\n"
]
}
],
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
@@ -211,39 +165,23 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"a[2] = 10.000000 (equivalent to a[(2,0)])\n",
"a[9] = 6.000000 (equivalent to a[(1,1)])\n",
"a[2,0] = 10.000000\n",
"a[2,4] = 14.000000\n",
"a[(2,4)] = 14.000000\n",
"a[2,3] = 100.000000\n",
"a[(2,4)] = 101.000000\n",
"tensor([[ 0., 1., 2., 3., 4.],\n",
" [ 5., 6., 7., 8., 9.],\n",
" [ 10., 11., 12., 100., 101.],\n",
" [ 15., 16., 17., 18., 19.],\n",
" [ 20., 21., 22., 23., 24.],\n",
" [ 25., 26., 27., 28., 29.],\n",
" [ 30., 31., 32., 33., 34.],\n",
" [ 35., 36., 37., 38., 39.]])\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def tensor_access_item(a: cute.Tensor):\n",
" # access data using linear index\n",
" cute.printf(\"a[2] = {} (equivalent to a[{}])\", a[2],\n",
" cute.make_identity_tensor(a.layout.shape)[2])\n",
" cute.printf(\"a[9] = {} (equivalent to a[{}])\", a[9],\n",
" cute.make_identity_tensor(a.layout.shape)[9])\n",
" cute.printf(\n",
" \"a[2] = {} (equivalent to a[{}])\",\n",
" a[2],\n",
" cute.make_identity_tensor(a.layout.shape)[2],\n",
" )\n",
" cute.printf(\n",
" \"a[9] = {} (equivalent to a[{}])\",\n",
" a[9],\n",
" cute.make_identity_tensor(a.layout.shape)[9],\n",
" )\n",
"\n",
" # access data using n-d coordinates, following two are equivalent\n",
" cute.printf(\"a[2,0] = {}\", a[2, 0])\n",
@@ -251,14 +189,14 @@
" cute.printf(\"a[(2,4)] = {}\", a[2, 4])\n",
"\n",
" # assign value to tensor@(2,4)\n",
" a[2,3] = 100.0\n",
" a[2,4] = 101.0\n",
" cute.printf(\"a[2,3] = {}\", a[2,3])\n",
" cute.printf(\"a[(2,4)] = {}\", a[(2,4)])\n",
" a[2, 3] = 100.0\n",
" a[2, 4] = 101.0\n",
" cute.printf(\"a[2,3] = {}\", a[2, 3])\n",
" cute.printf(\"a[(2,4)] = {}\", a[(2, 4)])\n",
"\n",
"\n",
"# Create a tensor with sequential data using torch\n",
"data = torch.arange(0, 8*5, dtype=torch.float32).reshape(8, 5)\n",
"data = torch.arange(0, 8 * 5, dtype=torch.float32).reshape(8, 5)\n",
"tensor_access_item(from_dlpack(data))\n",
"\n",
"print(data)"
@@ -287,14 +225,17 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### Coordinate Tensor\n",
"## Coordinate Tensors\n",
"\n",
"A coordinate tensor is a special type of tensor that maps coordinates to coordinates rather than to values. \n",
"The key distinction is that while regular tensors map coordinates to some value type (like numbers), \n",
"coordinate tensors map coordinates to other coordinates.\n",
"### Definition and Properties\n",
"\n",
"For example, given a shape (4,4), a coordinate tensor using row-major layout would appear as:\n",
"A coordinate tensor $T: Z^n → Z^m$ is a mathematical structure that establishes a mapping between coordinate spaces. Unlike standard tensors that map coordinates to scalar values, coordinate tensors map coordinates to other coordinates, forming a fundamental building block for tensor operations and transformations.\n",
"\n",
"### Examples\n",
"\n",
"Consider a `(4,4)` coordinate tensor:\n",
"\n",
"**Row-Major Layout (C-style):**\n",
"\\begin{bmatrix} \n",
"(0,0) & (0,1) & (0,2) & (0,3) \\\\\n",
"(1,0) & (1,1) & (1,2) & (1,3) \\\\\n",
@@ -302,8 +243,7 @@
"(3,0) & (3,1) & (3,2) & (3,3)\n",
"\\end{bmatrix}\n",
"\n",
"The same shape with a column-major layout would appear as:\n",
"\n",
"**Column-Major Layout (Fortran-style):**\n",
"\\begin{bmatrix}\n",
"(0,0) & (1,0) & (2,0) & (3,0) \\\\\n",
"(0,1) & (1,1) & (2,1) & (3,1) \\\\\n",
@@ -311,40 +251,50 @@
"(0,3) & (1,3) & (2,3) & (3,3)\n",
"\\end{bmatrix}\n",
"\n",
"The key points about coordinate tensors are:\n",
"- Each element in the tensor is itself a coordinate tuple (i,j) rather than a scalar value\n",
"- The coordinates map to themselves - so position (1,2) contains the coordinate (1,2)\n",
"- The layout (row-major vs column-major) determines how these coordinate tuples are arranged in memory\n",
"### Identity Tensor\n",
"\n",
"For example, coordinate tensors can be created using the `make_identity_tensor` utility:\n",
"An identity tensor $I$ is a special case of a coordinate tensor that implements the identity mapping function:\n",
"\n",
"**Definition:**\n",
"For a given shape $S = (s_1, s_2, ..., s_n)$, the identity tensor $I$ satisfies: $I(c) = c, \\forall c \\in \\prod_{i=1}^n [0, s_i)$\n",
"\n",
"**Properties:**\n",
"1. **Bijective Mapping**: The identity tensor establishes a one-to-one correspondence between coordinates.\n",
"2. **Layout Invariance**: The logical structure remains constant regardless of the underlying memory layout.\n",
"3. **Coordinate Preservation**: For any coordinate c, I(c) = c.\n",
"\n",
"\n",
"CuTe establishes an isomorphism between 1-D indices and N-D coordinates through lexicographical ordering. For a coordinate c = (c₁, c₂, ..., cₙ) in an identity tensor with shape S = (s₁, s₂, ..., sₙ):\n",
"\n",
"**Linear Index Formula:**\n",
"$\\text{idx} = c_1 + \\sum_{i=2}^{n} \\left(c_i \\prod_{j=1}^{i-1} s_j\\right)$\n",
"\n",
"**Example:**\n",
"```python\n",
"# Create an identity tensor from a given shape\n",
"coord_tensor = make_identity_tensor(layout.shape())\n",
"\n",
"# Access coordinate using linear index\n",
"coord = coord_tensor[linear_idx] # Returns the N-D coordinate\n",
"```\n",
"\n",
"This creates a tensor that maps each coordinate to itself, providing a reference point for understanding how other layouts transform these coordinates."
"This bidirectional mapping enables efficient conversion from linear indices to N-dimensional coordinates, facilitating tensor operations and memory access patterns."
]
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor<(0,0) o (8,4):(1@0,1@1)>\n"
]
}
],
"outputs": [],
"source": [
"@cute.jit\n",
"def print_tensor_coord(a: cute.Tensor):\n",
" coord_tensor = cute.make_identity_tensor(a.layout.shape)\n",
" print(coord_tensor)\n",
" cute.print_tensor(coord_tensor)\n",
"\n",
"a = torch.randn(8,4, dtype=torch_dtype(cutlass.Float32))\n",
"\n",
"a = torch.randn(8, 4, dtype=torch_dtype(cutlass.Float32))\n",
"print_tensor_coord(from_dlpack(a))"
]
}
examples/python/CuTeDSL/notebooks/tensorssa.ipynb
+37 -42
View File
@@ -10,8 +10,7 @@
"import cutlass.cute as cute\n",
"from cutlass.cute.runtime import from_dlpack\n",
"\n",
"import numpy as np\n",
"import torch"
"import numpy as np"
]
},
{
@@ -55,12 +54,13 @@
" :param b: The source tensor to be loaded.\n",
" \"\"\"\n",
" a_vec = a.load()\n",
" print(f\"a_vec: {a_vec}\") # prints `a_vec: vector<12xf32> o (3, 4)`\n",
" print(f\"a_vec: {a_vec}\") # prints `a_vec: vector<12xf32> o (3, 4)`\n",
" b_vec = b.load()\n",
" print(f\"b_vec: {b_vec}\") # prints `b_vec: vector<12xf32> o (3, 4)`\n",
" print(f\"b_vec: {b_vec}\") # prints `b_vec: vector<12xf32> o (3, 4)`\n",
" res.store(a_vec + b_vec)\n",
" cute.print_tensor(res)\n",
"\n",
"\n",
"a = np.ones(12).reshape((3, 4)).astype(np.float32)\n",
"b = np.ones(12).reshape((3, 4)).astype(np.float32)\n",
"c = np.zeros(12).reshape((3, 4)).astype(np.float32)\n",
@@ -101,6 +101,7 @@
" dst[0] = dst_vec\n",
" cute.print_tensor(dst)\n",
"\n",
"\n",
"def slice_1():\n",
" src_shape = (4, 2, 3)\n",
" dst_shape = (4, 3)\n",
@@ -124,6 +125,7 @@
" dst = np.random.randn(*dst_shape).astype(np.float32)\n",
" apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
"\n",
"\n",
"slice_1()"
]
},
@@ -141,6 +143,7 @@
" dst = np.random.randn(*dst_shape).astype(np.float32)\n",
" apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
"\n",
"\n",
"slice_2()"
]
},
@@ -169,22 +172,22 @@
" b_vec = b.load()\n",
"\n",
" add_res = a_vec + b_vec\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
"\n",
" sub_res = a_vec - b_vec\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
"\n",
" mul_res = a_vec * b_vec\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" div_res = a_vec / b_vec\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n",
" floor_div_res = a_vec // b_vec\n",
" cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n",
" cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n",
" mod_res = a_vec % b_vec\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
"\n",
"\n",
"a = np.empty((3,), dtype=np.float32)\n",
@@ -206,22 +209,23 @@
" a_vec = a.load()\n",
"\n",
" add_res = a_vec + c\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
" cute.print_tensor(add_res) # prints [3.000000, 3.000000, 3.000000]\n",
"\n",
" sub_res = a_vec - c\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
" cute.print_tensor(sub_res) # prints [-1.000000, -1.000000, -1.000000]\n",
"\n",
" mul_res = a_vec * c\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
" cute.print_tensor(mul_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" div_res = a_vec / c\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n",
" floor_div_res = a_vec // c\n",
" cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n",
" cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n",
" mod_res = a_vec % c\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
"\n",
"\n",
"a = np.empty((3,), dtype=np.float32)\n",
"a.fill(1.0)\n",
@@ -251,11 +255,12 @@
" eq_res = a_ == b_ # [False, False, False]\n",
" \"\"\"\n",
"\n",
"\n",
"a = np.array([1, 2, 3], dtype=np.float32)\n",
"b = np.array([2, 1, 4], dtype=np.float32)\n",
"res = np.empty((3,), dtype=np.bool_)\n",
"binary_op_3(from_dlpack(res), from_dlpack(a), from_dlpack(b))\n",
"print(res) # prints [False, True, False]\n"
"print(res) # prints [False, True, False]"
]
},
{
@@ -278,11 +283,12 @@
" # and_res = a_vec & b_vec\n",
" # res.store(and_res) # prints [0, 2, 0]\n",
"\n",
"\n",
"a = np.array([1, 2, 3], dtype=np.int32)\n",
"b = np.array([2, 2, 4], dtype=np.int32)\n",
"res = np.empty((3,), dtype=np.int32)\n",
"binary_op_4(from_dlpack(res), from_dlpack(a), from_dlpack(b))\n",
"print(res) # prints [3, 0, 7]"
"print(res) # prints [3, 0, 7]"
]
},
{
@@ -303,14 +309,15 @@
" a_vec = a.load()\n",
"\n",
" sqrt_res = cute.math.sqrt(a_vec)\n",
" cute.print_tensor(sqrt_res) # prints [2.000000, 2.000000, 2.000000]\n",
" cute.print_tensor(sqrt_res) # prints [2.000000, 2.000000, 2.000000]\n",
"\n",
" sin_res = cute.math.sin(a_vec)\n",
" res.store(sin_res)\n",
" cute.print_tensor(sin_res) # prints [-0.756802, -0.756802, -0.756802]\n",
" cute.print_tensor(sin_res) # prints [-0.756802, -0.756802, -0.756802]\n",
"\n",
" exp2_res = cute.math.exp2(a_vec)\n",
" cute.print_tensor(exp2_res) # prints [16.000000, 16.000000, 16.000000]\n",
" cute.print_tensor(exp2_res) # prints [16.000000, 16.000000, 16.000000]\n",
"\n",
"\n",
"a = np.array([4.0, 4.0, 4.0], dtype=np.float32)\n",
"res = np.empty((3,), dtype=np.float32)\n",
@@ -344,26 +351,14 @@
" :param src: The source tensor to be reduced.\n",
" \"\"\"\n",
" a_vec = a.load()\n",
" red_res = a_vec.reduce(\n",
" cute.ReductionOp.ADD,\n",
" 0.0,\n",
" reduction_profile=0\n",
" )\n",
" cute.printf(red_res) # prints 21.000000\n",
" red_res = a_vec.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=0)\n",
" cute.printf(red_res) # prints 21.000000\n",
"\n",
" red_res = a_vec.reduce(\n",
" cute.ReductionOp.ADD,\n",
" 0.0,\n",
" reduction_profile=(None, 1)\n",
" )\n",
" cute.print_tensor(red_res) # prints [6.000000, 15.000000]\n",
" red_res = a_vec.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=(None, 1))\n",
" cute.print_tensor(red_res) # prints [6.000000, 15.000000]\n",
"\n",
" red_res = a_vec.reduce(\n",
" cute.ReductionOp.ADD,\n",
" 1.0,\n",
" reduction_profile=(1, None)\n",
" )\n",
" cute.print_tensor(red_res) # prints [6.000000, 8.000000, 10.000000]\n",
" red_res = a_vec.reduce(cute.ReductionOp.ADD, 1.0, reduction_profile=(1, None))\n",
" cute.print_tensor(red_res) # prints [6.000000, 8.000000, 10.000000]\n",
"\n",
"\n",
"a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)\n",
@@ -399,7 +394,7 @@
"\n",
"@cute.jit\n",
"def broadcast_examples():\n",
" a = cute.make_fragment((1,3), dtype=cutlass.Float32)\n",
" a = cute.make_rmem_tensor((1, 3), dtype=cutlass.Float32)\n",
" a[0] = 0.0\n",
" a[1] = 1.0\n",
" a[2] = 2.0\n",
@@ -411,7 +406,7 @@
" # [ 0.000000, 1.000000, 2.000000, ],\n",
" # [ 0.000000, 1.000000, 2.000000, ]])\n",
"\n",
" c = cute.make_fragment((4,1), dtype=cutlass.Float32)\n",
" c = cute.make_rmem_tensor((4, 1), dtype=cutlass.Float32)\n",
" c[0] = 0.0\n",
" c[1] = 1.0\n",
" c[2] = 2.0\n",
@@ -494,7 +489,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.10"
"version": "3.12.11"
}
},
"nbformat": 4,
examples/python/CuTeDSL/utils/__init__.py
View File
examples/python/CuTeDSL/utils/fmha_helpers.py
+975
View File
@@ -0,0 +1,975 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
import enum
from typing import Tuple, Optional
import cutlass
from cutlass.cute.typing import Boolean
from cutlass.cutlass_dsl import (
Int32,
Float32,
min,
extract_mlir_values,
new_from_mlir_values,
)
from cutlass.utils.hardware_info import HardwareInfo
from cutlass.utils import WorkTileInfo
import cutlass.cute as cute
##############################################################################
# Fmha static tile scheduler
##############################################################################
class FmhaStaticTileSchedulerParams:
"""A class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler.
This class holds the configuration parameters needed to initialize and configure
the tile scheduler for FMHA operations.
:ivar is_persistent: Whether to use persistent kernel mode.
:type is_persistent: bool
:ivar problem_shape_mbh: Problem shape in (M, B, H) format.
:type problem_shape_mbh: cute.Shape
"""
def __init__(
self,
is_persistent: bool,
problem_shape_mbh: cute.Shape,
*,
loc=None,
ip=None,
):
"""
Initializes the FmhaStaticTileSchedulerParams with the given parameters.
:param is_persistent: Whether to use persistent kernel mode.
:type is_persistent: bool
:param problem_shape_mbh: Problem shape in (M, B, H) format.
:type problem_shape_mbh: cute.Shape
"""
self.is_persistent = is_persistent
self.problem_shape_mbh = problem_shape_mbh
self._loc = loc
self._ip = ip
def __extract_mlir_values__(self):
values, self._values_pos = [], []
for obj in [self.problem_shape_mbh]:
obj_values = extract_mlir_values(obj)
values += obj_values
self._values_pos.append(len(obj_values))
return values
def __new_from_mlir_values__(self, values):
obj_list = []
for obj, n_items in zip([self.problem_shape_mbh], self._values_pos):
obj_list.append(new_from_mlir_values(obj, values[:n_items]))
values = values[n_items:]
return FmhaStaticTileSchedulerParams(
self.is_persistent, *(tuple(obj_list)), loc=self._loc
)
class FmhaStaticTileScheduler:
"""A static tile scheduler for FMHA (Fused Multi-Head Attention) operations.
This class manages the scheduling of work tiles for FMHA kernels, supporting
both persistent and non-persistent kernel modes. It tracks the current work
position and advances through the problem space efficiently.
:ivar _params: Scheduler parameters.
:type _params: FmhaStaticTileSchedulerParams
:ivar _blk_coord: Block coordinates.
:type _blk_coord: cute.Coord
:ivar _grid_shape: Grid shape for the kernel.
:type _grid_shape: cute.Shape
:ivar _is_persistent: Whether to use persistent kernel mode.
:type _is_persistent: bool
:ivar _current_work_linear_idx: Current linear work index.
:type _current_work_linear_idx: Int32
:ivar _problem_shape_mbh: Problem shape in (M, B, H) format.
:type _problem_shape_mbh: cute.Layout
:ivar _num_blocks: Number of blocks in the problem.
:type _num_blocks: Int32
:ivar _is_first_block: Whether this is the first block.
:type _is_first_block: bool
:ivar num_persistent_sm: Number of persistent SMs.
:type num_persistent_sm: Int32
"""
def __init__(
self,
params: FmhaStaticTileSchedulerParams,
current_work_linear_idx: Int32,
blk_coord: cute.Coord,
grid_shape: cute.Shape,
*,
loc=None,
ip=None,
):
"""
Initializes the FmhaStaticTileScheduler with the given parameters.
:param params: Scheduler parameters.
:type params: FmhaStaticTileSchedulerParams
:param current_work_linear_idx: Current linear work index.
:type current_work_linear_idx: Int32
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param grid_shape: Grid shape for the kernel.
:type grid_shape: cute.Shape
"""
self._params = params
self._blk_coord = blk_coord
self._grid_shape = grid_shape
self._is_persistent = params.is_persistent
self._current_work_linear_idx = current_work_linear_idx
self._problem_shape_mbh = cute.make_layout(
params.problem_shape_mbh, loc=loc, ip=ip
)
self._num_blocks = cute.size(self._problem_shape_mbh, loc=loc, ip=ip)
self._is_first_block = True
self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip)
self._loc = loc
self._ip = ip
# called by host
@staticmethod
def get_grid_shape(
params: FmhaStaticTileSchedulerParams,
*,
loc=None,
ip=None,
) -> cute.Shape:
"""
Determine the grid shape for the FMHA kernel.
For persistent kernels, the grid shape is limited by the number of SMs
(Streaming Multiprocessors) available on the device. For non-persistent
kernels, the grid shape matches the problem shape.
:param params: Scheduler parameters.
:type params: FmhaStaticTileSchedulerParams
:return: Grid shape as (M, B, H) tuple.
:rtype: cute.Shape
"""
if params.is_persistent:
hardware_info = HardwareInfo()
sm_count = hardware_info.get_device_multiprocessor_count()
return (
min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)),
1,
1,
)
else:
return params.problem_shape_mbh
@staticmethod
def check_valid_work_for_seqlen_q(
q_tiler: int,
current_idx: Int32,
seqlen_q: Int32,
) -> Boolean:
"""
Check if the current work index is valid for the given query sequence length.
This method verifies that the current work tile index multiplied by the
query tiler size is within the bounds of the query sequence length.
:param q_tiler: Query tiler size.
:type q_tiler: int
:param current_idx: Current work index.
:type current_idx: Int32
:param seqlen_q: Query sequence length.
:type seqlen_q: Int32
:return: True if the work is valid, False otherwise.
:rtype: Boolean
"""
return current_idx * q_tiler < seqlen_q
def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
"""
Get information about the current work tile.
Determines if the current work is valid and computes the tile coordinates
based on whether the kernel is persistent or non-persistent.
:return: WorkTileInfo containing tile coordinates and validity flag.
:rtype: WorkTileInfo
"""
is_valid = (
self._current_work_linear_idx < self._num_blocks
if self._is_persistent
else self._is_first_block
)
blk_coord = (0, 0, 0)
if self._is_persistent:
blk_coord = self._problem_shape_mbh.get_hier_coord(
self._current_work_linear_idx, loc=loc, ip=ip
)
else:
blk_coord = self._blk_coord
# cur_tile_coord is (mid, 0, (bid, hid))
cur_tile_coord = (
blk_coord[0],
0,
(blk_coord[1], blk_coord[2]),
)
return WorkTileInfo(cur_tile_coord, is_valid)
def initial_work_tile_info(self, *, loc=None, ip=None):
"""
Get the initial work tile information.
:return: Initial WorkTileInfo.
:rtype: WorkTileInfo
"""
return self.get_current_work(loc=loc, ip=ip)
def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None):
"""
Advance to the next work tile.
For persistent kernels, advances by the number of persistent SMs.
For non-persistent kernels, marks that the first block has been processed.
:param advance_count: Number of steps to advance (default: 1).
:type advance_count: int
"""
if self._is_persistent:
self._current_work_linear_idx += advance_count * self.num_persistent_sm
self._is_first_block = False
def __extract_mlir_values__(self):
values = extract_mlir_values(self._params)
values.extend(extract_mlir_values(self._current_work_linear_idx))
values.extend(extract_mlir_values(self._blk_coord))
values.extend(extract_mlir_values(self._grid_shape))
return values
def __new_from_mlir_values__(self, values):
assert len(values) == 10
new_params = new_from_mlir_values(self._params, values[0:3])
new_current_work_linear_idx = new_from_mlir_values(
self._current_work_linear_idx, [values[3]]
)
new_blk_coord = new_from_mlir_values(self._blk_coord, values[4:7])
new_grid_shape = new_from_mlir_values(self._grid_shape, values[7:])
return FmhaStaticTileScheduler(
new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape
)
def create_fmha_static_tile_scheduler(
params: FmhaStaticTileSchedulerParams,
blk_coord: cute.Coord,
grid_shape: cute.Shape,
) -> FmhaStaticTileScheduler:
"""
Create a new FMHA static tile scheduler.
:param params: Scheduler parameters.
:type params: FmhaStaticTileSchedulerParams
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param grid_shape: Grid shape.
:type grid_shape: cute.Shape
:return: New FmhaStaticTileScheduler instance.
:rtype: FmhaStaticTileScheduler
"""
return FmhaStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape)
def create_fmha_static_tile_scheduler_params(
is_persistent: bool,
problem_shape_mbh: cute.Shape,
) -> FmhaStaticTileSchedulerParams:
"""
Create FMHA static tile scheduler parameters.
:param is_persistent: Whether to use persistent kernel mode.
:type is_persistent: bool
:param problem_shape_mbh: Problem shape in (M, B, H) format.
:type problem_shape_mbh: cute.Shape
:return: New FmhaStaticTileSchedulerParams instance.
:rtype: FmhaStaticTileSchedulerParams
"""
return FmhaStaticTileSchedulerParams(is_persistent, problem_shape_mbh)
def compute_grid(
o_shape: cute.Shape,
cta_tiler: Tuple[int, int, int],
is_persistent: bool,
) -> Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]:
"""
Compute grid parameters for FMHA operation.
This function calculates the appropriate grid shape and scheduler parameters
based on the output tensor shape, CTA (Cooperative Thread Array) tiler,
and whether to use persistent kernel mode.
The output tensor o has shape (s, d, ((h_r, h_k), b)) where:
- s: sequence length
- d: head dimension
- h_r: number of heads for query
- h_k: number of heads for key
- b: batch size
:param o_shape: Output tensor shape for grid computation.
:type o_shape: cute.Shape
:param cta_tiler: CTA tiler dimensions (M, N, K).
:type cta_tiler: Tuple[int, int, int]
:param is_persistent: Whether to use persistent kernel mode.
:type is_persistent: bool
:return: Tuple of (scheduler_params, grid_shape).
:rtype: Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]
"""
tile_sched_params = create_fmha_static_tile_scheduler_params(
is_persistent,
(
cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]),
cute.size(o_shape[2][0]),
cute.size(o_shape[2][1]),
),
)
grid = FmhaStaticTileScheduler.get_grid_shape(tile_sched_params)
return tile_sched_params, grid
##############################################################################
# Fused Mask
##############################################################################
class MaskEnum(enum.Enum):
"""Enumeration of mask types for FMHA operations.
- RESIDUAL_MASK: Residual mask for handling variable sequence lengths
- WINDOW_MASK: Window mask for attention which also includes causal and no mask
- WINDOW_MASK_INFERENCE: Same as the window mask, but has the limitation that the end of q is aligned with the end of k
- WINDOW_MASK_BWD: Window mask for backward pass
- WINDOW_MASK_BWD_INFERENCE: Same as the window mask for backward pass, but has the limitation that the end of q is aligned with the end of k
"""
RESIDUAL_MASK = enum.auto()
RESIDUAL_MASK_BWD = enum.auto()
WINDOW_MASK = enum.auto()
WINDOW_MASK_INFERENCE = enum.auto()
WINDOW_MASK_BWD = enum.auto()
WINDOW_MASK_BWD_INFERENCE = enum.auto()
class FusedMask:
"""A fused mask implementation for FMHA operations.
This class handles different types of attention masks including no mask,
residual mask for variable sequence lengths, and causal mask for
autoregressive attention patterns.
The class provides methods to:
- Calculate trip counts for different mask types
- Apply masks to attention scores
- Handle masked and unmasked trip calculations
"""
def get_trip_count(
mask_type: MaskEnum,
blk_coord: cute.Coord,
tile_shape: cute.Shape,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[Int32] = None,
window_size_right: Optional[Int32] = None,
) -> Int32:
"""
Calculate the number of trips needed for the current block.
The trip count depends on the mask type and the block coordinates.
For causal masks, it considers the autoregressive constraint.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param tile_shape: Shape of the tile.
:type tile_shape: cute.Shape
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Int32
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[Int32]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[Int32]
:return: Number of trips needed.
:rtype: Int32
"""
result = 0
offset = 0
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
offset = seqlen_k - seqlen_q
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
offset = seqlen_q - seqlen_k
if cutlass.const_expr(mask_type == MaskEnum.RESIDUAL_MASK):
result = cute.ceil_div(seqlen_k, tile_shape[1])
if cutlass.const_expr(mask_type is MaskEnum.RESIDUAL_MASK_BWD):
result = cute.ceil_div(seqlen_q, tile_shape[0])
if cutlass.const_expr(
mask_type == MaskEnum.WINDOW_MASK
or mask_type == MaskEnum.WINDOW_MASK_INFERENCE
):
if cutlass.const_expr(window_size_right is None):
result = cute.ceil_div(seqlen_k, tile_shape[1])
else:
max_idx_q = (blk_coord[0] + 1) * tile_shape[0]
idx_k = max_idx_q + offset + window_size_right
tmp_blocks_k = cute.ceil_div(idx_k, tile_shape[1])
max_blocks_k = cute.ceil_div(seqlen_k, tile_shape[1])
result = min(max_blocks_k, tmp_blocks_k)
if cutlass.const_expr(
mask_type == MaskEnum.WINDOW_MASK_BWD
or mask_type == MaskEnum.WINDOW_MASK_BWD_INFERENCE
):
if cutlass.const_expr(window_size_left is None):
result = cute.ceil_div(seqlen_q, tile_shape[0])
else:
max_idx_k = (blk_coord[1] + 1) * tile_shape[1]
idx_k = max_idx_k + offset + window_size_left
tmp_blocks_q = cute.ceil_div(idx_k, tile_shape[0])
max_blocks_q = cute.ceil_div(seqlen_q, tile_shape[0])
result = min(max_blocks_q, tmp_blocks_q)
start_block = FusedMask.get_trip_start(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
result = result - start_block
return result
@cute.jit
def get_trip_start(
mask_type: MaskEnum,
blk_coord: cute.Coord,
tile_shape: cute.Shape,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[Int32] = None,
window_size_right: Optional[Int32] = None,
) -> Int32:
"""
Get the start of the trip for the current block.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param tile_shape: Shape of the tile.
:type tile_shape: cute.Shape
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Int32
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[Int32]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[Int32]
"""
result = 0
offset = 0
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
offset = seqlen_k - seqlen_q
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
offset = seqlen_q - seqlen_k
if cutlass.const_expr(
mask_type is MaskEnum.WINDOW_MASK
or mask_type is MaskEnum.WINDOW_MASK_INFERENCE
):
if cutlass.const_expr(window_size_left is not None):
min_idx_q = blk_coord[0] * tile_shape[0]
idx_k = min_idx_q + offset - window_size_left
tmp_blocks_k = idx_k // tile_shape[1]
result = max(tmp_blocks_k, result)
if cutlass.const_expr(
mask_type is MaskEnum.WINDOW_MASK_BWD
or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE
):
if cutlass.const_expr(window_size_right is not None):
min_idx_k = blk_coord[1] * tile_shape[1]
idx_q = min_idx_k + offset - window_size_right
tmp_blocks_q = idx_q // tile_shape[0]
result = max(tmp_blocks_q, result)
return result
@cute.jit
def get_leading_mask_id(
mask_type: MaskEnum,
blk_coord: cute.Coord,
tile_shape: cute.Shape,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[Int32] = None,
window_size_right: Optional[Int32] = None,
) -> Tuple[Int32, Int32]:
"""
Get the begin and end tile idx for the leading mask.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param tile_shape: Shape of the tile.
:type tile_shape: cute.Shape
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Int32
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[Int32]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[Int32]
:return: Tuple of (begin, end) tile idx for the leading mask.
:rtype: Tuple[Int32, Int32]
"""
offset = 0
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
offset = seqlen_k - seqlen_q
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
offset = seqlen_q - seqlen_k
leading_mask_begin = FusedMask.get_trip_start(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
trip_count = FusedMask.get_trip_count(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
leading_mask_end = leading_mask_begin
if cutlass.const_expr(
mask_type is MaskEnum.WINDOW_MASK
or mask_type is MaskEnum.WINDOW_MASK_INFERENCE
):
if cutlass.const_expr(window_size_left is not None):
min_idx_q = (
(blk_coord[0] + 1) * tile_shape[0] + offset - window_size_left
)
leading_mask_end = min(
cute.ceil_div(min_idx_q, tile_shape[1]) - 1,
trip_count + leading_mask_begin - 1,
)
else:
leading_mask_end = leading_mask_begin - 1
elif cutlass.const_expr(
mask_type is MaskEnum.WINDOW_MASK_BWD
or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE
):
if cutlass.const_expr(window_size_right is not None):
min_idx_k = (
(blk_coord[1] + 1) * tile_shape[1] + offset - window_size_right
)
leading_mask_end = cute.ceil_div(min_idx_k, tile_shape[0]) - 1
else:
leading_mask_end = leading_mask_begin - 1
return leading_mask_begin, leading_mask_end
@cute.jit
def get_trailing_mask_id(
mask_type: MaskEnum,
blk_coord: cute.Coord,
tile_shape: cute.Shape,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[Int32] = None,
window_size_right: Optional[Int32] = None,
) -> Tuple[Optional[Int32], Optional[Int32]]:
"""
Get the begin and end tile idx for the trailing mask.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param tile_shape: Shape of the tile.
:type tile_shape: cute.Shape
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Int32
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[Int32]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[Int32]
:return: Tuple of (begin, end) tile idx for the trailing mask.
:rtype: Tuple[Int32, Int32]
"""
offset = 0
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
offset = seqlen_k - seqlen_q
if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
offset = seqlen_q - seqlen_k
trip_start = FusedMask.get_trip_start(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
trip_count = FusedMask.get_trip_count(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
trailing_mask_begin, trailing_mask_end = None, None
if cutlass.const_expr(
mask_type is MaskEnum.WINDOW_MASK
or mask_type is MaskEnum.WINDOW_MASK_INFERENCE
):
if cutlass.const_expr(window_size_right is not None):
min_idx_q = blk_coord[0] * tile_shape[0] + offset + window_size_right
trailing_mask_begin = min(
min_idx_q // tile_shape[1], trip_count + trip_start - 1
)
trailing_mask_end = trip_count + trip_start - 1
else:
# last tile, we always apply mask on it regardless whether it's a residual tile
trailing_mask_begin = trip_count + trip_start - 1
trailing_mask_end = trip_count + trip_start - 1
else:
if cutlass.const_expr(window_size_left is not None):
min_idx_k = blk_coord[1] * tile_shape[1] + offset + window_size_left + 1
max_idx_k = (
(blk_coord[1] + 1) * tile_shape[1] + offset + window_size_left
)
trailing_mask_begin = min(
cute.ceil_div(min_idx_k, tile_shape[0]) - 1,
trip_count + trip_start - 1,
)
trailing_mask_end = min(
cute.ceil_div(max_idx_k, tile_shape[0]) - 1,
trip_count + trip_start - 1,
)
else:
# last tile, we always apply mask on it regardless whether it's a residual tile
trailing_mask_begin = trip_count + trip_start - 1
trailing_mask_end = trip_count + trip_start - 1
return trailing_mask_begin, trailing_mask_end
@cute.jit
def get_masked_leading_count(
mask_type: MaskEnum,
blk_coord: cute.Coord,
tile_shape: cute.Shape,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[Int32] = None,
window_size_right: Optional[Int32] = None,
) -> Int32:
"""
Calculate the number of masked trips for the leading mask.
This is used for blocks that need special handling due to masking.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param tile_shape: Shape of the tile.
:type tile_shape: cute.Shape
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Int32
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[Int32]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[Int32]
:return: Number of masked trips.
:rtype: Int32
"""
result = 0
if cutlass.const_expr(
mask_type is not MaskEnum.RESIDUAL_MASK
and mask_type is not MaskEnum.RESIDUAL_MASK_BWD
):
if cutlass.const_expr(
window_size_left is not None or window_size_right is not None
):
leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
result = max(leading_mask_end - leading_mask_begin + 1, 0)
return result
@cute.jit
def get_masked_trailing_count(
mask_type: MaskEnum,
blk_coord: cute.Coord,
tile_shape: cute.Shape,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[Int32] = None,
window_size_right: Optional[Int32] = None,
rem_count: Optional[Int32] = 0,
) -> Int32:
"""
Calculate the number of masked trips for the trailing mask.
This is used for blocks that need special handling due to masking.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param tile_shape: Shape of the tile.
:type tile_shape: cute.Shape
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Int32
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[Int32]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[Int32]
:param rem_count: Remaining count from previous calculations.
:type rem_count: Int32
:return: Number of masked trips.
:rtype: Int32
"""
result = 0
if cutlass.const_expr(
mask_type is not MaskEnum.RESIDUAL_MASK
and mask_type is not MaskEnum.RESIDUAL_MASK_BWD
):
if cutlass.const_expr(
window_size_left is not None or window_size_right is not None
):
trailing_mask_begin, trailing_mask_end = FusedMask.get_trailing_mask_id(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
if cutlass.const_expr(
trailing_mask_begin is not None and trailing_mask_end is not None
):
if trailing_mask_begin <= leading_mask_end:
result = max(trailing_mask_end - leading_mask_end, 0)
else:
result = max(trailing_mask_end - trailing_mask_begin + 1, 0)
else:
if seqlen_k % tile_shape[1] != 0:
result = 1
else:
result = 0
return result + rem_count
@cute.jit
def get_unmasked_trip_count(
mask_type: MaskEnum,
blk_coord: cute.Coord,
tile_shape: cute.Shape,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[Int32] = None,
window_size_right: Optional[Int32] = None,
) -> Int32:
"""
Calculate the number of unmasked trips for the current block.
This represents the number of trips that don't require special
masking treatment.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param blk_coord: Block coordinates.
:type blk_coord: cute.Coord
:param tile_shape: Shape of the tile.
:type tile_shape: cute.Shape
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Int32
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[Int32]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[Int32]
:return: Number of unmasked trips.
:rtype: Int32
"""
result = (
FusedMask.get_trip_count(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
- FusedMask.get_masked_leading_count(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
)
- FusedMask.get_masked_trailing_count(
mask_type,
blk_coord,
tile_shape,
seqlen_q,
seqlen_k,
window_size_left,
window_size_right,
0,
)
)
return result
@cute.jit
def apply_mask(
mask_type: MaskEnum,
acc_qk: cute.Tensor,
index_qk: cute.Tensor,
seqlen_q: Int32,
seqlen_k: Int32,
window_size_left: Optional[int] = None,
window_size_right: Optional[int] = None,
index_transform: cutlass.Constexpr = lambda index_q, index_k: (
index_q,
index_k,
),
):
"""
Apply the appropriate mask to the attention scores.
This method modifies the attention scores (acc_qk) based on the mask type
and the positions in the index tensor.
:param mask_type: Type of mask to use
:type mask_type: utils.MaskEnum
:param acc_qk: Accumulated QK attention scores tensor.
:type acc_qk: cute.Tensor
:param index_qk: Index tensor containing position information.
:type index_qk: cute.Tensor
:param seqlen_k: Key sequence length for attention computation.
:type seqlen_k: Int32
:param seqlen_q: Query sequence length for attention computation.
:type seqlen_q: Optional[int]
:param window_size_left: Left-side sliding window size for attention masking.
:type window_size_left: Optional[int]
:param window_size_right: Right-side sliding window size for attention masking.
:type window_size_right: Optional[int]
"""
tidx, tidy, tidx = cute.arch.thread_idx()
offset = 0
offset = (
seqlen_k - seqlen_q
if cutlass.const_expr(
mask_type is MaskEnum.WINDOW_MASK_INFERENCE
or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE
)
else 0
)
for i in cutlass.range_constexpr(cute.size(acc_qk)):
index_q, index_k = index_transform(*index_qk[i])
if cutlass.const_expr(
window_size_left is not None or window_size_right is not None
):
if cutlass.const_expr(window_size_left is None):
if index_q + offset + window_size_right < index_k:
acc_qk[i] = -Float32.inf
if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask
acc_qk[i] = -Float32.inf
elif cutlass.const_expr(window_size_right is None):
if index_q + offset - window_size_left > index_k:
acc_qk[i] = -Float32.inf
if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask
acc_qk[i] = -Float32.inf
else:
max_K_index = min(index_q + offset + window_size_right, seqlen_k)
min_K_index = max(0, index_q + offset - window_size_left)
if index_k > max_K_index or index_k < min_K_index:
acc_qk[i] = -Float32.inf
if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask
acc_qk[i] = -Float32.inf
if cutlass.const_expr(
mask_type == MaskEnum.RESIDUAL_MASK
or mask_type == MaskEnum.RESIDUAL_MASK_BWD
):
if index_k >= seqlen_k or index_q >= seqlen_q:
acc_qk[i] = -Float32.inf
examples/python/CuTeDSL/utils/sparse_utils.py
+457
View File
@@ -0,0 +1,457 @@
import numpy as np
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
import torch
@cute.jit
def print_tensor_dlpack(src: cute.Tensor):
print(src)
cute.print_tensor(src)
# Sparse emulation
class SparseEmulation:
def __init__(self, M: int, N: int, K: int, L: int):
self.M = M
self.N = N
self.K = K
self.L = L
@cute.jit
def __call__(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor):
"""Sparse emulation"""
num_threads = 128
grid = (cute.ceil_div(self.M, num_threads), 1, 1)
block = (num_threads, 1, 1)
self.kernel(a, b, d, e).launch(grid=grid, block=block)
return
@cute.kernel
def kernel(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor):
"""CUDA kernel to emulate sparse tensor core"""
tidx, tidy, tidz = cute.arch.thread_idx()
bidx, bidy, bidz = cute.arch.block_idx()
row_idx = tidx + bidx * self.M
meta_idx = self.K // 4 // 8
if row_idx < self.M:
# each thread process 1 row
for col in range(self.N):
# each meta_idx stands for 32 elements
for e_idx in range(meta_idx):
meta_val = e[(row_idx, e_idx)]
for k in range(8):
# each k stands for 4 elements
meta_row = (meta_val >> (k * 4)) & 0xF
idx0 = meta_row & 0x3
idx1 = (meta_row >> 2) & 0x3
# calculate the idx in b tensor which has value in A tensor
km = e_idx * 16 + k * 2
km_1 = km + 1
kn = e_idx * 32 + k * 4 + idx0
kn_1 = e_idx * 32 + k * 4 + idx1
d[row_idx, col] += a[row_idx, km] * b[col, kn]
d[row_idx, col] += a[row_idx, km_1] * b[col, kn_1]
return
# Compressor
# compress a sparse tensor to a dense tensor && generate metadata
class Compressor:
def __init__(self, M: int, K: int, L: int):
self.M = M
self.K = K
self.L = L
self.pos_map = {
0x4: [0, 1],
0x8: [0, 2],
0xC: [0, 3],
0x9: [1, 2],
0xD: [1, 3],
0xE: [2, 3],
}
@cute.jit
def _init__(self, a: cute.Tensor):
self.__init__(a.shape[0], a.shape[1], a.shape[2])
def compress(self, a, a_compressed, meta, run_on_cpu: bool):
if run_on_cpu:
if a.device.type != "cpu":
raise ValueError("a must be on cpu")
return self.__compress_on_cpu(a, a_compressed, meta)
else:
if a.device.type != "cuda":
raise ValueError("a must be on cuda")
return self.__compress_on_cuda(a, a_compressed, meta)
def __compress_on_cpu(self, a, a_compressed, meta):
"""
compress the tensor on cpu
# Convert to 4-bit metadata value
# The metadata value represents which 2 elements are non-zero
# 0x4: [1,1,0,0] - first two elements are non-zero
# 0x8: [1,0,1,0] - first and third elements are non-zero
# 0xC: [1,0,0,1] - first and fourth elements are non-zero
# 0x9: [0,1,1,0] - second and third elements are non-zero
# 0xD: [0,1,0,1] - second and fourth elements are non-zero
# 0xE: [0,0,1,1] - third and fourth elements are non-zero
# special case:
# [0,0,0,0] == [0,0,1,1]
# [1,0,0,0] == [1,0,0,1]
# [0,1,0,0] == [0,1,0,1]
# [0,0,1,0] == [0,0,1,1]
# [0,0,0,1] == [0,0,1,1]
"""
M, K = a.shape
assert a_compressed.shape == (
M,
K // 2,
), f"Expected a_compressed shape {(M, K // 2)}, got {a_compressed.shape}"
assert meta.shape == (
M,
K // 4 // 8,
), f"Expected meta shape {(M, K // 4 // 8)}, got {meta.shape}"
for m in range(M):
k_meta = 0
for k in range(0, K, 4):
chunk = a[m, k : k + 4]
non_zero_indices = torch.nonzero(chunk).squeeze()
meta_val = 0xE
if torch.equal(non_zero_indices, torch.tensor([0, 1])):
meta_val = 0x4
elif torch.equal(non_zero_indices, torch.tensor([0, 2])):
meta_val = 0x8
elif torch.equal(non_zero_indices, torch.tensor([0, 3])) or torch.equal(
non_zero_indices, torch.tensor(0)
):
meta_val = 0xC
elif torch.equal(non_zero_indices, torch.tensor([1, 2])):
meta_val = 0x9
elif torch.equal(non_zero_indices, torch.tensor([1, 3])) or torch.equal(
non_zero_indices, torch.tensor(1)
):
meta_val = 0xD
elif torch.equal(non_zero_indices, torch.tensor([2, 3])) or torch.equal(
non_zero_indices, torch.tensor(2)
):
meta_val = 0xE
elif torch.equal(non_zero_indices, torch.tensor([])) or torch.equal(
non_zero_indices, torch.tensor(3)
):
meta_val = 0xE
else:
raise ValueError(f"Invalid non-zero pattern: {non_zero_indices}")
meta_idx = k // 4 // 8
meta_bit_pos = (k // 4) % 8
if k_meta == meta_idx:
k_meta = meta_idx + 1
meta[m, meta_idx] = 0
meta[m, meta_idx] |= meta_val << (meta_bit_pos * 4)
compressed_idx = k // 2
index = self.pos_map[meta_val]
a_compressed[m, compressed_idx] = chunk[index[0]]
a_compressed[m, compressed_idx + 1] = chunk[index[1]]
def __compress_on_cuda(self, a, a_compressed, meta):
"""
compress the tensor on cuda
"""
a_tensor = from_dlpack(a)
a_compressed_tensor = from_dlpack(a_compressed)
meta_tensor = from_dlpack(meta)
self.compress_on_cuda_impl(a_tensor, a_compressed_tensor, meta_tensor)
return
@cute.jit
def compress_on_cuda_impl(
self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor
):
"""Compress the input tensor using the metadata"""
num_threads = 128
grid = (cute.ceil_div(self.M, num_threads), 1, 1)
block = (num_threads, 1, 1)
self.compressor_impl(a, a_compressed, meta).launch(grid=grid, block=block)
@cute.kernel
def compressor_impl(
self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor
):
"""CUDA kernel to compress the tensor"""
tidx, tidy, tidz = cute.arch.thread_idx()
bidx, bidy, bidz = cute.arch.block_idx()
m = a.shape[0]
k = a.shape[1]
# each thread process 1 row
row_idx = tidx + bidx * self.M
meta_idx = self.K // 4 // 8
if row_idx < self.M:
# each meta_idx stands for 32 elements
for i in range(meta_idx):
meta[row_idx, i] = 0
# each k stands for 4 elements
for j in range(8):
val = a[row_idx, i * 32 + j * 4]
val_1 = a[row_idx, i * 32 + j * 4 + 1]
val_2 = a[row_idx, i * 32 + j * 4 + 2]
val_3 = a[row_idx, i * 32 + j * 4 + 3]
value_idx = 0
value_idx_1 = 0
value_idx_2 = 0
value_idx_3 = 0
pos0 = 0
pos1 = 0
if val != 0:
value_idx = 1
pos0 = 0
if val_1 != 0:
value_idx_1 = 1
if val_2 != 0:
value_idx_2 = 1
if val_3 != 0:
value_idx_3 = 1
pos = [value_idx, value_idx_1, value_idx_2, value_idx_3]
tmp = 0
if pos == [0, 0, 0, 0]:
tmp = 0xE
pos0 = 2
pos1 = 3
elif pos == [1, 0, 0, 0]:
tmp = 0xC
pos0 = 0
pos1 = 3
elif pos == [0, 1, 0, 0]:
tmp = 0xD
pos0 = 1
pos1 = 3
elif pos == [0, 0, 1, 0]:
tmp = 0xE
pos0 = 2
pos1 = 3
elif pos == [0, 0, 0, 1]:
tmp = 0xE
pos0 = 2
pos1 = 3
elif pos == [1, 1, 0, 0]:
tmp = 0x4
pos0 = 0
pos1 = 1
elif pos == [1, 0, 1, 0]:
tmp = 0x8
pos0 = 0
pos1 = 2
elif pos == [1, 0, 0, 1]:
tmp = 0xC
pos0 = 0
pos1 = 3
elif pos == [0, 1, 1, 0]:
tmp = 0x9
pos0 = 1
pos1 = 2
elif pos == [0, 1, 0, 1]:
tmp = 0xD
pos0 = 1
pos1 = 3
elif pos == [0, 0, 1, 1]:
tmp = 0xE
pos0 = 2
pos1 = 3
# cute.printf(row_idx, cutlass.Float32(val), cutlass.Float32(val_1), cutlass.Float32(val_2), cutlass.Float32(val_3), tmp)
meta[row_idx, i] |= tmp << (j * 4)
a_compressed[row_idx, i * 16 + j * 2] = a[
row_idx, i * 32 + j * 4 + pos0
]
a_compressed[row_idx, i * 16 + j * 2 + 1] = a[
row_idx, i * 32 + j * 4 + pos1
]
return
# SparseUtils is used to generate sparse tensor
# format torch.Tensor
class SparseUtils:
#!brief: SparseUtils is used to generate sparse tensor
#!param: M: int, K: int, L: int, dtype: cutlass.DataType
def __init__(self, M: int, K: int, L: int, dtype):
self.M = M
self.K = K
self.L = L
self.dtype = dtype
self.meta_data = self._generate_meta_data_4_2()
self._use_specific_meta_data = False
#!brief: cast cutlass.DataType to torch.Tensor
def _get_type(self):
if self.dtype == cutlass.Float16:
return torch.float16
elif self.dtype == cutlass.Float32:
return torch.float32
elif self.dtype == cutlass.Int8:
return torch.int8
else:
raise ValueError(f"Unsupported dtype: {self.dtype}")
def _generate_meta_data_4_2(self):
# metadata for 4:2 sparse will in range( 4,8,9,c,d,e)
# represents
# 0: [1,1,0,0] no zero pos 00,01 -> 0100 = 4
# 1: [1,0,1,0] no zero pos 00,10 -> 1000 = 8
# 2: [1,0,0,1] no zero pos 00,11 -> 1100 = c
# 3: [0,1,1,0] no zero pos 01,10 -> 1001 = 9
# 4: [0,1,0,1] no zero pos 01,11 -> 1101 = d
# 5: [0,0,1,1] no zero pos 10,11 -> 1011 = e
meta_value = [0x4, 0x8, 0x9, 0xC, 0xD, 0xE]
# 4:2 sparse, so each chunk is 4 elements, map to 4 bits
K_NumChunk = self.K // 4
meta_data = np.random.choice(
meta_value, size=(self.M, K_NumChunk), replace=True
)
meta_data = torch.from_numpy(
np.array(meta_data).astype(np.uint8).reshape(self.M, K_NumChunk)
)
return meta_data
#!brief: pack meta data
def _pack_meta_data(self):
tmp = []
K_NumChunk = self.K // 4
for i in range(self.M):
for j in range(K_NumChunk // 8):
v = 0
for k in range(8):
vv = int(self.meta_data[i, j * 8 + k] & 0xF)
tt = vv << (k * 4)
v = v | tt
tmp.append(v)
# debug print
# print([hex(vt) for vt in tmp])
result = torch.from_numpy(
np.array(tmp).astype(np.uint32).reshape(self.M, K_NumChunk // 8)
)
return result
#!brief: use specific meta data
def use_specific_meta_data(self, meta_data: torch.Tensor = None):
if meta_data is not None:
self.meta_data = meta_data
self._use_specific_meta_data = True
#!brief: generate sparse tensor with tensor
#!param: a: torch.Tensor
#!param: run_on_cpu: bool
#!return: torch.Tensor
def generate_sparse_4_2_tensor_with_tensor(self, a, run_on_cpu):
if run_on_cpu:
if a.device.type != "cpu":
raise ValueError("a must be on cpu")
return self.__generate_sparse_tensor_cpu(a)
else:
if a.device.type != "cuda":
raise ValueError("a must be on cuda")
a_tensor = from_dlpack(a)
packed_meta_data = self._pack_meta_data()
meta_tensor = from_dlpack(packed_meta_data.cuda())
self.__generate_sparse_tensor_cuda(a_tensor, meta_tensor)
return a
#!brief: generate sparse tensor
#!param: run_on_cpu: bool
#!return: torch.Tensor
def generate_4_2_sparse_tensor(self, run_on_cpu):
dtype = self._get_type()
a = torch.empty(self.M, self.K).random_(-5, 5).to(dtype)
if run_on_cpu:
return self.generate_sparse_4_2_tensor_with_tensor(a, run_on_cpu)
else:
return self.generate_sparse_4_2_tensor_with_tensor(a.cuda(), run_on_cpu)
#!brief: generate sparse tensor on cpu
#!param: a: torch.Tensor
#!return: torch.Tensor
def __generate_sparse_tensor_cpu(self, a):
if not self._use_specific_meta_data:
for m in range(self.M):
for k in range(0, self.K, 4):
# random choose 2 zero positions
zero_indices = torch.randperm(4)[:2]
a[m, k + zero_indices[0]] = 0
a[m, k + zero_indices[1]] = 0
return a
else:
# use specific meta data
tensor_mask = []
for i in range(self.M):
for j in range(self.K // 4):
meta_val = self.meta_data[i, j]
tmp = []
if meta_val == 0x4:
tmp = [1, 1, 0, 0]
elif meta_val == 0x8:
tmp = [1, 0, 1, 0]
elif meta_val == 0xC:
tmp = [1, 0, 0, 1]
elif meta_val == 0x9:
tmp = [0, 1, 1, 0]
elif meta_val == 0xD:
tmp = [0, 1, 0, 1]
elif meta_val == 0xE:
tmp = [0, 0, 1, 1]
tensor_mask.extend(tmp)
a = torch.reshape(a, (-1,))
mask = torch.tensor(tensor_mask)
a = a * mask
a = torch.reshape(a, (self.M, self.K))
return a
@cute.jit
def __generate_sparse_tensor_cuda(self, a: cute.Tensor, meta: cute.Tensor):
"""Generate a sparse tensor from a dense tensor using metadata"""
assert a.shape[0] == self.M and a.shape[1] == self.K
assert meta.shape[0] == self.M and meta.shape[1] == self.K // 4 // 8
num_threads = 128
grid = (cute.ceil_div(self.M, num_threads), 1, 1)
block = (num_threads, 1, 1)
self.kernel(a, meta).launch(grid=grid, block=block)
@cute.kernel
def kernel(self, a: cute.Tensor, meta: cute.Tensor):
"""Apply sparsity mask to input tensor using metadata"""
tidx, tidy, tidz = cute.arch.thread_idx()
bidx, bidy, bidz = cute.arch.block_idx()
# each thread process 1 ro
row_idx = tidx + bidx * self.M
meta_idx = self.K // 4 // 8
# each thread process 1 row
if row_idx < self.M:
# iterate over each chunk(32 elements)
for i in range(meta_idx):
meta_val = meta[(row_idx, i)]
# iterate over each sparse pattern(4 elements)
for j in range(8):
meta_row = (meta_val >> (j * 4)) & 0xF
idx0 = meta_row & 0x3
idx1 = (meta_row >> 2) & 0x3
r_id0 = 0
r_id1 = 0
# r_id is the idx that value is 0
if idx0 >= 2 and idx1 >= 2:
r_id0 = 0
r_id1 = 1
elif idx0 <= 1 and idx1 <= 1:
r_id0 = 2
r_id1 = 3
else:
r_id0 = idx0 ^ 0b1
r_id1 = idx1 ^ 0b1
row_id0 = r_id0 + i * 32 + j * 4
row_id1 = r_id1 + i * 32 + j * 4
a[row_idx, row_id0] = self.dtype(0.0)
a[row_idx, row_id1] = self.dtype(0.0)
return
examples/python/CuTeDSL/utils/test_sparse_utils.py
+104
View File
@@ -0,0 +1,104 @@
import sparse_utils as su
import cutlass
import torch
from cutlass.cute.runtime import from_dlpack
import numpy as np
import pytest
@pytest.mark.L0
def test_sparse_cpu():
M = 128
N = 32
K = 32
L = 1
debug = False
# generate sparse tensor
a = torch.empty(M, K).random_(-5, 5).to(torch.float16)
sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16)
if debug:
sparse_utils.use_specific_meta_data()
a_gen_from_cpu = sparse_utils.generate_sparse_4_2_tensor_with_tensor(a, True)
# print(a_gen_from_cpu)
# generate compressed tensor and meta data
a_compressed_cpu = torch.empty(M, K // 2).to(torch.float16)
meta_data_cpu = torch.empty(M, K // 4 // 8).to(torch.uint32)
compressor = su.Compressor(M, K, L)
compressor.compress(a_gen_from_cpu, a_compressed_cpu, meta_data_cpu, True)
# # test with gemm
b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda()
d = torch.empty(M, N).zero_().to(torch.float16).cuda()
b_tensor = from_dlpack(b)
d_tensor = from_dlpack(d)
a_compressed_cpu_tensor = from_dlpack(a_compressed_cpu.cuda())
meta_data_cpu_tensor = from_dlpack(meta_data_cpu.cuda())
sparse_emulation = su.SparseEmulation(M, N, K, 1)
sparse_emulation(a_compressed_cpu_tensor, b_tensor, d_tensor, meta_data_cpu_tensor)
ref = torch.einsum("mk,nk->mn", a_gen_from_cpu.cpu(), b.cpu())
if debug:
a_ori = a_gen_from_cpu.cpu().numpy()
np.savetxt("a.txt", a_ori, fmt="%f")
a_compressed_cpu_ori = a_compressed_cpu.cpu().numpy()
np.savetxt("a_compressed_cpu.txt", a_compressed_cpu_ori, fmt="%f")
meta_data_cpu_ori = meta_data_cpu.cpu().numpy()
np.savetxt("meta_data_cpu.txt", meta_data_cpu_ori, fmt="%f")
d_ori = d.cpu().numpy()
np.savetxt("d.txt", d_ori, fmt="%f")
ref_ori = ref.cpu().numpy()
np.savetxt("ref.txt", ref_ori, fmt="%f")
torch.testing.assert_close(d.cpu(), ref)
print("cpu d == ref")
@pytest.mark.L0
def test_sparse_cuda():
M = 128
N = 32
K = 32
L = 1
debug = False
sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16)
if debug:
sparse_utils.use_specific_meta_data()
# generate sparse tensor
a = torch.empty(M, K).random_(-5, 5).to(torch.float16).cuda()
a_gen_from_cuda = sparse_utils.generate_4_2_sparse_tensor(False)
# print(a_gen_from_cuda)
# generate compressed tensor and meta data
a_compressed_cuda = torch.empty(M, K // 2).to(torch.float16).cuda()
meta_data_cuda = torch.empty(M, K // 4 // 8).to(torch.uint32).cuda()
compressor = su.Compressor(M, K, L)
compressor.compress(a_gen_from_cuda, a_compressed_cuda, meta_data_cuda, False)
# test with gemm
b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda()
d = torch.empty(M, N).zero_().to(torch.float16).cuda()
b_tensor = from_dlpack(b)
d_tensor = from_dlpack(d)
a_compressed_cuda_tensor = from_dlpack(a_compressed_cuda)
meta_data_cuda_tensor = from_dlpack(meta_data_cuda)
sparse_emulation = su.SparseEmulation(M, N, K, 1)
sparse_emulation(
a_compressed_cuda_tensor, b_tensor, d_tensor, meta_data_cuda_tensor
)
ref = torch.einsum("mk,nk->mn", a_gen_from_cuda.cpu(), b.cpu())
if debug:
a_ori = a_gen_from_cuda.cpu().numpy()
np.savetxt("a.txt", a_ori, fmt="%f")
a_compressed_cuda_ori = a_compressed_cuda.cpu().numpy()
np.savetxt("a_compressed_cuda.txt", a_compressed_cuda_ori, fmt="%f")
meta_data_cuda_ori = meta_data_cuda.cpu().numpy()
np.savetxt("meta_data_cuda.txt", meta_data_cuda_ori, fmt="%f")
d_ori = d.cpu().numpy()
np.savetxt("d.txt", d_ori, fmt="%f")
ref_ori = ref.cpu().numpy()
np.savetxt("ref.txt", ref_ori, fmt="%f")
torch.testing.assert_close(d.cpu(), ref)
print("cuda d == ref")
if __name__ == "__main__":
cutlass.cuda.initialize_cuda_context()
test_sparse_cpu()
test_sparse_cuda()